/**
 * 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;
}