From 7976f68fd027d953e147abd9d15f61112803314a Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 6 Dec 2023 02:53:50 -0500 Subject: [PATCH] Save TTS results as media files --- .../inference_category.rs | 3 +- .../generic_synthetic_ids/id_category.rs | 14 +- .../media_file_origin_model_type.rs | 13 +- .../media_file_origin_product_category.rs | 13 +- .../enums/src/common/job_status_plus.rs | 11 + .../src/tokens/anonymous_visitor_tracking.rs | 2 +- .../src/tokens/generic_inference_jobs.rs | 2 +- crates/_types/tokens/src/tokens/users.rs | 2 +- .../web/get_inference_job_status.rs | 75 ++++++- .../insert_media_file_from_tacotron2.rs | 166 ++++++++++++++ .../src/queries/media_files/mod.rs | 1 + .../tts/tacotron2_v2_early_fakeyou/mod.rs | 1 + .../tacotron2_v2_early_fakeyou/process_job.rs | 90 ++------ .../tacotron2_dependencies.rs | 5 +- .../upload_results.rs | 202 ++++++++++++++++++ .../vall_e_x/process_job_inference_voice.rs | 2 +- crates/service/web/storyteller_web/Cargo.toml | 4 +- .../inference_job/get_inference_job_status.rs | 128 +++++++++-- .../tts/get_tts_inference_job_status.rs | 1 + 19 files changed, 627 insertions(+), 108 deletions(-) create mode 100644 crates/lib/mysql_queries/src/queries/media_files/insert_media_file_from_tacotron2.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/upload_results.rs diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs index faee8ff909..b38f181460 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs @@ -13,10 +13,11 @@ use strum::EnumIter; /// /// DO NOT CHANGE VALUES WITHOUT A MIGRATION STRATEGY. #[cfg_attr(test, derive(EnumIter, EnumCount))] -#[derive(Clone, Copy, Eq, PartialEq, Hash, Ord, PartialOrd, Deserialize, Serialize)] +#[derive(Clone, Copy, Eq, PartialEq, Hash, Ord, PartialOrd, Deserialize, Serialize, Default)] pub enum InferenceCategory { /// Eg. SadTalker and possibly Wav2Lip #[serde(rename = "lipsync_animation")] + #[default] LipsyncAnimation, #[serde(rename = "text_to_speech")] diff --git a/crates/_types/enums/src/by_table/generic_synthetic_ids/id_category.rs b/crates/_types/enums/src/by_table/generic_synthetic_ids/id_category.rs index ca1b502e14..eccff8cf4d 100644 --- a/crates/_types/enums/src/by_table/generic_synthetic_ids/id_category.rs +++ b/crates/_types/enums/src/by_table/generic_synthetic_ids/id_category.rs @@ -21,6 +21,11 @@ pub enum IdCategory { #[serde(rename = "lipsync_animation")] LipsyncAnimationResult, + /// Results from tacotron2 + /// Applies for RVC and SVC + #[serde(rename = "tts_result")] + TtsResult, + /// Results from voice conversion (which may live in the media_files table) /// Applies for RVC and SVC #[serde(rename = "voice_conversion")] @@ -52,6 +57,7 @@ impl IdCategory { match self { Self::MediaFile => "media_file", Self::LipsyncAnimationResult => "lipsync_animation", + Self::TtsResult => "tts_result", Self::VoiceConversionResult => "voice_conversion", Self::ZeroShotVoiceDataset => "zs_dataset", Self::ZeroShotVoiceEmbedding => "zs_voice", @@ -64,6 +70,7 @@ impl IdCategory { match value { "media_file" => Ok(Self::MediaFile), "lipsync_animation" => Ok(Self::LipsyncAnimationResult), + "tts_result" => Ok(Self::TtsResult), "voice_conversion" => Ok(Self::VoiceConversionResult), "zs_dataset" => Ok(Self::ZeroShotVoiceDataset), "zs_voice" => Ok(Self::ZeroShotVoiceEmbedding), @@ -79,6 +86,7 @@ impl IdCategory { BTreeSet::from([ Self::MediaFile, Self::LipsyncAnimationResult, + Self::TtsResult, Self::VoiceConversionResult, Self::ZeroShotTtsResult, Self::ZeroShotVoiceDataset, @@ -100,6 +108,7 @@ mod tests { fn test_serialization() { assert_serialization(IdCategory::MediaFile, "media_file"); assert_serialization(IdCategory::LipsyncAnimationResult, "lipsync_animation"); + assert_serialization(IdCategory::TtsResult, "tts_result"); assert_serialization(IdCategory::VoiceConversionResult, "voice_conversion"); assert_serialization(IdCategory::ZeroShotVoiceDataset, "zs_dataset"); assert_serialization(IdCategory::ZeroShotVoiceEmbedding, "zs_voice"); @@ -111,6 +120,7 @@ mod tests { fn to_str() { assert_eq!(IdCategory::MediaFile.to_str(), "media_file"); assert_eq!(IdCategory::LipsyncAnimationResult.to_str(), "lipsync_animation"); + assert_eq!(IdCategory::TtsResult.to_str(), "tts_result"); assert_eq!(IdCategory::VoiceConversionResult.to_str(), "voice_conversion"); assert_eq!(IdCategory::ZeroShotVoiceDataset.to_str(), "zs_dataset"); assert_eq!(IdCategory::ZeroShotVoiceEmbedding.to_str(), "zs_voice"); @@ -122,6 +132,7 @@ mod tests { fn from_str() { assert_eq!(IdCategory::from_str("media_file").unwrap(), IdCategory::MediaFile); assert_eq!(IdCategory::from_str("lipsync_animation").unwrap(), IdCategory::LipsyncAnimationResult); + assert_eq!(IdCategory::from_str("tts_result").unwrap(), IdCategory::TtsResult); assert_eq!(IdCategory::from_str("voice_conversion").unwrap(), IdCategory::VoiceConversionResult); assert_eq!(IdCategory::from_str("zs_dataset").unwrap(), IdCategory::ZeroShotVoiceDataset); assert_eq!(IdCategory::from_str("zs_voice").unwrap(), IdCategory::ZeroShotVoiceEmbedding); @@ -133,9 +144,10 @@ mod tests { fn all_variants() { // Static check let mut variants = IdCategory::all_variants(); - assert_eq!(variants.len(), 7); + assert_eq!(variants.len(), 8); assert_eq!(variants.pop_first(), Some(IdCategory::MediaFile)); assert_eq!(variants.pop_first(), Some(IdCategory::LipsyncAnimationResult)); + assert_eq!(variants.pop_first(), Some(IdCategory::TtsResult)); assert_eq!(variants.pop_first(), Some(IdCategory::VoiceConversionResult)); assert_eq!(variants.pop_first(), Some(IdCategory::ZeroShotTtsResult)); assert_eq!(variants.pop_first(), Some(IdCategory::ZeroShotVoiceDataset)); diff --git a/crates/_types/enums/src/by_table/media_files/media_file_origin_model_type.rs b/crates/_types/enums/src/by_table/media_files/media_file_origin_model_type.rs index b38f90a6dc..8cb9685bc7 100644 --- a/crates/_types/enums/src/by_table/media_files/media_file_origin_model_type.rs +++ b/crates/_types/enums/src/by_table/media_files/media_file_origin_model_type.rs @@ -25,6 +25,9 @@ pub enum MediaFileOriginModelType { #[serde(rename = "so_vits_svc")] SoVitsSvc, + #[serde(rename = "tacotron2")] + Tacotron2, + #[serde(rename = "vall_e_x")] VallEX, } @@ -41,6 +44,7 @@ impl MediaFileOriginModelType { Self::RvcV2 => "rvc_v2", Self::SadTalker => "sad_talker", Self::SoVitsSvc => "so_vits_svc", + Self::Tacotron2 => "tacotron2", Self::VallEX => "vall_e_x" } } @@ -50,6 +54,7 @@ impl MediaFileOriginModelType { "rvc_v2" => Ok(Self::RvcV2), "sad_talker" => Ok(Self::SadTalker), "so_vits_svc" => Ok(Self::SoVitsSvc), + "tacotron2" => Ok(Self::Tacotron2), "vall_e_x" => Ok(Self::VallEX), _ => Err(format!("invalid value: {:?}", value)), } @@ -62,6 +67,7 @@ impl MediaFileOriginModelType { Self::RvcV2, Self::SadTalker, Self::SoVitsSvc, + Self::Tacotron2, Self::VallEX ]) } @@ -85,6 +91,7 @@ mod tests { assert_eq!(MediaFileOriginModelType::RvcV2.to_str(), "rvc_v2"); assert_eq!(MediaFileOriginModelType::SadTalker.to_str(), "sad_talker"); assert_eq!(MediaFileOriginModelType::SoVitsSvc.to_str(), "so_vits_svc"); + assert_eq!(MediaFileOriginModelType::Tacotron2.to_str(), "tacotron2"); assert_eq!(MediaFileOriginModelType::VallEX.to_str(), "vall_e_x"); } @@ -93,17 +100,19 @@ mod tests { assert_eq!(MediaFileOriginModelType::from_str("rvc_v2").unwrap(), MediaFileOriginModelType::RvcV2); assert_eq!(MediaFileOriginModelType::from_str("sad_talker").unwrap(), MediaFileOriginModelType::SadTalker); assert_eq!(MediaFileOriginModelType::from_str("so_vits_svc").unwrap(), MediaFileOriginModelType::SoVitsSvc); + assert_eq!(MediaFileOriginModelType::from_str("tacotron2").unwrap(), MediaFileOriginModelType::Tacotron2); assert_eq!(MediaFileOriginModelType::from_str("vall_e_x").unwrap(), MediaFileOriginModelType::VallEX); - assert!(MediaFileOriginModelType::from_str("foo").is_err()); + assert!(MediaFileOriginModelType::from_str("foo").is_err()); } #[test] fn all_variants() { let mut variants = MediaFileOriginModelType::all_variants(); - assert_eq!(variants.len(), 4); + assert_eq!(variants.len(), 5); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::RvcV2)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::SadTalker)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::SoVitsSvc)); + assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::Tacotron2)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::VallEX)); assert_eq!(variants.pop_first(), None); } diff --git a/crates/_types/enums/src/by_table/media_files/media_file_origin_product_category.rs b/crates/_types/enums/src/by_table/media_files/media_file_origin_product_category.rs index 2d27b3824d..7b6f2a029a 100644 --- a/crates/_types/enums/src/by_table/media_files/media_file_origin_product_category.rs +++ b/crates/_types/enums/src/by_table/media_files/media_file_origin_product_category.rs @@ -21,6 +21,10 @@ pub enum MediaFileOriginProductCategory { #[serde(rename = "face_animator")] FaceAnimator, + /// Text to speech (Tacotron2, not voice designer / VallE-X) + #[serde(rename = "tts")] + TextToSpeech, + // TODO: This should be a temporary category until we migrate the DB to remove this default value /// Unknown which product is attached to the file (generated the file, the file was /// uploaded on behalf of, etc.) @@ -46,6 +50,7 @@ impl MediaFileOriginProductCategory { pub fn to_str(&self) -> &'static str { match self { Self::FaceAnimator => "face_animator", + Self::TextToSpeech => "tts", Self::Unknown => "unknown", Self::VoiceConversion => "voice_conversion", Self::ZeroShotVoice => "zs_voice", @@ -55,6 +60,7 @@ impl MediaFileOriginProductCategory { pub fn from_str(value: &str) -> Result { match value { "face_animator" => Ok(Self::FaceAnimator), + "tts" => Ok(Self::TextToSpeech), "unknown" => Ok(Self::Unknown), "voice_conversion" => Ok(Self::VoiceConversion), "zs_voice" => Ok(Self::ZeroShotVoice), @@ -67,6 +73,7 @@ impl MediaFileOriginProductCategory { // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. BTreeSet::from([ Self::FaceAnimator, + Self::TextToSpeech, Self::Unknown, Self::VoiceConversion, Self::ZeroShotVoice, @@ -85,6 +92,7 @@ mod tests { #[test] fn test_serialization() { assert_serialization(MediaFileOriginProductCategory::FaceAnimator, "face_animator"); + assert_serialization(MediaFileOriginProductCategory::TextToSpeech, "tts"); assert_serialization(MediaFileOriginProductCategory::Unknown, "unknown"); assert_serialization(MediaFileOriginProductCategory::VoiceConversion, "voice_conversion"); assert_serialization(MediaFileOriginProductCategory::ZeroShotVoice, "zs_voice"); @@ -93,6 +101,7 @@ mod tests { #[test] fn to_str() { assert_eq!(MediaFileOriginProductCategory::FaceAnimator.to_str(), "face_animator"); + assert_eq!(MediaFileOriginProductCategory::TextToSpeech.to_str(), "tts"); assert_eq!(MediaFileOriginProductCategory::Unknown.to_str(), "unknown"); assert_eq!(MediaFileOriginProductCategory::VoiceConversion.to_str(), "voice_conversion"); assert_eq!(MediaFileOriginProductCategory::ZeroShotVoice.to_str(), "zs_voice"); @@ -101,6 +110,7 @@ mod tests { #[test] fn from_str() { assert_eq!(MediaFileOriginProductCategory::from_str("face_animator").unwrap(), MediaFileOriginProductCategory::FaceAnimator); + assert_eq!(MediaFileOriginProductCategory::from_str("tts").unwrap(), MediaFileOriginProductCategory::TextToSpeech); assert_eq!(MediaFileOriginProductCategory::from_str("unknown").unwrap(), MediaFileOriginProductCategory::Unknown); assert_eq!(MediaFileOriginProductCategory::from_str("voice_conversion").unwrap(), MediaFileOriginProductCategory::VoiceConversion); assert_eq!(MediaFileOriginProductCategory::from_str("zs_voice").unwrap(), MediaFileOriginProductCategory::ZeroShotVoice); @@ -109,8 +119,9 @@ mod tests { #[test] fn all_variants() { let mut variants = MediaFileOriginProductCategory::all_variants(); - assert_eq!(variants.len(), 4); + assert_eq!(variants.len(), 5); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::FaceAnimator)); + assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::TextToSpeech)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::Unknown)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::VoiceConversion)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::ZeroShotVoice)); diff --git a/crates/_types/enums/src/common/job_status_plus.rs b/crates/_types/enums/src/common/job_status_plus.rs index 4f69359e73..be0ae8db25 100644 --- a/crates/_types/enums/src/common/job_status_plus.rs +++ b/crates/_types/enums/src/common/job_status_plus.rs @@ -34,6 +34,12 @@ impl_mysql_enum_coders!(JobStatusPlus); // NB: We can derive `sqlx::Type` instead of using `impl_mysql_enum_coders` +impl Default for JobStatusPlus { + fn default() -> Self { + Self::Pending + } +} + impl JobStatusPlus { pub fn to_str(&self) -> &'static str { match self { @@ -86,6 +92,11 @@ mod tests { mod explicit_checks { use super::*; + #[test] + fn test_default() { + assert_eq!(JobStatusPlus::default(), JobStatusPlus::Pending); + } + #[test] fn test_serialization() { assert_serialization(JobStatusPlus::Pending, "pending"); diff --git a/crates/_types/tokens/src/tokens/anonymous_visitor_tracking.rs b/crates/_types/tokens/src/tokens/anonymous_visitor_tracking.rs index 048d0f9a93..dbb1659ff9 100644 --- a/crates/_types/tokens/src/tokens/anonymous_visitor_tracking.rs +++ b/crates/_types/tokens/src/tokens/anonymous_visitor_tracking.rs @@ -8,7 +8,7 @@ use crate::prefixes::TokenPrefix; /// A rare instance that is *not* a primary key. /// These are generated as cookies to track users for a better logged out experience and for analytics. /// We use these as indices into several columns (ML result types, uploads, etc.) -#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, sqlx::Type, Debug, Serialize, Deserialize)] +#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, sqlx::Type, Debug, Serialize, Deserialize, Default)] #[sqlx(transparent)] pub struct AnonymousVisitorTrackingToken(pub String); diff --git a/crates/_types/tokens/src/tokens/generic_inference_jobs.rs b/crates/_types/tokens/src/tokens/generic_inference_jobs.rs index 48c3a78260..ddc5ca5c6a 100644 --- a/crates/_types/tokens/src/tokens/generic_inference_jobs.rs +++ b/crates/_types/tokens/src/tokens/generic_inference_jobs.rs @@ -7,7 +7,7 @@ use crate::prefixes::TokenPrefix; use utoipa::ToSchema; /// The primary key for "generic" inference jobs. -#[derive(Clone, PartialEq, Eq, sqlx::Type, Debug, Serialize, Deserialize,ToSchema)] +#[derive(Clone, PartialEq, Eq, sqlx::Type, Debug, Serialize, Deserialize, Default, ToSchema)] #[sqlx(transparent)] pub struct InferenceJobToken(String); diff --git a/crates/_types/tokens/src/tokens/users.rs b/crates/_types/tokens/src/tokens/users.rs index 2e28fca9fe..82438cc81c 100644 --- a/crates/_types/tokens/src/tokens/users.rs +++ b/crates/_types/tokens/src/tokens/users.rs @@ -6,7 +6,7 @@ use serde::Serialize; use crate::prefixes::TokenPrefix; /// The primary key for users. -#[derive(Clone, PartialEq, Eq, sqlx::Type, Debug, Serialize, Deserialize)] +#[derive(Clone, PartialEq, Eq, sqlx::Type, Debug, Serialize, Deserialize, Default)] #[sqlx(transparent)] pub struct UserToken(pub String); diff --git a/crates/lib/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs b/crates/lib/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs index 615ef68ca9..a59fd2e5d2 100644 --- a/crates/lib/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs +++ b/crates/lib/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs @@ -14,7 +14,7 @@ use tokens::tokens::users::UserToken; use crate::helpers::boolean_converters::i8_to_bool; -#[derive(Debug)] +#[derive(Debug, Default)] pub struct GenericInferenceJobStatus { pub job_token: InferenceJobToken, @@ -40,7 +40,7 @@ pub struct GenericInferenceJobStatus { /// Details about the user's original inference request /// (We may want to present it in the "pending" UI.) -#[derive(Debug)] +#[derive(Debug, Default)] pub struct RequestDetails { pub inference_category: InferenceCategory, pub maybe_model_type: Option, // TODO: Strongly type @@ -52,7 +52,7 @@ pub struct RequestDetails { } /// Details about the generated result -#[derive(Debug)] +#[derive(Debug, Default)] pub struct ResultDetails { pub entity_type: String, pub entity_token: String, @@ -71,7 +71,7 @@ pub struct ResultDetails { /// NB: DO NOT EXPOSE TO FRONTEND. /// This is used to gate access to job termination -#[derive(Debug)] +#[derive(Debug, Default)] pub struct UserDetails { pub maybe_creator_user_token: Option, pub maybe_creator_anonymous_visitor_token: Option, @@ -168,6 +168,11 @@ WHERE jobs.token = ? } }; + Ok(Some(raw_record_to_public_result(record))) +} + +/// Map the internal record type (since we query over several result tables) +fn raw_record_to_public_result(record: RawGenericInferenceJobStatus) -> GenericInferenceJobStatus { let maybe_model_title = match record.inference_category { InferenceCategory::LipsyncAnimation => Some("lipsync animation"), InferenceCategory::TextToSpeech => record.maybe_tts_model_title.as_deref(), @@ -213,7 +218,7 @@ WHERE jobs.token = ? }) }); - Ok(Some(GenericInferenceJobStatus { + GenericInferenceJobStatus { job_token: record.job_token, status: record.status, attempt_count: record.attempt_count, @@ -237,9 +242,10 @@ WHERE jobs.token = ? is_keepalive_required: i8_to_bool(record.is_keepalive_required), created_at: record.created_at, updated_at: record.updated_at, - })) + } } +#[derive(Debug, Default)] struct RawGenericInferenceJobStatus { pub job_token: InferenceJobToken, @@ -281,3 +287,60 @@ struct RawGenericInferenceJobStatus { pub maybe_successfully_completed_at: Option>, } +#[cfg(test)] +mod tests { + use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; + use crate::queries::generic_inference::web::get_inference_job_status::{raw_record_to_public_result, RawGenericInferenceJobStatus}; + + #[test] + fn text_to_speech_as_media_file() { + // NB(bt,2023-12-06): After tts migration to media files, these are the results the job system will return + let record = RawGenericInferenceJobStatus { + inference_category: InferenceCategory::TextToSpeech, + maybe_result_entity_type: Some("media_file".to_string()), // NB: Media file record + maybe_result_entity_token: Some("m_00af018cy75pxytpb8wbdx9jqgtp0p".to_string()), + maybe_media_file_public_bucket_directory_hash: Some("3vb91yq71z5zne56saazn4qntt52yter".to_string()), // NB: Media file + maybe_media_file_public_bucket_prefix: Some("fakeyou_".to_string()), + maybe_media_file_public_bucket_extension: Some(".wav".to_string()), + maybe_tts_public_bucket_path: None, // NB: Not a tts_result! + ..Default::default() + }; + + let record = raw_record_to_public_result(record); + + assert!(record.maybe_result_details.is_some()); + + let result_details = record.maybe_result_details.expect("should have result details"); + + assert_eq!(&result_details.entity_type, "media_file"); + assert_eq!(&result_details.entity_token, "m_00af018cy75pxytpb8wbdx9jqgtp0p"); + assert_eq!(&result_details.public_bucket_location_or_hash, "3vb91yq71z5zne56saazn4qntt52yter"); + assert_eq!(result_details.public_bucket_location_is_hash, true); + } + + #[test] + fn text_to_speech_as_tts_result() { + // NB(bt,2023-12-06): Prior to tts migration to media files, these are the results the job system will return + let record = RawGenericInferenceJobStatus { + inference_category: InferenceCategory::TextToSpeech, + maybe_result_entity_type: Some("text_to_speech".to_string()), // NB: TTS record + maybe_result_entity_token: Some("TR:wefhbk4j3yc8d6zwembedbdw4ad4s".to_string()), + maybe_tts_public_bucket_path: Some("/tts_inference_output/c/5/8/vocodes_c58aeaef-84df-4478-9e7f-c64280a852e8.wav".to_string()), // NB: tts_result! + maybe_media_file_public_bucket_directory_hash: None, // NB: Not a media file! + maybe_media_file_public_bucket_prefix: None, + maybe_media_file_public_bucket_extension: None, + ..Default::default() + }; + + let record = raw_record_to_public_result(record); + + assert!(record.maybe_result_details.is_some()); + + let result_details = record.maybe_result_details.expect("should have result details"); + + assert_eq!(&result_details.entity_type, "text_to_speech"); + assert_eq!(&result_details.entity_token, "TR:wefhbk4j3yc8d6zwembedbdw4ad4s"); + assert_eq!(&result_details.public_bucket_location_or_hash, "/tts_inference_output/c/5/8/vocodes_c58aeaef-84df-4478-9e7f-c64280a852e8.wav"); + assert_eq!(result_details.public_bucket_location_is_hash, false); + } +} diff --git a/crates/lib/mysql_queries/src/queries/media_files/insert_media_file_from_tacotron2.rs b/crates/lib/mysql_queries/src/queries/media_files/insert_media_file_from_tacotron2.rs new file mode 100644 index 0000000000..bc3010c1f8 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/media_files/insert_media_file_from_tacotron2.rs @@ -0,0 +1,166 @@ +use anyhow::anyhow; +use sqlx; +use sqlx::MySqlPool; + +use enums::by_table::generic_synthetic_ids::id_category::IdCategory; +use enums::by_table::media_files::media_file_origin_category::MediaFileOriginCategory; +use enums::by_table::media_files::media_file_origin_model_type::MediaFileOriginModelType; +use enums::by_table::media_files::media_file_origin_product_category::MediaFileOriginProductCategory; +use enums::by_table::media_files::media_file_type::MediaFileType; +use errors::AnyhowResult; +use tokens::tokens::media_files::MediaFileToken; +use tokens::tokens::users::UserToken; + +use crate::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; +use crate::queries::generic_synthetic_ids::transactional_increment_generic_synthetic_id::transactional_increment_generic_synthetic_id; + +pub struct InsertMediaFileArgs<'a> { + pub pool: &'a MySqlPool, + pub job: &'a AvailableInferenceJob, + + pub maybe_mime_type: Option<&'a str>, + pub file_size_bytes: u64, + pub duration_millis: u64, + pub sha256_checksum: &'a str, + + // TODO: Media duration. + //pub duration_millis: u64, + + pub public_bucket_directory_hash: &'a str, + pub maybe_public_bucket_prefix: Option<&'a str>, + pub maybe_public_bucket_extension: Option<&'a str>, + + pub is_on_prem: bool, + pub worker_hostname: &'a str, + pub worker_cluster: &'a str, +} + +pub async fn insert_media_file_from_tacotron2( + args: InsertMediaFileArgs<'_> +) -> AnyhowResult<(MediaFileToken, u64)> +{ + let result_token = MediaFileToken::generate(); + + let origin_model_type = MediaFileOriginModelType::Tacotron2; + + let mut maybe_creator_file_synthetic_id : Option = None; + let mut maybe_creator_category_synthetic_id : Option = None; + + let mut transaction = args.pool.begin().await?; + + if let Some(creator_user_token) = args.job.maybe_creator_user_token.as_deref() { + let user_token = UserToken::new_from_str(creator_user_token); + + let next_media_file_id = transactional_increment_generic_synthetic_id( + &user_token, + IdCategory::MediaFile, + &mut transaction + ).await?; + + let next_voice_conversion_id = transactional_increment_generic_synthetic_id( + &user_token, + IdCategory::TtsResult, + &mut transaction + ).await?; + + maybe_creator_file_synthetic_id = Some(next_media_file_id); + maybe_creator_category_synthetic_id = Some(next_voice_conversion_id); + } + + let vc_model_token = args.job.maybe_model_token.as_deref(); + let creator_ip_address = args.job.creator_ip_address.as_str(); + let creator_set_visibility = args.job.creator_set_visibility.clone(); + + const ORIGIN_CATEGORY : MediaFileOriginCategory = MediaFileOriginCategory::Inference; + const ORIGIN_PRODUCT_CATEGORY : MediaFileOriginProductCategory = MediaFileOriginProductCategory::TextToSpeech; + const MEDIA_TYPE : MediaFileType = MediaFileType::Audio; + + let record_id = { + let query_result = sqlx::query!( + r#" +INSERT INTO media_files +SET + token = ?, + + origin_category = ?, + origin_product_category = ?, + + maybe_origin_model_type = ?, + maybe_origin_model_token = ?, + + media_type = ?, + maybe_mime_type = ?, + file_size_bytes = ?, + + checksum_sha2 = ?, + + public_bucket_directory_hash = ?, + maybe_public_bucket_prefix = ?, + maybe_public_bucket_extension = ?, + + maybe_creator_user_token = ?, + maybe_creator_anonymous_visitor_token = ?, + creator_ip_address = ?, + + creator_set_visibility = ?, + + maybe_creator_file_synthetic_id = ?, + maybe_creator_category_synthetic_id = ?, + + is_generated_on_prem = ?, + generated_by_worker = ?, + generated_by_cluster = ? + + "#, + result_token.as_str(), + + ORIGIN_CATEGORY.to_str(), + ORIGIN_PRODUCT_CATEGORY.to_str(), + + origin_model_type.to_str(), + args.job.maybe_model_token, + + MEDIA_TYPE.to_str(), + args.maybe_mime_type, + args.file_size_bytes, + + args.sha256_checksum, + + args.public_bucket_directory_hash, + args.maybe_public_bucket_prefix, + args.maybe_public_bucket_extension, + + args.job.maybe_creator_user_token, + args.job.maybe_creator_anonymous_visitor_token, + args.job.creator_ip_address, + + args.job.creator_set_visibility.to_str(), + + maybe_creator_file_synthetic_id, + maybe_creator_category_synthetic_id, + + args.is_on_prem, + args.worker_hostname, + args.worker_cluster + ) + .execute(&mut *transaction) + .await; + + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + // TODO: handle better + //transaction.rollback().await?; + return Err(anyhow!("Mysql error: {:?}", err)); + } + }; + + record_id + }; + + transaction.commit().await?; + + Ok((result_token, record_id)) +} diff --git a/crates/lib/mysql_queries/src/queries/media_files/mod.rs b/crates/lib/mysql_queries/src/queries/media_files/mod.rs index 0590f82159..8930a3ae7c 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/mod.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/mod.rs @@ -2,6 +2,7 @@ pub mod delete_media_file; pub mod get_media_file; pub mod insert_media_file_from_cli_tool; pub mod insert_media_file_from_face_animation; +pub mod insert_media_file_from_tacotron2; pub mod insert_media_file_from_voice_conversion; pub mod insert_media_file_from_zero_shot_tts; pub mod list_media_files; diff --git a/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/mod.rs b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/mod.rs index a895808b59..c6d81ab3cf 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/mod.rs @@ -8,5 +8,6 @@ pub mod tacotron2_dependencies; pub mod tacotron2_inference_command; pub mod tacotron2_inference_sidecar_client; pub mod tacotron2_sidecar_health_check_client; +pub mod upload_results; pub mod virtual_lfu_cache; pub mod vocoder_option; diff --git a/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/process_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/process_job.rs index c6559a478a..0a3460beb4 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/process_job.rs @@ -27,6 +27,7 @@ use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::download_static_depe use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::health_check_trap::maybe_block_on_sidecar_health_check; use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::seconds_to_decoder_steps::seconds_to_decoder_steps; use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::tacotron2_inference_command::{InferenceArgs, MelMultiplyFactor}; +use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::upload_results::{upload_results, UploadResultArgs}; use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::vocoder_option::VocoderForInferenceOption; use crate::job_dependencies::JobDependencies; use crate::util::maybe_download_file_from_bucket::{maybe_download_file_from_bucket, MaybeDownloadArgs}; @@ -305,78 +306,27 @@ async fn process_job_with_cleanup( safe_delete_temp_file(&output_metadata_fs_path); - // ==================== UPLOAD AUDIO TO BUCKET ==================== // + // ==================== UPLOAD FILES AND SAVE RECORD ==================== // - job_progress_reporter.log_status("uploading result") - .map_err(|e| ProcessSingleJobError::Other(e))?; - - let audio_result_object_path = args.job_dependencies.buckets.bucket_path_unifier.tts_inference_wav_audio_output_path( - &job.uuid_idempotency_token); // TODO: Don't use this! - - info!("Audio destination bucket path: {:?}", &audio_result_object_path); - - info!("Uploading audio..."); - - args.job_dependencies.buckets.public_bucket_client.upload_filename_with_content_type( - &audio_result_object_path, - &output_audio_fs_path, - "audio/wav") - .await - .map_err(|e| ProcessSingleJobError::Other(e))?; - - safe_delete_temp_file(&output_audio_fs_path); - - // ==================== UPLOAD SPECTROGRAM TO BUCKETS ==================== // - - let spectrogram_result_object_path = args.job_dependencies.buckets.bucket_path_unifier.tts_inference_spectrogram_output_path( - &job.uuid_idempotency_token); // TODO: Don't use this! - - info!("Spectrogram destination bucket path: {:?}", &spectrogram_result_object_path); - - info!("Uploading spectrogram..."); - - args.job_dependencies.buckets.public_bucket_client.upload_filename_with_content_type( - &spectrogram_result_object_path, - &output_spectrogram_fs_path, - "application/json") - .await - .map_err(|e| ProcessSingleJobError::Other(e))?; - - safe_delete_temp_file(&output_spectrogram_fs_path); - - // ==================== DELETE DOWNLOADED FILE ==================== // - - // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... - safe_delete_temp_directory(&work_temp_dir.path()); - - // ==================== SAVE RECORDS ==================== // - - let text_hash = sha256_hash_string(&cleaned_inference_text) - .map_err(|e| ProcessSingleJobError::Other(e))?; - - info!("Saving tts inference record..."); - - let (id, inference_result_token) = insert_tts_result( - &args.job_dependencies.db.mysql_pool, - JobType::GenericInferenceJob(&job), - &text_hash, - Some(pretrained_vocoder), - &audio_result_object_path, - &spectrogram_result_object_path, - file_metadata.file_size_bytes, - file_metadata.duration_millis.unwrap_or(0), - args.job_dependencies.job.info.container.is_on_prem, - &args.job_dependencies.job.info.container.hostname, - args.job_dependencies.job.system.is_debug_worker) - .await - .map_err(|e| ProcessSingleJobError::Other(e))?; + let inference_result = upload_results(UploadResultArgs { + job_dependencies: args.job_dependencies, + job_progress_reporter: &mut job_progress_reporter, + job, + cleaned_inference_text: &cleaned_inference_text, + work_temp_dir, + pretrained_vocoder, + file_metadata: &file_metadata, + output_audio_fs_path: &output_audio_fs_path, + output_spectrogram_fs_path: &output_spectrogram_fs_path, + upload_as_media_file: model_dependencies.upload_as_media_file, + }).await?; info!("TTS Done. Original text was: {:?}", &job.maybe_raw_inference_text); args.job_dependencies.clients.firehose_publisher.tts_inference_finished( job.maybe_creator_user_token.as_deref(), tts_model.model_token.as_str(), - &inference_result_token) + &inference_result.entity_token) .await .map_err(|e| { error!("error publishing event: {:?}", e); @@ -386,20 +336,20 @@ async fn process_job_with_cleanup( job_progress_reporter.log_status("done") .map_err(|e| ProcessSingleJobError::Other(e))?; - info!("Job {:?} complete success! Downloaded, ran inference, and uploaded. Saved model record: {}, Result Token: {}", - job.id, id, &inference_result_token); + info!("Job {:?} complete success! Downloaded, ran inference, and uploaded. Result Token: {}", + job.id, &inference_result.entity_token); Ok(JobSuccessResult { maybe_result_entity: Some(ResultEntity { - entity_type: InferenceResultType::TextToSpeech, - entity_token: inference_result_token + entity_type: inference_result.entity_type, + entity_token: inference_result.entity_token, }), inference_duration, }) } #[derive(Deserialize, Default)] -struct FileMetadata { +pub struct FileMetadata { pub duration_millis: Option, pub mimetype: Option, pub file_size_bytes: u64, diff --git a/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/tacotron2_dependencies.rs b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/tacotron2_dependencies.rs index baf20bda9c..7b6df49855 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/tacotron2_dependencies.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/tacotron2_dependencies.rs @@ -20,6 +20,8 @@ pub struct Tacotron2Dependencies { /// Dependencies that are for the sidecar version of TT2 pub sidecar: SidecarDeps, + + pub upload_as_media_file: bool, } pub struct SidecarDeps { @@ -71,7 +73,8 @@ impl Tacotron2Dependencies { health_check_client, health_check_state: HealthCheckState::new(), virtual_lfu_cache, - } + }, + upload_as_media_file: easyenv::get_env_bool_or_default("TT2_AS_MEDIA_FILES", false), }) } } diff --git a/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/upload_results.rs b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/upload_results.rs new file mode 100644 index 0000000000..cd5df7ff23 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/upload_results.rs @@ -0,0 +1,202 @@ +use std::path::Path; +use anyhow::anyhow; + +use log::info; +use tempdir::TempDir; +use buckets::public::media_files::original_file::MediaFileBucketPath; + +use cloud_storage::bucket_path_unifier::BucketPathUnifier; +use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; +use filesys::safe_delete_temp_directory::safe_delete_temp_directory; +use filesys::safe_delete_temp_file::safe_delete_temp_file; +use hashing::sha256::sha256_hash_file::sha256_hash_file; +use hashing::sha256::sha256_hash_string::sha256_hash_string; +use jobs_common::job_progress_reporter::job_progress_reporter::JobProgressReporter; +use mysql_queries::column_types::vocoder_type::VocoderType; +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; +use mysql_queries::queries::media_files::insert_media_file_from_tacotron2::{insert_media_file_from_tacotron2, InsertMediaFileArgs}; +use mysql_queries::queries::tts::tts_results::insert_tts_result::JobType; +use mysql_queries::queries::tts::tts_results::insert_tts_result::insert_tts_result; + +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::process_job::FileMetadata; +use crate::job_dependencies::JobDependencies; + +const BUCKET_FILE_PREFIX: &str = "fakeyou_"; +const BUCKET_FILE_EXTENSION: &str = ".wav"; +const MIME_TYPE: &str = "audio/wav"; + +pub struct UploadResultArgs<'a> { + // System + pub job_dependencies: &'a JobDependencies, + pub job_progress_reporter: &'a mut Box, + + // Job + pub job: &'a AvailableInferenceJob, + pub cleaned_inference_text: &'a str, + pub work_temp_dir: &'a TempDir, + pub pretrained_vocoder: VocoderType, + + // Outputs + pub file_metadata: &'a FileMetadata, + pub output_audio_fs_path: &'a Path, + pub output_spectrogram_fs_path: &'a Path, + + // Flags + pub upload_as_media_file: bool, +} + +pub struct ResultDetails { + pub entity_type: InferenceResultType, + pub entity_token: String, +} + +pub async fn upload_results(args: UploadResultArgs<'_>) -> Result{ + if args.upload_as_media_file { + upload_as_media_file(args).await + } else { + upload_as_legacy_tts_result(args).await + } +} + +async fn upload_as_media_file(args: UploadResultArgs<'_>) -> Result{ + + // ==================== UPLOAD AUDIO TO BUCKET ==================== // + + args.job_progress_reporter.log_status("uploading result") + .map_err(|e| ProcessSingleJobError::Other(e))?; + + let result_bucket_location = MediaFileBucketPath::generate_new( + Some(BUCKET_FILE_PREFIX), + Some(BUCKET_FILE_EXTENSION) + ); + + let result_bucket_object_pathbuf = result_bucket_location.to_full_object_pathbuf(); + + info!("Audio destination bucket path: {:?}", &result_bucket_location.to_full_object_pathbuf()); + + info!("Uploading audio..."); + + args.job_dependencies.buckets.public_bucket_client.upload_filename_with_content_type( + &result_bucket_object_pathbuf, + &args.output_audio_fs_path, + "audio/wav") + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + let file_checksum = sha256_hash_file(&args.output_audio_fs_path).map_err(|err| { + ProcessSingleJobError::Other(anyhow!("Error hashing file: {:?}", err)) + })?; + + safe_delete_temp_file(&args.output_audio_fs_path); + + // ==================== UPLOAD SPECTROGRAM TO BUCKETS ==================== // + + // (SKIP FOR NOW) + + // ==================== DELETE WORK DIRECTORY ==================== // + + // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... + safe_delete_temp_directory(&args.work_temp_dir.path()); + + // ==================== SAVE RECORDS ==================== // + + info!("Saving tts inference record..."); + + let (inference_result_token, _id) = insert_media_file_from_tacotron2(InsertMediaFileArgs { + pool: &args.job_dependencies.db.mysql_pool, + job: &args.job, + maybe_mime_type: Some(&MIME_TYPE), + file_size_bytes: args.file_metadata.file_size_bytes, + duration_millis: args.file_metadata.duration_millis.unwrap_or(0), + sha256_checksum: &file_checksum, + public_bucket_directory_hash: result_bucket_location.get_object_hash(), + maybe_public_bucket_prefix: Some(BUCKET_FILE_PREFIX), + maybe_public_bucket_extension: Some(BUCKET_FILE_EXTENSION), + is_on_prem: args.job_dependencies.job.info.container.is_on_prem, + worker_hostname: &args.job_dependencies.job.info.container.hostname, + worker_cluster: &args.job_dependencies.job.info.container.cluster_name, + }) + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + Ok(ResultDetails { + entity_type: InferenceResultType::MediaFile, + entity_token: inference_result_token.to_string(), + }) +} + + +async fn upload_as_legacy_tts_result(args: UploadResultArgs<'_>) -> Result{ + + // ==================== UPLOAD AUDIO TO BUCKET ==================== // + + args.job_progress_reporter.log_status("uploading result") + .map_err(|e| ProcessSingleJobError::Other(e))?; + + let audio_result_object_path = args.job_dependencies.buckets.bucket_path_unifier.tts_inference_wav_audio_output_path( + &args.job.uuid_idempotency_token); // TODO: Don't use this! + + info!("Audio destination bucket path: {:?}", &audio_result_object_path); + + info!("Uploading audio..."); + + args.job_dependencies.buckets.public_bucket_client.upload_filename_with_content_type( + &audio_result_object_path, + &args.output_audio_fs_path, + "audio/wav") + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + safe_delete_temp_file(&args.output_audio_fs_path); + + // ==================== UPLOAD SPECTROGRAM TO BUCKETS ==================== // + + let spectrogram_result_object_path = args.job_dependencies.buckets.bucket_path_unifier.tts_inference_spectrogram_output_path( + &args.job.uuid_idempotency_token); // TODO: Don't use this! + + info!("Spectrogram destination bucket path: {:?}", &spectrogram_result_object_path); + + info!("Uploading spectrogram..."); + + args.job_dependencies.buckets.public_bucket_client.upload_filename_with_content_type( + &spectrogram_result_object_path, + &args.output_spectrogram_fs_path, + "application/json") + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + safe_delete_temp_file(&args.output_spectrogram_fs_path); + + // ==================== DELETE WORK DIRECTORY ==================== // + + // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... + safe_delete_temp_directory(&args.work_temp_dir.path()); + + // ==================== SAVE RECORDS ==================== // + + let text_hash = sha256_hash_string(&args.cleaned_inference_text) + .map_err(|e| ProcessSingleJobError::Other(e))?; + + info!("Saving tts inference record..."); + + let (_id, inference_result_token) = insert_tts_result( + &args.job_dependencies.db.mysql_pool, + JobType::GenericInferenceJob(&args.job), + &text_hash, + Some(args.pretrained_vocoder), + &audio_result_object_path, + &spectrogram_result_object_path, + args.file_metadata.file_size_bytes, + args.file_metadata.duration_millis.unwrap_or(0), + args.job_dependencies.job.info.container.is_on_prem, + &args.job_dependencies.job.info.container.hostname, + args.job_dependencies.job.system.is_debug_worker) + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + Ok(ResultDetails { + entity_type: InferenceResultType::TextToSpeech, + entity_token: inference_result_token, + }) +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_inference_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_inference_voice.rs index cb121025b2..20359f1e0d 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_inference_voice.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_inference_voice.rs @@ -23,7 +23,6 @@ use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; use crate::job::job_types::tts::vall_e_x::process_job::VALLEXProcessJobArgs; use crate::job::job_types::tts::vall_e_x::vall_e_x_inference_command::InferenceArgs; -// Clearify what this is for ? const BUCKET_FILE_PREFIX: &str = "fakeyou_"; const BUCKET_FILE_EXTENSION: &str = ".wav"; const MIME_TYPE: &str = "audio/wav"; @@ -208,6 +207,7 @@ pub async fn process_inference_voice( .map_err(|e| ProcessSingleJobError::Other(e))?; // ==================== UPLOAD AUDIO TO BUCKET ==================== + info!("Calculating sha256..."); let file_checksum = sha256_hash_file(&finished_file).map_err(|err| { diff --git a/crates/service/web/storyteller_web/Cargo.toml b/crates/service/web/storyteller_web/Cargo.toml index 669066a316..df7cab336c 100644 --- a/crates/service/web/storyteller_web/Cargo.toml +++ b/crates/service/web/storyteller_web/Cargo.toml @@ -54,6 +54,8 @@ user_input_common = { path = "../../../lib/user_input_common" } billing_component = { path = "../../plugins/billing" } users_component = { path = "../../plugins/users" } +# Workspace hack (faster compile times) +my-workspace-hack = { version = "0.1", path = "../../../my-workspace-hack" } # External anyhow = "1.0.75" @@ -69,6 +71,7 @@ log = "0.4.14" once_cell = "1.18.0" rand = "0.8.5" regex = "1.5.4" + #symphonia = { version = "0.5.2", features = ["aac", "alac", "flac", "isomp4", "mkv", "mp3", "ogg", "opus", "vorbis", "wav"] } # Opus support still hasn't landed in Symphonia, so we have to use an older version (~0.5) off-master. # Opus is the browser's default (only) recording codec: https://blog.addpipe.com/audio-constraints-getusermedia/ @@ -159,7 +162,6 @@ pulldown-cmark = "0.9.3" # The initial round of stabilizations for the async/await language feature in Rust 1.39 did not # include support for async fn in traits. async-trait = "0.1.58" -my-workspace-hack = { version = "0.1", path = "../../../my-workspace-hack" } utoipa = { version = "4", features = ["actix_extras"] } utoipa-swagger-ui = { version = "5", features = ["actix-web"] } diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs index b81b67e1b5..96be5ed7d7 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs @@ -14,7 +14,7 @@ use buckets::public::voice_conversion_results::original_file::VoiceConversionRes use enums::by_table::generic_inference_jobs::frontend_failure_category::FrontendFailureCategory; use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; use enums::common::job_status_plus::JobStatusPlus; -use mysql_queries::queries::generic_inference::web::get_inference_job_status::get_inference_job_status; +use mysql_queries::queries::generic_inference::web::get_inference_job_status::{GenericInferenceJobStatus, get_inference_job_status}; use redis_common::redis_keys::RedisKeys; use tokens::tokens::generic_inference_jobs::InferenceJobToken; use crate::http_server::responses::filter_model_name::maybe_filter_model_name; @@ -140,12 +140,15 @@ pub async fn get_inference_job_status_handler( if path.token.as_str().trim() == "None" { // NB: A bunch of Python clients use our API and can fail in this manner. // This was a large traffic driver during the 2023-03-08 outage. + // Presumably it's this client: https://github.com/shards-7/fakeyou.py return Err(GetInferenceJobStatusError::NotFound); } - // NB: Lookup failure is Err(RowNotFound). // NB: Since this is publicly exposed, we don't query sensitive data. - let maybe_status = get_inference_job_status(&path.token, &server_state.mysql_pool).await; + let maybe_status = get_inference_job_status( + &path.token, + &server_state.mysql_pool + ).await; let record = match maybe_status { Ok(Some(record)) => record, @@ -198,9 +201,32 @@ pub async fn get_inference_job_status_handler( }; } + let record_for_response = record_to_payload( + record, maybe_extra_status_description); + + let response = GetInferenceJobStatusSuccessResponse { + success: true, + state: record_for_response, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| { + error!("error returning response: {:?}", e); + GetInferenceJobStatusError::ServerError + })?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} + +fn record_to_payload( + record: GenericInferenceJobStatus, + maybe_extra_status_description: Option, +) -> InferenceJobStatusResponsePayload { let inference_category = record.request_details.inference_category.clone(); - let record_for_response = InferenceJobStatusResponsePayload { + InferenceJobStatusResponsePayload { job_token: record.job_token, request: RequestDetailsResponse { inference_category: record.request_details.inference_category, @@ -226,7 +252,7 @@ pub async fn get_inference_job_status_handler( MediaFileBucketPath::from_object_hash( &result_details.public_bucket_location_or_hash, result_details.maybe_media_file_public_bucket_prefix.as_deref(), - result_details.maybe_media_file_public_bucket_extension.as_deref()) + result_details.maybe_media_file_public_bucket_extension.as_deref()) .get_full_object_path_str() .to_string() } @@ -279,20 +305,80 @@ pub async fn get_inference_job_status_handler( }), created_at: record.created_at, updated_at: record.updated_at, - }; - - let response = GetInferenceJobStatusSuccessResponse { - success: true, - state: record_for_response, - }; - - let body = serde_json::to_string(&response) - .map_err(|e| { - error!("error returning response: {:?}", e); - GetInferenceJobStatusError::ServerError - })?; - - Ok(HttpResponse::Ok() - .content_type("application/json") - .body(body)) + } +} + +#[cfg(test)] +mod tests { + use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; + use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; + use mysql_queries::queries::generic_inference::web::get_inference_job_status::{GenericInferenceJobStatus, RequestDetails, ResultDetails}; + use crate::http_server::endpoints::inference_job::get_inference_job_status::record_to_payload; + + #[test] + fn text_to_speech_as_media_file() { + let status = GenericInferenceJobStatus { + request_details: RequestDetails { + inference_category: InferenceCategory::TextToSpeech, + maybe_model_type: Some("tacotron2".to_string()), + ..Default::default() + }, + maybe_result_details: Some(ResultDetails { + entity_type: "media_file".to_string(), + entity_token: "m_024pk3p3xrmeyk83chps90jeymnygk".to_string(), + public_bucket_location_or_hash: "tpk848b5s5zwhnwrph75jhyfyja3j42v".to_string(), + maybe_media_file_public_bucket_prefix: Some("fakeyou_".to_string()), + maybe_media_file_public_bucket_extension: Some(".wav".to_string()), + public_bucket_location_is_hash: true, + ..Default::default() + }), + ..Default::default() + }; + + let payload = + record_to_payload(status, None); + + assert!(payload.maybe_result.is_some()); + + let result = payload.maybe_result.expect("should have result"); + + assert_eq!( + Some("/media/t/p/k/8/4/tpk848b5s5zwhnwrph75jhyfyja3j42v/fakeyou_tpk848b5s5zwhnwrph75jhyfyja3j42v.wav"), + result.maybe_public_bucket_media_path.as_deref() + ); + } + + #[test] + fn text_to_speech_as_tts_result() { + let status = GenericInferenceJobStatus { + request_details: RequestDetails { + inference_category: InferenceCategory::TextToSpeech, + maybe_model_type: Some("tacotron2".to_string()), + ..Default::default() + }, + maybe_result_details: Some(ResultDetails { + entity_type: "text_to_speech".to_string(), + entity_token: "TR:qvgdvngw4y54t7rmfepwv8fd965eh".to_string(), + public_bucket_location_or_hash: + "/tts_inference_output/9/9/7/vocodes_997e710c-d733-406c-9b6e-5e1d0c471816.wav".to_string(), + maybe_media_file_public_bucket_prefix: None, + maybe_media_file_public_bucket_extension: None, + public_bucket_location_is_hash: false, + ..Default::default() + }), + ..Default::default() + }; + + let payload = + record_to_payload(status, None); + + assert!(payload.maybe_result.is_some()); + + let result = payload.maybe_result.expect("should have result"); + + assert_eq!( + Some("/tts_inference_output/9/9/7/vocodes_997e710c-d733-406c-9b6e-5e1d0c471816.wav"), + result.maybe_public_bucket_media_path.as_deref(), + ); + } } diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/get_tts_inference_job_status.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/get_tts_inference_job_status.rs index c1055c026f..be43c33d8a 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/get_tts_inference_job_status.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/get_tts_inference_job_status.rs @@ -150,6 +150,7 @@ async fn legacy_lookup(job_token: &str, server_state: &ServerState) GetTtsInferenceStatusError::ServerError })?; + // TODO(bt,2023-05-21): Make async. let extra_status_key = RedisKeys::tts_inference_extra_status_info(&job_token); let maybe_extra_status_description : Option = match redis.get(&extra_status_key) { Ok(Some(status)) => {