From 1614cba2f86b8a308a19a21ef76d80460a3b9f4e Mon Sep 17 00:00:00 2001
From: ArEnSc <michael.chung@databites.ca>
Date: Wed, 3 Jan 2024 23:16:00 -0500
Subject: [PATCH] properties added

---
 .../image_generation_payload.rs               | 31 ++++++++++++++++---
 .../image_gen/enqueue_image_generation.rs     | 10 +++++-
 2 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs
index 9a94b3a7e9..9ef9da3481 100644
--- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs
+++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs
@@ -20,9 +20,6 @@ pub enum StableDiffusionSource {
 
 #[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
 pub struct StableDiffusionArgs {
-    #[serde(rename = "is")]
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub maybe_image_source: Option<MediaFileToken>, 
 
     #[serde(rename = "sd")]
     #[serde(skip_serializing_if = "Option::is_none")]
@@ -32,6 +29,18 @@ pub struct StableDiffusionArgs {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub maybe_lora_model_token: Option<ModelWeightToken>,
 
+    #[serde(rename = "w")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub maybe_width: Option<i32>,
+
+    #[serde(rename = "h")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub maybe_height: Option<i32>,
+
+    #[serde(rename = "s")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub maybe_sampler: Option<i32>,
+
     #[serde(rename = "p")]
     #[serde(skip_serializing_if = "Option::is_none")]
     pub maybe_prompt: Option<String>,
@@ -51,13 +60,27 @@ pub struct StableDiffusionArgs {
     #[serde(rename = "mu")]
     #[serde(skip_serializing_if = "Option::is_none")]
     pub maybe_upload_path: Option<String>,
+
+    #[serde(rename = "cf")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub maybe_cfg_scale: Option<i32>,
     
     #[serde(rename = "lu")]
     #[serde(skip_serializing_if = "Option::is_none")]
     pub maybe_lora_upload_path: Option<String>,
 
+    #[serde(rename = "sa")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub maybe_number_of_samples: Option<i64>,
+
+    #[serde(rename = "bz")]
+    pub maybe_batch_size: Option<i32>,
+
+    #[serde(rename = "bc")]
+    pub maybe_batch_count: Option<i32>,
+
     #[serde(rename = "t")]
-    pub inference_type: String,
+    pub type_of_inference: String,
 }
 
 
diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs
index 6b03b384fd..c6960a7eb4 100644
--- a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs
+++ b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs
@@ -67,6 +67,7 @@ impl Display for TypeOfInference {
     }
 }
 
+
 #[derive(Deserialize,ToSchema)]
 pub struct EnqueueImageGenRequest {
     uuid_idempotency_token: String,
@@ -78,9 +79,15 @@ pub struct EnqueueImageGenRequest {
     maybe_a_prompt: Option<String>,
     maybe_n_prompt: Option<String>,
     maybe_seed: Option<i64>,
+    maybe_width: Option<i32>,
+    maybe_height: Option<i32>,
+    maybe_sampler: Option<String>,
     maybe_upload_path: Option<String>,
     maybe_lora_upload_path: Option<String>,
-    maybe_cfg_scale: Option<f32>
+    maybe_cfg_scale: Option<i32>,
+    maybe_number_of_samples: Option<i64>,
+    maybe_batch_size: Option<i32>,
+    maybe_batch_count: Option<i32>,
 }
 
 #[derive(Serialize,ToSchema)]
@@ -236,6 +243,7 @@ pub async fn enqueue_image_generation_request(
     let image_source_token = MediaFileToken(request.maybe_image_source.clone().unwrap_or_default());
     let sd_weight_token = ModelWeightToken(request.maybe_sd_model_token.clone().unwrap_or_default());
     let lora_token = ModelWeightToken(request.maybe_lora_model_token.clone().unwrap_or_default());
+    
 
     let a_prompt = request.maybe_a_prompt.clone().unwrap_or_default();
     let n_prompt = request.maybe_n_prompt.clone().unwrap_or_default();