Merge pull request #303 from storytold/feature/comfy-replication

save job args to file modification info

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,

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(),
+                },
             ],
         }
     }

crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs

@@ -149,7 +149,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);
                 }
             }
@@ -168,31 +168,21 @@ 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 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?;
+    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);
-
-    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))
-        })?;
+    // 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
         .as_ref()
@@ -216,14 +206,13 @@ 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();
         serde_json::to_writer(prompt_file, &prompt_json).unwrap();
     }
 
-
     // ==================== QUERY AND DOWNLOAD FILES ==================== //
 
     // Download SD model if specified
@@ -512,6 +501,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 {
@@ -527,6 +519,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))?;

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,