From edd521f2fddb07c1269cc68516be7938362d6b9f Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Sat, 4 Jan 2020 05:02:17 -0500
Subject: [PATCH] exact audio tensor match with python tacotron-melgan (wav
 audio still off)

---
 tacotron_melgan/src/melgan.rs | 21 +++++++++++----------
 tacotron_melgan/src/model.rs  |  8 ++++----
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/tacotron_melgan/src/melgan.rs b/tacotron_melgan/src/melgan.rs
index 97aa7d6d4c..38d9f202b6 100644
--- a/tacotron_melgan/src/melgan.rs
+++ b/tacotron_melgan/src/melgan.rs
@@ -9,19 +9,20 @@ use model::{
   load_model_file,
   load_wrapped_tensor_file,
 };
+use hound::Sample;
 
 const WRAPPED_MODEL_PATH : &'static str = "/home/bt/dev/voder/tacotron_melgan/melgan_jit_model_voder_c0cac635.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.containerized.pt";
 
-pub const MAX_WAV_VALUE : f32 = 32768.0f32;
-pub const HALF_WAV_VALUE : f32 = MAX_WAV_VALUE * 0.5f32;
+//pub const MAX_WAV_VALUE : f32 = 32768.0f32;
+//pub const HALF_WAV_VALUE : f32 = MAX_WAV_VALUE * 0.5f32;
 
 // TODO: This is an hparam and should be dynamic.
-pub const HOP_LENGTH : i64 = 256;
+//pub const HOP_LENGTH : i64 = 256;
 
-pub fn audio_tensor_to_audio_signal(mel: Tensor) -> Vec<f32> {
+pub fn audio_tensor_to_audio_signal(mel: Tensor) -> Vec<i16> {
   let mut flat_audio_tensor = mel.squeeze();
 
   println!("Sqeueezed tensor: {:?}, dim: {}",
@@ -37,7 +38,7 @@ pub fn audio_tensor_to_audio_signal(mel: Tensor) -> Vec<f32> {
   //println!("New size: {}", new_size);
   //flat_audio_tensor = flat_audio_tensor.resize_(&[new_size]);
 
-  flat_audio_tensor = flat_audio_tensor * MAX_WAV_VALUE as f64;
+  //flat_audio_tensor = flat_audio_tensor * MAX_WAV_VALUE as f64;
 
   let mut data : Vec<f32> = Vec::with_capacity(length);
 
@@ -47,7 +48,7 @@ pub fn audio_tensor_to_audio_signal(mel: Tensor) -> Vec<f32> {
 
   flat_audio_tensor.copy_data(data.as_mut_slice(), length as usize);
 
-  data
+  data.iter().map(|x| *x as i16).collect()
 }
 
 fn debug_print_sample(audio: &Vec<f32>, num_samples: usize) {
@@ -76,17 +77,17 @@ pub fn run_melgan_network() {
 
   let audio = audio_tensor_to_audio_signal(output);
 
-  debug_print_sample(&audio, 10);
+  //debug_print_sample(&audio, 10);
 
   write_audio_file(audio, "melgan_output.wav");
 }
 
-pub fn write_audio_file(audio_signal: Vec<f32>, filename: &str) {
+pub fn write_audio_file(audio_signal: Vec<i16>, filename: &str) {
   let spec = hound::WavSpec {
     channels: 1,
     sample_rate: 16000,
-    bits_per_sample: 32,
-    sample_format: hound::SampleFormat::Float,
+    bits_per_sample: 16,
+    sample_format: hound::SampleFormat::Int,
   };
 
   let mut writer = hound::WavWriter::create(filename, spec).unwrap();
diff --git a/tacotron_melgan/src/model.rs b/tacotron_melgan/src/model.rs
index 8f0d4754eb..5a50c8a7e0 100644
--- a/tacotron_melgan/src/model.rs
+++ b/tacotron_melgan/src/model.rs
@@ -55,7 +55,7 @@ impl TacoMelModel {
     }
   }
 
-  pub fn run_tts(&self, text: &str) -> Vec<f32> {
+  pub fn run_tts(&self, text: &str) -> Vec<i16> {
     println!("Text : {:?}", text);
     let copied = text.to_string().to_ascii_lowercase();
     let mut text_buffer : Vec<i64> = Vec::new();
@@ -105,12 +105,12 @@ impl TacoMelModel {
   }
 }
 
-pub fn audio_signal_to_wav_bytes(audio_signal: Vec<f32>) -> Vec<u8> {
+pub fn audio_signal_to_wav_bytes(audio_signal: Vec<i16>) -> Vec<u8> {
   let spec = WavSpec {
     channels: 1,
     sample_rate: 16000,
-    bits_per_sample: 32,
-    sample_format: SampleFormat::Float,
+    bits_per_sample: 16,
+    sample_format: SampleFormat::Int,
   };
   let bytes: Vec<u8> = Vec::new();
   let seek: Cursor<Vec<u8>> = Cursor::new(bytes);