From 89d08a449ef38f2c4b82ee42b64e071966c912ef Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Mon, 30 Dec 2019 03:48:24 -0500
Subject: [PATCH] Rust Tacotron2 model runs, albeit SLOWLY

---
 tacotron_melgan/src/main.rs     |  3 ++-
 tacotron_melgan/src/tacotron.rs | 35 ++++++++++++++++++++++++++++++---
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/tacotron_melgan/src/main.rs b/tacotron_melgan/src/main.rs
index 58ab2c6957..5e93c6c48d 100644
--- a/tacotron_melgan/src/main.rs
+++ b/tacotron_melgan/src/main.rs
@@ -12,5 +12,6 @@ pub mod tacotron;
 
 pub fn main() {
   println!("Tacotron2 + MelGan");
-  melgan::run_melgan();
+  //melgan::run_melgan();
+  tacotron::run_tacotron();
 }
\ No newline at end of file
diff --git a/tacotron_melgan/src/tacotron.rs b/tacotron_melgan/src/tacotron.rs
index 3b7a67b6b4..c8fe6ff540 100644
--- a/tacotron_melgan/src/tacotron.rs
+++ b/tacotron_melgan/src/tacotron.rs
@@ -1,12 +1,41 @@
+use tch::Tensor;
+use tch::CModule;
+use tch;
+use tch::nn::Module;
 
-const WRAPPED_MODEL_PATH : &'static str = "/home/bt/dev/voder/...";
 
-pub fn load_tacotron_model() {
+const WRAPPED_MODEL_PATH : &'static str = "/home/bt/dev/voder/tacotron_melgan/tacotron_container.pt";
+//const EXAMPLE_TEXT_SEQUENCE : &'static str = "/home/bt/dev/voder/data/text/tacotron_text_sequence.pt";
+const EXAMPLE_TEXT_SEQUENCE : &'static str = "/home/bt/dev/voder/data/text/tacotron_text_sequence.pt.containerized.pt";
+
+pub fn load_tacotron_model(filename: &str) -> CModule {
   // TODO
+  println!("Loading model: {}", filename);
+  tch::CModule::load(filename).unwrap()
 }
 
+pub fn load_wrapped_text(filename: &str) -> Tensor {
+  println!("Loading wrapped text sequence 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)
+}
+
+
 pub fn run_tacotron() {
-  // TODO
+  let mut text_sequence = load_wrapped_text(EXAMPLE_TEXT_SEQUENCE);
+  println!("Got text_sequence: {:?} of dim {}", text_sequence, text_sequence.dim());
+
+  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);
+
+  println!("Result tensor: {:?}", output);
 }
 
 /*println!("Tacotron2 + MelGan");