diff --git a/tacotron_melgan/src/main.rs b/tacotron_melgan/src/main.rs index e3b5232d23..cd6f1a6395 100644 --- a/tacotron_melgan/src/main.rs +++ b/tacotron_melgan/src/main.rs @@ -5,87 +5,10 @@ use tch::vision::resnet; use tch::nn::ModuleT; use tch::Tensor; - -//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"; - -// TODO: This isn't working because I need mels in the format of a torch tensor. Not python pickle file -const EXAMPLE_MEL_PATH : &'static str = "/home/bt/dev/2nd/melgan/temp/output/trump_2018_02_15-001.mel"; +pub mod melgan; +pub mod tacotron; pub fn main() { println!("Tacotron2 + MelGan"); - - // Create the model and load the weights from the file. - let mut vs = tch::nn::VarStore::new(tch::Device::Cpu); - - println!("Loading Wrapped model: {}", WRAPPED_MODEL_PATH); - - println!("loaded"); - - /* Just gotta do this... - model = Generator(hp.audio.n_mel_channels).cuda() - model.load_state_dict(checkpoint['model_g']) - model.eval(inference=False) - - with torch.no_grad(): - for melpath in tqdm.tqdm(glob.glob(os.path.join(args.input_folder, '*.mel'))): - mel = torch.load(melpath) - if len(mel.shape) == 2: - mel = mel.unsqueeze(0) - - filename = melpath.replace('.mel', '_reconstructed_epoch%04d.png' % checkpoint['epoch']) - #render_histogram(mel, filename) - mel = mel.cuda() - - audio = model.inference(mel) - audio = audio.cpu().detach().numpy() - - out_path = melpath.replace('.mel', '_reconstructed_epoch%04d.wav' % checkpoint['epoch']) - write(out_path, hp.audio.sampling_rate, audio) - */ - - //vs.load(WRAPPED_MODEL_PATH).unwrap(); - - println!("Loading trained model..."); - let model = tch::CModule::load(WRAPPED_MODEL_PATH).unwrap(); - - // NB: This is just random data for now. Kind of mel shaped since it's an image. - //let path = "/home/bt/Downloads/virtual-studio-design.jpg"; - //let image = imagenet::load_image_and_resize(path, 100, 100).unwrap(); - //let image = imagenet::load_image(EXAMPLE_MEL_PATH).unwrap(); - - println!("Loading mel file..."); - let mel_file = Tensor::load(EXAMPLE_MEL_PATH).unwrap(); - - println!("Forwarding mel to model..."); - let output = model.forward_ts(&[mel_file.unsqueeze(0)]).unwrap(); - - - /*// TODO: Now how do I evaluate the model? - //let resnet18 = tch::vision::resnet::resnet18(&vs.root(), imagenet::CLASS_COUNT); - - let path = "/home/bt/Downloads/virtual-studio-design.jpg"; - let image = imagenet::load_image_and_resize(path, 100, 100).unwrap(); - - // NB: This works - let net = Box::new(resnet::resnet18(&vs.root(), imagenet::CLASS_COUNT)); - - for _ in 0..10 { - let output = net - .forward_t(&image.unsqueeze(0), /*train=*/ false) - .softmax(-1, tch::Kind::Float); // Convert to probability. - - for (probability, class) in imagenet::top(&output, 5).iter() { - println!("{:50} {:5.2}%", class, 100.0 * probability) - } - }*/ - - /*// Apply the forward pass of the model to get the logits. - let output = net - .forward_t(&image.unsqueeze(0), /*train=*/ false) - .softmax(-1, tch::Kind::Float); // Convert to probability. - - // Print the top 5 categories for this image. - Ok(())*/ + melgan::run_melgan(); } \ No newline at end of file diff --git a/tacotron_melgan/src/melgan.rs b/tacotron_melgan/src/melgan.rs new file mode 100644 index 0000000000..83a8e4eed0 --- /dev/null +++ b/tacotron_melgan/src/melgan.rs @@ -0,0 +1,105 @@ +use tch::Tensor; +use tch::CModule; +use tch; + +//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"; + +// TODO: This isn't working because I need mels in the format of a torch tensor. Not python pickle file +const EXAMPLE_MEL_PATH : &'static str = "/home/bt/dev/2nd/voder/data/mels/trump_2018_02_15-001.mel"; + +pub fn load_melgan_model(filename: &str) -> CModule { + println!("Loading model: {}", filename); + tch::CModule::load(WRAPPED_MODEL_PATH).unwrap() +} + +pub fn load_mel(filename: &str) -> Tensor { + println!("Loading mel file: {}", filename); + Tensor::load(filename).unwrap() +} + +pub fn run_melgan() { + let mel = load_mel(EXAMPLE_MEL_PATH); + + 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!("Tacotron2 + MelGan"); + + // Create the model and load the weights from the file. + let mut vs = tch::nn::VarStore::new(tch::Device::Cpu); + + println!("Loading Wrapped model: {}", WRAPPED_MODEL_PATH); + + println!("loaded"); + + /* Just gotta do this... + model = Generator(hp.audio.n_mel_channels).cuda() + model.load_state_dict(checkpoint['model_g']) + model.eval(inference=False) + + with torch.no_grad(): + for melpath in tqdm.tqdm(glob.glob(os.path.join(args.input_folder, '*.mel'))): + mel = torch.load(melpath) + if len(mel.shape) == 2: + mel = mel.unsqueeze(0) + + filename = melpath.replace('.mel', '_reconstructed_epoch%04d.png' % checkpoint['epoch']) + #render_histogram(mel, filename) + mel = mel.cuda() + + audio = model.inference(mel) + audio = audio.cpu().detach().numpy() + + out_path = melpath.replace('.mel', '_reconstructed_epoch%04d.wav' % checkpoint['epoch']) + write(out_path, hp.audio.sampling_rate, audio) + */ + + //vs.load(WRAPPED_MODEL_PATH).unwrap(); + + println!("Loading trained model..."); + let model = tch::CModule::load(WRAPPED_MODEL_PATH).unwrap(); + + // NB: This is just random data for now. Kind of mel shaped since it's an image. + //let path = "/home/bt/Downloads/virtual-studio-design.jpg"; + //let image = imagenet::load_image_and_resize(path, 100, 100).unwrap(); + //let image = imagenet::load_image(EXAMPLE_MEL_PATH).unwrap(); + + println!("Loading mel file..."); + let mel_file = Tensor::load(EXAMPLE_MEL_PATH).unwrap(); + + println!("Forwarding mel to model..."); + let output = model.forward_ts(&[mel_file.unsqueeze(0)]).unwrap(); + + + /*// TODO: Now how do I evaluate the model? + //let resnet18 = tch::vision::resnet::resnet18(&vs.root(), imagenet::CLASS_COUNT); + + let path = "/home/bt/Downloads/virtual-studio-design.jpg"; + let image = imagenet::load_image_and_resize(path, 100, 100).unwrap(); + + // NB: This works + let net = Box::new(resnet::resnet18(&vs.root(), imagenet::CLASS_COUNT)); + + for _ in 0..10 { + let output = net + .forward_t(&image.unsqueeze(0), /*train=*/ false) + .softmax(-1, tch::Kind::Float); // Convert to probability. + + for (probability, class) in imagenet::top(&output, 5).iter() { + println!("{:50} {:5.2}%", class, 100.0 * probability) + } + }*/ + + /*// Apply the forward pass of the model to get the logits. + let output = net + .forward_t(&image.unsqueeze(0), /*train=*/ false) + .softmax(-1, tch::Kind::Float); // Convert to probability. + + // Print the top 5 categories for this image. + Ok(())*/ + */ +} diff --git a/tacotron_melgan/src/tacotron.rs b/tacotron_melgan/src/tacotron.rs new file mode 100644 index 0000000000..116fb679f5 --- /dev/null +++ b/tacotron_melgan/src/tacotron.rs @@ -0,0 +1,8 @@ + +pub fn load_tacotron_model() { + // TODO +} + +pub fn run_tacotron() { + // TODO +} \ No newline at end of file