make it work 🔥

Cargo.lock

@@ -3084,6 +3084,7 @@ dependencies = [
  "hyper",
  "images",
  "jobs_common",
+ "jsonpath_lib",
  "lfu",
  "log",
  "media",
@@ -3225,6 +3226,17 @@ dependencies = [
  "wasm-bindgen",
 ]
 
+[[package]]
+name = "jsonpath_lib"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "eaa63191d68230cccb81c5aa23abd53ed64d83337cacbb25a7b8c7979523774f"
+dependencies = [
+ "log",
+ "serde",
+ "serde_json",
+]
+
 [[package]]
 name = "jwt"
 version = "0.13.0"
@@ -5401,6 +5413,7 @@ version = "1.0.108"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3d1c7e3eac408d115102c4c24ad393e0821bb3a5df4d506a80f85f7a742a526b"
 dependencies = [
+ "indexmap",
  "itoa 1.0.9",
  "ryu",
  "serde",

crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs

@@ -65,6 +65,7 @@ impl InferenceCategory {
       "video_filter" => Ok(Self::VideoFilter),
       "image_generation" => Ok(Self::ImageGeneration),
       "mocap" => Ok(Self::Mocap),
+      "workflow" => Ok(Self::Workflow),
       _ => Err(format!("invalid value: {:?}", value)),
     }
   }

crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs

@@ -96,6 +96,7 @@ impl InferenceModelType {
       InferenceModelType::StableDiffusion,
       InferenceModelType::MocapNet,
       InferenceModelType::StyleTTS2,
+      InferenceModelType::ComfyUi,
     ])
   }
 }
@@ -154,7 +155,7 @@ mod tests {
   fn all_variants() {
     // Static check
     let mut variants = InferenceModelType::all_variants();
-    assert_eq!(variants.len(), 10);
+    assert_eq!(variants.len(), 11);
     assert_eq!(variants.pop_first(), Some(InferenceModelType::RvcV2));
     assert_eq!(variants.pop_first(), Some(InferenceModelType::SadTalker));
     assert_eq!(variants.pop_first(), Some(InferenceModelType::SoVitsSvc));

crates/cli/dev-database-seed/src/seeding/model_weights.rs

@@ -684,14 +684,22 @@ 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_token4 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously();
+    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 mut path_to_comfy = get_seed_tool_data_root();
-    path_to_comfy.push("models/workflows/comfyui/test-workflow.json");
+    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 metadata4 = remote_cloud_file_client.upload_file(comfy_weights_descriptor,path_to_comfy.as_path().to_str().unwrap()).await?;
-    let weights4 = CreateModelWeightsArgs {
-        token: &model_weight_token4, // replace with actual ModelWeightToken
+    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 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(),
@@ -703,15 +711,38 @@ pub async fn seed_workflows_for_testing_inference(mysql_pool: &Pool<MySql>, user
         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: 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),
+        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
     };
-    create_weight(weights4).await?;
+    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
+    };
+
+    create_weight(weights1).await?;
+    create_weight(weights2).await?;
 
     Ok(())
 }

crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs

@@ -40,6 +40,7 @@ impl RemoteCloudBucketDetails {
                 }
             },
             "svc" => Box::new(weights_descriptor::WeightsSVCDescriptor {}),
+            "workflow" => Box::new(weights_descriptor::WeightsWorkflowDescriptor {}),
             _ => panic!("Unknown prefix: {}", self.prefix)
         }
     }

crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs

@@ -113,8 +113,8 @@ pub enum PolymorphicInferenceArgs {
   // Mocap (Short name to save space when serializing.)
   Mc(MocapArgs),
 
-  // Workflow (Short name to save space when serializing.)
-  Wf(WorkflowArgs)
+  // ComfyUI (Short name to save space when serializing.)
+  Cu(WorkflowArgs)
 }
 
 impl GenericInferenceArgs {

crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs

@@ -20,14 +20,6 @@ pub struct WorkflowArgs {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub maybe_lora_model: Option<ModelWeightToken>,
 
-    #[serde(rename = "prompt")]
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub maybe_prompt: Option<String>,
-
-    #[serde(rename = "negative_prompt")]
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub maybe_negative_prompt: Option<String>,
-
     #[serde(rename = "workflow_config")]
     #[serde(skip_serializing_if = "Option::is_none")]
     pub maybe_workflow_config: Option<ModelWeightToken>,

crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_comfy_ui.rs

@@ -31,12 +31,10 @@ pub struct InsertArgs<'a> {
     pub is_on_prem: bool,
     pub worker_hostname: &'a str,
     pub worker_cluster: &'a str,
+    pub media_file_type: MediaFileType,
 }
 
-pub async fn insert_media_file_from_comfy_ui(
-    args: InsertArgs<'_>,
-    media_file_type: MediaFileType,
-) -> AnyhowResult<(MediaFileToken, u64)>
+pub async fn insert_media_file_from_comfy_ui(args: InsertArgs<'_>) -> AnyhowResult<(MediaFileToken, u64)>
 {
     let result_token = MediaFileToken::generate();
 
@@ -109,7 +107,7 @@ SET
       ORIGIN_PRODUCT_CATEGORY.to_str(),
       ORIGIN_MODEL_TYPE.to_str(),
 
-      media_file_type.to_str(),
+      args.media_file_type.to_str(),
       args.maybe_mime_type,
       args.file_size_bytes,
 

crates/lib/mysql_queries/src/queries/media_files/create/mod.rs

@@ -6,4 +6,4 @@ pub mod insert_media_file_from_tacotron2;
 pub mod insert_media_file_from_voice_conversion;
 pub mod insert_media_file_from_zero_shot_tts;
 pub mod insert_media_file_from_mocapnet;
-mod insert_media_file_from_comfy_ui;
+pub mod insert_media_file_from_comfy_ui;

crates/service/job/inference_job/Cargo.toml

@@ -82,3 +82,4 @@ 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"

crates/service/job/inference_job/config/inference-job.development.env

@@ -124,8 +124,10 @@ MOCAPNET_INFERENCE_EXECUTABLE_OR_COMMAND="pythonVirtualEnvironment/bin/python -m
 #            ComfyUI                #
 #                                   #
 # ================================= #
-COMFY_INFERENCE_ROOT_DIRECTORY="/home/salt/clone/storyteller-ml/image_generation/ComfyLauncher"
-COMFY_INFERENCE_EXECUTABLE_OR_COMMAND="../ComfyUI/venv/bin/python ComfyLauncher.py"
+COMFY_INFERENCE_ROOT_DIRECTORY="/home/salt/clone/storyteller-ml/image_generation/ComfyUI"
+COMFY_INFERENCE_EXECUTABLE_OR_COMMAND="venv/bin/python ../ComfyLauncher/ComfyRunner.py"
+STYLE_TTS2_INFERENCE_ROOT_DIRECTORY="/home/salt/clone/storyteller-ml/tts/style-tts2"
+STYLE_TTS2_INFERENCE_COMMAND="venv/bin/python3.8 main.py"
 
 # ================================= #
 #                                   #

crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs

@@ -262,7 +262,7 @@ pub async fn process_job(args: RvcV2ProcessJobArgs<'_>) -> Result<JobSuccessResu
         PolymorphicInferenceArgs::Rr(_) => None,
         PolymorphicInferenceArgs::Ig(_) => None,
         PolymorphicInferenceArgs::Mc(_) => None,
-        PolymorphicInferenceArgs::Wf(_) => None,
+        PolymorphicInferenceArgs::Cu(_) => None,
       })
       .flatten();
 
@@ -274,7 +274,7 @@ pub async fn process_job(args: RvcV2ProcessJobArgs<'_>) -> Result<JobSuccessResu
         PolymorphicInferenceArgs::Rr(_) => None,
         PolymorphicInferenceArgs::Ig(_) => None,
         PolymorphicInferenceArgs::Mc(_) => None,
-        PolymorphicInferenceArgs::Wf(_) => None,
+        PolymorphicInferenceArgs::Cu(_) => None,
       })
       .flatten();
 

crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs

@@ -182,7 +182,7 @@ pub async fn process_job(args: SoVitsSvcProcessJobArgs<'_>) -> Result<JobSuccess
         PolymorphicInferenceArgs::Rr(_) => None,
         PolymorphicInferenceArgs::Ig(_) => None,
         PolymorphicInferenceArgs::Mc(_) => None,
-        PolymorphicInferenceArgs::Wf(_) => None,
+        PolymorphicInferenceArgs::Cu(_) => None,
       })
       .flatten()
       .unwrap_or(false);
@@ -195,7 +195,7 @@ pub async fn process_job(args: SoVitsSvcProcessJobArgs<'_>) -> Result<JobSuccess
         PolymorphicInferenceArgs::Rr(_) => None,
         PolymorphicInferenceArgs::Ig(_) => None,
         PolymorphicInferenceArgs::Mc(_) => None,
-        PolymorphicInferenceArgs::Wf(_) => None,
+        PolymorphicInferenceArgs::Cu(_) => None,
       })
       .flatten();
 

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

@@ -34,7 +34,7 @@ static IGNORED_ENVIRONMENT_VARS : Lazy<HashSet<String>> = Lazy::new(|| {
 #[derive(Clone)]
 pub struct ComfyInferenceCommand {
     /// Where the code lives
-    comfy_root_code_directory: PathBuf,
+    pub(crate) comfy_root_code_directory: PathBuf,
 
     /// A single executable script or a much larger bash command.
     executable_or_command: ExecutableOrCommand,

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

@@ -0,0 +1,66 @@
+use std::path::PathBuf;
+
+use anyhow::anyhow;
+use tempdir::TempDir;
+
+use buckets::public::media_files::bucket_file_path::MediaFileBucketPath;
+use cloud_storage::bucket_client::BucketClient;
+use jobs_common::job_progress_reporter::job_progress_reporter::JobProgressReporter;
+use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob;
+use mysql_queries::queries::model_weights::get_weight::RetrievedModelWeight;
+
+use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
+use crate::util::maybe_download_file_from_bucket::{maybe_download_file_from_bucket, MaybeDownloadArgs};
+use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator;
+
+pub struct ModelFile {
+    pub filesystem_path: PathBuf,
+}
+// TODO replace this with remote cloud file manager.
+pub async fn download_model_file(
+    model_record: &Option<RetrievedModelWeight>,
+    public_bucket_client: &BucketClient,
+    job_progress_reporter: &mut Box<dyn JobProgressReporter>,
+    job: &AvailableInferenceJob,
+    temp_dir_creator: &ScopedTempDirCreator,
+    download_path: &PathBuf,
+) -> Result<ModelFile, ProcessSingleJobError> {
+
+    let bucket_object_path;
+    let model_filename;
+    match model_record {
+        Some(model_record) => {
+            let model_bucket_path =
+                MediaFileBucketPath::from_object_hash(
+                    &model_record.public_bucket_hash,
+                    model_record.maybe_public_bucket_prefix.as_deref(),
+                    model_record.maybe_public_bucket_extension.as_deref(),
+                );
+            let base_name = &model_record.public_bucket_hash;
+            let extension = model_record.maybe_public_bucket_extension.as_deref().unwrap_or("bin");
+            model_filename = format!("{base_name}.{extension}");
+            bucket_object_path = model_bucket_path.to_full_object_pathbuf();
+        }
+        None => {
+            return Err(ProcessSingleJobError::from_anyhow_error(
+                anyhow!("could not find model file")))
+        }
+    };
+    let downloaded_filesystem_path = download_path.join(&model_filename);
+
+    maybe_download_file_from_bucket(MaybeDownloadArgs {
+        name_or_description_of_file: "model file",
+        final_filesystem_file_path: &downloaded_filesystem_path,
+        bucket_object_path: &bucket_object_path,
+        bucket_client: public_bucket_client,
+        job_progress_reporter,
+        job_progress_update_description: "downloading",
+        job_id: job.id.0,
+        scoped_tempdir_creator: &temp_dir_creator,
+        maybe_existing_file_minimum_size_required: None,
+    }).await?;
+
+    Ok(ModelFile {
+        filesystem_path: downloaded_filesystem_path,
+    })
+}

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

@@ -2,4 +2,5 @@ mod validate_job;
 mod categorize_error;
 mod comfy_ui_inference_command;
 pub(crate) mod comfy_ui_dependencies;
-pub(crate) mod process_job;
+pub(crate) mod process_job;
+mod download_model_file;

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

@@ -1,31 +1,40 @@
 use std::fs::read_to_string;
+use std::path::PathBuf;
 use std::time::Instant;
 
-use anyhow::{anyhow, Error};
+use anyhow::{anyhow};
 use log::{error, info, warn};
-use serde::{Deserialize, Serialize};
+use jsonpath_lib::replace_with;
+use serde::de::Error;
+
+use serde_json::Value;
 
 use buckets::public::media_files::bucket_file_path::MediaFileBucketPath;
+use cloud_storage::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails;
+use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient;
 use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType;
+use enums::by_table::media_files::media_file_type::MediaFileType;
 use filesys::check_file_exists::check_file_exists;
 use filesys::file_size::file_size;
 use filesys::safe_delete_temp_directory::safe_delete_temp_directory;
 use filesys::safe_delete_temp_file::safe_delete_temp_file;
 use hashing::sha256::sha256_hash_file::sha256_hash_file;
 use mimetypes::mimetype_for_file::get_mimetype_for_file;
-use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs::Rr;
+use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs::{Cu};
+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_rerender::{insert_media_file_from_rerender, InsertArgs};
+use mysql_queries::queries::media_files::create::insert_media_file_from_comfy_ui::{insert_media_file_from_comfy_ui, InsertArgs};
+use mysql_queries::queries::model_weights::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_inference_command::InferenceArgs;
+use crate::job::job_types::workflow::comfy_ui::download_model_file::download_model_file;
 use crate::job::job_types::workflow::comfy_ui::validate_job::validate_job;
 use crate::job_dependencies::JobDependencies;
 
 
 const BUCKET_FILE_PREFIX: &str = "fakeyou_";
-// const BUCKET_FILE_EXTENSION: &str = ".mp4";
 
 pub struct ComfyProcessJobArgs<'a> {
     pub job_dependencies: &'a JobDependencies,
@@ -56,7 +65,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
 
     // ==================== TEMP DIR ==================== //
 
-    let work_temp_dir = format!("temp_rerender_inference_{}", job.id.0);
+    let work_temp_dir = format!("temp_comfy_inference{}", job.id.0);
 
     // NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem.
     let work_temp_dir = args.job_dependencies
@@ -67,20 +76,113 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
 
 
     // ==================== QUERY AND DOWNLOAD FILES ==================== //
+    let root_comfy_path = model_dependencies.inference_command.comfy_root_code_directory.clone();
 
-    //TODO: Download all deps
+    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")));
+        }
+    };
 
-    // let video_path = download_video_file(
-    //     &job_args.video_source,
-    //     &args.job_dependencies.buckets.public_bucket_client,
-    //     &mut job_progress_reporter,
-    //     job,
-    //     &args.job_dependencies.fs.scoped_temp_dir_creator_for_work,
-    //     &work_temp_dir,
-    //     &deps.db.mysql_pool
-    // ).await?;
-    //
-    // info!("Downloaded video file: {:?}", video_path.filesystem_path);
+    // Download SD model if specified
+    let mut maybe_sd_path: Option<PathBuf> = None;
+    match job_args.maybe_sd_model {
+        Some(sd_model) => {
+            let sd_dir = root_comfy_path.join("models").join("checkpoints");
+            let retrieved_sd_record =  get_weight_by_token(
+                sd_model,
+                false,
+                &deps.db.mysql_pool
+            ).await?.unwrap();
+
+            let bucket_details = RemoteCloudBucketDetails {
+                object_hash: retrieved_sd_record.public_bucket_hash,
+                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_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());
+            info!("Downloaded SD model to {:?}", sd_path);
+        }
+        None => {}
+    }
+    // Download Lora model if specified
+    let mut maybe_lora_path: Option<PathBuf> = None;
+    match job_args.maybe_lora_model {
+        Some(lora_model) => {
+            let lora_dir = root_comfy_path.join("models").join("checkpoints");
+            let retrieved_lora_record =  get_weight_by_token(
+                lora_model,
+                false,
+                &deps.db.mysql_pool
+            ).await?.unwrap();
+            let bucket_details = RemoteCloudBucketDetails {
+                object_hash: retrieved_lora_record.public_bucket_hash,
+                prefix: retrieved_lora_record.maybe_public_bucket_prefix.unwrap(),
+                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_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());
+            info!("Downloaded Lora model to {:?}", lora_path);
+        }
+        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 ==================== //
 
@@ -89,78 +191,8 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
     job_progress_reporter.log_status("running inference")
         .map_err(|e| ProcessSingleJobError::Other(e))?;
 
-    let output_video_fs_path = work_temp_dir.path().join("blend.mp4");
-
     info!("Running ComfyUI inference...");
 
-    // info!("Expected output video filename: {:?}", &output_video_fs_path);
-
-    // TODO: Limit output length for non-premium (???)
-
-    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!("Rerender args not found"))),
-        Some(args) => args,
-    };
-
-    let rr_args = match poly_args {
-        Rr(args) => args,
-        _ => return Err(ProcessSingleJobError::Other(anyhow!("Rerender args not found"))),
-    };
-
-
-    // =================== LOAD MODELS AND LORAS ==================== //
-
-    // let sd_model_token = rr_args.maybe_sd_model_token.as_ref().ok_or(ProcessSingleJobError::Other(anyhow!("SD model not found")))?;
-    //
-    // let retrieved_sd_record =  get_weight_by_token(
-    //     sd_model_token,
-    //     false,
-    //     &deps.db.mysql_pool
-    // ).await?;
-    //
-    // let model_path = download_model_file(
-    //     &retrieved_sd_record,
-    //     &args.job_dependencies.buckets.public_bucket_client,
-    //     &mut job_progress_reporter,
-    //     job,
-    //     &args.job_dependencies.fs.scoped_temp_dir_creator_for_work,
-    //     &work_temp_dir).await?;
-    //
-    // info!("Downloaded model file: {:?}", model_path.filesystem_path);
-    // let sd_model = model_path.filesystem_path.to_str().unwrap().to_string();
-    //
-    // let lora_path = match rr_args.maybe_lora_model_token.clone() {
-    //     Some(lora_model_token) => {
-    //         let retrieved_lora_record = get_weight_by_token(
-    //             &lora_model_token,
-    //             false,
-    //             &deps.db.mysql_pool
-    //         ).await?;
-    //
-    //         let lora_path = download_model_file(
-    //             &retrieved_lora_record,
-    //             &args.job_dependencies.buckets.public_bucket_client,
-    //             &mut job_progress_reporter,
-    //             job,
-    //             &args.job_dependencies.fs.scoped_temp_dir_creator_for_work,
-    //             &work_temp_dir).await?;
-    //
-    //         info!("Downloaded lora file: {:?}", lora_path.filesystem_path);
-    //
-    //         Some(lora_path.filesystem_path)
-    //     }
-    //     None => None
-    // };
-    // let sd_lora = match lora_path {
-    //     Some(lora_path) => Some(lora_path.to_str().unwrap().to_string()),
-    //     None => None
-    // };
-
     // ==================== RUN INFERENCE SCRIPT ==================== //
     let stderr_output_file = work_temp_dir.path().join("stderr.txt");
     let inference_start_time = Instant::now();
@@ -184,26 +216,39 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
             warn!("Captured stderr output: {}", contents);
         }
 
-        safe_delete_temp_file(&output_video_fs_path);
         safe_delete_temp_file(&stderr_output_file);
         safe_delete_temp_directory(&work_temp_dir);
+        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);
+        }
 
         return Err(error);
     }
 
-    // ==================== CHECK NON-WATERMARKED RESULT ==================== //
+    // ==================== GET OUTPUT FILE ======================== //
 
-    info!("Checking that output file exists: {:?} ...", output_video_fs_path);
+    // 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();
 
-    check_file_exists(&output_video_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?;
+    entries.sort_by_key(|path| std::fs::metadata(path).unwrap().modified().unwrap());
 
-    // ==================== OPTIONAL WATERMARK ==================== //
-
-    let finished_file = output_video_fs_path.clone();
 
     // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== //
 
-    info!("Checking that output watermark file exists: {:?} ...", finished_file);
+    // 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!("Interrogating result file size ...");
@@ -226,6 +271,14 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         _ => return Err(ProcessSingleJobError::Other(anyhow!("Mimetype not supported: {}", mimetype))),
     };
 
+    // determine media type from mime type
+    let media_type = match mimetype.as_str() {
+        "video/mp4" => MediaFileType::Video,
+        "image/png" => MediaFileType::Image,
+        "image/jpeg" => MediaFileType::Image,
+        _ => return Err(ProcessSingleJobError::Other(anyhow!("Mimetype not supported: {}", mimetype))),
+    };
+
     info!("Calculating sha256...");
 
     let file_checksum = sha256_hash_file(&finished_file)
@@ -244,7 +297,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
 
     let result_bucket_object_pathbuf = result_bucket_location.to_full_object_pathbuf();
 
-    info!("Video destination bucket path: {:?}", &result_bucket_object_pathbuf);
+    info!("Output file destination bucket path: {:?}", &result_bucket_object_pathbuf);
 
     info!("Uploading media ...");
 
@@ -257,9 +310,17 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
 
     // ==================== DELETE TEMP FILES ==================== //
 
-    safe_delete_temp_file(&output_video_fs_path);
     safe_delete_temp_file(&stderr_output_file);
-    safe_delete_temp_directory(&work_temp_dir);
+
+    // 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);
+    }
 
     // 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);
@@ -268,7 +329,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
 
     info!("Saving ComfyUI result (media_files table record) ...");
 
-    let (media_file_token, id) = insert_media_file_from_rerender(InsertArgs {
+    let (media_file_token, id) = insert_media_file_from_comfy_ui(InsertArgs {
         pool: &args.job_dependencies.db.mysql_pool,
         job: &job,
         maybe_mime_type: Some(&mimetype),
@@ -280,11 +341,12 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         is_on_prem: args.job_dependencies.job.info.container.is_on_prem,
         worker_hostname: &args.job_dependencies.job.info.container.hostname,
         worker_cluster: &args.job_dependencies.job.info.container.cluster_name,
+        media_file_type: media_type,
     })
         .await
         .map_err(|e| ProcessSingleJobError::Other(e))?;
 
-    info!("Rerender Done.");
+    info!("ComfyUI Done.");
 
     job_progress_reporter.log_status("done")
         .map_err(|e| ProcessSingleJobError::Other(e))?;
@@ -299,4 +361,17 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         }),
         inference_duration,
     })
+}
+
+fn replace_json_value(json: Value, path: &str, new_value: NewValue) -> Result<Value, serde_json::Error> {
+    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::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))
+    })
 }

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

@@ -11,6 +11,9 @@ use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
 pub struct JobArgs<'a> {
     pub workflow_source: &'a ModelWeightToken,
     pub maybe_json_modifications: &'a Option<HashMap<String, NewValue>>,
+    pub maybe_sd_model: &'a Option<ModelWeightToken>,
+    pub maybe_lora_model: &'a Option<ModelWeightToken>,
+
 }
 
 pub fn validate_job(job: &AvailableInferenceJob) -> Result<JobArgs, ProcessSingleJobError> {
@@ -42,7 +45,7 @@ pub fn validate_job(job: &AvailableInferenceJob) -> Result<JobArgs, ProcessSingl
     };
 
     let inference_args = match inference_args {
-        PolymorphicInferenceArgs::Wf(inference_args) => inference_args,
+        PolymorphicInferenceArgs::Cu(inference_args) => inference_args,
         _ => {
             return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!")));
         }
@@ -58,6 +61,8 @@ pub fn validate_job(job: &AvailableInferenceJob) -> Result<JobArgs, ProcessSingl
 
     Ok(JobArgs {
         workflow_source,
+        maybe_sd_model: &inference_args.maybe_sd_model,
+        maybe_lora_model: &inference_args.maybe_lora_model,
         maybe_json_modifications: &inference_args.maybe_json_modifications,
     })
 }

crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs

@@ -176,8 +176,6 @@ pub async fn enqueue_comfy_ui_handler(
     let inference_args = WorkflowArgs {
         maybe_sd_model: request.maybe_sd_model.clone(),
         maybe_lora_model: request.maybe_lora_model.clone(),
-        maybe_prompt: request.maybe_prompt.clone(),
-        maybe_negative_prompt: request.maybe_negative_prompt.clone(),
         maybe_json_modifications: request.maybe_json_modifications.clone(),
         maybe_workflow_config: request.maybe_workflow_config.clone(),
     };
@@ -195,7 +193,7 @@ pub async fn enqueue_comfy_ui_handler(
         maybe_max_duration_seconds: None,
         maybe_inference_args: Some(GenericInferenceArgs {
             inference_category: Some(InferenceCategoryAbbreviated::Workflow),
-            args: Some(PolymorphicInferenceArgs::Wf(inference_args)),
+            args: Some(PolymorphicInferenceArgs::Cu(inference_args)),
         }),
         maybe_creator_user_token: maybe_user_token.as_ref(),
         maybe_avt_token: maybe_avt_token.as_ref(),