From 2bee82f0ae11d01e2ebde0e81c48da9e94cd2aae Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 16 Jul 2024 03:45:24 -0400 Subject: [PATCH] connect the last few pieces --- .../src/job/job_loop/process_single_job.rs | 120 +++++++++++------- .../workflow/process_single_workflow_job.rs | 20 ++- .../enqueue_live_portrait_workflow_handler.rs | 2 +- 3 files changed, 93 insertions(+), 49 deletions(-) diff --git a/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs b/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs index 719165afcd..cf40cbed64 100644 --- a/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs +++ b/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs @@ -13,6 +13,7 @@ use mysql_queries::queries::generic_inference::job::mark_generic_inference_job_s use redis_common::redis_keys::RedisKeys; use crate::job::job_loop::determine_dependency_status::determine_dependency_status; +use crate::job::job_loop::job_success_result::JobSuccessResult; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; use crate::job::job_loop::process_single_job_success_case::ProcessSingleJobSuccessCase; use crate::job::job_types::format_conversion::process_single_format_conversion_job::process_single_format_conversion_job; @@ -182,50 +183,10 @@ async fn do_process_single_job( // ==================== HANDLE DIFFERENT INFERENCE TYPES ==================== // - // TODO(bt,2024-07-16): This logic needs to be cleaned up and better coordinated with the setup env vars. - - let job_success_result = match job.inference_category { - InferenceCategory::LipsyncAnimation => { - process_single_lipsync_job(job_dependencies, job).await? - } - InferenceCategory::TextToSpeech => { - process_single_tts_job(job_dependencies, job).await? - } - InferenceCategory::VoiceConversion => { - process_single_vc_job(job_dependencies, job).await? - } - InferenceCategory::VideoFilter => { - process_single_vf_job(job_dependencies, job).await? - } - InferenceCategory::ImageGeneration => { - process_single_ig_job(job_dependencies, job).await? - } - InferenceCategory::Mocap => { - process_single_mc_job(job_dependencies, job).await? - } - InferenceCategory::FormatConversion => { - process_single_format_conversion_job(job_dependencies, job).await? - } - InferenceCategory::ConvertBvhToWorkflow => { - process_single_render_engine_scene_job(job_dependencies, job).await? - } - InferenceCategory::Workflow => { - process_single_workflow_job(job_dependencies, job).await? - } - InferenceCategory::DeprecatedField => { - match job.job_type { - InferenceJobType::VideoRender - | InferenceJobType::LivePortrait - | InferenceJobType::ComfyUi => { - // NB: These are all comfy workflow jobs too - process_single_workflow_job(job_dependencies, job).await? - } - _ => { - return Err(ProcessSingleJobError::InvalidJob( - anyhow!("invalid job for dispatch. type: {:?}", job.job_type))) - } - } - } + let job_success_result = if can_use_new_dispatch(job) { + new_dispatch(job_dependencies, job).await? + } else { + old_dispatch(job_dependencies, job).await? }; let maybe_entity_type = job_success_result.maybe_result_entity @@ -288,3 +249,74 @@ async fn do_process_single_job( Ok(ProcessSingleJobSuccessCase::JobCompleted) } + +fn can_use_new_dispatch(job: &AvailableInferenceJob) -> bool { + match job.job_type { + InferenceJobType::VideoRender => true, + InferenceJobType::LivePortrait => true, + _ => false, + } +} + +async fn new_dispatch( + job_dependencies: &JobDependencies, + job: &AvailableInferenceJob, +) -> Result { + + let job_success_result = match job.job_type { + InferenceJobType::VideoRender + | InferenceJobType::LivePortrait + | InferenceJobType::ComfyUi => { + // NB: These are all comfy workflow jobs too + process_single_workflow_job(job_dependencies, job).await? + } + _ => { + return Err(ProcessSingleJobError::InvalidJob( + anyhow!("invalid job type for dispatch: {:?}", job.job_type))) + } + }; + + Ok(job_success_result) +} + +async fn old_dispatch( + job_dependencies: &JobDependencies, + job: &AvailableInferenceJob, +) -> Result { + + let job_success_result = match job.inference_category { + InferenceCategory::LipsyncAnimation => { + process_single_lipsync_job(job_dependencies, job).await? + } + InferenceCategory::TextToSpeech => { + process_single_tts_job(job_dependencies, job).await? + } + InferenceCategory::VoiceConversion => { + process_single_vc_job(job_dependencies, job).await? + } + InferenceCategory::VideoFilter => { + process_single_vf_job(job_dependencies, job).await? + } + InferenceCategory::ImageGeneration => { + process_single_ig_job(job_dependencies, job).await? + } + InferenceCategory::Mocap => { + process_single_mc_job(job_dependencies, job).await? + } + InferenceCategory::FormatConversion => { + process_single_format_conversion_job(job_dependencies, job).await? + } + InferenceCategory::ConvertBvhToWorkflow => { + process_single_render_engine_scene_job(job_dependencies, job).await? + } + InferenceCategory::Workflow => { + process_single_workflow_job(job_dependencies, job).await? + } + InferenceCategory::DeprecatedField => { + return Err(ProcessSingleJobError::InvalidJob( + anyhow!("invalid job category for dispatch: {:?}", job.inference_category))) + } + }; + + Ok(job_success_result) +} diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/process_single_workflow_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/process_single_workflow_job.rs index 8963529253..c54973ceb3 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/process_single_workflow_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/process_single_workflow_job.rs @@ -1,10 +1,12 @@ +use enums::by_table::generic_inference_jobs::inference_job_type::InferenceJobType; use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; use crate::job::job_loop::job_success_result::JobSuccessResult; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; -use crate::job::job_types::workflow::upload_workflow; -use crate::job::job_types::workflow::video_style_transfer; +use crate::job::job_types::workflow::live_portrait::process_live_portrait_job::process_live_portrait_job; +use crate::job::job_types::workflow::upload_workflow::process_upload_workflow_job::process_upload_workflow_job; use crate::job::job_types::workflow::video_style_transfer::extract_vst_workflow_payload_from_job::extract_vst_workflow_payload_from_job; +use crate::job::job_types::workflow::video_style_transfer::process_video_style_transfer_job::process_video_style_transfer_job; use crate::job_dependencies::JobDependencies; pub async fn process_single_workflow_job( @@ -12,14 +14,24 @@ pub async fn process_single_workflow_job( job: &AvailableInferenceJob ) -> Result { + // First try to dispatch with the newer "job type". + match job.job_type { + InferenceJobType::LivePortrait => { + let job_success_result = process_live_portrait_job(job_dependencies, job).await?; + return Ok(job_success_result); + } + _ => {} // Fall through + } + + // If we couldn't dispatch with "job type", fall back to older heuristics. let workflow_args = extract_vst_workflow_payload_from_job(&job)?; let job_success_result = match workflow_args.maybe_google_drive_link { Some(_link) => { - upload_workflow::process_upload_workflow_job::process_upload_workflow_job(job_dependencies, job).await? + process_upload_workflow_job(job_dependencies, job).await? } None => { - video_style_transfer::process_video_style_transfer_job::process_video_style_transfer_job(job_dependencies, job).await? + process_video_style_transfer_job(job_dependencies, job).await? } }; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_live_portrait_workflow_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_live_portrait_workflow_handler.rs index 7b4e49b730..60b4f7469a 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_live_portrait_workflow_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_live_portrait_workflow_handler.rs @@ -208,7 +208,7 @@ pub async fn enqueue_live_portrait_workflow_handler( uuid_idempotency_token: &request.uuid_idempotency_token, job_type: InferenceJobType::LivePortrait, inference_category: InferenceCategory::DeprecatedField, - maybe_model_type: None, + maybe_model_type: Some(InferenceModelType::ComfyUi), // TODO(bt): Remove! maybe_model_token: None, maybe_input_source_token: None, maybe_input_source_token_type: None,