diff --git a/Cargo.lock b/Cargo.lock index daaded6471..4b1fff70f1 100644 --- a/Cargo.lock +++ b/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]] diff --git a/crates/cli/dev-database-seed/src/main.rs b/crates/cli/dev-database-seed/src/main.rs index 6ab4c2710c..e8ba82d32c 100644 --- a/crates/cli/dev-database-seed/src/main.rs +++ b/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?; diff --git a/crates/cli/dev-database-seed/src/seeding/media_files.rs b/crates/cli/dev-database-seed/src/seeding/media_files.rs index d9df7b7c3a..9d0caad288 100644 --- a/crates/cli/dev-database-seed/src/seeding/media_files.rs +++ b/crates/cli/dev-database-seed/src/seeding/media_files.rs @@ -43,7 +43,7 @@ pub async fn seed_test_videos(mysql_pool: &Pool, user_token: UserToken) - } pub async fn seed_media_seedtool(mysql_pool: &Pool) -> AnyhowResult<()> { - info!("Seeding weights..."); + info!("Seeding media..."); let user_token = match get_user_token_by_username(HANASHI_USERNAME, mysql_pool).await? { None => { diff --git a/crates/cli/dev-database-seed/src/seeding/model_weights.rs b/crates/cli/dev-database-seed/src/seeding/model_weights.rs index c77fe219c1..9d1ee6b494 100644 --- a/crates/cli/dev-database-seed/src/seeding/model_weights.rs +++ b/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, 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("

Description

".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, 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("

Description

".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("

Description

".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("

Description

".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("

Description

".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) -> 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(()) } diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs index 077b0265fc..3cbfe469c8 100644 --- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs +++ b/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" => { diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs index 0bbbfac9ce..2b3c3d4d92 100644 --- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs +++ b/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 } -} \ No newline at end of file +} diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs index 8f662a22b3..d60c9b1432 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs +++ b/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>, + + #[serde(rename = "in")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_input_file: Option, + + #[serde(rename = "out")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_output_path: Option, } diff --git a/crates/service/job/inference_job/Cargo.toml b/crates/service/job/inference_job/Cargo.toml index bd1e10467f..aeb2f5c193 100644 --- a/crates/service/job/inference_job/Cargo.toml +++ b/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" \ No newline at end of file +jsonpath_lib = "0.3.0" +walkdir = "2.4.0" +reqwest = "0.11.22" \ No newline at end of file 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 c9b0ebfd6a..f2234988f6 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 @@ -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 { 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 +} + +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(), + }, + ], + } + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs index d8fadd09dd..c878aa5a75 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs +++ b/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); 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 8b7d90070c..feef4cc8a1 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 @@ -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 { let job = args.job; let deps = args.job_dependencies; @@ -71,8 +81,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result) -> Result = 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 = 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 = 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 = 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 = 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) -> Result) -> Result = std::fs::read_dir(output_dir).unwrap() - .map(|res| res.map(|e| e.path())) - .collect::, 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) -> Result) -> Result) -> Result = std::fs::read_dir(output_dir).unwrap() - .map(|res| res.map(|e| e.path())) - .collect::, 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) -> Result Result { - replace_with(json, path, &mut |_| { +fn replace_json_value(json: Value, path: &str, new_value: NewValue) -> anyhow::Result { + // 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) }) -} \ No newline at end of file +} diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/validate_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/validate_job.rs index 989a07e813..a82046ef20 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/validate_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/validate_job.rs @@ -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>, pub maybe_sd_model: &'a Option, pub maybe_lora_model: &'a Option, - + pub maybe_input_file: &'a Option, } pub fn validate_job(job: &AvailableInferenceJob) -> Result { @@ -58,11 +60,20 @@ pub fn validate_job(job: &AvailableInferenceJob) -> Result 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, }) } diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs index 02d8694a46..a796f1b2d4 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs +++ b/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, maybe_json_modifications: Option>, maybe_workflow_config: Option, + maybe_input_file: Option, + maybe_output_path: Option, creator_set_visibility: Option, } @@ -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...");