From ec0b475457afcb3ea06ae3dd2adf63e2cb6534e2 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 1 Feb 2024 03:25:25 -0500 Subject: [PATCH] other job types --- .idea/dictionaries/bt.xml | 1 + .../inference_job_type.rs | 93 ++++++++++++++++++- .../animation/enqueue_face_animation.rs | 2 +- .../image_gen/enqueue_image_generation.rs | 2 +- .../endpoints/mocap/enqueue_mocapnet.rs | 2 +- .../enqueue_infer_tts_handler.rs | 2 +- .../enqueue_voice_conversion_inference.rs | 12 ++- .../inference/enqueue_tts_request.rs | 2 +- .../voice_designer/voices/create_voice.rs | 2 +- .../endpoints/workflows/enqueue_comfy_ui.rs | 2 +- 10 files changed, 109 insertions(+), 11 deletions(-) diff --git a/.idea/dictionaries/bt.xml b/.idea/dictionaries/bt.xml index 1cce46b73d..f638c1a06f 100644 --- a/.idea/dictionaries/bt.xml +++ b/.idea/dictionaries/bt.xml @@ -26,6 +26,7 @@ hparams huggingface lipsync + mocap mochi nestedness pathbuf diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_job_type.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_job_type.rs index c7ca70ff18..0445c145a8 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_job_type.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_job_type.rs @@ -13,16 +13,45 @@ use strum::EnumIter; /// YOU CAN ADD NEW VALUES, BUT DO NOT CHANGE EXISTING VALUES WITHOUT A MIGRATION STRATEGY. #[cfg_attr(test, derive(EnumIter, EnumCount))] #[derive(Clone, Copy, Eq, PartialEq, Hash, Ord, PartialOrd, Deserialize, Serialize, Default)] +#[serde(rename_all = "snake_case")] pub enum InferenceJobType { + /// Jobs that run ComfyUI workflows + ComfyUi, + /// A job that turns "FBX" game engine files into "GLTF" files (Bevy-compatible). #[serde(rename = "convert_fbx_gltf")] ConvertFbxToGltf, - #[serde(rename = "rerender_a_video")] + /// Process a video into BVH mocap animation data for game engines + MocapNet, + + /// Re-render a video is a video style transfer algorithm. We developed code + /// around it, but chose to develop AnimateDiff / ComfyUI support instead. RerenderAVideo, + /// RVC is a voice conversion model. RVCv2 is the most popular such model currently. + #[serde(rename = "rvc_v2")] + RvcV2, + + /// SadTalker does image-to-video lip-syncing when given an audio file and image. + SadTalker, + + /// so-vits-svc voice conversion. This predates RVCv2. + SoVitsSvc, + + /// Stable diffusion image generation + StableDiffusion, + + /// StyleTTS2 is a zero shot multi-speaker TTS model. + /// This job type should handle both speaker vector encoding and inference. + #[serde(rename = "styletts2")] + StyleTTS2, + + /// TT2 Text to speech + Tacotron2, + /// A value we may use in the future for historical jobs - #[serde(rename = "unknown")] + /// (i.e. when we backfill the database column and make it non-nullable) #[default] Unknown, } @@ -35,15 +64,32 @@ impl_mysql_enum_coders!(InferenceJobType); impl InferenceJobType { pub fn to_str(&self) -> &'static str { match self { + Self::ComfyUi => "comfy_ui", Self::ConvertFbxToGltf => "convert_fbx_gltf", + Self::MocapNet => "mocap_net", Self::RerenderAVideo => "rerender_a_video", + Self::RvcV2 => "rvc_v2", + Self::SadTalker => "sad_talker", + Self::SoVitsSvc => "so_vits_svc", + Self::StableDiffusion => "stable_diffusion", + Self::StyleTTS2 => "styletts2", + Self::Tacotron2 => "tacotron2", Self::Unknown => "unknown", } } pub fn from_str(value: &str) -> Result { match value { + "comfy_ui" => Ok(Self::ComfyUi), "convert_fbx_gltf" => Ok(Self::ConvertFbxToGltf), + "mocap_net" => Ok(Self::MocapNet), + "rerender_a_video" => Ok(Self::RerenderAVideo), + "rvc_v2" => Ok(Self::RvcV2), + "sad_talker" => Ok(Self::SadTalker), + "so_vits_svc" => Ok(Self::SoVitsSvc), + "stable_diffusion" => Ok(Self::StableDiffusion), + "styletts2" => Ok(Self::StyleTTS2), + "tacotron2" => Ok(Self::Tacotron2), "unknown" => Ok(Self::Unknown), _ => Err(format!("invalid value: {:?}", value)), } @@ -53,7 +99,16 @@ impl InferenceJobType { // NB: BTreeSet is sorted // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. BTreeSet::from([ + Self::ComfyUi, Self::ConvertFbxToGltf, + Self::MocapNet, + Self::RerenderAVideo, + Self::RvcV2, + Self::SadTalker, + Self::SoVitsSvc, + Self::StableDiffusion, + Self::StyleTTS2, + Self::Tacotron2, Self::Unknown, ]) } @@ -74,22 +129,46 @@ mod tests { #[test] fn test_serialization() { + assert_serialization(InferenceJobType::ComfyUi, "comfy_ui"); assert_serialization(InferenceJobType::ConvertFbxToGltf, "convert_fbx_gltf"); + assert_serialization(InferenceJobType::MocapNet, "mocap_net"); assert_serialization(InferenceJobType::RerenderAVideo, "rerender_a_video"); + assert_serialization(InferenceJobType::RvcV2, "rvc_v2"); + assert_serialization(InferenceJobType::SadTalker, "sad_talker"); + assert_serialization(InferenceJobType::SoVitsSvc, "so_vits_svc"); + assert_serialization(InferenceJobType::StableDiffusion, "stable_diffusion"); + assert_serialization(InferenceJobType::StyleTTS2, "styletts2"); + assert_serialization(InferenceJobType::Tacotron2, "tacotron2"); assert_serialization(InferenceJobType::Unknown, "unknown"); } #[test] fn to_str() { + assert_eq!(InferenceJobType::ComfyUi.to_str(), "comfy_ui"); assert_eq!(InferenceJobType::ConvertFbxToGltf.to_str(), "convert_fbx_gltf"); + assert_eq!(InferenceJobType::MocapNet.to_str(), "mocap_net"); assert_eq!(InferenceJobType::RerenderAVideo.to_str(), "rerender_a_video"); + assert_eq!(InferenceJobType::RvcV2.to_str(), "rvc_v2"); + assert_eq!(InferenceJobType::SadTalker.to_str(), "sad_talker"); + assert_eq!(InferenceJobType::SoVitsSvc.to_str(), "so_vits_svc"); + assert_eq!(InferenceJobType::StableDiffusion.to_str(), "stable_diffusion"); + assert_eq!(InferenceJobType::StyleTTS2.to_str(), "styletts2"); + assert_eq!(InferenceJobType::Tacotron2.to_str(), "tacotron2"); assert_eq!(InferenceJobType::Unknown.to_str(), "unknown"); } #[test] fn from_str() { + assert_eq!(InferenceJobType::from_str("comfy_ui").unwrap(), InferenceJobType::ComfyUi); assert_eq!(InferenceJobType::from_str("convert_fbx_gltf").unwrap(), InferenceJobType::ConvertFbxToGltf); + assert_eq!(InferenceJobType::from_str("mocap_net").unwrap(), InferenceJobType::MocapNet); assert_eq!(InferenceJobType::from_str("rerender_a_video").unwrap(), InferenceJobType::RerenderAVideo); + assert_eq!(InferenceJobType::from_str("rvc_v2").unwrap(), InferenceJobType::RvcV2); + assert_eq!(InferenceJobType::from_str("sad_talker").unwrap(), InferenceJobType::SadTalker); + assert_eq!(InferenceJobType::from_str("so_vits_svc").unwrap(), InferenceJobType::SoVitsSvc); + assert_eq!(InferenceJobType::from_str("stable_diffusion").unwrap(), InferenceJobType::StableDiffusion); + assert_eq!(InferenceJobType::from_str("styletts2").unwrap(), InferenceJobType::StyleTTS2); + assert_eq!(InferenceJobType::from_str("tacotron2").unwrap(), InferenceJobType::Tacotron2); assert_eq!(InferenceJobType::from_str("unknown").unwrap(), InferenceJobType::Unknown); } @@ -97,9 +176,17 @@ mod tests { fn all_variants() { // Static check let mut variants = InferenceJobType::all_variants(); - assert_eq!(variants.len(), 3); + assert_eq!(variants.len(), 11); + assert_eq!(variants.pop_first(), Some(InferenceJobType::ComfyUi)); assert_eq!(variants.pop_first(), Some(InferenceJobType::ConvertFbxToGltf)); + assert_eq!(variants.pop_first(), Some(InferenceJobType::MocapNet)); assert_eq!(variants.pop_first(), Some(InferenceJobType::RerenderAVideo)); + assert_eq!(variants.pop_first(), Some(InferenceJobType::RvcV2)); + assert_eq!(variants.pop_first(), Some(InferenceJobType::SadTalker)); + assert_eq!(variants.pop_first(), Some(InferenceJobType::SoVitsSvc)); + assert_eq!(variants.pop_first(), Some(InferenceJobType::StableDiffusion)); + assert_eq!(variants.pop_first(), Some(InferenceJobType::StyleTTS2)); + assert_eq!(variants.pop_first(), Some(InferenceJobType::Tacotron2)); assert_eq!(variants.pop_first(), Some(InferenceJobType::Unknown)); assert_eq!(variants.pop_first(), None); diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/animation/enqueue_face_animation.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/animation/enqueue_face_animation.rs index ea2274fdd6..95bcc1e867 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/animation/enqueue_face_animation.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/animation/enqueue_face_animation.rs @@ -308,7 +308,7 @@ pub async fn enqueue_face_animation_handler( let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { uuid_idempotency_token: &request.uuid_idempotency_token, - job_type: InferenceJobType::Unknown, + job_type: InferenceJobType::SadTalker, inference_category: InferenceCategory::LipsyncAnimation, maybe_model_type: Some(InferenceModelType::SadTalker), // NB: Model is static during inference maybe_model_token: None, // NB: Model is static during inference diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs index 5cb51cdf5f..13e4a06a1a 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs @@ -407,7 +407,7 @@ pub async fn enqueue_image_generation_request( // create the job record here! let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { uuid_idempotency_token: &request.uuid_idempotency_token, - job_type: InferenceJobType::Unknown, + job_type: InferenceJobType::StableDiffusion, inference_category: InferenceCategory::ImageGeneration, maybe_model_type: Some(InferenceModelType::StableDiffusion), // NB: Model is static during inference maybe_model_token: None, // NB: Model is static during inference diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/mocap/enqueue_mocapnet.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/mocap/enqueue_mocapnet.rs index 99d4175f3e..0fe5fd8fb5 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/mocap/enqueue_mocapnet.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/mocap/enqueue_mocapnet.rs @@ -188,7 +188,7 @@ pub async fn enqueue_mocapnet_handler( let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { uuid_idempotency_token: &request.uuid_idempotency_token, - job_type: InferenceJobType::Unknown, + job_type: InferenceJobType::MocapNet, inference_category: InferenceCategory::Mocap, maybe_model_type: Some(InferenceModelType::MocapNet), // NB: Model is static during inference maybe_model_token: None, // NB: Model is static during inference diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/enqueue_infer_tts_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/enqueue_infer_tts_handler.rs index d924baf630..553b5f3225 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/enqueue_infer_tts_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/enqueue_infer_tts_handler.rs @@ -398,7 +398,7 @@ pub async fn enqueue_infer_tts_handler( let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { uuid_idempotency_token: &request.uuid_idempotency_token, - job_type: InferenceJobType::Unknown, + job_type: InferenceJobType::Tacotron2, inference_category: InferenceCategory::TextToSpeech, maybe_model_type: Some(model_type), maybe_model_token: Some(request.tts_model_token.as_str()), diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_conversion/inference/enqueue_voice_conversion_inference.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_conversion/inference/enqueue_voice_conversion_inference.rs index f75b8a3225..1f3fd5db47 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_conversion/inference/enqueue_voice_conversion_inference.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_conversion/inference/enqueue_voice_conversion_inference.rs @@ -290,9 +290,19 @@ pub async fn enqueue_voice_conversion_inference_handler( }); } + let job_type = match model_inference_info.job_model_type { + InferenceModelType::SoVitsSvc => InferenceJobType::SoVitsSvc, + InferenceModelType::RvcV2 => InferenceJobType::RvcV2, + _ => { + // In theory, this shouldn't catch anything. + error!("wrong model type for voice conversion: {:?}", model_inference_info.job_model_type); + return Err(EnqueueVoiceConversionInferenceError::ServerError) + } + }; + let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { uuid_idempotency_token: &request.uuid_idempotency_token, - job_type: InferenceJobType::Unknown, + job_type, inference_category: InferenceCategory::VoiceConversion, maybe_model_type: Some(model_inference_info.job_model_type), maybe_model_token: Some(&model_token), diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/inference/enqueue_tts_request.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/inference/enqueue_tts_request.rs index 9ad344136a..f435a6b2d8 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/inference/enqueue_tts_request.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/inference/enqueue_tts_request.rs @@ -190,7 +190,7 @@ pub async fn enqueue_tts_request( // create the job record here! let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { uuid_idempotency_token: &request.uuid_idempotency_token, - job_type: InferenceJobType::Unknown, + job_type: InferenceJobType::StyleTTS2, inference_category: InferenceCategory::TextToSpeech, maybe_model_type: Some(InferenceModelType::StyleTTS2), // NB: Model is static during inference maybe_model_token: None, // NB: Model is static during inference diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/create_voice.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/create_voice.rs index b552bfc572..cb4af090bf 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/create_voice.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/create_voice.rs @@ -177,7 +177,7 @@ pub async fn create_voice_handler( // create the job record here! explore the table insert the inference args in here as json! keep it short let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { uuid_idempotency_token: &request.uuid_idempotency_token, - job_type: InferenceJobType::Unknown, + job_type: InferenceJobType::StyleTTS2, inference_category: InferenceCategory::TextToSpeech, maybe_model_type: Some(InferenceModelType::StyleTTS2), // NB: Model is static during inference maybe_model_token: None, // NB: Model is static during inference diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs index 899130c160..8bcd9b7edb 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs @@ -190,7 +190,7 @@ pub async fn enqueue_comfy_ui_handler( let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { uuid_idempotency_token: &request.uuid_idempotency_token, - job_type: InferenceJobType::Unknown, + job_type: InferenceJobType::ComfyUi, inference_category: InferenceCategory::Workflow, maybe_model_type: Some(InferenceModelType::ComfyUi), // NB: Model is static during inference maybe_model_token: None, // NB: Model is static during inference