From 55367e8723c4cfdd1120ffae1e9d27b30622011d Mon Sep 17 00:00:00 2001
From: saltacc <alotofcatz@gmail.com>
Date: Mon, 22 Jan 2024 10:37:19 -0800
Subject: [PATCH 1/5] Resolved merge conflict by keeping my changes in
 process_job.rs

---
 .../src/seeding/model_weights.rs              |  9 +++++----
 .../remote_file_manager/media_descriptor.rs   | 20 +++++++++++++++++++
 .../remote_cloud_bucket_details.rs            | 15 ++++++++++++++
 .../remote_file_manager/weights_descriptor.rs |  2 +-
 .../workflow_payload.rs                       |  9 +++++++++
 .../workflow/comfy_ui/validate_job.rs         | 13 +++++++++++-
 .../endpoints/workflows/enqueue_comfy_ui.rs   |  5 +++++
 7 files changed, 67 insertions(+), 6 deletions(-)

diff --git a/crates/cli/dev-database-seed/src/seeding/model_weights.rs b/crates/cli/dev-database-seed/src/seeding/model_weights.rs
index 171d3a5942..987fdea2bf 100644
--- a/crates/cli/dev-database-seed/src/seeding/model_weights.rs
+++ b/crates/cli/dev-database-seed/src/seeding/model_weights.rs
@@ -686,6 +686,7 @@ Noosphere by skumerz + dalcefoPainting + 饭特稀V08 by zhazhahui345 + GhostMix
 pub async fn seed_workflows_for_testing_inference(mysql_pool: &Pool<MySql>, user_token: UserToken) -> AnyhowResult<()>{
     let model_weight_token1 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously();
     let model_weight_token2 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously();
+    let model_weight_token3 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously();
 
     let mut path_to_comfy = get_seed_tool_data_root();
     path_to_comfy.push("models/workflows/comfyui/workflow_api.json");
@@ -1061,9 +1062,9 @@ pub async fn seed_weights(mysql_pool: &Pool<MySql>) -> AnyhowResult<()> {
     // original_seed_weights(mysql_pool,user_token).await?;
     // seed_weights_for_user_token(mysql_pool, user_token).await?;
     // seed_weights_for_paging(mysql_pool,user_token).await?;
-    seed_weights_for_testing_inference(mysql_pool,user_token.clone()).await?;
-    println!("TESTING DOWLOAD");
-    test_seed_weights_files().await?;
-    seed_workflows_for_testing_inference(mysql_pool,user_token.clone()).await?;
+    // seed_weights_for_testing_inference(mysql_pool,user_token.clone()).await?;
+    // println!("TESTING DOWLOAD");
+    // test_seed_weights_files().await?;
+    // seed_workflows_for_testing_inference(mysql_pool,user_token.clone()).await?;
     Ok(())
 }
diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/media_descriptor.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/media_descriptor.rs
index d969193271..3e0c949642 100644
--- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/media_descriptor.rs
+++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/media_descriptor.rs
@@ -1,4 +1,5 @@
 use crate::remote_file_manager::file_descriptor::FileDescriptor;
+
 // TODO ASK BRANDON WHAT path should this be?
 const MEDIA_FILE_DIRECTORY: &str = "/media";
 
@@ -20,3 +21,22 @@ impl FileDescriptor for MediaImagePngDescriptor {
         true
     }
 }
+
+pub struct MediaVideoMp4Descriptor;
+impl FileDescriptor for MediaVideoMp4Descriptor {
+    fn remote_directory_path(&self) -> &str {
+        MEDIA_FILE_DIRECTORY
+    }
+
+    fn get_suffix(&self) -> String {
+        "mp4".to_string()
+    }
+
+    fn get_prefix(&self) -> String {
+        "video".to_string()
+    }
+
+    fn is_public(&self) -> bool {
+        true
+    }
+}
diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs
index 9ea88495a7..c99b7391a8 100644
--- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs
+++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs
@@ -1,3 +1,4 @@
+use crate::remote_file_manager::media_descriptor;
 use crate::remote_file_manager::weights_descriptor::{self};
 use super::file_descriptor::FileDescriptor;
 
@@ -28,6 +29,7 @@ impl RemoteCloudBucketDetails {
 
     pub fn file_descriptor_from_bucket_details(&self) -> Box<dyn FileDescriptor> {
         match self.prefix.as_str() {
+            // Weights
             "loRA" => Box::new(weights_descriptor::WeightsLoRADescriptor {}),
             "sd15" => Box::new(weights_descriptor::WeightsSD15Descriptor {}),
             "sdxl" => Box::new(weights_descriptor::WeightsSDXLDescriptor {}),
@@ -41,6 +43,19 @@ impl RemoteCloudBucketDetails {
             },
             "svc" => Box::new(weights_descriptor::WeightsSVCDescriptor {}),
             "workflow" => Box::new(weights_descriptor::WeightsWorkflowDescriptor {}),
+            // Media
+            "video" => {
+                match self.suffix.as_str() {
+                    "mp4" => Box::new(media_descriptor::MediaVideoMp4Descriptor {}),
+                    _ => panic!("Unknown suffix: {}",self.suffix)
+                }
+            },
+            "image" => {
+                match self.suffix.as_str() {
+                    "png" => Box::new(media_descriptor::MediaImagePngDescriptor {}),
+                    _ => panic!("Unknown suffix: {}",self.suffix)
+                }
+            },
             _ => panic!("Unknown prefix: {}", self.prefix)
         }
     }
diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs
index 0bbbfac9ce..b5072ffa96 100644
--- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs
+++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs
@@ -189,4 +189,4 @@ impl FileDescriptor for WeightsWorkflowDescriptor {
     fn is_public(&self) -> bool {
         true
     }
-}
\ No newline at end of file
+}
diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs
index 8f662a22b3..d60c9b1432 100644
--- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs
+++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs
@@ -1,4 +1,5 @@
 use std::collections::HashMap;
+use tokens::tokens::media_files::MediaFileToken;
 use tokens::tokens::model_weights::ModelWeightToken;
 
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
@@ -27,4 +28,12 @@ pub struct WorkflowArgs {
     #[serde(rename = "json_modifications")]
     #[serde(skip_serializing_if = "Option::is_none")]
     pub maybe_json_modifications: Option<HashMap<String, NewValue>>,
+
+    #[serde(rename = "in")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub maybe_input_file: Option<MediaFileToken>,
+
+    #[serde(rename = "out")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub maybe_output_path: Option<String>,
 }
diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/validate_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/validate_job.rs
index 989a07e813..a82046ef20 100644
--- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/validate_job.rs
+++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/validate_job.rs
@@ -4,16 +4,18 @@ use anyhow::anyhow;
 use mysql_queries::payloads::generic_inference_args::generic_inference_args::{InferenceCategoryAbbreviated, PolymorphicInferenceArgs};
 use mysql_queries::payloads::generic_inference_args::workflow_payload::{NewValue};
 use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob;
+use tokens::tokens::media_files::MediaFileToken;
 use tokens::tokens::model_weights::ModelWeightToken;
 
 use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
 
 pub struct JobArgs<'a> {
     pub workflow_source: &'a ModelWeightToken,
+    pub output_path: &'a String,
     pub maybe_json_modifications: &'a Option<HashMap<String, NewValue>>,
     pub maybe_sd_model: &'a Option<ModelWeightToken>,
     pub maybe_lora_model: &'a Option<ModelWeightToken>,
-
+    pub maybe_input_file: &'a Option<MediaFileToken>,
 }
 
 pub fn validate_job(job: &AvailableInferenceJob) -> Result<JobArgs, ProcessSingleJobError> {
@@ -58,11 +60,20 @@ pub fn validate_job(job: &AvailableInferenceJob) -> Result<JobArgs, ProcessSingl
         }
     };
 
+    let output_path = match &inference_args.maybe_output_path {
+        Some(args) => args,
+        None => {
+            return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("No output file provided!")));
+        }
+    };
+
 
     Ok(JobArgs {
         workflow_source,
+        output_path,
         maybe_sd_model: &inference_args.maybe_sd_model,
         maybe_lora_model: &inference_args.maybe_lora_model,
         maybe_json_modifications: &inference_args.maybe_json_modifications,
+        maybe_input_file: &inference_args.maybe_input_file,
     })
 }
diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs
index 02d8694a46..a796f1b2d4 100644
--- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs
+++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs
@@ -19,6 +19,7 @@ use mysql_queries::payloads::generic_inference_args::generic_inference_args::{Ge
 use mysql_queries::payloads::generic_inference_args::workflow_payload::{NewValue, WorkflowArgs};
 use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{insert_generic_inference_job, InsertGenericInferenceArgs};
 use tokens::tokens::generic_inference_jobs::InferenceJobToken;
+use tokens::tokens::media_files::MediaFileToken;
 use tokens::tokens::model_weights::ModelWeightToken;
 use tokens::tokens::users::UserToken;
 
@@ -44,6 +45,8 @@ pub struct EnqueueComfyRequest {
     maybe_negative_prompt: Option<String>,
     maybe_json_modifications: Option<HashMap<String, NewValue>>,
     maybe_workflow_config: Option<ModelWeightToken>,
+    maybe_input_file: Option<MediaFileToken>,
+    maybe_output_path: Option<String>,
 
     creator_set_visibility: Option<Visibility>,
 }
@@ -178,6 +181,8 @@ pub async fn enqueue_comfy_ui_handler(
         maybe_lora_model: request.maybe_lora_model.clone(),
         maybe_json_modifications: request.maybe_json_modifications.clone(),
         maybe_workflow_config: request.maybe_workflow_config.clone(),
+        maybe_input_file: request.maybe_input_file.clone(),
+        maybe_output_path: request.maybe_output_path.clone(),
     };
 
     info!("Creating ComfyUI job record...");

From 84a2aa610806f65d83232e411fc60932190de524 Mon Sep 17 00:00:00 2001
From: saltacc <alotofcatz@gmail.com>
Date: Tue, 23 Jan 2024 12:58:09 -0800
Subject: [PATCH 2/5] use input file, fix ext and prefix

---
 .../remote_cloud_bucket_details.rs            |  1 +
 .../workflow/comfy_ui/process_job.rs          | 94 +++++++++++++------
 2 files changed, 64 insertions(+), 31 deletions(-)

diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs
index 49bbc02682..077b0265fc 100644
--- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs
+++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs
@@ -1,3 +1,4 @@
+use crate::remote_file_manager::media_descriptor;
 use crate::remote_file_manager::weights_descriptor::{self};
 use super::file_descriptor::FileDescriptor;
 
diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs
index 8b7d90070c..65dcfeda83 100644
--- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs
+++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs
@@ -23,6 +23,7 @@ use mysql_queries::payloads::generic_inference_args::generic_inference_args::Pol
 use mysql_queries::payloads::generic_inference_args::workflow_payload::NewValue;
 use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob;
 use mysql_queries::queries::media_files::create::insert_media_file_from_comfy_ui::{insert_media_file_from_comfy_ui, InsertArgs};
+use mysql_queries::queries::media_files::get_media_file::get_media_file;
 use mysql_queries::queries::model_weights::get::get_weight::get_weight_by_token;
 
 use crate::job::job_loop::job_success_result::{JobSuccessResult, ResultEntity};
@@ -31,8 +32,6 @@ use crate::job::job_types::workflow::comfy_ui::comfy_ui_inference_command::Infer
 use crate::job::job_types::workflow::comfy_ui::validate_job::validate_job;
 use crate::job_dependencies::JobDependencies;
 
-const BUCKET_FILE_PREFIX: &str = "fakeyou_";
-
 pub struct ComfyProcessJobArgs<'a> {
     pub job_dependencies: &'a JobDependencies,
     pub job: &'a AvailableInferenceJob,
@@ -94,7 +93,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
                 sd_model,
                 false,
                 &deps.db.mysql_pool
-            ).await?.unwrap();
+            ).await?.ok_or_else(|| ProcessSingleJobError::Other(anyhow!("SD model not found")))?;
 
             let bucket_details = RemoteCloudBucketDetails {
                 object_hash: retrieved_sd_record.public_bucket_hash,
@@ -106,6 +105,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
             remote_cloud_file_client.download_file(bucket_details, sd_path.clone()).await?;
             maybe_sd_path = Some(sd_path.parse().unwrap());
             info!("Downloaded SD model to {:?}", sd_path);
+            maybe_sd_path = Some(sd_path.parse().unwrap());
         }
         None => {}
     }
@@ -118,7 +118,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
                 lora_model,
                 false,
                 &deps.db.mysql_pool
-            ).await?.unwrap();
+            ).await?.ok_or_else(|| ProcessSingleJobError::Other(anyhow!("Lora model not found")))?;
             let bucket_details = RemoteCloudBucketDetails {
                 object_hash: retrieved_lora_record.public_bucket_hash,
                 prefix: retrieved_lora_record.maybe_public_bucket_prefix.unwrap(),
@@ -130,6 +130,31 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
             remote_cloud_file_client.download_file(bucket_details, lora_path.clone()).await?;
             maybe_lora_path = Some(lora_path.parse().unwrap());
             info!("Downloaded Lora model to {:?}", lora_path);
+            maybe_lora_path = Some(lora_path.parse().unwrap());
+        }
+        None => {}
+    }
+    // Download Input file if specified
+    let mut maybe_input_path: Option<PathBuf> = None;
+    match job_args.maybe_input_file {
+        Some(input_file) => {
+            let input_dir = root_comfy_path.join("input");
+            let retrieved_input_record =  get_media_file(
+                input_file,
+                false,
+                &deps.db.mysql_pool
+            ).await?.ok_or_else(|| ProcessSingleJobError::Other(anyhow!("Input file not found")))?;
+            let bucket_details = RemoteCloudBucketDetails {
+                object_hash: retrieved_input_record.public_bucket_directory_hash,
+                prefix: retrieved_input_record.maybe_public_bucket_prefix.unwrap(),
+                suffix: retrieved_input_record.maybe_public_bucket_extension.clone().unwrap(),
+            };
+            // Download to "input.EXTENSION"
+            let input_filename = format!("input.{}", retrieved_input_record.maybe_public_bucket_extension.unwrap());
+            let input_path = input_dir.join(input_filename).to_str().unwrap().to_string();
+            remote_cloud_file_client.download_file(bucket_details, input_path.clone()).await?;
+            info!("Downloaded input file to {:?}", input_path);
+            maybe_input_path = Some(input_path.parse().unwrap());
         }
         None => {}
     }
@@ -222,49 +247,46 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         if let Some(lora_path) = maybe_lora_path {
             safe_delete_temp_file(&lora_path);
         }
+        if let Some(input_path) = maybe_input_path {
+            safe_delete_temp_file(&input_path);
+        }
 
         return Err(error);
     }
 
     // ==================== GET OUTPUT FILE ======================== //
-
-    // take the latest file in the output directory
-    let output_dir = root_comfy_path.join("output");
-    let mut entries: Vec<PathBuf> = std::fs::read_dir(output_dir).unwrap()
-        .map(|res| res.map(|e| e.path()))
-        .collect::<Result<Vec<_>, std::io::Error>>()
-        .unwrap();
-
-    entries.sort_by_key(|path| std::fs::metadata(path).unwrap().modified().unwrap());
-
+    let mut output_file = root_comfy_path.join("output");
+    output_file = output_file.join(job_args.output_path);
 
     // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== //
-
-    // check if entries is empty
-    if entries.is_empty() {
-        return Err(ProcessSingleJobError::Other(anyhow!("No output files found")));
-    }
-    let finished_file = entries.last().unwrap();
-    info!("Checking that file exists: {:?} ...", finished_file);
-    check_file_exists(&finished_file).map_err(|e| ProcessSingleJobError::Other(e))?;
+    info!("Checking that file exists: {:?} ...", output_file);
+    check_file_exists(&output_file).map_err(|e| ProcessSingleJobError::Other(e))?;
 
     info!("Interrogating result file size ...");
 
-    let file_size_bytes = file_size(&finished_file)
+    let file_size_bytes = file_size(&output_file)
         .map_err(|err| ProcessSingleJobError::Other(err))?;
 
     info!("Interrogating result mimetype ...");
 
-    let mimetype = get_mimetype_for_file(&finished_file)
+    let mimetype = get_mimetype_for_file(&output_file)
         .map_err(|err| ProcessSingleJobError::from_io_error(err))?
         .map(|mime| mime.to_string())
         .ok_or(ProcessSingleJobError::Other(anyhow!("Mimetype could not be determined")))?;
 
     // create ext from mimetype
     let ext = match mimetype.as_str() {
-        "video/mp4" => ".mp4",
-        "image/png" => ".png",
-        "image/jpeg" => ".jpg",
+        "video/mp4" => "mp4",
+        "image/png" => "png",
+        "image/jpeg" => "jpg",
+        _ => return Err(ProcessSingleJobError::Other(anyhow!("Mimetype not supported: {}", mimetype))),
+    };
+
+    // create prefix from mimetype
+    let prefix = match mimetype.as_str() {
+        "video/mp4" => "video",
+        "image/png" => "image",
+        "image/jpeg" => "image",
         _ => return Err(ProcessSingleJobError::Other(anyhow!("Mimetype not supported: {}", mimetype))),
     };
 
@@ -278,7 +300,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
 
     info!("Calculating sha256...");
 
-    let file_checksum = sha256_hash_file(&finished_file)
+    let file_checksum = sha256_hash_file(&output_file)
         .map_err(|err| {
             ProcessSingleJobError::Other(anyhow!("Error hashing file: {:?}", err))
         })?;
@@ -289,7 +311,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         .map_err(|e| ProcessSingleJobError::Other(e))?;
 
     let result_bucket_location = MediaFileBucketPath::generate_new(
-        Some(BUCKET_FILE_PREFIX),
+        Some(prefix),
         Some(ext));
 
     let result_bucket_object_pathbuf = result_bucket_location.to_full_object_pathbuf();
@@ -300,7 +322,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
 
     args.job_dependencies.buckets.public_bucket_client.upload_filename_with_content_type(
         &result_bucket_object_pathbuf,
-        &finished_file,
+        &output_file,
         &mimetype) // TODO: We should check the mimetype to make sure bad payloads can't get uploaded
         .await
         .map_err(|e| ProcessSingleJobError::Other(e))?;
@@ -308,6 +330,16 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
     // ==================== DELETE TEMP FILES ==================== //
 
     safe_delete_temp_file(&stderr_output_file);
+    safe_delete_temp_file(&workflow_path);
+    if let Some(sd_path) = maybe_sd_path {
+        safe_delete_temp_file(&sd_path);
+    }
+    if let Some(lora_path) = maybe_lora_path {
+        safe_delete_temp_file(&lora_path);
+    }
+    if let Some(input_path) = maybe_input_path {
+        safe_delete_temp_file(&input_path);
+    }
 
     // delete all files in output directory
     let output_dir = root_comfy_path.join("output");
@@ -333,7 +365,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         file_size_bytes,
         sha256_checksum: &file_checksum,
         public_bucket_directory_hash: result_bucket_location.get_object_hash(),
-        maybe_public_bucket_prefix: Some(BUCKET_FILE_PREFIX),
+        maybe_public_bucket_prefix: Some(prefix),
         maybe_public_bucket_extension: Some(ext),
         is_on_prem: args.job_dependencies.job.info.container.is_on_prem,
         worker_hostname: &args.job_dependencies.job.info.container.hostname,

From aa719ba56ffdacce986c49a459564b50d5954bd7 Mon Sep 17 00:00:00 2001
From: saltacc <alotofcatz@gmail.com>
Date: Tue, 30 Jan 2024 11:55:30 -0800
Subject: [PATCH 3/5] download requirements, fixes

---
 Cargo.lock                                    |   1 +
 crates/cli/dev-database-seed/src/main.rs      |  15 +-
 .../src/seeding/media_files.rs                |   2 +-
 .../src/seeding/model_weights.rs              | 141 ++++++++++++------
 .../remote_cloud_bucket_details.rs            |   7 +-
 .../remote_file_manager/weights_descriptor.rs |  24 +++
 crates/service/job/inference_job/Cargo.toml   |   3 +-
 .../comfy_ui/comfy_ui_dependencies.rs         |  30 ++++
 .../comfy_ui/comfy_ui_inference_command.rs    |   3 +
 .../workflow/comfy_ui/process_job.rs          |  85 ++++++++---
 10 files changed, 238 insertions(+), 73 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index d31846707b..4c7fb5d4f4 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3110,6 +3110,7 @@ dependencies = [
  "tokens",
  "tokio",
  "tts_common",
+ "walkdir",
 ]
 
 [[package]]
diff --git a/crates/cli/dev-database-seed/src/main.rs b/crates/cli/dev-database-seed/src/main.rs
index 6ab4c2710c..e8ba82d32c 100644
--- a/crates/cli/dev-database-seed/src/main.rs
+++ b/crates/cli/dev-database-seed/src/main.rs
@@ -8,11 +8,6 @@ use errors::AnyhowResult;
 use crate::bucket_clients::get_bucket_clients;
 use crate::cli_args::parse_cli_args;
 use crate::seeding::model_weights::seed_weights;
-use crate::seeding::media_files::seed_media_seedtool;
-use crate::seeding::seed_media_files::seed_media_files;
-use crate::seeding::tts_tacotron2::seed_tts_tacotron2;
-use crate::seeding::users::seed_user_accounts;
-use crate::seeding::voice_conversion::seed_voice_conversion;
 
 pub mod bucket_clients;
 
@@ -52,13 +47,13 @@ pub async fn main() -> AnyhowResult<()> {
 //    maybe_elasticsearch = Some(get_elasticsearch_client());
 //  }
 
-  seed_user_accounts(&pool).await?;
-  seed_media_files(&pool, maybe_bucket_clients.as_ref()).await?;
+  // seed_user_accounts(&pool).await?;
+  // seed_media_files(&pool, maybe_bucket_clients.as_ref()).await?;
   //seed_zero_shot_tts(&pool, maybe_bucket_clients.as_ref()).await?;
-  seed_voice_conversion(&pool).await?;
+  // seed_voice_conversion(&pool).await?;
   seed_weights(&pool).await?;
-  seed_media_seedtool(&pool).await?;
-  seed_tts_tacotron2(&pool, maybe_bucket_clients.as_ref()).await?;
+  // seed_media_seedtool(&pool).await?;
+  // seed_tts_tacotron2(&pool, maybe_bucket_clients.as_ref()).await?;
   
   // should seed the weights with a few files for hanashi
   //seed_weights_files(&pool, maybe_bucket_clients.as_ref()).await?;
diff --git a/crates/cli/dev-database-seed/src/seeding/media_files.rs b/crates/cli/dev-database-seed/src/seeding/media_files.rs
index d9df7b7c3a..9d0caad288 100644
--- a/crates/cli/dev-database-seed/src/seeding/media_files.rs
+++ b/crates/cli/dev-database-seed/src/seeding/media_files.rs
@@ -43,7 +43,7 @@ pub async fn seed_test_videos(mysql_pool: &Pool<MySql>, user_token: UserToken) -
 }
 
 pub async fn seed_media_seedtool(mysql_pool: &Pool<MySql>) -> AnyhowResult<()> {
-    info!("Seeding weights...");
+    info!("Seeding media...");
 
     let user_token = match get_user_token_by_username(HANASHI_USERNAME, mysql_pool).await? {
         None => {
diff --git a/crates/cli/dev-database-seed/src/seeding/model_weights.rs b/crates/cli/dev-database-seed/src/seeding/model_weights.rs
index a935cc245f..d8938a2ad2 100644
--- a/crates/cli/dev-database-seed/src/seeding/model_weights.rs
+++ b/crates/cli/dev-database-seed/src/seeding/model_weights.rs
@@ -6,7 +6,7 @@ use sqlx::{MySql, Pool};
 
 use cloud_storage::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails;
 use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient;
-use cloud_storage::remote_file_manager::weights_descriptor::{WeightsLoRADescriptor, WeightsSD15Descriptor, WeightsWorkflowDescriptor};
+use cloud_storage::remote_file_manager::weights_descriptor::{WeightsLoRADescriptor, WeightsSD15Descriptor, WeightsSD15CkptDescriptor};
 use enums::by_table::model_weights::{
     weights_category::WeightsCategory,
     weights_types::WeightsType,
@@ -687,63 +687,120 @@ pub async fn seed_workflows_for_testing_inference(mysql_pool: &Pool<MySql>, user
     let model_weight_token1 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously();
     let model_weight_token2 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously();
     let model_weight_token3 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously();
+    let model_weight_token4 = ModelWeightToken::new_from_str("weight_n8sz47gmfw2zx02snrbz88ns9");
 
     let mut path_to_comfy = get_seed_tool_data_root();
     path_to_comfy.push("models/workflows/comfyui/workflow_api.json");
     let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?;
-    let comfy_weights_descriptor = Box::new(WeightsWorkflowDescriptor {});
-    let metadata1 = remote_cloud_file_client.upload_file(comfy_weights_descriptor, path_to_comfy.as_path().to_str().unwrap()).await?;
+    // let comfy_weights_descriptor = Box::new(WeightsWorkflowDescriptor {});
+    // let metadata1 = remote_cloud_file_client.upload_file(comfy_weights_descriptor, path_to_comfy.as_path().to_str().unwrap()).await?;
 
-    let mut path_to_comfy2 = get_seed_tool_data_root();
-    path_to_comfy2.push("models/workflows/comfyui/majicmixRealistic_v7.safetensors");
-    let comfy_weights_descriptor2 = Box::new(WeightsSD15Descriptor {});
-    let metadata2 = remote_cloud_file_client.upload_file(comfy_weights_descriptor2, path_to_comfy2.as_path().to_str().unwrap()).await?;
+    // let mut path_to_comfy2 = get_seed_tool_data_root();
+    // path_to_comfy2.push("models/workflows/comfyui/majicmixRealistic_v7.safetensors");
+    // let comfy_weights_descriptor2 = Box::new(WeightsSD15Descriptor {});
+    // let metadata2 = remote_cloud_file_client.upload_file(comfy_weights_descriptor2, path_to_comfy2.as_path().to_str().unwrap()).await?;
 
-    let weights1 = CreateModelWeightsArgs {
-        token: &model_weight_token1, // replace with actual ModelWeightToken
+    // let mut path_to_comfy3 = get_seed_tool_data_root();
+    // path_to_comfy3.push("models/workflows/comfyui/yae_video_test.json");
+    // let comfy_weights_descriptor3 = Box::new(WeightsWorkflowDescriptor {});
+    // let metadata3 = remote_cloud_file_client.upload_file(comfy_weights_descriptor3, path_to_comfy3.as_path().to_str().unwrap()).await?;
+
+    let mut path_to_comfy4 = get_seed_tool_data_root();
+    path_to_comfy4.push("models/workflows/comfyui/v3_sd15_mm.ckpt");
+    let comfy_weights_descriptor4 = Box::new(WeightsSD15CkptDescriptor {});
+    let metadata4 = remote_cloud_file_client.upload_file(comfy_weights_descriptor4, path_to_comfy4.as_path().to_str().unwrap()).await?;
+
+    // let weights1 = CreateModelWeightsArgs {
+    //     token: &model_weight_token1, // replace with actual ModelWeightToken
+    //     weights_type: WeightsType::ComfyUi, // replace with actual WeightsType
+    //     weights_category: WeightsCategory::WorkflowConfig, // replace with actual WeightsCategory
+    //     title: "comfy-test-workflow".to_string(),
+    //     maybe_description_markdown: Some("Test workflow for ComfyUI".to_string()),
+    //     maybe_description_rendered_html: Some("<p>Description</p>".to_string()),
+    //     creator_user_token: Some(&user_token), // replace with actual UserToken
+    //     creator_ip_address: "192.168.1.1",
+    //     creator_set_visibility: Visibility::Public,
+    //     maybe_last_update_user_token: Some("Last Update User Token".to_string()),
+    //     original_download_url: Some("https://github.com/comfyanonymous/ComfyUI".to_string()),
+    //     original_filename: Some("test-workflow.json".to_string()),
+    //     file_size_bytes: metadata1.file_size_bytes as i32,
+    //     file_checksum_sha2: metadata1.sha256_checksum.to_string(),
+    //     public_bucket_hash: metadata1.bucket_details.clone().unwrap().object_hash,
+    //     maybe_public_bucket_prefix: Some(metadata1.bucket_details.clone().unwrap().prefix),
+    //     maybe_public_bucket_extension: Some(metadata1.bucket_details.clone().unwrap().suffix),
+    //     version: 1,
+    //     mysql_pool: &mysql_pool, // replace with actual MySqlPool
+    // };
+    // let weights2 = CreateModelWeightsArgs {
+    //     token: &model_weight_token2, // replace with actual ModelWeightToken
+    //     weights_type: WeightsType::ComfyUi, // replace with actual WeightsType
+    //     weights_category: WeightsCategory::WorkflowConfig, // replace with actual WeightsCategory
+    //     title: "v1-5-pruned-emaonly".to_string(),
+    //     maybe_description_markdown: Some("Test model for ComfyUI".to_string()),
+    //     maybe_description_rendered_html: Some("<p>Description</p>".to_string()),
+    //     creator_user_token: Some(&user_token), // replace with actual UserToken
+    //     creator_ip_address: "192.168.1.1",
+    //     creator_set_visibility: Visibility::Public,
+    //     maybe_last_update_user_token: Some("Last Update User Token".to_string()),
+    //     original_download_url: Some("https://huggingface.co/runwayml/stable-diffusion-v1-5".to_string()),
+    //     original_filename: Some("v1-5-pruned-emaonly.ckpt".to_string()),
+    //     file_size_bytes: metadata2.file_size_bytes as i32,
+    //     file_checksum_sha2: metadata2.sha256_checksum.to_string(),
+    //     public_bucket_hash: metadata2.bucket_details.clone().unwrap().object_hash,
+    //     maybe_public_bucket_prefix: Some(metadata2.bucket_details.clone().unwrap().prefix),
+    //     maybe_public_bucket_extension: Some(metadata2.bucket_details.clone().unwrap().suffix),
+    //     version: 1,
+    //     mysql_pool: &mysql_pool, // replace with actual MySqlPool
+    // };
+
+    // let weights3 = CreateModelWeightsArgs {
+    //     token: &model_weight_token3, // replace with actual ModelWeightToken
+    //     weights_type: WeightsType::ComfyUi, // replace with actual WeightsType
+    //     weights_category: WeightsCategory::WorkflowConfig, // replace with actual WeightsCategory
+    //     title: "yae-video-test".to_string(),
+    //     maybe_description_markdown: Some("Test workflow for ComfyUI".to_string()),
+    //     maybe_description_rendered_html: Some("<p>Description</p>".to_string()),
+    //     creator_user_token: Some(&user_token), // replace with actual UserToken
+    //     creator_ip_address: "192.168.1.1",
+    //     creator_set_visibility: Visibility::Public,
+    //     maybe_last_update_user_token: Some("Last Update User Token".to_string()),
+    //     original_download_url: Some("https://github.com/comfyanonymous/ComfyUI".to_string()),
+    //     original_filename: Some("yae-video-test.json".to_string()),
+    //     file_size_bytes: metadata3.file_size_bytes as i32,
+    //     file_checksum_sha2: metadata3.sha256_checksum.to_string(),
+    //     public_bucket_hash: metadata3.bucket_details.clone().unwrap().object_hash,
+    //     maybe_public_bucket_prefix: Some(metadata3.bucket_details.clone().unwrap().prefix),
+    //     maybe_public_bucket_extension: Some(metadata3.bucket_details.clone().unwrap().suffix),
+    //     version: 1,
+    //     mysql_pool: &mysql_pool, // replace with actual MySqlPool
+    // };
+
+    let weights4 = CreateModelWeightsArgs {
+        token: &model_weight_token4, // replace with actual ModelWeightToken
         weights_type: WeightsType::ComfyUi, // replace with actual WeightsType
         weights_category: WeightsCategory::WorkflowConfig, // replace with actual WeightsCategory
-        title: "comfy-test-workflow".to_string(),
-        maybe_description_markdown: Some("Test workflow for ComfyUI".to_string()),
+        title: "v3_sd15_mm".to_string(),
+        maybe_description_markdown: Some("ComfyUI motion module".to_string()),
         maybe_description_rendered_html: Some("<p>Description</p>".to_string()),
         creator_user_token: Some(&user_token), // replace with actual UserToken
         creator_ip_address: "192.168.1.1",
         creator_set_visibility: Visibility::Public,
         maybe_last_update_user_token: Some("Last Update User Token".to_string()),
-        original_download_url: Some("https://github.com/comfyanonymous/ComfyUI".to_string()),
-        original_filename: Some("test-workflow.json".to_string()),
-        file_size_bytes: metadata1.file_size_bytes as i32,
-        file_checksum_sha2: metadata1.sha256_checksum.to_string(),
-        public_bucket_hash: metadata1.bucket_details.clone().unwrap().object_hash,
-        maybe_public_bucket_prefix: Some(metadata1.bucket_details.clone().unwrap().prefix),
-        maybe_public_bucket_extension: Some(metadata1.bucket_details.clone().unwrap().suffix),
-        version: 1,
-        mysql_pool: &mysql_pool, // replace with actual MySqlPool
-    };
-    let weights2 = CreateModelWeightsArgs {
-        token: &model_weight_token2, // replace with actual ModelWeightToken
-        weights_type: WeightsType::ComfyUi, // replace with actual WeightsType
-        weights_category: WeightsCategory::WorkflowConfig, // replace with actual WeightsCategory
-        title: "v1-5-pruned-emaonly".to_string(),
-        maybe_description_markdown: Some("Test model for ComfyUI".to_string()),
-        maybe_description_rendered_html: Some("<p>Description</p>".to_string()),
-        creator_user_token: Some(&user_token), // replace with actual UserToken
-        creator_ip_address: "192.168.1.1",
-        creator_set_visibility: Visibility::Public,
-        maybe_last_update_user_token: Some("Last Update User Token".to_string()),
-        original_download_url: Some("https://huggingface.co/runwayml/stable-diffusion-v1-5".to_string()),
-        original_filename: Some("v1-5-pruned-emaonly.ckpt".to_string()),
-        file_size_bytes: metadata2.file_size_bytes as i32,
-        file_checksum_sha2: metadata2.sha256_checksum.to_string(),
-        public_bucket_hash: metadata2.bucket_details.clone().unwrap().object_hash,
-        maybe_public_bucket_prefix: Some(metadata2.bucket_details.clone().unwrap().prefix),
-        maybe_public_bucket_extension: Some(metadata2.bucket_details.clone().unwrap().suffix),
+        original_download_url: Some("https://huggingface.co/guoyww/animatediff/blob/main/v3_sd15_mm.ckpt".to_string()),
+        original_filename: Some("v3_sd15_mm.ckpt".to_string()),
+        file_size_bytes: metadata4.file_size_bytes as i32,
+        file_checksum_sha2: metadata4.sha256_checksum.to_string(),
+        public_bucket_hash: metadata4.bucket_details.clone().unwrap().object_hash,
+        maybe_public_bucket_prefix: Some(metadata4.bucket_details.clone().unwrap().prefix),
+        maybe_public_bucket_extension: Some(metadata4.bucket_details.clone().unwrap().suffix),
         version: 1,
         mysql_pool: &mysql_pool, // replace with actual MySqlPool
     };
 
-    create_weight(weights1).await?;
-    create_weight(weights2).await?;
+    // create_weight(weights1).await?;
+    // create_weight(weights2).await?;
+    // create_weight(weights3).await?;
+    create_weight(weights4).await?;
 
     Ok(())
 }
@@ -1065,6 +1122,6 @@ pub async fn seed_weights(mysql_pool: &Pool<MySql>) -> AnyhowResult<()> {
     // seed_weights_for_testing_inference(mysql_pool,user_token.clone()).await?;
     // println!("TESTING DOWLOAD");
     // test_seed_weights_files().await?;
-    // seed_workflows_for_testing_inference(mysql_pool,user_token.clone()).await?;
+    seed_workflows_for_testing_inference(mysql_pool,user_token.clone()).await?;
     Ok(())
 }
diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs
index 077b0265fc..3cbfe469c8 100644
--- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs
+++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs
@@ -31,7 +31,12 @@ impl RemoteCloudBucketDetails {
         match self.prefix.as_str() {
             // Weights
             "loRA" => Box::new(weights_descriptor::WeightsLoRADescriptor {}),
-            "sd15" => Box::new(weights_descriptor::WeightsSD15Descriptor {}),
+            "sd15" =>
+                match self.suffix.as_str() {
+                    "safetensors" => Box::new(weights_descriptor::WeightsSD15Descriptor {}),
+                    "ckpt" => Box::new(weights_descriptor::WeightsSD15CkptDescriptor {}),
+                    _ => panic!("Unknown suffix: {}",self.suffix)
+                },
             "sdxl" => Box::new(weights_descriptor::WeightsSDXLDescriptor {}),
             "valle_prompt" => Box::new(weights_descriptor::WeightsVallePromptDescriptor {}),
             "rvc" => {
diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs
index b5072ffa96..2b3c3d4d92 100644
--- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs
+++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs
@@ -50,6 +50,30 @@ impl FileDescriptor for WeightsSD15Descriptor {
     }
 }
 
+pub struct WeightsSD15CkptDescriptor;
+
+impl FileDescriptor for WeightsSD15CkptDescriptor {
+    fn remote_directory_path(&self) -> &str {
+        REMOTE_FILE_DIRECTORY
+    }
+    // this will be the type of file peroid is handled by the file formatter
+    // e.g safetensors bin jpg
+    fn get_suffix(&self)->String {
+       "ckpt".to_string()
+    }
+    // This will be the prefix of the media type or the weights type.
+    // name of the weights or the name of the media type
+    // vall-e_prompt, SD15, sd15, sdxl when implmenting add to the end
+    fn get_prefix(&self)->String {
+        "sd15".to_string()
+    }
+
+    // This will be ensure that the right bucket is picked
+    fn is_public(&self) -> bool {
+      true
+    }
+}
+
 
 pub struct WeightsSDXLDescriptor;
 
diff --git a/crates/service/job/inference_job/Cargo.toml b/crates/service/job/inference_job/Cargo.toml
index bd1e10467f..e4fcf4d680 100644
--- a/crates/service/job/inference_job/Cargo.toml
+++ b/crates/service/job/inference_job/Cargo.toml
@@ -82,4 +82,5 @@ my-workspace-hack = { version = "0.1", path = "../../../my-workspace-hack" }
 time = "0.3.30"
 mime = "0.3.17"
 rand = { version = "0.8.5", features = [] }
-jsonpath_lib = "0.3.0"
\ No newline at end of file
+jsonpath_lib = "0.3.0"
+walkdir = "2.4.0"
\ No newline at end of file
diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs
index c9b0ebfd6a..386abe78d5 100644
--- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs
+++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs
@@ -1,15 +1,45 @@
+use std::path::PathBuf;
 use errors::AnyhowResult;
 
 use crate::job::job_types::workflow::comfy_ui::comfy_ui_inference_command::ComfyInferenceCommand;
 
 pub struct ComfyDependencies {
     pub inference_command: ComfyInferenceCommand,
+    pub dependency_tokens: RequiredModels,
 }
 
 impl ComfyDependencies {
     pub fn setup() -> AnyhowResult<Self> {
         Ok(Self {
             inference_command: ComfyInferenceCommand::from_env()?,
+            dependency_tokens: RequiredModels::init(),
         })
     }
 }
+
+
+pub struct ComfyDependency {
+    pub(crate) location: PathBuf,
+    pub(crate) prefix: String,
+    pub(crate) hash: String,
+    pub(crate) extension: String,
+}
+
+pub struct RequiredModels {
+    pub(crate) comfy: Vec<ComfyDependency>
+}
+
+impl RequiredModels {
+    pub fn init() -> Self {
+        Self {
+            comfy: vec![
+                ComfyDependency {
+                    location: PathBuf::from("custom_nodes/ComfyUI-AnimateDiff-Evolved/models/v3_sd15_mm.ckpt"),
+                    prefix: "sd15".to_string(),
+                    hash: "d07ac2v29ak90cttyx4v846yhept2fbr".to_string(),
+                    extension: "ckpt".to_string(),
+                },
+            ],
+        }
+    }
+}
diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs
index da134c50e3..d3444d8890 100644
--- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs
+++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs
@@ -64,6 +64,7 @@ pub enum ExecutableOrCommand {
 
 pub struct InferenceArgs<'s> {
     pub stderr_output_file: &'s Path,
+    pub stdout_output_file: &'s Path,
 }
 
 impl ComfyInferenceCommand {
@@ -180,7 +181,9 @@ impl ComfyInferenceCommand {
         info!("stderr will be written to file: {:?}", args.stderr_output_file.as_os_str());
 
         let stderr_file = File::create(&args.stderr_output_file)?;
+        let stdout_file = File::create(&args.stdout_output_file)?;
         config.stderr = Redirection::File(stderr_file);
+        config.stdout = Redirection::File(stdout_file);
 
         if !env_vars.is_empty() {
             config.env = Some(env_vars);
diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs
index 65dcfeda83..e578bd4d5a 100644
--- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs
+++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs
@@ -1,5 +1,5 @@
 use std::fs::read_to_string;
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use std::time::Instant;
 
 use anyhow::anyhow;
@@ -7,6 +7,7 @@ use jsonpath_lib::replace_with;
 use log::{error, info, warn};
 use serde::de::Error;
 use serde_json::Value;
+use walkdir::WalkDir;
 
 use buckets::public::media_files::bucket_file_path::MediaFileBucketPath;
 use cloud_storage::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails;
@@ -37,6 +38,16 @@ pub struct ComfyProcessJobArgs<'a> {
     pub job: &'a AvailableInferenceJob,
 }
 
+fn recursively_delete_files_in(path: &Path) -> std::io::Result<()> {
+    for entry in WalkDir::new(path) {
+        let entry = entry?;
+        if entry.path().is_file() {
+            safe_delete_temp_file(entry.path());
+        }
+    }
+    Ok(())
+}
+
 pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResult, ProcessSingleJobError> {
     let job = args.job;
     let deps = args.job_dependencies;
@@ -70,6 +81,40 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         .new_tempdir(&work_temp_dir)
         .map_err(|e| ProcessSingleJobError::from_io_error(e))?;
 
+    // ===================== DOWNLOAD REQUIRED MODELS IF NOT EXIST ===================== //
+
+    let all_models = &args.job_dependencies.job.job_specific_dependencies.maybe_comfy_ui_dependencies;
+    match all_models {
+        Some(models) => {
+            for model in &models.dependency_tokens.comfy {
+                let mut comfy_dir = model_dependencies.inference_command.comfy_root_code_directory.clone();
+                comfy_dir = comfy_dir.join(model.location.clone());
+                if !comfy_dir.exists() {
+                    let bucket_details = RemoteCloudBucketDetails {
+                        object_hash: model.hash.clone(),
+                        prefix: model.prefix.clone(),
+                        suffix: model.extension.clone(),
+                    };
+                    let comfy_path = comfy_dir.to_str().unwrap().to_string();
+                    let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await;
+                    let remote_cloud_file_client = match remote_cloud_file_client {
+                        Ok(res) => {
+                            res
+                        }
+                        Err(_) => {
+                            return Err(ProcessSingleJobError::from(anyhow!("failed to get remote cloud file client")));
+                        }
+                    };
+                    remote_cloud_file_client.download_file(bucket_details, comfy_path.clone()).await?;
+                    info!("Downloaded model to {:?}", comfy_path);
+                }
+            }
+        }
+        None => {
+            info!("No models specified to download")
+        }
+    }
+
 
     // ==================== QUERY AND DOWNLOAD FILES ==================== //
     let root_comfy_path = model_dependencies.inference_command.comfy_root_code_directory.clone();
@@ -217,21 +262,36 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
 
     // ==================== RUN INFERENCE SCRIPT ==================== //
     let stderr_output_file = work_temp_dir.path().join("stderr.txt");
+    let stdout_output_file = work_temp_dir.path().join("stdout.txt");
+
     let inference_start_time = Instant::now();
 
     let command_exit_status = model_dependencies
         .inference_command
         .execute_inference(InferenceArgs {
             stderr_output_file: &stderr_output_file,
+            stdout_output_file: &stdout_output_file,
         });
 
     let inference_duration = Instant::now().duration_since(inference_start_time);
 
     info!("Inference took duration to complete: {:?}", &inference_duration);
 
-    if !command_exit_status.is_success() {
+    // ==================== GET OUTPUT FILE ======================== //
+    let mut output_file = root_comfy_path.join("output");
+    output_file = output_file.join(job_args.output_path);
+
+    // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== //
+
+    // check stdout for success and check if file exists
+    let stdout_output = read_to_string(&stdout_output_file).unwrap();
+    // check for "Prompt executed" in stdout (comfyui only outputs this for success)
+    if !stdout_output.contains("Prompt executed") || check_file_exists(&output_file).is_err() {
         error!("Inference failed: {:?}", command_exit_status);
 
+        error!("Captured stdout output: {}", stdout_output);
+        error!("Captured stderr output: {}", read_to_string(&stderr_output_file).unwrap());
+
         let error = ProcessSingleJobError::Other(anyhow!("CommandExitStatus: {:?}", command_exit_status));
 
         if let Ok(contents) = read_to_string(&stderr_output_file) {
@@ -239,6 +299,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         }
 
         safe_delete_temp_file(&stderr_output_file);
+        safe_delete_temp_file(&stdout_output_file);
         safe_delete_temp_directory(&work_temp_dir);
         safe_delete_temp_file(&workflow_path);
         if let Some(sd_path) = maybe_sd_path {
@@ -250,18 +311,12 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         if let Some(input_path) = maybe_input_path {
             safe_delete_temp_file(&input_path);
         }
+        let output_dir = root_comfy_path.join("output");
+        recursively_delete_files_in(&output_dir).unwrap();
 
         return Err(error);
     }
 
-    // ==================== GET OUTPUT FILE ======================== //
-    let mut output_file = root_comfy_path.join("output");
-    output_file = output_file.join(job_args.output_path);
-
-    // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== //
-    info!("Checking that file exists: {:?} ...", output_file);
-    check_file_exists(&output_file).map_err(|e| ProcessSingleJobError::Other(e))?;
-
     info!("Interrogating result file size ...");
 
     let file_size_bytes = file_size(&output_file)
@@ -330,6 +385,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
     // ==================== DELETE TEMP FILES ==================== //
 
     safe_delete_temp_file(&stderr_output_file);
+    safe_delete_temp_file(&stdout_output_file);
     safe_delete_temp_file(&workflow_path);
     if let Some(sd_path) = maybe_sd_path {
         safe_delete_temp_file(&sd_path);
@@ -341,15 +397,8 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         safe_delete_temp_file(&input_path);
     }
 
-    // delete all files in output directory
     let output_dir = root_comfy_path.join("output");
-    let entries: Vec<PathBuf> = std::fs::read_dir(output_dir).unwrap()
-        .map(|res| res.map(|e| e.path()))
-        .collect::<Result<Vec<_>, std::io::Error>>()
-        .unwrap();
-    for entry in entries {
-        safe_delete_temp_file(&entry);
-    }
+    recursively_delete_files_in(&output_dir).unwrap();
 
     // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk...
     safe_delete_temp_directory(&work_temp_dir);

From f0f2788e1bdf9ddb216e5f52b1f7d7f0a3323ac2 Mon Sep 17 00:00:00 2001
From: saltacc <alotofcatz@gmail.com>
Date: Wed, 31 Jan 2024 17:17:13 -0800
Subject: [PATCH 4/5] make things work with new stuff :/

---
 Cargo.lock                                    |   1 +
 .../src/seeding/model_weights.rs              | 108 +++++++++---------
 crates/service/job/inference_job/Cargo.toml   |   3 +-
 .../comfy_ui/comfy_ui_dependencies.rs         |  64 ++++++++++-
 .../workflow/comfy_ui/process_job.rs          |  61 ++++++----
 5 files changed, 153 insertions(+), 84 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 37ac76bd68..4b1fff70f1 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3099,6 +3099,7 @@ dependencies = [
  "r2d2_redis",
  "rand 0.8.5",
  "redis_common",
+ "reqwest",
  "serde",
  "serde_derive",
  "serde_json",
diff --git a/crates/cli/dev-database-seed/src/seeding/model_weights.rs b/crates/cli/dev-database-seed/src/seeding/model_weights.rs
index d8938a2ad2..9d1ee6b494 100644
--- a/crates/cli/dev-database-seed/src/seeding/model_weights.rs
+++ b/crates/cli/dev-database-seed/src/seeding/model_weights.rs
@@ -6,7 +6,7 @@ use sqlx::{MySql, Pool};
 
 use cloud_storage::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails;
 use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient;
-use cloud_storage::remote_file_manager::weights_descriptor::{WeightsLoRADescriptor, WeightsSD15Descriptor, WeightsSD15CkptDescriptor};
+use cloud_storage::remote_file_manager::weights_descriptor::{WeightsLoRADescriptor, WeightsSD15Descriptor, WeightsSD15CkptDescriptor, WeightsWorkflowDescriptor};
 use enums::by_table::model_weights::{
     weights_category::WeightsCategory,
     weights_types::WeightsType,
@@ -363,7 +363,7 @@ pub async fn seed_weights_for_paging(mysql_pool: &Pool<MySql>, user_token: UserT
                 //cached_user_ratings_positive_count,
                 //cached_user_ratings_negative_count,
                 //maybe_cached_user_ratings_ratio: Some(cached_user_ratings_ratio as f32),
-                version,
+                version: version as u32,
                 mysql_pool: &mysql_pool, // replace with actual MySqlPool
             };
         
@@ -623,7 +623,7 @@ Noosphere by skumerz + dalcefoPainting + 饭特稀V08 by zhazhahui345 + GhostMix
         maybe_last_update_user_token: Some("Last Update User Token".to_string()),
         original_download_url: Some("https://civitai.com/models/41865/majicmix-fantasy".to_string()),
         original_filename: Some("majicmixFantasy_v30Vae.safetensors".to_string()),
-        file_size_bytes: metadata1.file_size_bytes as i32,
+        file_size_bytes: metadata1.file_size_bytes,
         file_checksum_sha2: metadata1.sha256_checksum.to_string(),
         public_bucket_hash: metadata1.bucket_details.clone().unwrap().object_hash,
         maybe_public_bucket_prefix: Some(metadata1.bucket_details.clone().unwrap().prefix),
@@ -645,7 +645,7 @@ Noosphere by skumerz + dalcefoPainting + 饭特稀V08 by zhazhahui345 + GhostMix
         maybe_last_update_user_token: Some("<p> Honkai <p>".to_string()),
         original_download_url: Some("https://civitai.com/models/95243/mihoyo-collection-honkai-impact-3rd-or-honkai-star-rail-or-genshin-impact-or-zenless-zone-zero".to_string()),
         original_filename: Some("xiawolei-v100-000019.safetensors".to_string()),
-        file_size_bytes: metadata2.file_size_bytes as i32,
+        file_size_bytes: metadata2.file_size_bytes,
         file_checksum_sha2: metadata2.sha256_checksum.to_string(),
         public_bucket_hash: metadata2.bucket_details.clone().unwrap().object_hash.clone(),
         maybe_public_bucket_prefix: Some(metadata2.bucket_details.clone().unwrap().prefix),
@@ -667,7 +667,7 @@ Noosphere by skumerz + dalcefoPainting + 饭特稀V08 by zhazhahui345 + GhostMix
         maybe_last_update_user_token: Some("Last Update User Token".to_string()),
         original_download_url: Some("https://civitai.com/models/97653/z-vae".to_string()),
         original_filename: Some("zVae_v20.safetensors".to_string()),
-        file_size_bytes: metadata3.file_size_bytes as i32,
+        file_size_bytes: metadata3.file_size_bytes,
         file_checksum_sha2: metadata3.sha256_checksum.to_string(),
         public_bucket_hash: metadata3.bucket_details.clone().unwrap().object_hash,
         maybe_public_bucket_prefix: Some(metadata3.bucket_details.clone().unwrap().prefix),
@@ -690,10 +690,10 @@ pub async fn seed_workflows_for_testing_inference(mysql_pool: &Pool<MySql>, user
     let model_weight_token4 = ModelWeightToken::new_from_str("weight_n8sz47gmfw2zx02snrbz88ns9");
 
     let mut path_to_comfy = get_seed_tool_data_root();
-    path_to_comfy.push("models/workflows/comfyui/workflow_api.json");
+    path_to_comfy.push("models/workflows/comfyui/yae_video_prod.json");
     let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?;
-    // let comfy_weights_descriptor = Box::new(WeightsWorkflowDescriptor {});
-    // let metadata1 = remote_cloud_file_client.upload_file(comfy_weights_descriptor, path_to_comfy.as_path().to_str().unwrap()).await?;
+    let comfy_weights_descriptor = Box::new(WeightsWorkflowDescriptor {});
+    let metadata1 = remote_cloud_file_client.upload_file(comfy_weights_descriptor, path_to_comfy.as_path().to_str().unwrap()).await?;
 
     // let mut path_to_comfy2 = get_seed_tool_data_root();
     // path_to_comfy2.push("models/workflows/comfyui/majicmixRealistic_v7.safetensors");
@@ -705,32 +705,32 @@ pub async fn seed_workflows_for_testing_inference(mysql_pool: &Pool<MySql>, user
     // let comfy_weights_descriptor3 = Box::new(WeightsWorkflowDescriptor {});
     // let metadata3 = remote_cloud_file_client.upload_file(comfy_weights_descriptor3, path_to_comfy3.as_path().to_str().unwrap()).await?;
 
-    let mut path_to_comfy4 = get_seed_tool_data_root();
+/*    let mut path_to_comfy4 = get_seed_tool_data_root();
     path_to_comfy4.push("models/workflows/comfyui/v3_sd15_mm.ckpt");
     let comfy_weights_descriptor4 = Box::new(WeightsSD15CkptDescriptor {});
-    let metadata4 = remote_cloud_file_client.upload_file(comfy_weights_descriptor4, path_to_comfy4.as_path().to_str().unwrap()).await?;
+    let metadata4 = remote_cloud_file_client.upload_file(comfy_weights_descriptor4, path_to_comfy4.as_path().to_str().unwrap()).await?;*/
 
-    // let weights1 = CreateModelWeightsArgs {
-    //     token: &model_weight_token1, // replace with actual ModelWeightToken
-    //     weights_type: WeightsType::ComfyUi, // replace with actual WeightsType
-    //     weights_category: WeightsCategory::WorkflowConfig, // replace with actual WeightsCategory
-    //     title: "comfy-test-workflow".to_string(),
-    //     maybe_description_markdown: Some("Test workflow for ComfyUI".to_string()),
-    //     maybe_description_rendered_html: Some("<p>Description</p>".to_string()),
-    //     creator_user_token: Some(&user_token), // replace with actual UserToken
-    //     creator_ip_address: "192.168.1.1",
-    //     creator_set_visibility: Visibility::Public,
-    //     maybe_last_update_user_token: Some("Last Update User Token".to_string()),
-    //     original_download_url: Some("https://github.com/comfyanonymous/ComfyUI".to_string()),
-    //     original_filename: Some("test-workflow.json".to_string()),
-    //     file_size_bytes: metadata1.file_size_bytes as i32,
-    //     file_checksum_sha2: metadata1.sha256_checksum.to_string(),
-    //     public_bucket_hash: metadata1.bucket_details.clone().unwrap().object_hash,
-    //     maybe_public_bucket_prefix: Some(metadata1.bucket_details.clone().unwrap().prefix),
-    //     maybe_public_bucket_extension: Some(metadata1.bucket_details.clone().unwrap().suffix),
-    //     version: 1,
-    //     mysql_pool: &mysql_pool, // replace with actual MySqlPool
-    // };
+    let weights1 = CreateModelWeightsArgs {
+        token: &model_weight_token1, // replace with actual ModelWeightToken
+        weights_type: WeightsType::ComfyUi, // replace with actual WeightsType
+        weights_category: WeightsCategory::WorkflowConfig, // replace with actual WeightsCategory
+        title: "yae-video-prod".to_string(),
+        maybe_description_markdown: Some("ComfyUI Workflow".to_string()),
+        maybe_description_rendered_html: Some("<p>Description</p>".to_string()),
+        creator_user_token: Some(&user_token), // replace with actual UserToken
+        creator_ip_address: "192.168.1.1",
+        creator_set_visibility: Visibility::Public,
+        maybe_last_update_user_token: Some("Last Update User Token".to_string()),
+        original_download_url: Some("https://github.com/comfyanonymous/ComfyUI".to_string()),
+        original_filename: Some("yae_video_prod.json".to_string()),
+        file_size_bytes: metadata1.file_size_bytes,
+        file_checksum_sha2: metadata1.sha256_checksum.to_string(),
+        public_bucket_hash: metadata1.bucket_details.clone().unwrap().object_hash,
+        maybe_public_bucket_prefix: Some(metadata1.bucket_details.clone().unwrap().prefix),
+        maybe_public_bucket_extension: Some(metadata1.bucket_details.clone().unwrap().suffix),
+        version: 1,
+        mysql_pool: &mysql_pool, // replace with actual MySqlPool
+    };
     // let weights2 = CreateModelWeightsArgs {
     //     token: &model_weight_token2, // replace with actual ModelWeightToken
     //     weights_type: WeightsType::ComfyUi, // replace with actual WeightsType
@@ -775,32 +775,32 @@ pub async fn seed_workflows_for_testing_inference(mysql_pool: &Pool<MySql>, user
     //     mysql_pool: &mysql_pool, // replace with actual MySqlPool
     // };
 
-    let weights4 = CreateModelWeightsArgs {
-        token: &model_weight_token4, // replace with actual ModelWeightToken
-        weights_type: WeightsType::ComfyUi, // replace with actual WeightsType
-        weights_category: WeightsCategory::WorkflowConfig, // replace with actual WeightsCategory
-        title: "v3_sd15_mm".to_string(),
-        maybe_description_markdown: Some("ComfyUI motion module".to_string()),
-        maybe_description_rendered_html: Some("<p>Description</p>".to_string()),
-        creator_user_token: Some(&user_token), // replace with actual UserToken
-        creator_ip_address: "192.168.1.1",
-        creator_set_visibility: Visibility::Public,
-        maybe_last_update_user_token: Some("Last Update User Token".to_string()),
-        original_download_url: Some("https://huggingface.co/guoyww/animatediff/blob/main/v3_sd15_mm.ckpt".to_string()),
-        original_filename: Some("v3_sd15_mm.ckpt".to_string()),
-        file_size_bytes: metadata4.file_size_bytes as i32,
-        file_checksum_sha2: metadata4.sha256_checksum.to_string(),
-        public_bucket_hash: metadata4.bucket_details.clone().unwrap().object_hash,
-        maybe_public_bucket_prefix: Some(metadata4.bucket_details.clone().unwrap().prefix),
-        maybe_public_bucket_extension: Some(metadata4.bucket_details.clone().unwrap().suffix),
-        version: 1,
-        mysql_pool: &mysql_pool, // replace with actual MySqlPool
-    };
+    // let weights4 = CreateModelWeightsArgs {
+    //     token: &model_weight_token4, // replace with actual ModelWeightToken
+    //     weights_type: WeightsType::ComfyUi, // replace with actual WeightsType
+    //     weights_category: WeightsCategory::WorkflowConfig, // replace with actual WeightsCategory
+    //     title: "v3_sd15_mm".to_string(),
+    //     maybe_description_markdown: Some("ComfyUI motion module".to_string()),
+    //     maybe_description_rendered_html: Some("<p>Description</p>".to_string()),
+    //     creator_user_token: Some(&user_token), // replace with actual UserToken
+    //     creator_ip_address: "192.168.1.1",
+    //     creator_set_visibility: Visibility::Public,
+    //     maybe_last_update_user_token: Some("Last Update User Token".to_string()),
+    //     original_download_url: Some("https://huggingface.co/guoyww/animatediff/blob/main/v3_sd15_mm.ckpt".to_string()),
+    //     original_filename: Some("v3_sd15_mm.ckpt".to_string()),
+    //     file_size_bytes: metadata4.file_size_bytes as i32,
+    //     file_checksum_sha2: metadata4.sha256_checksum.to_string(),
+    //     public_bucket_hash: metadata4.bucket_details.clone().unwrap().object_hash,
+    //     maybe_public_bucket_prefix: Some(metadata4.bucket_details.clone().unwrap().prefix),
+    //     maybe_public_bucket_extension: Some(metadata4.bucket_details.clone().unwrap().suffix),
+    //     version: 1,
+    //     mysql_pool: &mysql_pool, // replace with actual MySqlPool
+    // };
 
-    // create_weight(weights1).await?;
+    create_weight(weights1).await?;
     // create_weight(weights2).await?;
     // create_weight(weights3).await?;
-    create_weight(weights4).await?;
+    // create_weight(weights4).await?;
 
     Ok(())
 }
diff --git a/crates/service/job/inference_job/Cargo.toml b/crates/service/job/inference_job/Cargo.toml
index e4fcf4d680..aeb2f5c193 100644
--- a/crates/service/job/inference_job/Cargo.toml
+++ b/crates/service/job/inference_job/Cargo.toml
@@ -83,4 +83,5 @@ time = "0.3.30"
 mime = "0.3.17"
 rand = { version = "0.8.5", features = [] }
 jsonpath_lib = "0.3.0"
-walkdir = "2.4.0"
\ No newline at end of file
+walkdir = "2.4.0"
+reqwest = "0.11.22"
\ No newline at end of file
diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs
index 386abe78d5..f2234988f6 100644
--- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs
+++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs
@@ -20,9 +20,7 @@ impl ComfyDependencies {
 
 pub struct ComfyDependency {
     pub(crate) location: PathBuf,
-    pub(crate) prefix: String,
-    pub(crate) hash: String,
-    pub(crate) extension: String,
+    pub(crate) url: String,
 }
 
 pub struct RequiredModels {
@@ -33,11 +31,65 @@ impl RequiredModels {
     pub fn init() -> Self {
         Self {
             comfy: vec![
+                ComfyDependency {
+                    location: PathBuf::from("models/loras/v3_sd15_adapter.ckpt"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/v3_sd15_adapter.ckpt".to_string(),
+                },
                 ComfyDependency {
                     location: PathBuf::from("custom_nodes/ComfyUI-AnimateDiff-Evolved/models/v3_sd15_mm.ckpt"),
-                    prefix: "sd15".to_string(),
-                    hash: "d07ac2v29ak90cttyx4v846yhept2fbr".to_string(),
-                    extension: "ckpt".to_string(),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/v3_sd15_mm.ckpt".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("custom_nodes/ComfyUI-AnimateDiff-Evolved/models/v3_sd15_mm.ckpt"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/v3_sd15_mm.ckpt".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("models/controlnet/v3_sd15_sparsectrl_scribble.ckpt"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/v3_sd15_sparsectrl_scribble.ckpt".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("models/controlnet/control_v11f1e_sd15_tile.bin"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/control_v11f1e_sd15_tile.bin".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("models/controlnet/control_v11p_sd15_softedge.safetensors"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/control_v11p_sd15_softedge.safetensors".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("models/controlnet/control_v11p_sd15s2_lineart_anime.safetensors"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/control_v11p_sd15s2_lineart_anime.safetensors".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("models/controlnet/control_v11p_sd15_lineart.safetensors"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/control_v11p_sd15_lineart.safetensors".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("models/controlnet/control_v11f1p_sd15_depth.safetensors"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/control_v11f1p_sd15_depth.safetensors".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("models/controlnet/control_v11p_sd15_openpose.safetensors"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/control_v11p_sd15_openpose.safetensors".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("models/controlnet/animatediff_controlnet.ckpt"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/animatediff_controlnet.ckpt".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("models/controlnet/control_v2p_sd15_mediapipe_face.safetensors"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/control_v2p_sd15_mediapipe_face.safetensors".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("models/controlnet/ControlNetHED.pth"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/ControlNetHED.pth".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("models/ipadapter/ip-adapter_sd15_light.safetensors"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/ip-adapter_sd15_light.safetensors".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("models/clip_vision/CLIP-ViT-H-14-laion2B-s32B-b79K.safetensors"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/model.safetensors".to_string(),
                 },
             ],
         }
diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs
index e578bd4d5a..1494df93a3 100644
--- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs
+++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs
@@ -1,12 +1,14 @@
-use std::fs::read_to_string;
+use std::fs::{File, read_to_string};
 use std::path::{Path, PathBuf};
 use std::time::Instant;
 
 use anyhow::anyhow;
 use jsonpath_lib::replace_with;
 use log::{error, info, warn};
+use reqwest::Url;
 use serde::de::Error;
 use serde_json::Value;
+use tokio::io::AsyncWriteExt;
 use walkdir::WalkDir;
 
 use buckets::public::media_files::bucket_file_path::MediaFileBucketPath;
@@ -83,30 +85,43 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
 
     // ===================== DOWNLOAD REQUIRED MODELS IF NOT EXIST ===================== //
 
+    async fn download_file(file_url: String, dep_path: PathBuf) -> Result<(), anyhow::Error> {
+        // Send a GET request to the file_url
+        let response = reqwest::get(&file_url).await?;
+
+        // Ensure the request was successful
+        if response.status().is_success() {
+            // Get the byte stream of the file
+            let bytes = response.bytes().await?;
+
+            // Ensure the parent directory exists
+            if let Some(parent) = dep_path.parent() {
+                // Skip errors (for example if the directory already exists)
+                let _ = tokio::fs::create_dir_all(&parent).await;
+            }
+
+            // Create or open the file at dep_path asynchronously
+            let mut file = tokio::fs::File::create(&dep_path).await?;
+
+            // Write the bytes to the file asynchronously
+            file.write_all(&bytes).await?;
+            println!("File downloaded successfully.");
+        } else {
+            println!("Failed to download the file.");
+        }
+
+        Ok(())
+    }
+
     let all_models = &args.job_dependencies.job.job_specific_dependencies.maybe_comfy_ui_dependencies;
     match all_models {
         Some(models) => {
             for model in &models.dependency_tokens.comfy {
-                let mut comfy_dir = model_dependencies.inference_command.comfy_root_code_directory.clone();
-                comfy_dir = comfy_dir.join(model.location.clone());
-                if !comfy_dir.exists() {
-                    let bucket_details = RemoteCloudBucketDetails {
-                        object_hash: model.hash.clone(),
-                        prefix: model.prefix.clone(),
-                        suffix: model.extension.clone(),
-                    };
-                    let comfy_path = comfy_dir.to_str().unwrap().to_string();
-                    let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await;
-                    let remote_cloud_file_client = match remote_cloud_file_client {
-                        Ok(res) => {
-                            res
-                        }
-                        Err(_) => {
-                            return Err(ProcessSingleJobError::from(anyhow!("failed to get remote cloud file client")));
-                        }
-                    };
-                    remote_cloud_file_client.download_file(bucket_details, comfy_path.clone()).await?;
-                    info!("Downloaded model to {:?}", comfy_path);
+                let mut dep_path = model_dependencies.inference_command.comfy_root_code_directory.clone();
+                dep_path = dep_path.join(model.location.clone());
+                if !dep_path.exists() {
+                    download_file(model.url.clone(), dep_path.clone()).await.map_err(|e| ProcessSingleJobError::Other(e))?;
+                    info!("Downloaded model to {:?}", dep_path);
                 }
             }
         }
@@ -145,7 +160,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
                 prefix: retrieved_sd_record.maybe_public_bucket_prefix.unwrap(),
                 suffix: retrieved_sd_record.maybe_public_bucket_extension.unwrap(),
             };
-            let sd_filename = retrieved_sd_record.original_filename.unwrap_or("sd_model.safetensors".to_string());
+            let sd_filename = "model.safetensors";
             let sd_path = sd_dir.join(sd_filename).to_str().unwrap().to_string();
             remote_cloud_file_client.download_file(bucket_details, sd_path.clone()).await?;
             maybe_sd_path = Some(sd_path.parse().unwrap());
@@ -170,7 +185,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
                 suffix: retrieved_lora_record.maybe_public_bucket_extension.unwrap(),
             };
 
-            let lora_filename = retrieved_lora_record.original_filename.unwrap_or("lora_model.safetensors".to_string());
+            let lora_filename = "lora.safetensors";
             let lora_path = lora_dir.join(lora_filename).to_str().unwrap().to_string();
             remote_cloud_file_client.download_file(bucket_details, lora_path.clone()).await?;
             maybe_lora_path = Some(lora_path.parse().unwrap());

From bfebf60289c8ddd802b40ecf30fa4ff2f0b31d79 Mon Sep 17 00:00:00 2001
From: saltacc <alotofcatz@gmail.com>
Date: Wed, 31 Jan 2024 17:38:56 -0800
Subject: [PATCH 5/5] properly handle invalid jsonpaths

---
 .../workflow/comfy_ui/process_job.rs          | 140 ++++++++++--------
 1 file changed, 75 insertions(+), 65 deletions(-)

diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs
index 1494df93a3..feef4cc8a1 100644
--- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs
+++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs
@@ -3,9 +3,7 @@ use std::path::{Path, PathBuf};
 use std::time::Instant;
 
 use anyhow::anyhow;
-use jsonpath_lib::replace_with;
 use log::{error, info, warn};
-use reqwest::Url;
 use serde::de::Error;
 use serde_json::Value;
 use tokio::io::AsyncWriteExt;
@@ -84,6 +82,17 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         .map_err(|e| ProcessSingleJobError::from_io_error(e))?;
 
     // ===================== DOWNLOAD REQUIRED MODELS IF NOT EXIST ===================== //
+    let root_comfy_path = model_dependencies.inference_command.comfy_root_code_directory.clone();
+
+    let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await;
+    let remote_cloud_file_client = match remote_cloud_file_client {
+        Ok(res) => {
+            res
+        }
+        Err(_) => {
+            return Err(ProcessSingleJobError::from(anyhow!("failed to get remote cloud file client")));
+        }
+    };
 
     async fn download_file(file_url: String, dep_path: PathBuf) -> Result<(), anyhow::Error> {
         // Send a GET request to the file_url
@@ -130,19 +139,53 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         }
     }
 
+    // Download workflow to ComfyRunner
+    let workflow_dir = root_comfy_path.join("../ComfyLauncher");
+    let retrieved_workflow_record =  get_weight_by_token(
+        job_args.workflow_source,
+        false,
+        &deps.db.mysql_pool
+    ).await?.unwrap();
+
+    let bucket_details = RemoteCloudBucketDetails {
+        object_hash: retrieved_workflow_record.public_bucket_hash,
+        prefix: retrieved_workflow_record.maybe_public_bucket_prefix.unwrap(),
+        suffix: retrieved_workflow_record.maybe_public_bucket_extension.unwrap(),
+    };
+    let workflow_path = workflow_dir.join("prompt.json").to_str().unwrap().to_string();
+    remote_cloud_file_client.download_file(bucket_details, workflow_path.clone()).await?;
+
+    let maybe_args = job.maybe_inference_args
+        .as_ref()
+        .map(|args| args.args.as_ref())
+        .flatten();
+
+    let poly_args = match maybe_args {
+        None => return Err(ProcessSingleJobError::Other(anyhow!("ComfyUi args not found"))),
+        Some(args) => args,
+    };
+
+    let comfy_args = match poly_args {
+        Cu(args) => args,
+        _ => return Err(ProcessSingleJobError::Other(anyhow!("ComfyUi args not found"))),
+    };
+
+    // Apply modifications if they exist
+    if let Some(modifications) = comfy_args.maybe_json_modifications.clone() {
+        // Load prompt.json
+        let prompt_file = File::open(&workflow_path).unwrap();
+        let mut prompt_json: Value = serde_json::from_reader(prompt_file).unwrap();
+        // Modify json
+        for (path, new_value) in modifications {
+            prompt_json = replace_json_value(prompt_json, &path, new_value).map_err(|e| ProcessSingleJobError::Other(e))?;
+        }
+        // Save prompt.json
+        let prompt_file = File::create(&workflow_path).unwrap();
+        serde_json::to_writer(prompt_file, &prompt_json).unwrap();
+    }
+
 
     // ==================== QUERY AND DOWNLOAD FILES ==================== //
-    let root_comfy_path = model_dependencies.inference_command.comfy_root_code_directory.clone();
-
-    let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await;
-    let remote_cloud_file_client = match remote_cloud_file_client {
-        Ok(res) => {
-            res
-        }
-        Err(_) => {
-            return Err(ProcessSingleJobError::from(anyhow!("failed to get remote cloud file client")));
-        }
-    };
 
     // Download SD model if specified
     let mut maybe_sd_path: Option<PathBuf> = None;
@@ -219,53 +262,6 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         None => {}
     }
 
-    // Download workflow to ComfyRunner
-    let workflow_dir = root_comfy_path.join("../ComfyLauncher");
-    let retrieved_workflow_record =  get_weight_by_token(
-        job_args.workflow_source,
-        false,
-        &deps.db.mysql_pool
-    ).await?.unwrap();
-
-    let bucket_details = RemoteCloudBucketDetails {
-        object_hash: retrieved_workflow_record.public_bucket_hash,
-        prefix: retrieved_workflow_record.maybe_public_bucket_prefix.unwrap(),
-        suffix: retrieved_workflow_record.maybe_public_bucket_extension.unwrap(),
-    };
-    let workflow_path = workflow_dir.join("prompt.json").to_str().unwrap().to_string();
-    remote_cloud_file_client.download_file(bucket_details, workflow_path.clone()).await?;
-
-    let maybe_args = job.maybe_inference_args
-        .as_ref()
-        .map(|args| args.args.as_ref())
-        .flatten();
-
-    let poly_args = match maybe_args {
-        None => return Err(ProcessSingleJobError::Other(anyhow!("ComfyUi args not found"))),
-        Some(args) => args,
-    };
-
-    let comfy_args = match poly_args {
-        Cu(args) => args,
-        _ => return Err(ProcessSingleJobError::Other(anyhow!("ComfyUi args not found"))),
-    };
-
-    // Apply modifications if they exist
-    if let Some(modifications) = comfy_args.maybe_json_modifications.clone() {
-        // Load prompt.json
-        let prompt_file = std::fs::File::open(&workflow_path).unwrap();
-        let mut prompt_json: Value = serde_json::from_reader(prompt_file).unwrap();
-        // Modify json
-        for (path, new_value) in modifications {
-            prompt_json = replace_json_value(prompt_json, &path, new_value).expect(
-                format!("Failed to replace json value at path {}", path).as_str()
-            );
-        }
-        // Save prompt.json
-        let prompt_file = std::fs::File::create(&workflow_path).unwrap();
-        serde_json::to_writer(prompt_file, &prompt_json).unwrap();
-    }
-
     // ==================== SETUP FOR INFERENCE ==================== //
 
     info!("Ready for ComfyUI inference...");
@@ -456,15 +452,29 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
     })
 }
 
-fn replace_json_value(json: Value, path: &str, new_value: NewValue) -> Result<Value, serde_json::Error> {
-    replace_with(json, path, &mut |_| {
+fn replace_json_value(json: Value, path: &str, new_value: NewValue) -> anyhow::Result<Value> {
+    // First, attempt to read the value at the specified path
+    let results = jsonpath_lib::select(&json, path).map_err(|err| {
+        anyhow!("Invalid jsonpath '{}': {:?}", path, err)
+    })?;
+
+    // If the path does not exist or returns no results, return an error
+    if results.is_empty() {
+        return Err(anyhow!("Path '{}' does not exist in the provided JSON.", path));
+    }
+
+    // If the path exists, proceed with the replacement
+    // Assuming replace_with returns a Result, handle it appropriately
+    jsonpath_lib::replace_with(json, path, &mut |_| {
         match &new_value {
             NewValue::String(s) => Some(Value::String(s.clone())),
-            NewValue::Float(f) => Some(Value::Number(serde_json::Number::from_f64(*f as f64).unwrap())),
+            NewValue::Float(f) => serde_json::Number::from_f64(*f as f64)
+                .map(Value::Number) // Convert Option to Some(Value::Number) if Some, else None
+                .or_else(|| Some(Value::Null)), // If None, use Value::Null instead
             NewValue::Int(i) => Some(Value::Number(serde_json::Number::from(*i))),
             NewValue::Bool(b) => Some(Value::Bool(*b)),
         }
     }).map_err(|err| {
-        serde_json::Error::custom(format!("Failed to replace json value at path {}: {:?}", path, err))
+        anyhow!("Failed to replace json value at path '{}': {:?}", path, err)
     })
-}
\ No newline at end of file
+}