diff --git a/Cargo.lock b/Cargo.lock index da61fa765a..adb739c3c5 100644 --- a/Cargo.lock +++ b/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)", ] diff --git a/tacotron_melgan/Cargo.toml b/tacotron_melgan/Cargo.toml index d7fb0baa2e..49fb92fe0d 100644 --- a/tacotron_melgan/Cargo.toml +++ b/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) diff --git a/tacotron_melgan/src/main.rs b/tacotron_melgan/src/main.rs index cd6f1a6395..72cc562516 100644 --- a/tacotron_melgan/src/main.rs +++ b/tacotron_melgan/src/main.rs @@ -1,4 +1,5 @@ extern crate tch; +extern crate hound; use tch::vision::imagenet; use tch::vision::resnet; diff --git a/tacotron_melgan/src/melgan.rs b/tacotron_melgan/src/melgan.rs index 704839e753..d9d99d7ee9 100644 --- a/tacotron_melgan/src/melgan.rs +++ b/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 = 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"); diff --git a/tacotron_melgan/src/tacotron.rs b/tacotron_melgan/src/tacotron.rs index 116fb679f5..b8ed3106e0 100644 --- a/tacotron_melgan/src/tacotron.rs +++ b/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 } diff --git a/tacotron_melgan/wrap_torch_model.py b/tacotron_melgan/wrap_torch_model.py index c167f7e081..02530845c6 100755 --- a/tacotron_melgan/wrap_torch_model.py +++ b/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'],