diff --git a/tacotron_melgan/src/main.rs b/tacotron_melgan/src/main.rs index 5e93c6c48d..8afd2714e1 100644 --- a/tacotron_melgan/src/main.rs +++ b/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."); } \ No newline at end of file diff --git a/tacotron_melgan/src/melgan.rs b/tacotron_melgan/src/melgan.rs index 74f7e1aa3f..ef19e77c74 100644 --- a/tacotron_melgan/src/melgan.rs +++ b/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 { data } +pub fn write_audio_file(audio_signal: Vec, 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()); diff --git a/tacotron_melgan/src/model.rs b/tacotron_melgan/src/model.rs new file mode 100644 index 0000000000..2605060977 --- /dev/null +++ b/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) +} diff --git a/tacotron_melgan/src/tacotron.rs b/tacotron_melgan/src/tacotron.rs index c8fe6ff540..e8f1d703ec 100644 --- a/tacotron_melgan/src/tacotron.rs +++ b/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");