From 550a11c493151ef7afa49da60968c8f68f5982db Mon Sep 17 00:00:00 2001 From: saltacc Date: Fri, 19 Jan 2024 12:40:19 -0800 Subject: [PATCH] make it work :fire: --- Cargo.lock | 13 + .../inference_category.rs | 1 + .../inference_model_type.rs | 3 +- .../src/seeding/model_weights.rs | 53 +++- .../remote_cloud_bucket_details.rs | 1 + .../generic_inference_args.rs | 4 +- .../workflow_payload.rs | 8 - .../create/insert_media_file_from_comfy_ui.rs | 8 +- .../src/queries/media_files/create/mod.rs | 2 +- crates/service/job/inference_job/Cargo.toml | 1 + .../config/inference-job.development.env | 6 +- .../job/job_types/vc/rvc_v2/process_job.rs | 4 +- .../job_types/vc/so_vits_svc/process_job.rs | 4 +- .../comfy_ui/comfy_ui_inference_command.rs | 2 +- .../workflow/comfy_ui/download_model_file.rs | 66 +++++ .../job/job_types/workflow/comfy_ui/mod.rs | 3 +- .../workflow/comfy_ui/process_job.rs | 277 +++++++++++------- .../workflow/comfy_ui/validate_job.rs | 7 +- .../endpoints/workflows/enqueue_comfy_ui.rs | 4 +- 19 files changed, 326 insertions(+), 141 deletions(-) create mode 100644 crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/download_model_file.rs diff --git a/Cargo.lock b/Cargo.lock index fcc378de3..61f301f0b 100644 --- a/Cargo.lock +++ b/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", diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs index 641dab68c..2f591b463 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs +++ b/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)), } } diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs index cd139eb49..3e2f5c4de 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs +++ b/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)); 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 497d996af..fa9edaad8 100644 --- a/crates/cli/dev-database-seed/src/seeding/model_weights.rs +++ b/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, 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, 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("

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 + }; + + create_weight(weights1).await?; + create_weight(weights2).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 68627337f..9ea88495a 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 @@ -40,6 +40,7 @@ impl RemoteCloudBucketDetails { } }, "svc" => Box::new(weights_descriptor::WeightsSVCDescriptor {}), + "workflow" => Box::new(weights_descriptor::WeightsWorkflowDescriptor {}), _ => panic!("Unknown prefix: {}", self.prefix) } } diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs index 10547fca3..962d5a93c 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs +++ b/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 { 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 2aadc215d..8f662a22b 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 @@ -20,14 +20,6 @@ pub struct WorkflowArgs { #[serde(skip_serializing_if = "Option::is_none")] pub maybe_lora_model: Option, - #[serde(rename = "prompt")] - #[serde(skip_serializing_if = "Option::is_none")] - pub maybe_prompt: Option, - - #[serde(rename = "negative_prompt")] - #[serde(skip_serializing_if = "Option::is_none")] - pub maybe_negative_prompt: Option, - #[serde(rename = "workflow_config")] #[serde(skip_serializing_if = "Option::is_none")] pub maybe_workflow_config: Option, diff --git a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_comfy_ui.rs b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_comfy_ui.rs index 3db08a99f..178e83f43 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_comfy_ui.rs +++ b/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, diff --git a/crates/lib/mysql_queries/src/queries/media_files/create/mod.rs b/crates/lib/mysql_queries/src/queries/media_files/create/mod.rs index b705c6d45..7d1e1cc8a 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/create/mod.rs +++ b/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; \ No newline at end of file diff --git a/crates/service/job/inference_job/Cargo.toml b/crates/service/job/inference_job/Cargo.toml index 80897d083..bd1e10467 100644 --- a/crates/service/job/inference_job/Cargo.toml +++ b/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" \ No newline at end of file diff --git a/crates/service/job/inference_job/config/inference-job.development.env b/crates/service/job/inference_job/config/inference-job.development.env index 34dd79737..36231d9bc 100644 --- a/crates/service/job/inference_job/config/inference-job.development.env +++ b/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" # ================================= # # # diff --git a/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs b/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs index 81bc88948..ab33c2048 100644 --- a/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs +++ b/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 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 None, PolymorphicInferenceArgs::Ig(_) => None, PolymorphicInferenceArgs::Mc(_) => None, - PolymorphicInferenceArgs::Wf(_) => None, + PolymorphicInferenceArgs::Cu(_) => None, }) .flatten(); diff --git a/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs b/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs index 110d3f9a1..b5f519f2c 100644 --- a/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs +++ b/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 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 None, PolymorphicInferenceArgs::Ig(_) => None, PolymorphicInferenceArgs::Mc(_) => None, - PolymorphicInferenceArgs::Wf(_) => None, + PolymorphicInferenceArgs::Cu(_) => None, }) .flatten(); 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 760439af0..da134c50e 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 @@ -34,7 +34,7 @@ static IGNORED_ENVIRONMENT_VARS : Lazy> = 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, diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/download_model_file.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/download_model_file.rs new file mode 100644 index 000000000..1c32e96b2 --- /dev/null +++ b/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, + public_bucket_client: &BucketClient, + job_progress_reporter: &mut Box, + job: &AvailableInferenceJob, + temp_dir_creator: &ScopedTempDirCreator, + download_path: &PathBuf, +) -> Result { + + 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, + }) +} diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/mod.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/mod.rs index 60a9e9ede..141ca8230 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/mod.rs +++ b/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; \ No newline at end of file +pub(crate) mod process_job; +mod download_model_file; \ No newline at end of file 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 237729afd..04863ca5d 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,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) -> Result { + 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 = 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 = 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 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 = std::fs::read_dir(output_dir).unwrap() + .map(|res| res.map(|e| e.path())) + .collect::, 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 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) -> 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); + } // 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) -> Result) -> Result Result { + 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)) + }) } \ 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 ed4309799..989a07e81 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 @@ -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>, + pub maybe_sd_model: &'a Option, + pub maybe_lora_model: &'a Option, + } pub fn validate_job(job: &AvailableInferenceJob) -> Result { @@ -42,7 +45,7 @@ pub fn validate_job(job: &AvailableInferenceJob) -> Result 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