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 fc8bfa78bb..e192d352b2 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 @@ -3,6 +3,7 @@ use errors::AnyhowResult; use crate::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; use crate::payloads::generic_inference_args::lipsync_payload::LipsyncArgs; +use crate::payloads::generic_inference_args::live_portrait_payload::LivePortraitPayload; use crate::payloads::generic_inference_args::mocap_payload::MocapArgs; use crate::payloads::generic_inference_args::render_engine_scene_to_video_payload::RenderEngineSceneToVideoArgs; use crate::payloads::generic_inference_args::tts_payload::TTSArgs; @@ -56,6 +57,10 @@ pub enum InferenceCategoryAbbreviated { #[serde(alias = "workflow")] Workflow, + #[serde(rename = "lp")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + #[serde(alias = "live_portrait")] + LivePortrait, + #[serde(rename = "fc")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. #[serde(alias = "format_conversion")] FormatConversion, @@ -125,6 +130,9 @@ pub enum PolymorphicInferenceArgs { /// ComfyUI (Short name to save space when serializing.) Cu(WorkflowArgs), + /// ComfyUI (Short name to save space when serializing.) + Lp(LivePortraitPayload), + /// Render engine scene to video args Es(RenderEngineSceneToVideoArgs), } @@ -164,6 +172,7 @@ impl InferenceCategoryAbbreviated { Self::ImageGeneration =>InferenceCategory::ImageGeneration, Self::Mocap => InferenceCategory::Mocap, Self::Workflow => InferenceCategory::Workflow, + Self::LivePortrait => InferenceCategory::Workflow, // TODO(bt,2024-07-16): Not bidirectional/reflective Self::FormatConversion => InferenceCategory::FormatConversion, Self::ConvertBvhToWorkflow => InferenceCategory::ConvertBvhToWorkflow, } diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/live_portrait_payload.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/live_portrait_payload.rs new file mode 100644 index 0000000000..9101499d80 --- /dev/null +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/live_portrait_payload.rs @@ -0,0 +1,23 @@ +use enums::common::visibility::Visibility; +use tokens::tokens::model_weights::ModelWeightToken; + +#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)] +pub struct LivePortraitPayload { + /// Either an image or video. + #[serde(rename = "p")] + #[serde(skip_serializing_if = "Option::is_none")] + pub portrait_media_file_token: Option, + + /// A video that drives the face animation. + #[serde(rename = "d")] + #[serde(skip_serializing_if = "Option::is_none")] + pub driver_media_file_token: Option, + + #[serde(rename = "rm")] + #[serde(skip_serializing_if = "Option::is_none")] + pub remove_watermark: Option, + + #[serde(rename = "cv")] + #[serde(skip_serializing_if = "Option::is_none")] + pub creator_visibility: Option, +} diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs index 17a1fd86bb..45478e0100 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs @@ -1,6 +1,7 @@ pub mod generic_inference_args; pub mod image_generation_payload; pub mod lipsync_payload; +pub mod live_portrait_payload; pub mod mocap_payload; pub mod render_engine_scene_to_video_payload; pub mod tts_payload; 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 e54020cfab..5eddbd949a 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 @@ -259,6 +259,7 @@ pub async fn process_job(args: RvcV2ProcessJobArgs<'_>) -> Result None, PolymorphicInferenceArgs::Cu(_) => None, PolymorphicInferenceArgs::Es(_) => None, + PolymorphicInferenceArgs::Lp(_) => None, }) .flatten(); @@ -272,6 +273,7 @@ pub async fn process_job(args: RvcV2ProcessJobArgs<'_>) -> Result None, PolymorphicInferenceArgs::Cu(_) => None, PolymorphicInferenceArgs::Es(_) => None, + PolymorphicInferenceArgs::Lp(_) => 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 55a10ab80e..8df79621cc 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 @@ -180,6 +180,7 @@ pub async fn process_job(args: SoVitsSvcProcessJobArgs<'_>) -> Result None, PolymorphicInferenceArgs::Cu(_) => None, PolymorphicInferenceArgs::Es(_) => None, + PolymorphicInferenceArgs::Lp(_) => None, }) .flatten() .unwrap_or(false); @@ -194,6 +195,7 @@ pub async fn process_job(args: SoVitsSvcProcessJobArgs<'_>) -> Result None, PolymorphicInferenceArgs::Cu(_) => None, PolymorphicInferenceArgs::Es(_) => None, + PolymorphicInferenceArgs::Lp(_) => None, }) .flatten(); diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/extract_live_portrait_payload_from_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/extract_live_portrait_payload_from_job.rs new file mode 100644 index 0000000000..98c1aee44b --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/extract_live_portrait_payload_from_job.rs @@ -0,0 +1,26 @@ +use anyhow::anyhow; + +use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; +use mysql_queries::payloads::generic_inference_args::live_portrait_payload::LivePortraitPayload; +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; + +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::util::get_polymorphic_args_from_job::get_polymorphic_args_from_job; + +pub fn extract_live_portrait_payload_from_job( + job: &AvailableInferenceJob, +) -> Result { + + let polymorphic_args = get_polymorphic_args_from_job(&job)?; + + let some_args = match polymorphic_args { + PolymorphicInferenceArgs::Lp(payload) => payload.clone(), + _ => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!")) + ); + } + }; + + Ok(some_args) +} diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/mod.rs b/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/mod.rs index 3d43f2a9b1..d8cfc2e0bf 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/live_portrait/mod.rs @@ -1 +1,2 @@ +pub mod extract_live_portrait_payload_from_job; pub mod process_live_portrait_job; \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/mod.rs b/crates/service/job/inference_job/src/job/job_types/workflow/mod.rs index b51166592e..4ce2c05051 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/mod.rs @@ -1,5 +1,4 @@ pub mod comfy_ui_dependencies; -pub mod get_workflow_args_from_job; pub mod live_portrait; pub mod process_single_workflow_job; pub mod upload_workflow; 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 967e633589..8963529253 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 @@ -2,9 +2,9 @@ use mysql_queries::queries::generic_inference::job::list_available_generic_infer 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::get_workflow_args_from_job::get_workflow_args_from_job; use crate::job::job_types::workflow::upload_workflow; use crate::job::job_types::workflow::video_style_transfer; +use crate::job::job_types::workflow::video_style_transfer::extract_vst_workflow_payload_from_job::extract_vst_workflow_payload_from_job; use crate::job_dependencies::JobDependencies; pub async fn process_single_workflow_job( @@ -12,7 +12,7 @@ pub async fn process_single_workflow_job( job: &AvailableInferenceJob ) -> Result { - let workflow_args = get_workflow_args_from_job(&job)?; + let workflow_args = extract_vst_workflow_payload_from_job(&job)?; let job_success_result = match workflow_args.maybe_google_drive_link { Some(_link) => { diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/upload_workflow/process_upload_workflow_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/upload_workflow/process_upload_workflow_job.rs index b9864be1b1..36327af12c 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/upload_workflow/process_upload_workflow_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/upload_workflow/process_upload_workflow_job.rs @@ -19,13 +19,13 @@ use tokens::tokens::users::UserToken; 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::get_workflow_args_from_job::get_workflow_args_from_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_dependencies::JobDependencies; pub async fn process_upload_workflow_job(deps: &JobDependencies, job: &AvailableInferenceJob) -> Result{ let mysql_pool = &deps.db.mysql_pool; - let wf_args = get_workflow_args_from_job(&job)?; + let wf_args = extract_vst_workflow_payload_from_job(&job)?; let title = match wf_args.maybe_title { Some(val) => { diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/get_workflow_args_from_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/extract_vst_workflow_payload_from_job.rs similarity index 95% rename from crates/service/job/inference_job/src/job/job_types/workflow/get_workflow_args_from_job.rs rename to crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/extract_vst_workflow_payload_from_job.rs index 4e1b0112a3..21e938d0a0 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/get_workflow_args_from_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/extract_vst_workflow_payload_from_job.rs @@ -5,7 +5,7 @@ use mysql_queries::queries::generic_inference::job::list_available_generic_infer use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; use crate::util::get_polymorphic_args_from_job::get_polymorphic_args_from_job; -pub fn get_workflow_args_from_job( +pub fn extract_vst_workflow_payload_from_job( job: &AvailableInferenceJob, ) -> Result { diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/mod.rs b/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/mod.rs index 586d57ded9..231ff3b9ca 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/mod.rs @@ -1,3 +1,4 @@ +pub mod extract_vst_workflow_payload_from_job; pub mod process_video_style_transfer_job; pub mod steps; pub mod util; 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 ae78e07df5..85f2730445 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 @@ -44,11 +44,11 @@ pub struct EnqueueLivePortraitWorkflowRequest { /// Source media token /// This can be an image or a video media file token - portrait_media_file: MediaFileToken, + portrait_media_file_token: MediaFileToken, /// Driving media token /// This must be a video media file token. - driver_media_file: MediaFileToken, + driver_media_file_token: MediaFileToken, /// Remove watermark from the output /// Only for premium accounts