From 7a4c035960f80d0fe0b0cc9c551fbcaec4ce54ce Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 16 Nov 2023 20:32:19 -0500 Subject: [PATCH] upload of weights works --- .../buckets/src/public/zs_voices/directory.rs | 41 ++++++++----- .../buckets/src/public/zs_voices/file.rs | 34 ++++------- .../inference_result_type.rs | 12 +++- .../tts/vall_e_x/process_job_create_voice.rs | 57 +++++++++---------- 4 files changed, 78 insertions(+), 66 deletions(-) diff --git a/crates/_types/buckets/src/public/zs_voices/directory.rs b/crates/_types/buckets/src/public/zs_voices/directory.rs index a4a5300d80..1cfc740769 100644 --- a/crates/_types/buckets/src/public/zs_voices/directory.rs +++ b/crates/_types/buckets/src/public/zs_voices/directory.rs @@ -6,7 +6,7 @@ use crate::util::hashed_directory_path_long_string::hashed_directory_path_long_s // TODO: Generate these from a macro. // TODO: Use a central path registry for quick reference -const DIRECTORY: &str = "/user/zero_shot_embeddings"; +const DIRECTORY: &str = "/user/zs_embeddings"; #[derive(Debug, Clone, Copy, Eq, PartialEq)] pub enum ModelCategory { @@ -23,11 +23,25 @@ impl ModelCategory { } } +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +pub enum ModelType { + VallEx, +} + +impl ModelType { + pub fn to_str(&self) -> &'static str { + match self { + ModelType::VallEx => "vallex", + } + } +} + /// Directory for zero shot voice embeddings /// Each embedding gets its own directory so we can store new versions alongside the original. #[derive(Clone)] pub struct ZeroShotVoiceEmbeddingBucketDirectory { model_category: ModelCategory, + model_type: ModelType, object_hash: String, directory: String, } @@ -36,18 +50,19 @@ impl PrivatePath for ZeroShotVoiceEmbeddingBucketDirectory {} impl ZeroShotVoiceEmbeddingBucketDirectory { - pub fn generate_new(model_category: ModelCategory) -> Self { + pub fn generate_new(model_category: ModelCategory, model_type: ModelType) -> Self { let entropy = crockford_entropy_lower(32); - Self::from_object_hash(model_category, &entropy) + Self::from_object_hash(model_category, model_type, &entropy) } - pub fn from_object_hash(model_category: ModelCategory, object_hash: &str) -> Self { + pub fn from_object_hash(model_category: ModelCategory, model_type: ModelType, object_hash: &str) -> Self { // TODO: Path construction could be cleaner. - let subdir = model_category.to_str(); + let subdirs = format!("{}/{}", model_category.to_str(), model_type.to_str()); let middle = hashed_directory_path_long_string(object_hash); - let directory = format!("{}/{}/{}{}", DIRECTORY, subdir, middle, object_hash); + let directory = format!("{}/{}/{}{}", DIRECTORY, subdirs, middle, object_hash); Self { model_category, + model_type, object_hash: object_hash.to_string(), directory, } @@ -64,30 +79,30 @@ impl ZeroShotVoiceEmbeddingBucketDirectory { #[cfg(test)] mod tests { - use crate::public::zs_voices::directory::{ModelCategory, ZeroShotVoiceEmbeddingBucketDirectory}; + use crate::public::zs_voices::directory::{ModelCategory, ModelType, ZeroShotVoiceEmbeddingBucketDirectory}; #[test] pub fn generate_new_entropy() { - let directory = ZeroShotVoiceEmbeddingBucketDirectory::generate_new(ModelCategory::Tts); + let directory = ZeroShotVoiceEmbeddingBucketDirectory::generate_new(ModelCategory::Tts, ModelType::VallEx); assert_eq!(directory.get_object_hash().len(), 32); } #[test] pub fn get_directory_path_str() { - let directory = ZeroShotVoiceEmbeddingBucketDirectory::from_object_hash(ModelCategory::Tts, "abcdefghijk"); - assert_eq!(directory.get_directory_path_str(), "/user/zero_shot_embeddings/tts/a/b/c/d/e/abcdefghijk"); + let directory = ZeroShotVoiceEmbeddingBucketDirectory::from_object_hash(ModelCategory::Tts, ModelType::VallEx, "abcdefghijk"); + assert_eq!(directory.get_directory_path_str(), "/user/zs_embeddings/tts/vallex/a/b/c/d/e/abcdefghijk"); } #[test] pub fn get_directory_path_str_short_name() { - let directory = ZeroShotVoiceEmbeddingBucketDirectory::from_object_hash(ModelCategory::Tts, "foo"); - assert_eq!(directory.get_directory_path_str(), "/user/zero_shot_embeddings/tts/f/o/foo"); + let directory = ZeroShotVoiceEmbeddingBucketDirectory::from_object_hash(ModelCategory::Tts, ModelType::VallEx, "foo"); + assert_eq!(directory.get_directory_path_str(), "/user/zs_embeddings/tts/vallex/f/o/foo"); } #[test] pub fn get_object_hash() { let hash = "abcdefghijk"; - let directory = ZeroShotVoiceEmbeddingBucketDirectory::from_object_hash(ModelCategory::Tts, hash); + let directory = ZeroShotVoiceEmbeddingBucketDirectory::from_object_hash(ModelCategory::Tts, ModelType::VallEx, hash); assert_eq!(directory.get_object_hash(), hash); } } diff --git a/crates/_types/buckets/src/public/zs_voices/file.rs b/crates/_types/buckets/src/public/zs_voices/file.rs index 7bb6aeb9b2..6428936c7c 100644 --- a/crates/_types/buckets/src/public/zs_voices/file.rs +++ b/crates/_types/buckets/src/public/zs_voices/file.rs @@ -3,24 +3,11 @@ use std::path::PathBuf; use crockford::crockford_entropy_lower; use crate::private::private_path::PrivatePath; -use crate::public::zs_voices::directory::{ModelCategory, ZeroShotVoiceEmbeddingBucketDirectory}; +use crate::public::zs_voices::directory::{ModelCategory, ModelType, ZeroShotVoiceEmbeddingBucketDirectory}; // TODO: Generate these from a macro. -#[derive(Debug, Clone, Copy, Eq, PartialEq)] -pub enum ModelType { - VallEx, -} - -impl ModelType { - pub fn to_str(&self) -> &'static str { - match self { - ModelType::VallEx => "vallex", - } - } -} - /// The original user upload file. /// It may have derivative files (down samples, crops, etc.) that live alongside it. #[derive(Clone)] @@ -32,21 +19,21 @@ pub struct ZeroShotVoiceEmbeddingBucketPath { model_category: ModelCategory, model_type: ModelType, - model_version: u32, + model_version: u64, } impl PrivatePath for ZeroShotVoiceEmbeddingBucketPath {} impl ZeroShotVoiceEmbeddingBucketPath { - pub fn generate_new(model_category: ModelCategory, model_type: ModelType, model_version: u32) -> Self { + pub fn generate_new(model_category: ModelCategory, model_type: ModelType, model_version: u64) -> Self { let entropy = crockford_entropy_lower(32); Self::from_object_hash(&entropy, model_category, model_type, model_version) } - pub fn from_object_hash(hash: &str, model_category: ModelCategory, model_type: ModelType, model_version: u32) -> Self { + pub fn from_object_hash(hash: &str, model_category: ModelCategory, model_type: ModelType, model_version: u64) -> Self { // TODO: Path construction could be cleaner. - let directory = ZeroShotVoiceEmbeddingBucketDirectory::from_object_hash(model_category, hash); + let directory = ZeroShotVoiceEmbeddingBucketDirectory::from_object_hash(model_category, model_type, hash); let filename = format!("{}_{}_{}", model_type.to_str(), hash, model_version); @@ -104,7 +91,7 @@ mod tests { #[test] pub fn get_full_object_path_str() { let file = ZeroShotVoiceEmbeddingBucketPath::from_object_hash("abcdefghijk", ModelCategory::Tts, ModelType::VallEx, 0); - assert_eq!(file.get_full_object_path_str(), "/user/zero_shot_embeddings/tts/a/b/c/d/e/abcdefghijk/vallex_abcdefghijk_0"); + assert_eq!(file.get_full_object_path_str(), "/user/zs_embeddings/tts/vallex/a/b/c/d/e/abcdefghijk/vallex_abcdefghijk_0"); assert_eq!(file.get_basename(), "vallex_abcdefghijk_0"); } } @@ -112,6 +99,7 @@ mod tests { mod without_prefix_and_extension { use crate::public::zs_voices::directory::ModelCategory; use crate::public::zs_voices::file::{ModelType, ZeroShotVoiceEmbeddingBucketPath}; + use super::*; #[test] @@ -124,16 +112,16 @@ mod tests { #[test] pub fn from_object_hash_tts_vallex_version() { let file = ZeroShotVoiceEmbeddingBucketPath::from_object_hash("abcdefghijk", ModelCategory::Tts, ModelType::VallEx, 123); - assert_eq!(file.get_full_object_path_str(), "/user/zero_shot_embeddings/tts/a/b/c/d/e/abcdefghijk/vallex_abcdefghijk_123"); - assert_eq!(file.to_full_object_pathbuf(), PathBuf::from("/user/zero_shot_embeddings/tts/a/b/c/d/e/abcdefghijk/vallex_abcdefghijk_123")); + assert_eq!(file.get_full_object_path_str(), "/user/zs_embeddings/tts/vallex/a/b/c/d/e/abcdefghijk/vallex_abcdefghijk_123"); + assert_eq!(file.to_full_object_pathbuf(), PathBuf::from("/user/zs_embeddings/tts/vallex/a/b/c/d/e/abcdefghijk/vallex_abcdefghijk_123")); assert_eq!(file.get_basename(), "vallex_abcdefghijk_123"); } #[test] pub fn from_object_hash_vc_vallex_version() { let file = ZeroShotVoiceEmbeddingBucketPath::from_object_hash("abcdefghijk", ModelCategory::Vc, ModelType::VallEx, 0); - assert_eq!(file.get_full_object_path_str(), "/user/zero_shot_embeddings/vc/a/b/c/d/e/abcdefghijk/vallex_abcdefghijk_0"); - assert_eq!(file.to_full_object_pathbuf(), PathBuf::from("/user/zero_shot_embeddings/vc/a/b/c/d/e/abcdefghijk/vallex_abcdefghijk_0")); + assert_eq!(file.get_full_object_path_str(), "/user/zs_embeddings/vc/vallex/a/b/c/d/e/abcdefghijk/vallex_abcdefghijk_0"); + assert_eq!(file.to_full_object_pathbuf(), PathBuf::from("/user/zs_embeddings/vc/vallex/a/b/c/d/e/abcdefghijk/vallex_abcdefghijk_0")); assert_eq!(file.get_basename(), "vallex_abcdefghijk_0"); } diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs index c687124416..4a7d879c21 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs @@ -26,6 +26,9 @@ pub enum InferenceResultType { #[serde(rename = "voice_conversion")] VoiceConversion, + + #[serde(rename = "zs_voice_embedding")] + ZeroShotVoiceEmbedding, } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -39,6 +42,7 @@ impl InferenceResultType { Self::MediaFile => "media_file", Self::TextToSpeech => "text_to_speech", Self::VoiceConversion => "voice_conversion", + Self::ZeroShotVoiceEmbedding => "zs_voice_embedding", } } @@ -47,6 +51,7 @@ impl InferenceResultType { "media_file" => Ok(Self::MediaFile), "text_to_speech" => Ok(Self::TextToSpeech), "voice_conversion" => Ok(Self::VoiceConversion), + "zs_voice_embedding" => Ok(Self::ZeroShotVoiceEmbedding), _ => Err(format!("invalid value: {:?}", value)), } } @@ -58,6 +63,7 @@ impl InferenceResultType { InferenceResultType::MediaFile, InferenceResultType::TextToSpeech, InferenceResultType::VoiceConversion, + InferenceResultType::ZeroShotVoiceEmbedding, ]) } } @@ -72,6 +78,7 @@ mod tests { assert_serialization(InferenceResultType::MediaFile, "media_file"); assert_serialization(InferenceResultType::TextToSpeech, "text_to_speech"); assert_serialization(InferenceResultType::VoiceConversion, "voice_conversion"); + assert_serialization(InferenceResultType::ZeroShotVoiceEmbedding, "zs_voice_embedding"); } #[test] @@ -79,6 +86,7 @@ mod tests { assert_eq!(InferenceResultType::MediaFile.to_str(), "media_file"); assert_eq!(InferenceResultType::TextToSpeech.to_str(), "text_to_speech"); assert_eq!(InferenceResultType::VoiceConversion.to_str(), "voice_conversion"); + assert_eq!(InferenceResultType::ZeroShotVoiceEmbedding.to_str(), "zs_voice_embedding"); } #[test] @@ -86,16 +94,18 @@ mod tests { assert_eq!(InferenceResultType::from_str("media_file").unwrap(), InferenceResultType::MediaFile); assert_eq!(InferenceResultType::from_str("text_to_speech").unwrap(), InferenceResultType::TextToSpeech); assert_eq!(InferenceResultType::from_str("voice_conversion").unwrap(), InferenceResultType::VoiceConversion); + assert_eq!(InferenceResultType::from_str("zs_voice_embedding").unwrap(), InferenceResultType::ZeroShotVoiceEmbedding); } #[test] fn all_variants() { // Static check let mut variants = InferenceResultType::all_variants(); - assert_eq!(variants.len(), 3); + assert_eq!(variants.len(), 4); assert_eq!(variants.pop_first(), Some(InferenceResultType::MediaFile)); assert_eq!(variants.pop_first(), Some(InferenceResultType::TextToSpeech)); assert_eq!(variants.pop_first(), Some(InferenceResultType::VoiceConversion)); + assert_eq!(variants.pop_first(), Some(InferenceResultType::ZeroShotVoiceEmbedding)); assert_eq!(variants.pop_first(), None); // Generated check diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_create_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_create_voice.rs index bc85db7971..7a310be608 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_create_voice.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_create_voice.rs @@ -1,12 +1,13 @@ use std::fs::read_to_string; use std::path::PathBuf; -use std::thread; -use std::time::{Duration, Instant}; +use std::time::Instant; use anyhow::anyhow; use log::{error, info, warn}; use buckets::public::media_files::original_file::MediaFileBucketPath; +use buckets::public::zs_voices::directory::{ModelCategory, ModelType}; +use buckets::public::zs_voices::file::ZeroShotVoiceEmbeddingBucketPath; use cloud_storage::bucket_client::BucketClient; use cloud_storage::bucket_path_unifier::BucketPathUnifier; use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; @@ -20,7 +21,6 @@ use mysql_queries::queries::voice_designer::voice_samples::list_dataset_samples_ use mysql_queries::queries::voice_designer::voice_samples::list_dataset_samples_for_dataset_token::list_dataset_samples_for_dataset_token; use mysql_queries::queries::voice_designer::voices::create_voice::create_voice; use mysql_queries::queries::voice_designer::voices::create_voice::CreateVoiceArgs; -use tokens::tokens::media_files::MediaFileToken; use tokens::tokens::users::UserToken; use crate::job; @@ -33,6 +33,9 @@ const BUCKET_FILE_PREFIX_CREATE: &str = "fakeyou_"; const BUCKET_FILE_EXTENSION_CREATE: &str = ".bin"; const MIME_TYPE_CREATE: &str = "application/x-binary"; +/// We'll increment this if we change the underlying Vall-E code in an incompatible way. +const CURRENT_VALLE_MODEL_VERSION: u64 = 0; + pub async fn process_create_voice( args: VALLEXProcessJobArgs<'_>, dataset_token: String @@ -244,19 +247,22 @@ pub async fn process_create_voice( return Err(error); } - info!("Success; waiting..."); - thread::sleep(Duration::from_secs(300)); + info!("Inference success!"); + + //info!("Success; waiting..."); + //thread::sleep(Duration::from_secs(300)); // STEP 4. Download dataset each audio file info!("Uploading Media ..."); - let result_bucket_location: MediaFileBucketPath = MediaFileBucketPath::generate_new( - Some(BUCKET_FILE_PREFIX_CREATE), - Some(BUCKET_FILE_EXTENSION_CREATE) + let embedding_bucket_location = ZeroShotVoiceEmbeddingBucketPath::generate_new( + ModelCategory::Tts, + ModelType::VallEx, + CURRENT_VALLE_MODEL_VERSION ); - let result_bucket_object_pathbuf = result_bucket_location.to_full_object_pathbuf(); + let embedding_bucket_object_pathbuf = embedding_bucket_location.to_full_object_pathbuf(); // Get Finished File let mut finished_file = work_temp_dir.path().to_path_buf(); @@ -264,12 +270,12 @@ pub async fn process_create_voice( let output_bucket_file_name = String::from("temp.npz"); // use extension for bucket upload. finished_file.push(&output_bucket_file_name); - info!("Upload Bucket Path: {:?}", result_bucket_object_pathbuf); info!("Upload File Path: {:?}", finished_file); + info!("Upload Bucket Path: {:?}", embedding_bucket_object_pathbuf); args.job_dependencies.private_bucket_client .upload_filename_with_content_type( - &result_bucket_object_pathbuf, + &embedding_bucket_object_pathbuf, &finished_file, &MIME_TYPE_CREATE ) @@ -278,7 +284,6 @@ pub async fn process_create_voice( // CLEARIFY! these items // 1.Should this be object hash? - let bucket_hash = result_bucket_location.get_object_hash().clone(); // 2.As well as this what should the voice name be? let voice_name = single_dataset.title; @@ -287,30 +292,24 @@ pub async fn process_create_voice( dataset_token: &voice_dataset_token, model_category: ZsVoiceModelCategory::Tts, model_type: ZsVoiceModelType::VallEX, - model_version: 0, + model_version: CURRENT_VALLE_MODEL_VERSION, model_encoding_type: ZsVoiceEncodingType::Encodec, voice_title: &voice_name, - bucket_hash, + bucket_hash: embedding_bucket_location.get_object_hash().clone(), maybe_creator_user_token: Some(&creator_user_token), creator_ip_address: &creator_ip_address, creator_set_visibility: Visibility::Public, mysql_pool, - }).await; + }).await + .map_err(|e| ProcessSingleJobError::Other(e))?; - let media_file_token = MediaFileToken::generate(); - - match voice_token { - Ok(_value) => { - Ok(JobSuccessResult { - maybe_result_entity: Some(ResultEntity { - entity_type: InferenceResultType::MediaFile, - entity_token: media_file_token.to_string(), - }), - inference_duration, - }) - } - Err(e) => { Err(ProcessSingleJobError::Other(e)) } - } + Ok(JobSuccessResult { + maybe_result_entity: Some(ResultEntity { + entity_type: InferenceResultType::ZeroShotVoiceEmbedding, + entity_token: voice_token.to_string(), + }), + inference_duration, + }) } pub struct AudioFile {