Merge pull request #294 from storytold/feature/comfyui

Feature/comfyui

Cargo.lock

@@ -3099,6 +3099,7 @@ dependencies = [
  "r2d2_redis",
  "rand 0.8.5",
  "redis_common",
+ "reqwest",
  "serde",
  "serde_derive",
  "serde_json",
@@ -3110,6 +3111,7 @@ dependencies = [
  "tokens",
  "tokio",
  "tts_common",
+ "walkdir",
 ]
 
 [[package]]

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?;

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 => {

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, WeightsWorkflowDescriptor};
 use enums::by_table::model_weights::{
     weights_category::WeightsCategory,
     weights_types::WeightsType,
@@ -686,31 +686,43 @@ 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 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 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 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()),
+        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("test-workflow.json".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,
@@ -719,30 +731,76 @@ pub async fn seed_workflows_for_testing_inference(mysql_pool: &Pool<MySql>, user
         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,
-        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 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: "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(weights2).await?;
+    // create_weight(weights2).await?;
+    // create_weight(weights3).await?;
+    // create_weight(weights4).await?;
 
     Ok(())
 }
@@ -1061,9 +1119,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_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(())
 }

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" => {

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;
 
@@ -189,4 +213,4 @@ impl FileDescriptor for WeightsWorkflowDescriptor {
     fn is_public(&self) -> bool {
         true
     }
-}
+}

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>,
 }

crates/service/job/inference_job/Cargo.toml

@@ -82,4 +82,6 @@ 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"
+jsonpath_lib = "0.3.0"
+walkdir = "2.4.0"
+reqwest = "0.11.22"

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

@@ -1,15 +1,97 @@
+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) url: String,
+}
+
+pub struct RequiredModels {
+    pub(crate) comfy: Vec<ComfyDependency>
+}
+
+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"),
+                    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(),
+                },
+            ],
+        }
+    }
+}

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

@@ -63,6 +63,7 @@ pub enum ExecutableOrCommand {
 
 pub struct InferenceArgs<'s> {
     pub stderr_output_file: &'s Path,
+    pub stdout_output_file: &'s Path,
 }
 
 impl ComfyInferenceCommand {
@@ -179,7 +180,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);

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

@@ -1,12 +1,13 @@
-use std::fs::read_to_string;
-use std::path::PathBuf;
+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 serde::de::Error;
 use serde_json::Value;
+use tokio::io::AsyncWriteExt;
+use walkdir::WalkDir;
 
 use buckets::public::media_files::bucket_file_path::MediaFileBucketPath;
 use cloud_storage::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails;
@@ -23,6 +24,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,13 +33,21 @@ 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,
 }
 
+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;
@@ -71,8 +81,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         .new_tempdir(&work_temp_dir)
         .map_err(|e| ProcessSingleJobError::from_io_error(e))?;
 
-
-    // ==================== QUERY AND DOWNLOAD FILES ==================== //
+    // ===================== 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;
@@ -85,53 +94,49 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         }
     };
 
-    // 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();
+    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?;
 
-            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);
+        // 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.");
         }
-        None => {}
+
+        Ok(())
     }
-    // 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);
+    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 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);
+                }
+            }
+        }
+        None => {
+            info!("No models specified to download")
         }
-        None => {}
     }
 
     // Download workflow to ComfyRunner
@@ -168,19 +173,95 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
     // 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 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).expect(
-                format!("Failed to replace json value at path {}", path).as_str()
-            );
+            prompt_json = replace_json_value(prompt_json, &path, new_value).map_err(|e| ProcessSingleJobError::Other(e))?;
         }
         // Save prompt.json
-        let prompt_file = std::fs::File::create(&workflow_path).unwrap();
+        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
+    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?.ok_or_else(|| ProcessSingleJobError::Other(anyhow!("SD model not found")))?;
+
+            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 = "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());
+            info!("Downloaded SD model to {:?}", sd_path);
+            maybe_sd_path = Some(sd_path.parse().unwrap());
+        }
+        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?.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(),
+                suffix: retrieved_lora_record.maybe_public_bucket_extension.unwrap(),
+            };
+
+            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());
+            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 => {}
+    }
+
     // ==================== SETUP FOR INFERENCE ==================== //
 
     info!("Ready for ComfyUI inference...");
@@ -192,21 +273,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) {
@@ -214,6 +310,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 {
@@ -222,49 +319,40 @@ 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);
+        }
+        let output_dir = root_comfy_path.join("output");
+        recursively_delete_files_in(&output_dir).unwrap();
 
         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());
-
-
-    // ==================== 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!("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 +366,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 +377,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 +388,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,16 +396,20 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
     // ==================== DELETE TEMP FILES ==================== //
 
     safe_delete_temp_file(&stderr_output_file);
-
-    // 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);
+    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);
     }
+    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);
+    }
+
+    let output_dir = root_comfy_path.join("output");
+    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);
@@ -333,7 +425,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,
@@ -360,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)
     })
-}
+}

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,
     })
 }

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...");