Testing inference_job

crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/model_downloaders.rs

@@ -5,6 +5,7 @@ pub struct VallEXDownloaders {
     pub whisper_lg: WhisperLargeModelMapping,
     pub whisper_md: WhisperMediumModelMapping,
     pub vocos_encodec_24khz: VocosEncodecModelMapping,
+    pub vocos_config: VocosEncodecConfigMapping
 }
 
 crate::impl_model_downloader!(
@@ -29,18 +30,27 @@ crate::impl_model_downloader!(
     WhisperMediumModelMapping,
     "whisper Medium Model Mapping",
     "WHISPER_MEDIUM_TENSOR_BUCKET_PATH",
-    "/dependencies/zero_shot_tts/vall-e-x_1.0/whisper-medium/model.safetensors",
+    "/dependencies/zero_shot_tts/vall-e-x_1.0/whisper/medium.pt",
     "WHISPER_MEDIUM_TENSOR_FILESYSTEM_PATH",
-    "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/whisper-medium/model.safetensors"
+    "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/whisper/medium.pt"
 );
 
 crate::impl_model_downloader!(
     VocosEncodecModelMapping,
     "Vocos Encodec Model Mapping",
     "VOCOS_ENCODEC_TENSOR_BUCKET_PATH",
-    "/dependencies/zero_shot_tts/vall-e-x_1.0/encodec_pytorch_model.bin",
+    "/dependencies/zero_shot_tts/vall-e-x_1.0/vocos-encodec/encodec_pytorch_model.bin",
     "VOCOS_ENCODEC_TENSOR_FILESYSTEM_PATH",
-    "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/encodec_pytorch_model.bin"
+    "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vocos-encodec/encodec_pytorch_model.bin"
+);
+
+crate::impl_model_downloader!(
+    VocosEncodecConfigMapping,
+    "Vocos Config Mapping",
+    "VOCOS_ENCODEC_CONFIG_BUCKET_PATH",
+    "/dependencies/zero_shot_tts/vall-e-x_1.0/vocos-encodec/config.yaml",
+    "VOCOS_ENCODEC_CONFIG_FILESYSTEM_PATH",
+    "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vocos-encodec/config.yaml"
 );
 
 impl VallEXDownloaders {
@@ -50,11 +60,12 @@ impl VallEXDownloaders {
             whisper_lg: WhisperLargeModelMapping::from_env(),
             whisper_md: WhisperMediumModelMapping::from_env(),
             vocos_encodec_24khz: VocosEncodecModelMapping::from_env(),
+            vocos_config: VocosEncodecConfigMapping::from_env(),
         }
     }
 
     pub fn all_downloaders(&self) -> Vec<&dyn ModelDownloader> {
-        vec![&self.vall_e_x, &self.whisper_lg, &self.whisper_md, &self.vocos_encodec_24khz]
+        vec![&self.vall_e_x, &self.whisper_lg, &self.whisper_md, &self.vocos_encodec_24khz,&self.vocos_config]
     }
 }
 

crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs

@@ -21,17 +21,17 @@ pub struct VoiceFile {
 pub async fn download_voice_embedding_from_hash(
     bucket_hash: &str,
     name: &str,
-    private_bucket_client: BucketClient
+    private_bucket_client: &BucketClient,
+    path: &str
 ) -> Result<VoiceFile, ProcessSingleJobError> {
-
     let unifer = BucketPathUnifier::default_paths();
     let object_path = unifer.zero_shot_tts_speaker_encoding(bucket_hash, 0);
-    let filesystem_path = format!("{}_weights.npz", name);
-    
+    let filesystem_path = format!("{}/{}_weights.npz", path, name);
+
     let result = private_bucket_client.download_file_to_disk(object_path, &filesystem_path).await;
 
     let voice_file = VoiceFile {
-        filesystem_path: PathBuf::from(&filesystem_path), 
+        filesystem_path: PathBuf::from(&filesystem_path),
     };
 
     Ok(voice_file)
@@ -50,22 +50,24 @@ pub async fn process_job(
     let job = args.job;
     let deps = args.job_dependencies;
     let mysql_pool = &deps.mysql_pool;
-    
+
     // get some globals
-    let job_progress_reporter = deps.job_progress_reporter
+    let mut job_progress_reporter = deps.job_progress_reporter
         .new_generic_inference(job.inference_job_token.as_str())
         .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?;
 
     // get job args
     let text = match job.maybe_raw_inference_text.clone() {
-        Some(value) => { value },
-        None => { return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Text for Inference"))) }
+        Some(value) => { value }
+        None => {
+            return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Text for Inference")));
+        }
     };
 
     // get args token
     let jobArgs = validate_job(&job)?; // bubbles error up
     let voice_token = tokens::tokens::zs_voices::ZsVoiceToken(jobArgs.voice_token);
-    
+
     // Get voice bucket hash - from voice token
     let voice_lookup_result = get_voice_by_token(&voice_token, false, &mysql_pool).await;
 
@@ -73,7 +75,7 @@ pub async fn process_job(
         Ok(Some(voice)) => voice,
         Ok(None) => {
             warn!("Voice not found: {:?}", voice_token);
-            return Err(ProcessSingleJobError::Other(anyhow!("Voice not found: {:?}",voice_token)));
+            return Err(ProcessSingleJobError::Other(anyhow!("Voice not found: {:?}", voice_token)));
         }
         Err(err) => {
             warn!("Error looking up voice: {:?}", err);
@@ -81,10 +83,6 @@ pub async fn process_job(
         }
     };
 
-    //download_voice_embedding_from_hash( bucket_hash, name, private_bucket_client);
-
-
-
     // Need to download the models
     info!("Download models (if not present)...");
 
@@ -100,57 +98,61 @@ pub async fn process_job(
         }
     }
 
-    // Download embeddings file using embedding token
-    // Create a temp dir to download things to
-
-    // Run Inference
-    let command_exit_status =
-        args.job_dependencies.job_type_details.vall_e_x.inference_command.execute_inference(
-            InferenceArgs {
-                input_embedding: String::from(""), // name of the embedding.npz in the tmp dir
-                input_text: String::from(text), // name of the text
-                /// --result_file: path to final file output
-                output_file_name: String::from("output.wav"), // output file name in the output folder
-                stderr_output_file: String::from(""),
-            }
-        );
-
-    // upload audio to bucket
-
-    // deletetemp dir files
-
-    // save result ?
-
-    // ==================== UPLOAD AUDIO TO BUCKET ==================== //
-    // Download embeddings file using embedding token
-
     // run inference
     let work_temp_dir = format!("temp_zeroshot_inference_{}", job.id.0);
+    let work_temp_dir_string = work_temp_dir.clone();
     // NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem.
     let work_temp_dir = args.job_dependencies.fs.scoped_temp_dir_creator_for_work
         .new_tempdir(&work_temp_dir)
         .map_err(|e| ProcessSingleJobError::from_io_error(e))?;
 
-    // delete temp dir files
+    let voiceFile = download_voice_embedding_from_hash(
+        &voice.bucket_hash,
+        &voice.title,
+        &deps.private_bucket_client,
+        &work_temp_dir_string
+    ).await?;
+    // // Download embeddings file using embedding token
+    // // Create a temp dir to download things to
+    job_progress_reporter
+        .log_status("running inference")
+        .map_err(|e| ProcessSingleJobError::Other(e))?;
 
-    // ==================== UPLOAD AUDIO TO BUCKET ==================== //
-    let maybe_args = job.maybe_inference_args
-        .as_ref()
-        .map(|args| args.args.as_ref())
-        .flatten();
+    // Run Inference
+    let command_exit_status =
+        args.job_dependencies.job_type_details.vall_e_x.inference_command.execute_inference(
+            InferenceArgs {
+                input_embedding_path: format!("{}_weights.npz", &voice.title), // name of the embedding.npz in the tmp dir
+                input_text: String::from(text), // text
+                output_file_name: String::from("output.wav"), // output file name in the output folder
+                stderr_output_file: String::from("test.txt"),
+            }
+        );
+        
+    // // upload audio to bucket
+    // // deletetemp dir files
 
-    let workdir = work_temp_dir.path().to_path_buf();
-    let stderr_output_file = work_temp_dir.path().join("stderr.txt");
-    let inference_start_time = Instant::now();
+    // // ==================== UPLOAD AUDIO TO BUCKET ==================== //
+
+    // // delete temp dir files
+
+    // // ==================== UPLOAD AUDIO TO BUCKET ==================== //
+    // let maybe_args = job.maybe_inference_args
+    //     .as_ref()
+    //     .map(|args| args.args.as_ref())
+    //     .flatten();
+
+    // let workdir = work_temp_dir.path().to_path_buf();
+    // let stderr_output_file = work_temp_dir.path().join("stderr.txt");
+    // let inference_start_time = Instant::now();
 
     // upload audio to public bucket ( for voice )
     // return ok
     Err(ProcessSingleJobError::InvalidJob(anyhow!("this job flow is not yet implemented")))
 }
-
 mod tests {
     #[test]
     fn test_length_zero() {
         println!("Hello Test")
     }
-}
+}

crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs

@@ -43,7 +43,7 @@ static IGNORED_ENVIRONMENT_VARS : Lazy<HashSet<String>> = Lazy::new(|| {
 // so we have all the files downloaded for the inputs.
 pub struct InferenceArgs<P: AsRef<Path>> {
     /// --driven_audio: path to the input audio
-    pub input_embedding: P, // name of the embedding.npz in the tmp dir
+    pub input_embedding_path: P, // name of the embedding.npz in the work dir
     pub input_text: P, // name of the text
     /// --result_file: path to final file output
     pub output_file_name: P, // output file name in the output folder
@@ -97,7 +97,6 @@ impl VallEXInferenceCommand {
         let maybe_inference_command = easyenv::get_env_string_optional(
           "VALL_E_X_INFERENCE_COMMAND");
     
-
           // EVENTUALLY should remove the check
         // Optional, eg. `./infer.py`. Typically we'll use the command form instead.
         let maybe_inference_executable = easyenv::get_env_pathbuf_optional(
@@ -187,15 +186,30 @@ impl VallEXInferenceCommand {
     
         // ===== Begin Python Args =====
     
-        command.push_str(" -text ");
+        command.push_str(" --text ");
         command.push_str(&path_to_string(args.input_text));
     
-        command.push_str(" -use_prompt_name ");
-        command.push_str(&path_to_string(args.input_embedding));
+        command.push_str(" --prompt-name ");
+        command.push_str(&path_to_string(args.input_embedding_path));
     
-        command.push_str(" -output_file_name ");
+        command.push_str(" --audio-name ");
         command.push_str(&path_to_string(args.output_file_name));
 
+        command.push_str(" mode ");
+        command.push_str(&path_to_string("0"));
+        // TODO improve and use a better model for premium users
+        command.push_str(" --whisper-model ");
+        command.push_str(&path_to_string(Path::new("medium")));
+
+        command.push_str(" whisper-path ");
+        command.push_str(&path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/medium.pt")));
+
+        command.push_str(" --vocos-folder-path ");
+        command.push_str(&path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vocos-encodec/")));
+
+        command.push_str(" --vallex-path ");
+        command.push_str(&path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vallex-checkpoint.pt")));
+
         // ===== End Python Args =====
     
         if let Some(docker_options) = self.maybe_docker_options.as_ref() {

crates/service/web/storyteller_web/src/routes.rs

@@ -171,7 +171,7 @@ pub fn add_routes<T, B> (app: App<T>) -> App<T>
   app = add_subscription_routes(app); /* /v1/subscriptions/... */
 
   // TODO find a long term feature flag solution, since this code is likely deployed into production we don't want the route found.
-  let enable_voice_designer_route = false;
+  let enable_voice_designer_route = true;
   if enable_voice_designer_route {
     app = add_voice_designer_routes(app); /* /v1/voice_designer */
   }