From 5528993768aa51a0e2b5307ef7f625fb7dcea97a Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Sat, 3 Feb 2024 23:44:21 -0500
Subject: [PATCH] config for comfyui workflow

---
 .../comfy_ui/comfy_ui_dependencies.rs         |  2 +
 .../workflow/comfy_ui/process_job.rs          | 40 ++++++++++++++-----
 2 files changed, 31 insertions(+), 11 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 e85ca31fb3..60733a8421 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,6 +6,7 @@ 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 {
@@ -13,6 +14,7 @@ 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")?,
         })
     }
 }
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 7b4fad6890..a2cae5017d 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
@@ -29,6 +29,7 @@ use mysql_queries::queries::model_weights::get::get_weight::get_weight_by_token;
 
 use crate::job::job_loop::job_success_result::{JobSuccessResult, ResultEntity};
 use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
+use crate::job::job_types::workflow::comfy_ui::comfy_ui_dependencies::ComfyDependencies;
 use crate::job::job_types::workflow::comfy_ui::comfy_ui_inference_command::InferenceArgs;
 use crate::job::job_types::workflow::comfy_ui::validate_job::validate_job;
 use crate::job_dependencies::JobDependencies;
@@ -122,6 +123,11 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         Ok(())
     }
 
+    let comfy_deps = match args.job_dependencies.job.job_specific_dependencies.maybe_comfy_ui_dependencies {
+        Some(ref deps) => deps,
+        None => return Err(ProcessSingleJobError::from(anyhow!("no comfy deps"))),
+    };
+
     let all_models = &args.job_dependencies.job.job_specific_dependencies.maybe_comfy_ui_dependencies;
     match all_models {
         Some(models) => {
@@ -145,19 +151,31 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
     if !workflow_dir.exists() {
         std::fs::create_dir_all(&workflow_dir).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(),
-    };
     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 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?;
+
+    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 VAE: {:?}", err);
+            ProcessSingleJobError::from_anyhow_error(anyhow!("could not download VAE: {:?}", err))
+        })?;
 
     let maybe_args = job.maybe_inference_args
         .as_ref()