extractor

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,
     }

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<ModelWeightToken>,
+
+  /// A video that drives the face animation.
+  #[serde(rename = "d")]
+  #[serde(skip_serializing_if = "Option::is_none")]
+  pub driver_media_file_token: Option<ModelWeightToken>,
+
+  #[serde(rename = "rm")]
+  #[serde(skip_serializing_if = "Option::is_none")]
+  pub remove_watermark: Option<bool>,
+
+  #[serde(rename = "cv")]
+  #[serde(skip_serializing_if = "Option::is_none")]
+  pub creator_visibility: Option<Visibility>,
+}

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;

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<JobSuccessResu
         PolymorphicInferenceArgs::Mc(_) => None,
         PolymorphicInferenceArgs::Cu(_) => None,
         PolymorphicInferenceArgs::Es(_) => None,
+        PolymorphicInferenceArgs::Lp(_) => None,
       })
       .flatten();
 
@@ -272,6 +273,7 @@ pub async fn process_job(args: RvcV2ProcessJobArgs<'_>) -> Result<JobSuccessResu
         PolymorphicInferenceArgs::Mc(_) => None,
         PolymorphicInferenceArgs::Cu(_) => None,
         PolymorphicInferenceArgs::Es(_) => None,
+        PolymorphicInferenceArgs::Lp(_) => None,
       })
       .flatten();
 

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<JobSuccess
         PolymorphicInferenceArgs::Mc(_) => 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<JobSuccess
         PolymorphicInferenceArgs::Mc(_) => None,
         PolymorphicInferenceArgs::Cu(_) => None,
         PolymorphicInferenceArgs::Es(_) => None,
+        PolymorphicInferenceArgs::Lp(_) => None,
       })
       .flatten();
 

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<LivePortraitPayload, ProcessSingleJobError> {
+
+  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)
+}

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;

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;

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<JobSuccessResult, ProcessSingleJobError> {
 
-  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) => {

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<JobSuccessResult, ProcessSingleJobError>{
    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) => {

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<WorkflowArgs, ProcessSingleJobError> {
 

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;

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