upload of weights works

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

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

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

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 {