Document API, make the implementation match the API, and update the same.

Common/StableDiffusion.proto

@@ -1,63 +1,247 @@
+/**
+ * Stable Diffusion RPC service for Apple Platforms.
+ */
 syntax = "proto3";
 package gay.pizza.stable.diffusion;
 
+/**
+ * Utilize a prefix of 'Sd' for Swift.
+ */
 option swift_prefix = "Sd";
 
-message ModelInfo {
-    string name = 1;
-    string attention = 2;
-    bool is_loaded = 3;
+/**
+ * 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;
 }
 
-message Image {
-    bytes content = 1;
-}
-
-message ListModelsRequest {}
-message ListModelsResponse {
-    repeated ModelInfo models = 1;
-}
-
-message ReloadModelsRequest {}
-message ReloadModelsResponse {}
-
+/**
+ * Represents the schedulers that are used to sample images.
+ */
 enum Scheduler {
+    /**
+     * The PNDM (Pseudo numerical methods for diffusion models) scheduler.
+     */
     pndm = 0;
-    dpmSolverPlusPlus = 1;
+
+    /**
+     * 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;
 }
 
-message LoadModelRequest {
-    string model_name = 1;
-    ComputeUnits compute_units = 2;
-    bool reduce_memory = 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;
+}
+
+/**
+ * Represents a request to list the models available on the host.
+ */
+message ListModelsRequest {}
+
+/**
+ * Represents a response to listing the models available on the host.
+ */
+message ListModelsResponse {
+    /**
+     * The available models on the Stable Diffusion server.
+     */
+    repeated ModelInfo available_models = 1;
+}
+
+/**
+ * Represents a request to load a model into a specified compute unit.
+ */
+message LoadModelRequest {
+    /**
+     * The model name to load onto the compute unit.
+     */
+    string model_name = 1;
+
+    /**
+     * The compute units to load the model onto.
+     */
+    ComputeUnits compute_units = 2;
+}
+
+/**
+ * Represents a response to loading a model.
+ */
 message LoadModelResponse {}
 
+/**
+ * The model service, for management and loading of models.
+ */
 service ModelService {
+    /**
+     * Lists the available models on the host.
+     * This will return both models that are currently loaded, and models that are not yet loaded.
+     */
     rpc ListModels(ListModelsRequest) returns (ListModelsResponse);
-    rpc ReloadModels(ReloadModelsRequest) returns (ReloadModelsResponse);
+
+    /**
+     * Loads a model onto a compute unit.
+     */
     rpc LoadModel(LoadModelRequest) returns (LoadModelResponse);
 }
 
+/**
+ * Represents a request to generate images using a loaded model.
+ */
 message GenerateImagesRequest {
+    /**
+     * The model name to use for generation.
+     * The model must be already be loaded using ModelService.LoadModel RPC method.
+     */
     string model_name = 1;
-    uint32 image_count = 2;
-    string prompt = 3;
-    string negative_prompt = 4;
+
+    /**
+     * The output format for generated images.
+     */
+    ImageFormat output_image_format = 2;
+
+    /**
+     * The number of batches of images to generate.
+     */
+    uint32 batch_count = 3;
+
+    /**
+     * The number of images inside a single batch.
+     */
+    uint32 batch_size = 4;
+
+    /**
+     * The positive textual prompt for image generation.
+     */
+    string prompt = 5;
+
+    /**
+     * The negative prompt for image generation.
+     */
+    string negative_prompt = 6;
+    
+    /**
+     * The random seed to use.
+     * Zero indicates that the seed should be random.
+     */
+    uint32 seed = 7;
 }
 
+/**
+ * Represents the response from image generation.
+ */
 message GenerateImagesResponse {
+    /**
+     * The set of generated images by the Stable Diffusion pipeline.
+     */
     repeated Image images = 1;
+    
+    /**
+     * The seeds that were used to generate the images.
+     */
+    repeated uint32 seeds = 2;
 }
 
+/**
+ * The image generation service, for generating images from loaded models.
+ */
 service ImageGenerationService {
-    rpc GenerateImage(GenerateImagesRequest) returns (GenerateImagesResponse);
+    /**
+     * Generates images using a loaded model.
+     */
+    rpc GenerateImages(GenerateImagesRequest) returns (GenerateImagesResponse);
 }