Model outputs audio, but all zeroes right now

Cargo.lock

@@ -1152,6 +1152,7 @@ dependencies = [
 name = "tacotron_melgan"
 version = "0.0.1"
 dependencies = [
+ "hound 3.4.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "tch 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 

tacotron_melgan/Cargo.toml

@@ -9,4 +9,5 @@ name = "tacotron_melgan"
 path = "src/main.rs"
 
 [dependencies]
+hound = "3.4.0" # wav codec
 tch = "0.1.3" # libtorch bindings (pytorch)

tacotron_melgan/src/main.rs

@@ -1,4 +1,5 @@
 extern crate tch;
+extern crate hound;
 
 use tch::vision::imagenet;
 use tch::vision::resnet;

tacotron_melgan/src/melgan.rs

@@ -3,14 +3,14 @@ use tch::CModule;
 use tch;
 use tch::nn::Module;
 
+
 //const TACOTRON_MODEL_PATH : &'static str = "/home/bt/models/tacotron2-nvidia/tacotron2_statedict.pt";
 //const MELGAN_MODEL_PATH : &'static str = "/home/bt/models/melgan-swpark/firstgo_a7c2351_1100.pt";
-//const WRAPPED_MODEL_PATH : &'static str = "/home/bt/dev/voder/tacotron_melgan/container.pt";
-const WRAPPED_MODEL_PATH : &'static str = "/home/bt/dev/voder/tacotron_melgan/cpu_container.pt";
+const WRAPPED_MODEL_PATH : &'static str = "/home/bt/dev/voder/tacotron_melgan/container.pt";
+//const WRAPPED_MODEL_PATH : &'static str = "/home/bt/dev/voder/tacotron_melgan/cpu_container.pt";
 
-// TODO: This isn't working because I need mels in the format of a torch tensor. Not python pickle file
 //const EXAMPLE_MEL_1: &'static str = "/home/bt/dev/voder/data/mels/LJ002-0320.mel";
-const EXAMPLE_MEL_1: &'static str = "/home/bt/dev/voder/tacotron_melgan/mel_file.pt";
+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";
 
 pub fn load_melgan_model(filename: &str) -> CModule {
@@ -18,6 +18,14 @@ pub fn load_melgan_model(filename: &str) -> CModule {
   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.
 pub fn load_wrapped_mel(filename: &str) -> Tensor {
   println!("Loading wrapped mel file: {}", filename);
   let module = tch::CModule::load(filename).unwrap();
@@ -28,20 +36,50 @@ pub fn load_wrapped_mel(filename: &str) -> Tensor {
   module.forward(&temp)
 }
 
-pub fn load_mel(filename: &str) -> Tensor {
-  println!("Loading mel file: {}", filename);
-  // The file format is the same as the one used by the PyTorch C++ API.
-  Tensor::load(filename).unwrap()
-}
-
 pub fn run_melgan() {
-  let mel = load_wrapped_mel(EXAMPLE_MEL_1);
-  println!("Got mel: {:?}", mel);
+  let mut mel = load_wrapped_mel(EXAMPLE_MEL_1);
+  println!("Got mel: {:?} of dim {}", mel, mel.dim());
+
+  if mel.dim() == 2 {
+    println!("mel unsqeeze");
+    mel = mel.unsqueeze(0);
+  }
 
   let mut vs = tch::nn::VarStore::new(tch::Device::Cpu);
   let melgan_model = load_melgan_model(WRAPPED_MODEL_PATH);
 
-  let output = melgan_model.forward_ts(&[mel.unsqueeze(0)]).unwrap();
+  println!("Evaluating model...");
+  let output = melgan_model.forward(&mel);
+
+  println!("Result tensor: {:?}", output);
+
+  let flat = output.squeeze();
+
+  println!("Sqeueezed tensor: {:?}, dim: {}", flat, flat.dim());
+
+  let length = flat.size1().unwrap() as usize;
+  println!("Length: {}", length);
+
+  let mut data : Vec<f32> = Vec::with_capacity(length);
+  for i in 0 .. length {
+    data.push(0.0f32);
+  }
+
+  flat.copy_data(data.as_mut_slice(), length as usize);
+
+  let spec = hound::WavSpec {
+    channels: 1,
+    sample_rate: 16000,
+    bits_per_sample: 32,
+    sample_format: hound::SampleFormat::Float,
+  };
+
+  let mut writer = hound::WavWriter::create("output.wav", spec).unwrap();
+
+  for sample in data {
+    writer.write_sample(sample).unwrap();
+  }
+
 
   /*println!("Tacotron2 + MelGan");
 

tacotron_melgan/src/tacotron.rs

@@ -1,4 +1,6 @@
 
+const WRAPPED_MODEL_PATH : &'static str = "/home/bt/dev/voder/...";
+
 pub fn load_tacotron_model() {
   // TODO
 }

tacotron_melgan/wrap_torch_model.py

@@ -112,8 +112,9 @@ class Container(torch.nn.Module):
 
 
 print('Loading melgan model...')
-melgan_model_file = '/home/bt/models/melgan-swpark/firstgo_a7c2351_1100.pt'
-melgan_model = torch.load(melgan_model_file)
+#melgan_model_file = '/home/bt/models/melgan-swpark/firstgo_a7c2351_1100.pt'
+melgan_model_file = '/home/bt/models/firstgo_a7c2351_1100.pt'
+melgan_model = torch.load(melgan_model_file, map_location=torch.device('cpu'))
 
 def cuda_to_cpu(model):
     """Recursively make everything non-CUDA"""
@@ -141,7 +142,7 @@ graph = {
     # NB: You can't save a toplevel dictionary of dictionaries,
     # like so: 'melgan': melgan_model,
     # Per inference.py, it looks like only model_g is used.
-    #'melgan_model_g': melgan_model['model_g'], # NB: keeping this out to load below.
+    'melgan_model_g': melgan_model['model_g'], # NB: keeping this out to load below.
     #'melgan_model_d': melgan_model['model_d'],
     #'melgan_optim_g': melgan_model['optim_g'],
     #'melgan_optim_d': melgan_model['optim_d'],