This works end-to-end

tacotron_melgan/src/main.rs

@@ -2,16 +2,54 @@ extern crate hound;
 extern crate rand;
 extern crate tch;
 
-use tch::vision::imagenet;
-use tch::vision::resnet;
-use tch::nn::ModuleT;
+use tch::CModule;
 use tch::Tensor;
+use tch::nn::Module;
+use tch::nn::ModuleT;
 
 pub mod melgan;
+pub mod model;
 pub mod tacotron;
 
 pub fn main() {
   println!("Tacotron2 + MelGan");
+
   //melgan::run_melgan();
-  tacotron::run_tacotron();
+  //tacotron::run_tacotron();
+  run_end_to_end();
+}
+
+fn run_end_to_end() {
+  println!("Load Tacotron Model");
+  let tacotron_model =
+      model::load_model("/home/bt/dev/voder/tacotron_melgan/tacotron_container.pt");
+
+  println!("Load Melgan Model");
+  let melgan_model =
+      model::load_model("/home/bt/dev/voder/tacotron_melgan/container2.pt");
+
+  println!("Load Text Sequence Tensor");
+  let text_sequence =
+      model::load_wrapped_tensor("/home/bt/dev/voder/data/text/tacotron_text_sequence.pt.containerized.pt");
+
+  println!("Evaluating Tacotron Model");
+  let mut mel_tensor = tacotron_model.forward(&text_sequence);
+
+  println!("Result mel tensor: {:?}", mel_tensor);
+
+  if mel_tensor.dim() == 2 {
+    println!("mel unsqeeze");
+    mel_tensor = mel_tensor.unsqueeze(0);
+  }
+
+  let audio_tensor = melgan_model.forward(&mel_tensor);
+
+  println!("Result audio tensor: {:?}", audio_tensor);
+
+  let audio = melgan::process_raw_mel(audio_tensor);
+
+  println!("Writing audio file...");
+  melgan::write_audio_file(audio, "double_model_results.wav");
+
+  println!("Done.");
 }

tacotron_melgan/src/melgan.rs

@@ -14,10 +14,10 @@ const WRAPPED_MODEL_PATH : &'static str = "/home/bt/dev/voder/tacotron_melgan/co
 const EXAMPLE_MEL_1: &'static str = "/home/bt/dev/voder/data/mels/LJ002-0320.mel.containerized.pt";
 const EXAMPLE_MEL_2 : &'static str = "/home/bt/dev/voder/data/mels/trump_2018_02_15-001.mel.containerized.pt";
 
-const MAX_WAV_VALUE : f32 = 32768.0f32;
+pub const MAX_WAV_VALUE : f32 = 32768.0f32;
 
 // TODO: This is an hparam and should be dynamic.
-const HOP_LENGTH : i64 = 256;
+pub const HOP_LENGTH : i64 = 256;
 
 pub fn load_melgan_model(filename: &str) -> CModule {
   println!("Loading model: {}", filename);
@@ -118,6 +118,22 @@ pub fn process_raw_mel(mel: Tensor) -> Vec<f32> {
   data
 }
 
+pub fn write_audio_file(audio_signal: Vec<f32>, filename: &str) {
+  let spec = hound::WavSpec {
+    channels: 1,
+    sample_rate: 16000,
+    bits_per_sample: 32,
+    sample_format: hound::SampleFormat::Float,
+  };
+
+  let mut writer = hound::WavWriter::create(filename, spec).unwrap();
+
+  for sample in audio_signal {
+    let sample = MAX_WAV_VALUE * sample;
+    writer.write_sample(sample).unwrap();
+  }
+}
+
 pub fn run_melgan() {
   let mut mel = load_wrapped_mel(EXAMPLE_MEL_1);
   println!("Got mel: {:?} of dim {}", mel, mel.dim());

tacotron_melgan/src/model.rs

@@ -0,0 +1,30 @@
+use tch::Tensor;
+use tch::CModule;
+use tch;
+use tch::nn::Module;
+
+/// Load a Torch model
+pub fn load_model(filename: &str) -> CModule {
+  println!("Loading model: {}", filename);
+  tch::CModule::load(filename).unwrap()
+}
+
+/// This is a _HACK_ to load a single Tensor.
+///
+/// It seems currently impossible to load tensors into tch.rs (libtorch),
+/// (eg with 'Tensor::load(filename)') that are saved from pytorch. They
+/// use different serialization formats. Luckily, I can embed a tensor in
+/// a JIT module and unpack it from there instead. The most
+/// straightforward way of getting it out is to define a 'forward()'
+/// method that simply returns the wrapped tensor.
+///
+/// I wrote a script to wrap tensors: 'jit_containerize_tensor.py'.
+pub fn load_wrapped_tensor(filename: &str) -> Tensor {
+  println!("Loading wrapped tensor file: {}", filename);
+  let module = tch::CModule::load(filename).unwrap();
+  let mut temp = Tensor::zeros(
+    &[10, 10, 10],
+    (tch::Kind::Float, tch::Device::Cpu)
+  );
+  module.forward(&temp)
+}

tacotron_melgan/src/tacotron.rs

@@ -32,10 +32,12 @@ pub fn run_tacotron() {
   let mut vs = tch::nn::VarStore::new(tch::Device::Cpu);
   let tacotron_model = load_tacotron_model(WRAPPED_MODEL_PATH);
 
-  println!("Evaluating model...");
-  let output = tacotron_model.forward(&text_sequence);
+  for i in 0..10 {
+    println!("Evaluating model... {}", i);
+    let output = tacotron_model.forward(&text_sequence);
 
-  println!("Result tensor: {:?}", output);
+    println!("Result tensor: {:?}", output);
+  }
 }
 
 /*println!("Tacotron2 + MelGan");