From 48b1f6306b325c66d950358f14af1dd8fa71a83b Mon Sep 17 00:00:00 2001
From: salt <alotofcatz@gmail.com>
Date: Sat, 17 Feb 2024 06:21:51 -0800
Subject: [PATCH 1/3] save job args to file modification info

---
 .../create/insert_media_file_from_comfy_ui.rs |  4 ++++
 .../workflow/comfy_ui/process_job.rs          | 23 ++++++++++++-------
 .../workflow/comfy_ui/validate_job.rs         |  1 +
 3 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_comfy_ui.rs b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_comfy_ui.rs
index 178e83f438..284cd6894e 100644
--- a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_comfy_ui.rs
+++ b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_comfy_ui.rs
@@ -32,6 +32,8 @@ pub struct InsertArgs<'a> {
     pub worker_hostname: &'a str,
     pub worker_cluster: &'a str,
     pub media_file_type: MediaFileType,
+
+    pub extra_file_modification_info: Option<&'a str>,
 }
 
 pub async fn insert_media_file_from_comfy_ui(args: InsertArgs<'_>) -> AnyhowResult<(MediaFileToken, u64)>
@@ -92,6 +94,7 @@ SET
   creator_ip_address = ?,
 
   creator_set_visibility = ?,
+  extra_file_modification_info = ?,
 
   maybe_creator_file_synthetic_id = ?,
   maybe_creator_category_synthetic_id = ?,
@@ -122,6 +125,7 @@ SET
       args.job.creator_ip_address,
 
       args.job.creator_set_visibility.to_str(),
+      args.extra_file_modification_info,
 
       maybe_creator_file_synthetic_id,
       maybe_creator_category_synthetic_id,
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 f165fa6e5f..a9939fe586 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
@@ -136,7 +136,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
                 info!("Checking if path exists: {:?}", dep_path);
                 if !dep_path.exists() {
                     warn!("Path does not exist: {:?}", dep_path);
-                    download_file(model.url.clone(), dep_path.clone()).await.map_err(|e| ProcessSingleJobError::Other(e))?;
+                    download_file(model.url.clone(), dep_path.clone()).await.map_err(|e| ProcessSingleJobError::Other(e.into()))?;
                     info!("Downloaded model to {:?}", dep_path);
                 }
             }
@@ -155,22 +155,25 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
 
     let workflow_path = workflow_dir.join("prompt.json").to_str().unwrap().to_string();
 
-    //let retrieved_workflow_record =  get_weight_by_token(
+    // let retrieved_workflow_record =  get_weight_by_token(
     //    job_args.workflow_source,
     //    false,
     //    &deps.db.mysql_pool
-    //).await?.unwrap();
-
-    //let bucket_details = RemoteCloudBucketDetails {
+    // ).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(),
-    //};
-    //remote_cloud_file_client.download_file(bucket_details, workflow_path.clone()).await?;
+    // };
+    // remote_cloud_file_client.download_file(bucket_details, workflow_path.clone()).await?;
+    // info!("Downloaded workflow to {:?}", workflow_path);
 
     info!("Downloading workflow {:?}", &comfy_deps.workflow_bucket_path);
     info!("Downloading workflow to {:?}", workflow_path);
 
+    // download_file(comfy_deps.workflow_bucket_path.clone(), PathBuf::from(workflow_path.clone())).await.map_err(|e| ProcessSingleJobError::Other(e))?;
+
     args.job_dependencies
         .buckets
         .public_bucket_client
@@ -203,7 +206,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         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))?;
+            prompt_json = replace_json_value(prompt_json, &path, new_value).map_err(|e| ProcessSingleJobError::Other(e.into()))?;
         }
         // Save prompt.json
         let prompt_file = File::create(&workflow_path).unwrap();
@@ -475,6 +478,9 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
 
     // ==================== SAVE RECORDS ==================== //
 
+    // create a json detailing the args used to create the media
+    let args_json = serde_json::to_string(&job_args).map_err(|e| ProcessSingleJobError::Other(e.into()))?;
+
     info!("Saving ComfyUI result (media_files table record) ...");
 
     let (media_file_token, id) = insert_media_file_from_comfy_ui(InsertArgs {
@@ -490,6 +496,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         worker_hostname: &args.job_dependencies.job.info.container.hostname,
         worker_cluster: &args.job_dependencies.job.info.container.cluster_name,
         media_file_type: media_type,
+        extra_file_modification_info: Some(&args_json),
     })
         .await
         .map_err(|e| ProcessSingleJobError::Other(e))?;
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 a82046ef20..f6e283023b 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
@@ -9,6 +9,7 @@ use tokens::tokens::model_weights::ModelWeightToken;
 
 use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
 
+#[derive(Serialize, Debug)]
 pub struct JobArgs<'a> {
     pub workflow_source: &'a ModelWeightToken,
     pub output_path: &'a String,

From aa9e1c2fe47d949379d86db3a00c6fa15f207b2b Mon Sep 17 00:00:00 2001
From: salt <alotofcatz@gmail.com>
Date: Tue, 20 Feb 2024 07:41:05 -0800
Subject: [PATCH 2/3] new dependencies + download workflow from weight token

---
 .../comfy_ui/comfy_ui_dependencies.rs         | 10 ++++-
 .../workflow/comfy_ui/process_job.rs          | 38 +++++++------------
 2 files changed, 21 insertions(+), 27 deletions(-)

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 60733a8421..9b39c9877d 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
@@ -6,7 +6,6 @@ use crate::job::job_types::workflow::comfy_ui::comfy_ui_inference_command::Comfy
 pub struct ComfyDependencies {
     pub inference_command: ComfyInferenceCommand,
     pub dependency_tokens: RequiredModels,
-    pub workflow_bucket_path: String,
 }
 
 impl ComfyDependencies {
@@ -14,7 +13,6 @@ impl ComfyDependencies {
         Ok(Self {
             inference_command: ComfyInferenceCommand::from_env()?,
             dependency_tokens: RequiredModels::init(),
-            workflow_bucket_path: easyenv::get_env_string_required("COMFY_WORKFLOW_BUCKET_PATH")?,
         })
     }
 }
@@ -133,6 +131,14 @@ impl RequiredModels {
                     location: PathBuf::from("input/example.png"),
                     url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/example.png".to_string(),
                 },
+                ComfyDependency {
+                    location: PathBuf::from("custom_nodes/ComfyUI-AnimateDiff-Evolved/models/sd15_t2v_beta_motion.ckpt"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/sd15_t2v_beta_motion.ckpt".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("loras/LCM_LoRA_Weights_SD15.safetensors"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/LCM_LoRA_Weights_SD15.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 a9939fe586..621b517615 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
@@ -155,35 +155,24 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
 
     let workflow_path = workflow_dir.join("prompt.json").to_str().unwrap().to_string();
 
-    // 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(),
-    // };
-    // remote_cloud_file_client.download_file(bucket_details, workflow_path.clone()).await?;
-    // info!("Downloaded workflow to {:?}", workflow_path);
+    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(),
+    };
+    remote_cloud_file_client.download_file(bucket_details, workflow_path.clone()).await?;
+    info!("Downloaded workflow to {:?}", workflow_path);
 
-    info!("Downloading workflow {:?}", &comfy_deps.workflow_bucket_path);
     info!("Downloading workflow to {:?}", workflow_path);
 
     // download_file(comfy_deps.workflow_bucket_path.clone(), PathBuf::from(workflow_path.clone())).await.map_err(|e| ProcessSingleJobError::Other(e))?;
 
-    args.job_dependencies
-        .buckets
-        .public_bucket_client
-        .download_file_to_disk(&comfy_deps.workflow_bucket_path, &workflow_path)
-        .await
-        .map_err(|err| {
-            error!("could not download workflow: {:?}", err);
-            ProcessSingleJobError::from_anyhow_error(anyhow!("could not download workflow: {:?}", err))
-        })?;
-
     let maybe_args = job.maybe_inference_args
         .as_ref()
         .map(|args| args.args.as_ref())
@@ -213,7 +202,6 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         serde_json::to_writer(prompt_file, &prompt_json).unwrap();
     }
 
-
     // ==================== QUERY AND DOWNLOAD FILES ==================== //
 
     // Download SD model if specified

From 07d8125de9276798aad8c6de6a3c3c9632f10a7e Mon Sep 17 00:00:00 2001
From: salt <alotofcatz@gmail.com>
Date: Tue, 20 Feb 2024 07:42:27 -0800
Subject: [PATCH 3/3] remove excess message

---
 .../src/job/job_types/workflow/comfy_ui/process_job.rs          | 2 --
 1 file changed, 2 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 621b517615..74a246a5c6 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
@@ -169,8 +169,6 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
     remote_cloud_file_client.download_file(bucket_details, workflow_path.clone()).await?;
     info!("Downloaded workflow to {:?}", workflow_path);
 
-    info!("Downloading workflow to {:?}", workflow_path);
-
     // download_file(comfy_deps.workflow_bucket_path.clone(), PathBuf::from(workflow_path.clone())).await.map_err(|e| ProcessSingleJobError::Other(e))?;
 
     let maybe_args = job.maybe_inference_args