Job management and preparation for multi-hosting.

Common/shared.proto

@@ -0,0 +1,130 @@
+/**
+ * Shared messages for the Stable Diffusion RPC service.
+ */
+syntax = "proto3";
+package gay.pizza.stable.diffusion;
+
+/**
+ * Utilize a prefix of 'Sd' for Swift.
+ */
+option swift_prefix = "Sd";
+option java_multiple_files = true;
+
+/**
+ * Represents the model attention. Model attention has to do with how the model is encoded, and
+ * can determine what compute units are able to support a particular model.
+ */
+enum ModelAttention {
+  /**
+   * The model is an original attention type. It can be loaded only onto CPU & GPU compute units.
+   */
+  original = 0;
+
+  /**
+   * The model is a split-ein-sum attention type. It can be loaded onto all compute units,
+   * including the Apple Neural Engine.
+   */
+  split_ein_sum = 1;
+}
+
+/**
+ * Represents the schedulers that are used to sample images.
+ */
+enum Scheduler {
+  /**
+   * The PNDM (Pseudo numerical methods for diffusion models) scheduler.
+   */
+  pndm = 0;
+
+  /**
+   * The DPM-Solver++ scheduler.
+   */
+  dpm_solver_plus_plus = 1;
+}
+
+/**
+ * Represents a specifier for what compute units are available for ML tasks.
+ */
+enum ComputeUnits {
+  /**
+   * The CPU as a singular compute unit.
+   */
+  cpu = 0;
+
+  /**
+   * The CPU & GPU combined into a singular compute unit.
+   */
+  cpu_and_gpu = 1;
+
+  /**
+   * Allow the usage of all compute units. CoreML will decided where the model is loaded.
+   */
+  all = 2;
+
+  /**
+   * The CPU & Neural Engine combined into a singular compute unit.
+   */
+  cpu_and_neural_engine = 3;
+}
+
+/**
+ * Represents information about an available model.
+ * The primary key of a model is it's 'name' field.
+ */
+message ModelInfo {
+  /**
+   * The name of the available model. Note that within the context of a single RPC server,
+   * the name of a model is a unique identifier. This may not be true when utilizing a cluster or
+   * load balanced server, so keep that in mind.
+   */
+  string name = 1;
+
+  /**
+   * The attention of the model. Model attention determines what compute units can be used to
+   * load the model and make predictions.
+   */
+  ModelAttention attention = 2;
+
+  /**
+   * Whether the model is currently loaded onto an available compute unit.
+   */
+  bool is_loaded = 3;
+
+  /**
+   * The compute unit that the model is currently loaded into, if it is loaded to one at all.
+   * When is_loaded is false, the value of this field should be null.
+   */
+  ComputeUnits loaded_compute_units = 4;
+
+  /**
+   * The compute units that this model supports using.
+   */
+  repeated ComputeUnits supported_compute_units = 5;
+}
+
+/**
+ * Represents the format of an image.
+ */
+enum ImageFormat {
+  /**
+   * The PNG image format.
+   */
+  png = 0;
+}
+
+/**
+ * Represents an image within the Stable Diffusion context.
+ * This could be an input image for an image generation request, or it could be
+ * a generated image from the Stable Diffusion model.
+ */
+message Image {
+  /**
+   * The format of the image.
+   */
+  ImageFormat format = 1;
+
+  /**
+   * The raw data of the image, in the specified format.
+   */
+  bytes data = 2;
+}