From cd8ca0092ba2540f5a88cc9677128c438ef8bcb7 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 16 Nov 2023 20:07:14 -0500 Subject: [PATCH] modify bucket pathing --- crates/_types/buckets/src/public/mod.rs | 3 +- .../buckets/src/public/zs_voices/directory.rs | 93 +++++++++++ .../buckets/src/public/zs_voices/file.rs | 154 ++++++++++++++++++ .../buckets/src/public/zs_voices/mod.rs | 2 + .../job/job_types/tts/vall_e_x/process_job.rs | 3 - .../tts/vall_e_x/process_job_create_voice.rs | 16 +- .../vall_e_x/vall_e_x_inference_command.rs | 2 +- 7 files changed, 262 insertions(+), 11 deletions(-) create mode 100644 crates/_types/buckets/src/public/zs_voices/directory.rs create mode 100644 crates/_types/buckets/src/public/zs_voices/file.rs create mode 100644 crates/_types/buckets/src/public/zs_voices/mod.rs diff --git a/crates/_types/buckets/src/public/mod.rs b/crates/_types/buckets/src/public/mod.rs index c0879da257..d85de56720 100644 --- a/crates/_types/buckets/src/public/mod.rs +++ b/crates/_types/buckets/src/public/mod.rs @@ -1,4 +1,5 @@ +pub mod media_files; pub mod media_uploads; pub mod public_path; pub mod voice_conversion_results; -pub mod media_files; \ No newline at end of file +pub mod zs_voices; \ No newline at end of file diff --git a/crates/_types/buckets/src/public/zs_voices/directory.rs b/crates/_types/buckets/src/public/zs_voices/directory.rs new file mode 100644 index 0000000000..a4a5300d80 --- /dev/null +++ b/crates/_types/buckets/src/public/zs_voices/directory.rs @@ -0,0 +1,93 @@ +use crockford::crockford_entropy_lower; +use crate::private::private_path::PrivatePath; + +use crate::util::hashed_directory_path_long_string::hashed_directory_path_long_string; + +// TODO: Generate these from a macro. + +// TODO: Use a central path registry for quick reference +const DIRECTORY: &str = "/user/zero_shot_embeddings"; + +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +pub enum ModelCategory { + Tts, + Vc, +} + +impl ModelCategory { + pub fn to_str(&self) -> &'static str { + match self { + ModelCategory::Tts => "tts", + ModelCategory::Vc => "vc", + } + } +} + +/// 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, + object_hash: String, + directory: String, +} + +impl PrivatePath for ZeroShotVoiceEmbeddingBucketDirectory {} + +impl ZeroShotVoiceEmbeddingBucketDirectory { + + pub fn generate_new(model_category: ModelCategory) -> Self { + let entropy = crockford_entropy_lower(32); + Self::from_object_hash(model_category, &entropy) + } + + pub fn from_object_hash(model_category: ModelCategory, object_hash: &str) -> Self { + // TODO: Path construction could be cleaner. + let subdir = model_category.to_str(); + let middle = hashed_directory_path_long_string(object_hash); + let directory = format!("{}/{}/{}{}", DIRECTORY, subdir, middle, object_hash); + Self { + model_category, + object_hash: object_hash.to_string(), + directory, + } + } + + pub fn get_directory_path_str(&self) -> &str { + &self.directory + } + + pub fn get_object_hash(&self) -> &str { + &self.object_hash + } +} + +#[cfg(test)] +mod tests { + use crate::public::zs_voices::directory::{ModelCategory, ZeroShotVoiceEmbeddingBucketDirectory}; + + #[test] + pub fn generate_new_entropy() { + let directory = ZeroShotVoiceEmbeddingBucketDirectory::generate_new(ModelCategory::Tts); + 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"); + } + + #[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"); + } + + #[test] + pub fn get_object_hash() { + let hash = "abcdefghijk"; + let directory = ZeroShotVoiceEmbeddingBucketDirectory::from_object_hash(ModelCategory::Tts, 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 new file mode 100644 index 0000000000..7bb6aeb9b2 --- /dev/null +++ b/crates/_types/buckets/src/public/zs_voices/file.rs @@ -0,0 +1,154 @@ +use std::path::PathBuf; + +use crockford::crockford_entropy_lower; + +use crate::private::private_path::PrivatePath; +use crate::public::zs_voices::directory::{ModelCategory, 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)] +pub struct ZeroShotVoiceEmbeddingBucketPath { + directory: ZeroShotVoiceEmbeddingBucketDirectory, + + filename: String, + full_object_path: String, + + model_category: ModelCategory, + model_type: ModelType, + model_version: u32, +} + +impl PrivatePath for ZeroShotVoiceEmbeddingBucketPath {} + +impl ZeroShotVoiceEmbeddingBucketPath { + + pub fn generate_new(model_category: ModelCategory, model_type: ModelType, model_version: u32) -> 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 { + // TODO: Path construction could be cleaner. + let directory = ZeroShotVoiceEmbeddingBucketDirectory::from_object_hash(model_category, hash); + + let filename = format!("{}_{}_{}", model_type.to_str(), hash, model_version); + + let full_object_path = format!("{}/{}", directory.get_directory_path_str(), filename); + + Self { + directory, + filename, + full_object_path, + model_category, + model_type, + model_version, + } + } + + pub fn get_full_object_path_str(&self) -> &str { + &self.full_object_path + } + + pub fn to_full_object_pathbuf(&self) -> PathBuf { + PathBuf::from(&self.full_object_path) + } + + pub fn get_directory(&self) -> &ZeroShotVoiceEmbeddingBucketDirectory { + &self.directory + } + + pub fn get_object_hash(&self) -> &str { + &self.directory.get_object_hash() + } + + pub fn get_basename(&self) -> &str { + &self.filename + } +} + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + + use crate::public::media_files::original_file::MediaFileBucketPath; + + mod with_prefix_and_extension { + use crate::public::zs_voices::directory::ModelCategory; + use crate::public::zs_voices::file::{ModelType, ZeroShotVoiceEmbeddingBucketPath}; + + #[test] + pub fn generate_new_entropy() { + let file = ZeroShotVoiceEmbeddingBucketPath::generate_new(ModelCategory::Tts, ModelType::VallEx, 0); + assert_eq!(file.get_object_hash().len(), 32); + assert_eq!(file.get_object_hash().len(), 32); + assert_eq!(file.get_directory().get_object_hash().len(), 32); + } + + #[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_basename(), "vallex_abcdefghijk_0"); + } + } + + mod without_prefix_and_extension { + use crate::public::zs_voices::directory::ModelCategory; + use crate::public::zs_voices::file::{ModelType, ZeroShotVoiceEmbeddingBucketPath}; + use super::*; + + #[test] + pub fn generate_new_entropy() { + let file = ZeroShotVoiceEmbeddingBucketPath::generate_new(ModelCategory::Vc, ModelType::VallEx, 123); + assert_eq!(file.get_object_hash().len(), 32); + assert_eq!(file.get_directory().get_object_hash().len(), 32); + } + + #[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_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_basename(), "vallex_abcdefghijk_0"); + } + + #[test] + pub fn get_full_object_path_str_starts_with_directory() { + let file = ZeroShotVoiceEmbeddingBucketPath::generate_new(ModelCategory::Vc, ModelType::VallEx, 123); + let full_path = file.get_full_object_path_str(); + assert!(full_path.starts_with(file.get_directory().get_directory_path_str())) + } + + #[test] + pub fn get_object_hash() { + let hash = "abcdefghijk"; + let file = MediaFileBucketPath::from_object_hash(hash, None, None); + assert_eq!(file.get_object_hash(), hash); + } + } +} diff --git a/crates/_types/buckets/src/public/zs_voices/mod.rs b/crates/_types/buckets/src/public/zs_voices/mod.rs new file mode 100644 index 0000000000..a5c0d4bfab --- /dev/null +++ b/crates/_types/buckets/src/public/zs_voices/mod.rs @@ -0,0 +1,2 @@ +pub mod directory; +pub mod file; \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs index 9f9dfb2481..0dab146192 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs @@ -17,8 +17,6 @@ pub struct VALLEXProcessJobArgs<'a> { pub job: &'a AvailableInferenceJob, } - - // query using the token then grab the bucket hash pub async fn process_job( args: VALLEXProcessJobArgs<'_> @@ -45,4 +43,3 @@ pub async fn process_job( } } } - 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 0a95cfc206..bc85db7971 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,6 +1,7 @@ use std::fs::read_to_string; use std::path::PathBuf; -use std::time::Instant; +use std::thread; +use std::time::{Duration, Instant}; use anyhow::anyhow; use log::{error, info, warn}; @@ -186,24 +187,25 @@ pub async fn process_create_voice( .log_status("running inference") .map_err(|e| ProcessSingleJobError::Other(e))?; - let inference_start_time = Instant::now(); - // Command line arg for a list of paths to insert the container let audio_files = join_paths(downloaded_dataset); info!("Files to process: {:?}", audio_files); // Name of the output file - let output_file_name = String::from("temp"); // don't use the extension... for the inference since the container will add the extension. + // NB: don't use the extension... for the inference since the container will add the extension. + let output_file_name = PathBuf::from("temp"); let stderr_output_file = work_temp_dir.path().join("zero_shot_create_voice_err.txt"); + let inference_start_time = Instant::now(); + // Run Inference let command_exit_status = args.job_dependencies.job_type_details.vall_e_x.create_embedding_command.execute_inference( job::job_types::tts::vall_e_x::vall_e_x_inference_command::CreateVoiceInferenceArgs { output_embedding_path: &workdir, - output_embedding_name: output_file_name.clone(), + output_embedding_name: &output_file_name, audio_files, stderr_output_file: &stderr_output_file, } @@ -242,6 +244,9 @@ pub async fn process_create_voice( return Err(error); } + info!("Success; waiting..."); + thread::sleep(Duration::from_secs(300)); + // STEP 4. Download dataset each audio file info!("Uploading Media ..."); @@ -255,7 +260,6 @@ pub async fn process_create_voice( // Get Finished File let mut finished_file = work_temp_dir.path().to_path_buf(); - //let mut finished_file = workdir; let output_bucket_file_name = String::from("temp.npz"); // use extension for bucket upload. finished_file.push(&output_bucket_file_name); diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs index 6e295d522d..7d7aa8aa4c 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs @@ -48,7 +48,7 @@ pub struct InferenceArgs> { pub struct CreateVoiceInferenceArgs> { pub output_embedding_path: P, - pub output_embedding_name: String, + pub output_embedding_name: P, pub audio_files: String, pub stderr_output_file: P, }