diff --git a/Cargo.lock b/Cargo.lock index dc04c73647..882400c47d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3521,6 +3521,7 @@ dependencies = [ "opentelemetry", "opentelemetry-otlp", "opentelemetry_sdk", + "primitives", "r2d2_redis", "rand 0.8.5", "redis_common", diff --git a/crates/lib/data/primitives/src/bool_to_str.rs b/crates/lib/data/primitives/src/bool_to_str.rs new file mode 100644 index 0000000000..2bf36c6d83 --- /dev/null +++ b/crates/lib/data/primitives/src/bool_to_str.rs @@ -0,0 +1,24 @@ + +/// Convert a boolean to a lower case string, with values of either "true" or "false". +pub fn bool_to_str(value: bool) -> &'static str { + if value { + "true" + } else { + "false" + } +} + +#[cfg(test)] +mod tests { + use crate::bool_to_str::bool_to_str; + + #[test] + fn test_true() { + assert_eq!(bool_to_str(true), "true"); + } + + #[test] + fn test_false() { + assert_eq!(bool_to_str(false), "false"); + } +} diff --git a/crates/lib/data/primitives/src/lib.rs b/crates/lib/data/primitives/src/lib.rs index c2003b13d9..d9b38a254e 100644 --- a/crates/lib/data/primitives/src/lib.rs +++ b/crates/lib/data/primitives/src/lib.rs @@ -18,6 +18,7 @@ #![allow(dead_code)] #![allow(non_snake_case)] +pub mod bool_to_str; pub mod lazy_any_option_true; pub mod str_to_bool; pub mod traits; diff --git a/crates/service/job/inference_job/Cargo.toml b/crates/service/job/inference_job/Cargo.toml index c8a42c44a3..4149b4337a 100644 --- a/crates/service/job/inference_job/Cargo.toml +++ b/crates/service/job/inference_job/Cargo.toml @@ -39,6 +39,7 @@ memory_caching = { path = "../../../lib/caching/memory_caching" } migration = { path = "../../../lib/migration" } mimetypes = { path = "../../../lib/files/mimetypes" } mysql_queries = { path = "../../../lib/mysql_queries" } +primitives = { path = "../../../lib/data/primitives" } redis_common = { path = "../../../lib/redis_common" } subprocess_common = { path = "../../../lib/subprocess_common" } thumbnail_generator = { path = "../../../lib/api_clients/thumbnail_generator"} diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/command_args.rs b/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/command_args.rs index 30516a45d2..d4d8a85678 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/command_args.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/command_args.rs @@ -1,6 +1,7 @@ use std::path::Path; use filesys::path_to_string::path_to_string; +use primitives::bool_to_str::bool_to_str; use subprocess_common::command_runner::command_args::CommandArgs; #[derive(Debug)] @@ -9,6 +10,7 @@ pub struct LivePortraitCommandArgs<'a> { pub driver_file: &'a Path, pub tempdir: &'a Path, pub output_file: &'a Path, + pub input_is_image: bool, pub stderr_output_file: &'a Path, pub stdout_output_file: &'a Path, @@ -30,6 +32,9 @@ impl CommandArgs for LivePortraitCommandArgs<'_> { command.push_str(" --output "); command.push_str(&path_to_string(self.output_file)); + command.push_str(" --input-is-image "); + command.push_str(bool_to_str(self.input_is_image)); + command.push_str(" "); command diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/process_live_portrait_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/process_live_portrait_job.rs index 1cb7a78292..9d6f4e0fc4 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/process_live_portrait_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/process_live_portrait_job.rs @@ -122,7 +122,7 @@ pub async fn process_live_portrait_job( let portrait_media_token = job_payload.portrait_media_file_token.ok_or_else(|| anyhow!("no portrait media token"))?; let portrait_file_path = work_temp_dir.path().join("portrait.bin"); - download_media_file(DownloadMediaFileArgs { + let portrait = download_media_file(DownloadMediaFileArgs { mysql_pool: &deps.db.mysql_pool, remote_cloud_file_client: &remote_cloud_file_client, media_file_token: &portrait_media_token, @@ -134,13 +134,21 @@ pub async fn process_live_portrait_job( let driver_media_token = job_payload.driver_media_file_token.ok_or_else(|| anyhow!("no driver media token"))?; let driver_file_path = work_temp_dir.path().join("driver.bin"); - download_media_file(DownloadMediaFileArgs { + let driver = download_media_file(DownloadMediaFileArgs { mysql_pool: &deps.db.mysql_pool, remote_cloud_file_client: &remote_cloud_file_client, media_file_token: &driver_media_token, download_path: &driver_file_path, }).await?; + let input_is_image = match portrait.media_file.media_type { + MediaFileType::Image + | MediaFileType::Jpg + | MediaFileType::Png + | MediaFileType::Gif => true, + _ => false, + }; + // ==================== RUN COMFY INFERENCE ==================== // info!("Preparing for ComfyUI inference..."); @@ -155,9 +163,6 @@ pub async fn process_live_portrait_job( info!("Running ComfyUI inference..."); - // TODO TEMP - let media_file_token = MediaFileToken::generate(); - let command_exit_status = comfy_deps .inference_command // TODO(bt,2024-07-15): Move this to its own runner. Just hacking this quickly. @@ -168,6 +173,7 @@ pub async fn process_live_portrait_job( output_file: &output_file_path, stderr_output_file: &stderr_output_file, stdout_output_file: &stdout_output_file, + input_is_image, }); let inference_duration = Instant::now().duration_since(inference_start_time);