support for multi-speaker glow-tts models (broken in tensor dims)

deployment.yaml

@@ -86,7 +86,7 @@ spec:
         # LSpeech model that sounds 100% amazing
         # image: docker.pkg.github.com/echelon/voder/tts-service:6b5c22567614
         #
-        image: docker.pkg.github.com/ml-applications/voder/tts-service:0281e206db16
+        image: docker.pkg.github.com/ml-applications/voder/tts-service:42f51ecd18e2
         resources:
           limits:
             # 4 GiB (4096 MiB)

tts_service/models.toml

@@ -25,6 +25,14 @@ model_pipeline = "ArpabetTacotronMelgan"
 tacotron = "tacotron2_ljs_high_pitch_arpabet_2020-03-28.jit"
 melgan = "melgan_ljs-high-pitch_container.pt"
 
+[[speakers]]
+name = "Arnold Schwarzenegger"
+slug = "arnold"
+model_pipeline = "ArpabetGlowTtsMultiSpeakerMelgan"
+glow_tts_multi_speaker = "glow_multispeaker.torchjit"
+melgan = "melgan_trump_txlearn_2020-05-30_steno_chkpt9800.torchjit"
+speaker_id = 2
+
 # --------------- Model configs ---------------
 
 [[model_locations]]
@@ -40,6 +48,10 @@ base_directory = "/home/bt/dev/jit-models/glow-tts-jaywalnut"
 model_type = "Melgan"
 base_directory = "/home/bt/dev/tacotron-melgan/"
 
+[[model_locations]]
+model_type = "ArpabetGlowTtsMultiSpeaker"
+base_directory = "/home/bt/dev/jit-models/glow-tts-jaywalnut-multispeaker"
+
 # --------------- Tacotron ---------------
 
 # TODO: Get a high-quality LJS Tacotron
@@ -160,6 +172,12 @@ description = "First experiments (continued); modified with longer input to JIT"
 file_path = "glow_tts_ljs_txlearn_arpabet_2020-06-01_G_3614.torchjit"
 description = "First experiments (continued); modified with longer input to JIT"
 
+
+# ---------------- Glow TTS Multi-speaker ----------------
+[[glow_tts_multi_speaker]]
+file_path = "glow_multispeaker.torchjit"
+description = "Mutli-speaker garbage"
+
 # ---------------- Melgan ----------------
 
 # TODO: Get a high-quality LJS Melgan

tts_service/src/endpoints/legacy_speak.rs

@@ -217,6 +217,7 @@ pub async fn legacy_get_speak(request: HttpRequest,
           .set_header(header::CONTENT_RANGE, content_range_value)
           .body(wav_data))
     },
+    ModelPipeline::ArpabetGlowTtsMultiSpeakerMelgan=> unimplemented!(),
     ModelPipeline::RawTextTacotronMelgan => unimplemented!(),
   }
 }

tts_service/src/endpoints/speak.rs

@@ -15,7 +15,7 @@ use crate::text::arpabet::text_to_arpabet_encoding;
 use crate::model::model_config::ModelPipeline;
 use crate::model::old_model::TacoMelModel;
 use crate::text::cleaners::clean_text;
-use crate::model::pipelines::arpabet_glow_tts_melgan_pipeline;
+use crate::model::pipelines::{arpabet_glow_tts_melgan_pipeline, arpabet_glow_tts_multi_speaker_melgan_pipeline};
 use crate::database::model::NewSentence;
 
 #[derive(Deserialize)]
@@ -148,6 +148,45 @@ pub async fn post_speak(request: HttpRequest,
           .set_header(header::CONTENT_RANGE, content_range_value)
           .body(wav_data))
     },
+    ModelPipeline::ArpabetGlowTtsMultiSpeakerMelgan=> {
+      let glow_tts_multi_speaker_model = speaker.glow_tts_multi_speaker
+          .as_ref()
+          .map(|s| s.clone())
+          .expect("TODO ERROR HANDLING");
+
+      let melgan_model = speaker.melgan
+          .as_ref()
+          .map(|s| s.clone())
+          .expect("TODO ERROR HANDLING");
+
+      info!("Glow-TTS Multi-Speaker Model: {}", glow_tts_multi_speaker_model);
+      info!("Melgan Model: {}", melgan_model);
+      info!("Text: {}", text);
+
+      let speaker_id = speaker.speaker_id.expect("Should have speaker_id");
+
+      let glow_tts_multi_speaker = app_state
+          .model_cache
+          .get_or_load_arbabet_glow_tts_multi_speaker(&glow_tts_multi_speaker_model)
+          .expect(&format!("Couldn't load glow-tts multi-speaker model: {}", glow_tts_multi_speaker_model));
+
+      let melgan = app_state
+          .model_cache
+          .get_or_load_melgan(&melgan_model)
+          .expect(&format!("Couldn't load melgan model: {}", melgan_model));
+
+      let wav_data = arpabet_glow_tts_multi_speaker_melgan_pipeline(
+        &cleaned_text, speaker_id, &glow_tts_multi_speaker, &melgan);
+
+      // To make iOS Safari work, you need a Content-Range and Content-Length header:
+      // https://stackoverflow.com/a/17835399
+      let content_range_value = format!("bytes 0-{}/{}", wav_data.len(), wav_data.len());
+      Ok(HttpResponse::build(StatusCode::OK)
+          .content_type("audio/wav")
+          .set_header(header::CONTENT_DISPOSITION, "attachment; filename = \"generated.wav\"")
+          .set_header(header::CONTENT_RANGE, content_range_value)
+          .body(wav_data))
+    },
     ModelPipeline::RawTextTacotronMelgan => unimplemented!(),
   }
 }

tts_service/src/model/arpabet_glow_tts_multi_speaker_model.rs

@@ -0,0 +1,33 @@
+use crate::model::model_container::ModelContainer;
+use anyhow::{Result, Error};
+use std::path::Path;
+use std::thread;
+use std::time::Duration;
+use tch::Tensor;
+use crate::text::arpabet::text_to_arpabet_encoding_glow_tts;
+
+pub struct ArpabetGlowTtsMultiSpeakerModel {
+  model_container: ModelContainer,
+}
+
+impl ArpabetGlowTtsMultiSpeakerModel {
+  pub fn load(filename: &Path) -> Result<Self> {
+    let model_container = ModelContainer::load(filename)?;
+    Ok(Self {
+      model_container,
+    })
+  }
+
+  pub fn encoded_arpabet_to_mel(&self, arpabet_encodings: &Vec<i64>, speaker_id: i64) -> Tensor {
+    let text_tensor = Tensor::of_slice(&arpabet_encodings.as_slice());
+    let text_tensor_2d = text_tensor.unsqueeze(0); // Add another axis
+
+    let lengths = [arpabet_encodings.len() as i64];
+    let lengths_tensor = Tensor::of_slice(&lengths);
+
+    let speaker = [speaker_id];
+    let speaker_tensor = Tensor::of_slice(&speaker);
+
+    self.model_container.forward3(&text_tensor_2d, &lengths_tensor, &speaker_tensor)
+  }
+}

tts_service/src/model/mod.rs

@@ -1,4 +1,5 @@
 pub mod arpabet_glow_tts_model;
+pub mod arpabet_glow_tts_multi_speaker_model;
 pub mod arpabet_tacotron_model;
 pub mod melgan_model;
 pub mod model_cache;

tts_service/src/model/model_cache.rs

@@ -6,6 +6,7 @@ use std::path::PathBuf;
 use std::sync::Arc;
 use std::sync::Mutex;
 use crate::model::arpabet_glow_tts_model::ArpabetGlowTtsModel;
+use crate::model::arpabet_glow_tts_multi_speaker_model::ArpabetGlowTtsMultiSpeakerModel;
 
 pub struct ModelCache {
   /// Base directories for each model type(if configured)
@@ -13,6 +14,7 @@ pub struct ModelCache {
 
   arpabet_tacotron_models: Arc<Mutex<HashMap<String, Arc<ArpabetTacotronModel>>>>,
   arpabet_glow_tts_models: Arc<Mutex<HashMap<String, Arc<ArpabetGlowTtsModel>>>>,
+  arpabet_glow_tts_multi_speaker_models: Arc<Mutex<HashMap<String, Arc<ArpabetGlowTtsMultiSpeakerModel>>>>,
   melgan_models: Arc<Mutex<HashMap<String, Arc<MelganModel>>>>,
 }
 
@@ -31,6 +33,7 @@ impl ModelCache {
       base_directories,
       arpabet_tacotron_models: Arc::new(Mutex::new(HashMap::new())),
       arpabet_glow_tts_models: Arc::new(Mutex::new(HashMap::new())),
+      arpabet_glow_tts_multi_speaker_models: Arc::new(Mutex::new(HashMap::new())),
       melgan_models: Arc::new(Mutex::new(HashMap::new())),
     }
   }
@@ -95,6 +98,36 @@ impl ModelCache {
     }
   }
 
+  pub fn get_or_load_arbabet_glow_tts_multi_speaker(&self, filename: &str) -> Option<Arc<ArpabetGlowTtsMultiSpeakerModel>> {
+    let base_directory = self.get_base_directory(ModelType::ArpabetGlowTtsMultiSpeaker);
+
+    let mut file_path = PathBuf::from(filename);
+
+    if !file_path.is_absolute() {
+      file_path = base_directory.map(|base_dir| base_dir.join(&file_path))
+          .unwrap_or(file_path);
+    }
+
+    let mut lock = self.arpabet_glow_tts_multi_speaker_models.lock()
+        .expect("should unlock");
+
+    if let Some(model) = lock.get(filename) {
+      return Some(model.clone());
+    }
+
+    match ArpabetGlowTtsMultiSpeakerModel::load(&file_path) {
+      Ok(model) => {
+        let arc = Arc::new(model);
+        lock.insert(filename.to_string(), arc.clone());
+        Some(arc)
+      },
+      Err(e) => {
+        warn!("There was an error loading the model `{:?}`: {}", file_path, e);
+        None
+      },
+    }
+  }
+
   pub fn get_or_load_melgan(&self, filename: &str) -> Option<Arc<MelganModel>> {
     let base_directory = self.get_base_directory(ModelType::Melgan);
 

tts_service/src/model/model_config.rs

@@ -8,6 +8,7 @@ pub struct ModelConfigs {
   pub model_locations: Vec<ModelLocation>,
   tacotron: Vec<ModelDetails>,
   glow_tts: Vec<ModelDetails>,
+  glow_tts_multi_speaker: Vec<ModelDetails>,
   melgan: Vec<ModelDetails>,
 }
 
@@ -17,12 +18,16 @@ pub struct Speaker {
   pub name: String,
   /// URL slug / key
   pub slug: String,
+  /// Speaker Id, if used
+  pub speaker_id: Option<i64>,
   /// Model pipeline to use
   pub model_pipeline: ModelPipeline,
   /// Tacotron model, if used.
   pub tacotron: Option<String>,
   /// Glow-tts model, if used.
   pub glow_tts: Option<String>,
+  /// Glow-tts multi-speaker model, if used.
+  pub glow_tts_multi_speaker: Option<String>,
   /// Melgan model, if used.
   pub melgan: Option<String>,
 }
@@ -33,6 +38,7 @@ pub enum ModelType {
   RawTextTacotron,
   ArpabetTacotron,
   ArpabetGlowTts, // NB: This uses a different arpabet preprocessor
+  ArpabetGlowTtsMultiSpeaker, // NB: This uses a different arpabet preprocessor
   Melgan,
 }
 
@@ -51,6 +57,7 @@ pub enum ModelPipeline {
   ArpabetTacotronMelgan,
   RawTextTacotronMelgan,
   ArpabetGlowTtsMelgan,
+  ArpabetGlowTtsMultiSpeakerMelgan,
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]

tts_service/src/model/model_container.rs

@@ -40,4 +40,14 @@ impl ModelContainer {
 
     result.expect("SHOULD WORK")
   }
+
+  pub fn forward3(&self, arg1: &Tensor, arg2: &Tensor, arg3: &Tensor) -> Tensor {
+    let result = self.jit_model.forward_ts(&[arg1, arg2, arg3]);
+
+    if let Err(err) = result.as_ref() {
+      println!("error: {:?}", err);
+    }
+
+    result.expect("SHOULD WORK")
+  }
 }

tts_service/src/model/pipelines.rs

@@ -9,11 +9,29 @@ use hound::WavWriter;
 use std::io::{Cursor, BufWriter};
 use tch::Tensor;
 use crate::model::arpabet_glow_tts_model::ArpabetGlowTtsModel;
+use crate::model::arpabet_glow_tts_multi_speaker_model::ArpabetGlowTtsMultiSpeakerModel;
 
 // TODO: This might be useful to implement as a multi-stage
 //  state machine struct with functions, that way you can pull out
 //  intermediate pieces (eg. text to send to the db, or reject the request)
 
+pub fn arpabet_glow_tts_multi_speaker_melgan_pipeline(
+  cleaned_text: &str,
+  speaker_id: i64,
+  arpabet_glow_tts: &ArpabetGlowTtsMultiSpeakerModel,
+  melgan: &MelganModel) -> Vec<u8> {
+
+  let arpabet = Arpabet::load_cmudict(); // TODO: Inefficient.
+  let arpabet_encodings = text_to_arpabet_encoding_glow_tts(arpabet, &cleaned_text);
+
+  let mel_tensor = arpabet_glow_tts.encoded_arpabet_to_mel(&arpabet_encodings, speaker_id);
+
+  let audio_tensor = melgan.tacotron_mel_to_audio(&mel_tensor);
+  let audio_signal = mel_audio_tensor_to_audio_signal(audio_tensor);
+
+  audio_signal_to_wav_bytes(audio_signal)
+}
+
 pub fn arpabet_glow_tts_melgan_pipeline(
   cleaned_text: &str,
   arpabet_glow_tts: &ArpabetGlowTtsModel,