From 76775dec605fd9ea40e3d81d0226b60bc147290f Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Fri, 28 Apr 2023 15:54:53 -0700
Subject: [PATCH] hubert path config

---
 .../config/inference-job.development.env      |  3 +++
 .../so_vits_svc_inference_command.rs          | 26 ++++++++++++++++---
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/crates/service/job/inference_job/config/inference-job.development.env b/crates/service/job/inference_job/config/inference-job.development.env
index e70632e006..3fc1f22909 100644
--- a/crates/service/job/inference_job/config/inference-job.development.env
+++ b/crates/service/job/inference_job/config/inference-job.development.env
@@ -9,6 +9,9 @@ BUCKET_TIMEOUT_SECONDS=6000
 #                                   #
 # ================================= #
 
+# Cache hubert in a recoverable place so do don't constantly re-download it.
+HUBERT_PATH="/tmp/hubert/checkpoint_best_legacy_500.pt"
+
 SO_VITS_SVC_INFERENCE_ROOT_DIRECTORY="/models/voice_conversion/so-vits-svc/src"
 
 #SO_VITS_SVC_INFERENCE_MAYBE_DOCKER_IMAGE="abd0ded43f90" # Working as of 2023-04-28
diff --git a/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/so_vits_svc_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/so_vits_svc_inference_command.rs
index efa9792840..50af39f034 100644
--- a/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/so_vits_svc_inference_command.rs
+++ b/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/so_vits_svc_inference_command.rs
@@ -1,11 +1,11 @@
-use std::env;
-use std::ffi::OsString;
+use anyhow::anyhow;
 use container_common::anyhow_result::AnyhowResult;
 use filesys::path_to_string::path_to_string;
 use log::info;
+use std::env;
+use std::ffi::OsString;
 use std::fs::OpenOptions;
 use std::path::{Path, PathBuf};
-use anyhow::anyhow;
 use subprocess::{Popen, PopenConfig, Redirection};
 use subprocess_common::docker_options::{DockerEnvVar, DockerFilesystemMount, DockerGpu, DockerOptions};
 
@@ -28,6 +28,9 @@ pub struct SoVitsSvcInferenceCommand {
   /// If this is run under Docker (eg. in development), these are the options.
   maybe_docker_options: Option<DockerOptions>,
 
+  /// Where to cache hubert. We send this to the model via HUBERT_PATH env var.
+  maybe_hubert_path: Option<PathBuf>,
+
   maybe_huggingface_cache_dir: Option<PathBuf>,
   maybe_nltk_cache_dir: Option<PathBuf>,
 }
@@ -75,6 +78,7 @@ impl SoVitsSvcInferenceCommand {
     maybe_virtual_env_activation_command: Option<&str>,
     maybe_default_config_path: Option<P>,
     maybe_docker_options: Option<DockerOptions>,
+    maybe_hubert_path: Option<P>,
     maybe_huggingface_cache_dir: Option<P>,
     maybe_nltk_cache_dir: Option<P>,
   ) -> Self {
@@ -84,6 +88,7 @@ impl SoVitsSvcInferenceCommand {
       maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| s.to_string()),
       maybe_default_config_path: maybe_default_config_path.map(|p| p.as_ref().to_path_buf()),
       maybe_docker_options,
+      maybe_hubert_path: maybe_hubert_path.map(|s| s.as_ref().to_path_buf()),
       maybe_huggingface_cache_dir: maybe_huggingface_cache_dir.map(|s| s.as_ref().to_path_buf()),
       maybe_nltk_cache_dir: maybe_nltk_cache_dir.map(|s| s.as_ref().to_path_buf()),
     }
@@ -116,6 +121,9 @@ impl SoVitsSvcInferenceCommand {
     let maybe_default_config_path = easyenv::get_env_pathbuf_optional(
       "SO_VITS_SVC_INFERENCE_MAYBE_DEFAULT_CONFIG_PATH");
 
+    let maybe_hubert_path =
+        easyenv::get_env_pathbuf_optional("HUBERT_PATH");
+
     let maybe_huggingface_cache_dir =
         easyenv::get_env_pathbuf_optional("HF_DATASETS_CACHE");
 
@@ -156,6 +164,7 @@ impl SoVitsSvcInferenceCommand {
       maybe_virtual_env_activation_command,
       maybe_default_config_path,
       maybe_docker_options,
+      maybe_hubert_path,
       maybe_huggingface_cache_dir,
       maybe_nltk_cache_dir,
     })
@@ -264,6 +273,7 @@ impl SoVitsSvcInferenceCommand {
     }*/
 
     // Copy all environment variables from the parent process.
+    // This is necessary to send all the kubernetes settings for Nvidia / CUDA.
     for (env_key, env_value) in env::vars() {
       env_vars.push((
         OsString::from(env_key),
@@ -271,6 +281,16 @@ impl SoVitsSvcInferenceCommand {
       ));
     }
 
+    // In production / k8s, we should get this env var from the deployment and handle it
+    // more generally when copying over all environment variables, but in local development
+    // we're explicit since it may not be set outside of config files.
+    if let Some(hubert_path) = self.maybe_hubert_path.as_ref() {
+      env_vars.push((
+        OsString::from("HUBERT_PATH"),
+        OsString::from(hubert_path),
+      ));
+    }
+
     let mut config = PopenConfig::default();
 
     if !env_vars.is_empty() {