modify bucket pathing

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;
+pub mod zs_voices;

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);
+  }
+}

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);
+    }
+  }
+}

crates/_types/buckets/src/public/zs_voices/mod.rs

@@ -0,0 +1,2 @@
+pub mod directory;
+pub mod file;

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

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);

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<P: AsRef<Path>> {
 
 pub struct CreateVoiceInferenceArgs<P: AsRef<Path>> {
   pub output_embedding_path: P,
-  pub output_embedding_name: String,
+  pub output_embedding_name: P,
   pub audio_files: String,
   pub stderr_output_file: P,
 }