From 8d36d6c7099c3ab257701820b5ff8b985ee7fb11 Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Thu, 5 Oct 2023 21:43:36 -0400
Subject: [PATCH] Add some downloaders for vall-e-x

---
 .idea/dictionaries/bt.xml                     |  2 +
 .../inference_job/src/job/job_types/vc/mod.rs |  3 +-
 .../src/job/job_types/vc/vall_e_x_temp/mod.rs |  1 +
 .../vc/vall_e_x_temp/model_downloaders.rs     | 62 +++++++++++++++++++
 4 files changed, 67 insertions(+), 1 deletion(-)
 create mode 100644 crates/service/job/inference_job/src/job/job_types/vc/vall_e_x_temp/mod.rs
 create mode 100644 crates/service/job/inference_job/src/job/job_types/vc/vall_e_x_temp/model_downloaders.rs

diff --git a/.idea/dictionaries/bt.xml b/.idea/dictionaries/bt.xml
index a9bebf1881..96e5c2a274 100644
--- a/.idea/dictionaries/bt.xml
+++ b/.idea/dictionaries/bt.xml
@@ -12,7 +12,9 @@
       <w>dalle</w>
       <w>demuxer</w>
       <w>dereferencing</w>
+      <w>downloaders</w>
       <w>downvote</w>
+      <w>encodec</w>
       <w>fakeyou</w>
       <w>ffprobe</w>
       <w>greenlight</w>
diff --git a/crates/service/job/inference_job/src/job/job_types/vc/mod.rs b/crates/service/job/inference_job/src/job/job_types/vc/mod.rs
index 041ecf69a1..181429cc16 100644
--- a/crates/service/job/inference_job/src/job/job_types/vc/mod.rs
+++ b/crates/service/job/inference_job/src/job/job_types/vc/mod.rs
@@ -1,4 +1,5 @@
 pub mod process_single_vc_job;
 pub mod rvc_v2;
 pub mod so_vits_svc;
-pub mod softvc;
\ No newline at end of file
+pub mod softvc;
+pub mod vall_e_x_temp;
diff --git a/crates/service/job/inference_job/src/job/job_types/vc/vall_e_x_temp/mod.rs b/crates/service/job/inference_job/src/job/job_types/vc/vall_e_x_temp/mod.rs
new file mode 100644
index 0000000000..d1186b0820
--- /dev/null
+++ b/crates/service/job/inference_job/src/job/job_types/vc/vall_e_x_temp/mod.rs
@@ -0,0 +1 @@
+pub mod model_downloaders;
\ No newline at end of file
diff --git a/crates/service/job/inference_job/src/job/job_types/vc/vall_e_x_temp/model_downloaders.rs b/crates/service/job/inference_job/src/job/job_types/vc/vall_e_x_temp/model_downloaders.rs
new file mode 100644
index 0000000000..e04841fc72
--- /dev/null
+++ b/crates/service/job/inference_job/src/job/job_types/vc/vall_e_x_temp/model_downloaders.rs
@@ -0,0 +1,62 @@
+use crate::util::model_downloader::ModelDownloader;
+
+// Core models
+
+crate::impl_model_downloader!(
+  VallExEncodec,
+  "vall-e-x encodec model",
+  "VALL_E_X_CLOUD_PATH_ENCODEC",
+  "/dependencies/zero_shot_tts/vall-e-x_1.0/encodec_pytorch_model.bin",
+  "VALL_E_X_FS_PATH_ENCODEC",
+  "/tmp/downloads/vall_e_x/encodec_pytorch_model.bin"
+);
+
+crate::impl_model_downloader!(
+  VallExModel,
+  "vall-e-x model",
+  "VALL_E_X_CLOUD_PATH_MODEL",
+  "/dependencies/zero_shot_tts/vall-e-x_1.0/vallex-checkpoint.pt",
+  "VALL_E_X_FS_PATH_MODEL",
+  "/tmp/downloads/vall_e_x/vallex-checkpoint.pt"
+);
+
+// Whisper (Medium)
+
+crate::impl_model_downloader!(
+  VallExWhisperMediumFlax,
+  "vall-e-x whisper medium flax_model.msgpack",
+  "VALL_E_X_CLOUD_PATH_WM_FLAX",
+  "/dependencies/zero_shot_tts/vall-e-x_1.0/whisper-medium/flax_model.msgpack ",
+  "VALL_E_X_FS_PATH_WM_FLAX",
+  "/tmp/downloads/vall_e_x/whisper-medium/flax_model.msgpack "
+);
+
+// TODO (remaining files)
+
+// Whisper (Large)
+
+// TODO (remaining files)
+
+pub struct VallExDownloaders {
+  pub encodec: VallExEncodec,
+  pub model: VallExModel,
+  pub whisper_medium_flax: VallExWhisperMediumFlax,
+}
+
+impl VallExDownloaders {
+  pub fn build_all_from_env() -> Self {
+    Self {
+      encodec: VallExEncodec::from_env(),
+      model: VallExModel::from_env(),
+      whisper_medium_flax: VallExWhisperMediumFlax::from_env(),
+    }
+  }
+
+  pub fn all_downloaders(&self) -> Vec<&dyn ModelDownloader> {
+    vec![
+      &self.encodec,
+      &self.model,
+      &self.whisper_medium_flax,
+    ]
+  }
+}