Model picker! Updated frontend and backend

This allows the frontend user to choose which melgan and tacotron are
used at runtime without restarting the server. It could use UI
improvement, but it works.

The available models are configured with a toml config file loaded by
the backend.

Cargo.lock

@@ -311,6 +311,11 @@ dependencies = [
  "winapi 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 
+[[package]]
+name = "anyhow"
+version = "1.0.30"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+
 [[package]]
 name = "arc-swap"
 version = "0.4.4"
@@ -2266,6 +2271,14 @@ name = "toml"
 version = "0.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 
+[[package]]
+name = "toml"
+version = "0.5.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+dependencies = [
+ "serde 1.0.104 (registry+https://github.com/rust-lang/crates.io-index)",
+]
+
 [[package]]
 name = "torch-sys"
 version = "0.1.3"
@@ -2325,6 +2338,7 @@ dependencies = [
  "actix-rt 1.0.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "actix-utils 1.0.4 (registry+https://github.com/rust-lang/crates.io-index)",
  "actix-web 2.0.0 (registry+https://github.com/rust-lang/crates.io-index)",
+ "anyhow 1.0.30 (registry+https://github.com/rust-lang/crates.io-index)",
  "arpabet 1.0.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "hound 3.4.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "numerics 0.0.1 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -2332,6 +2346,7 @@ dependencies = [
  "serde 1.0.104 (registry+https://github.com/rust-lang/crates.io-index)",
  "serde_derive 1.0.104 (registry+https://github.com/rust-lang/crates.io-index)",
  "tch 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)",
+ "toml 0.5.6 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 
 [[package]]
@@ -2632,6 +2647,7 @@ dependencies = [
 "checksum aligned_alloc 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)" = "9dcebfb002ccde769c15bc841d0d5548a90e80fcd2ffed5131339e8074746f0a"
 "checksum alsa-sys 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)" = "b0edcbbf9ef68f15ae1b620f722180b82a98b6f0628d30baa6b8d2a5abc87d58"
 "checksum ansi_term 0.11.0 (registry+https://github.com/rust-lang/crates.io-index)" = "ee49baf6cb617b853aa8d93bf420db2383fab46d314482ca2803b40d5fde979b"
+"checksum anyhow 1.0.30 (registry+https://github.com/rust-lang/crates.io-index)" = "2494382e9ba43995f3c56359e518641f450f5c36feeb4632a75cde2ec297c867"
 "checksum arc-swap 0.4.4 (registry+https://github.com/rust-lang/crates.io-index)" = "d7b8a9123b8027467bce0099fe556c628a53c8d83df0507084c31e9ba2e39aff"
 "checksum arpabet 1.0.0 (registry+https://github.com/rust-lang/crates.io-index)" = "c76d18638cfaa15c3e7c521ef8a223614bee6b802433d788815347f28a0c4add"
 "checksum async-trait 0.1.22 (registry+https://github.com/rust-lang/crates.io-index)" = "c8df72488e87761e772f14ae0c2480396810e51b2c2ade912f97f0f7e5b95e3c"
@@ -2855,6 +2871,7 @@ dependencies = [
 "checksum tokio 0.2.6 (registry+https://github.com/rust-lang/crates.io-index)" = "0e1bef565a52394086ecac0a6fa3b8ace4cb3a138ee1d96bd2b93283b56824e3"
 "checksum tokio-util 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)" = "571da51182ec208780505a32528fc5512a8fe1443ab960b3f2f3ef093cd16930"
 "checksum toml 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)" = "736b60249cb25337bc196faa43ee12c705e426f3d55c214d73a4e7be06f92cb4"
+"checksum toml 0.5.6 (registry+https://github.com/rust-lang/crates.io-index)" = "ffc92d160b1eef40665be3a05630d003936a3bc7da7421277846c2613e92c71a"
 "checksum torch-sys 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)" = "5baf9a7559875b8903e30e5ca0f69154a4f2359a6a8bd304e6d29ddea3a7e0cd"
 "checksum trust-dns-proto 0.18.0-alpha.2 (registry+https://github.com/rust-lang/crates.io-index)" = "2a7f3a2ab8a919f5eca52a468866a67ed7d3efa265d48a652a9a3452272b413f"
 "checksum trust-dns-resolver 0.18.0-alpha.2 (registry+https://github.com/rust-lang/crates.io-index)" = "6f90b1502b226f8b2514c6d5b37bafa8c200d7ca4102d57dc36ee0f3b7a04a2f"

tacotron_melgan/wrap_tacotron2_model.py

@@ -11,82 +11,17 @@ from math import sqrt
 
 from tacotron_model import Tacotron2
 
-#tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_arpabet_ljs_checkpoint_22000'
-#tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_arpabet_ljs_checkpoint_54000'
-#tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_arpabet_ljs_checkpoint_93000'
-#tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_arpabet_ljs_checkpoint_128000'
-#tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_arpabet_txlearn_trump_checkpoint_3000'
-#tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_10000'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_9500'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_2000'
-# AWFUL
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_11500'
-# Start new... (bad)
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_500'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_1500'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_20500'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_29500'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_41000'
-# Great
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_5000'
-# Awful
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_13000'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_6000'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_12000'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_500'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_1500'
-# Distributed run with carefully tweaked validation set
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_4500'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_6000'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_7000'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_15500'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_16500'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_20500'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_500'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_1000'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_9500'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_11000'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_16000'
-# Really good:
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_26000'
-# So awful (loss skyrocketed for no reason):
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_32500'
-# This one was right before the uptick in loss:
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_30000'
-# Trying this again with more data
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_500'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_13000'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_28000'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_35500'
-# Trying with my voice (very limited data)
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_32500'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_3000'
-# Another training run
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_500'
-# Sounds awesome, but 100% spectral robot noise:
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_2000'
-# OH GOD, AWFUL. NOT EVEN WORDS:
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_7500'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_3000'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_10000'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_76000'
-# Modified LJS - pitch +700 reverb +2 +noise
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_4500'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_6500'
-# Got some spectral noise. Isn't clean. But English is good:
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_17500'
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_78500'
-# Training a larger transformed LJS dataset of the same construction:
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_9500'
-# Post-coronavirus outbreak training
-tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_17500'
-
 # Going back in time, this is an old trump txlearned model:
 tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_trump_txlearn_ljs_arpabet_ckpt30000_2020_02_06'
+# Starting again with new samples:
+tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_trump_txlearn_arpabet_2020-05-14_chkpt7000'
+tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_trump_txlearn_arpabet_2020-05-14_chkpt11500'
 
 
 #output_filename = 'tacotron_arpabet_1.pt'
 output_filename = 'tacotron2_trump_txlearn_ljs_arpabet_ckpt30000_2020_02_06.jit'
+output_filename = 'tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt7000.jit'
+output_filename = 'tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt11500.jit'
 
 # NB: Tacotron's original shape is [148, 512]. Arpabet encoding is [256, 512].
 NUM_SYMBOLS_ENCODED = 256 # 148 by default, 256 with my custom Arpabet encoding

tts_frontend/src/TtsService.ts

@@ -5,6 +5,10 @@ class TtsRequest {
   text: String
   speaker: String
 
+  // Which model to use
+  melgan_model?: Number
+  arpabet_tacotron_model?: Number
+
   constructor(text: String) {
     this.text = text;
     // LJSpeech dataset from Keith Ito & LibriVox, spoken by Linda Johnson.

tts_frontend/src/components/MainComponent.tsx

@@ -2,6 +2,7 @@ import React from 'react';
 import Howl from 'howler';
 import TextInput from './TextInput';
 import TrackList from './TrackList';
+import { ModelPickerDropdownComponent } from './model_picker/ModelPickerDropdownComponent';
 
 interface Props {
 }
@@ -26,7 +27,6 @@ class MainComponent extends React.Component<Props, State> {
   constructor(props: Props) {
     super(props);
     this.state = { utterances: new Array() };
-    //this.setState({utterances: new Array()});
   }
 
   /** Add a new audio track. */

tts_frontend/src/components/TextInput.tsx

@@ -1,6 +1,7 @@
 import React from 'react';
 import Howl from 'howler';
 import {TextAudioPair }from './MainComponent'
+import { ModelPickerDropdownComponent } from './model_picker/ModelPickerDropdownComponent';
 
 interface Props {
   appendUtteranceCallback: (utterance: TextAudioPair) => void
@@ -9,12 +10,16 @@ interface Props {
 interface State {
   text?: String,
   howl?: Howl,
+  arpabet_tacotron_model?: string,
+  melgan_model?: string,
 }
 
 /// Requests to the backend.
 class TtsRequest {
   text: String
   speaker: String
+  arpabet_tacotron_model?: string
+  melgan_model?: string
 
   constructor(text: String) {
     this.text = text;
@@ -42,6 +47,18 @@ class TextInput extends React.Component<Props, State> {
     this.setState({text: text});
   }
 
+  handleArpabetTacotronModelChange = (ev: React.FormEvent<HTMLSelectElement>) => {
+    console.log('callback A');
+    const model = (ev.target as HTMLInputElement).value;
+    this.setState({arpabet_tacotron_model: model});
+  }
+
+  handleMelganModelChange = (ev: React.FormEvent<HTMLSelectElement>) => {
+    console.log('callback B');
+    const model = (ev.target as HTMLInputElement).value;
+    this.setState({melgan_model: model});
+  }
+
   makeRequest = (ev: React.FormEvent<HTMLFormElement>) => {
     console.log("Form Submit");
 
@@ -50,8 +67,10 @@ class TextInput extends React.Component<Props, State> {
     }
 
     let request = new TtsRequest(this.state.text);
+    request.arpabet_tacotron_model = this.state.arpabet_tacotron_model || "";
+    request.melgan_model = this.state.melgan_model || "";
 
-    const url = `/tts`;
+    const url = '/tts';
     fetch(url, {
       method: 'POST',
       headers: {
@@ -63,8 +82,10 @@ class TextInput extends React.Component<Props, State> {
     .then(res => res.blob())
     .then(blob => {
       console.log(blob);
+
       const url = window.URL.createObjectURL(blob);
       console.log(url);
+
       const sound = new Howl.Howl({
         src: [url],
         format: 'wav',
@@ -76,10 +97,7 @@ class TextInput extends React.Component<Props, State> {
       this.props.appendUtteranceCallback(new TextAudioPair(request.text, sound));
 
       (window as any).sound = sound;
-    })
-
-    //.then(res => res.json())
-    //.then(...
+    });
 
     ev.preventDefault();
     return false;
@@ -89,6 +107,10 @@ class TextInput extends React.Component<Props, State> {
     return (
       <form onSubmit={this.makeRequest}>
         <input onChange={this.handleTextChange} />
+        <ModelPickerDropdownComponent 
+          changeArpabetTacotronCallback={this.handleArpabetTacotronModelChange}
+          changeMelganCallback={this.handleMelganModelChange}
+        />
         <button>Submit</button>
       </form>
     );

tts_frontend/src/components/model_picker/DropdownArpabetTacotron.tsx

@@ -0,0 +1,30 @@
+import React from 'react';
+import { ModelDetails } from './ModelResponse';
+
+interface Props {
+  callback: (ev: React.FormEvent<HTMLSelectElement>) => void,
+  model_details: ModelDetails[]
+}
+
+interface State {
+}
+
+class DropdownArpabetTacotron extends React.Component<Props, State> {
+
+  constructor(props: Props) {
+    super(props);
+  }
+
+  public render() {
+    return (
+      <select onChange={this.props.callback}>
+        {this.props.model_details.map(function(model_details: ModelDetails) {
+          let file : string = model_details.file || "unknown";
+          return (<option value={file}>{file}</option>)
+        })};
+      </select>
+    );
+  }
+}
+
+export {DropdownArpabetTacotron};

tts_frontend/src/components/model_picker/DropdownMelgan.tsx

@@ -0,0 +1,30 @@
+import React from 'react';
+import { ModelDetails } from './ModelResponse';
+
+interface Props {
+  callback: (ev: React.FormEvent<HTMLSelectElement>) => void,
+  model_details: ModelDetails[]
+}
+
+interface State {
+}
+
+class DropdownMelgan extends React.Component<Props, State> {
+
+  constructor(props: Props) {
+    super(props);
+  }
+
+  public render() {
+    return (
+      <select onChange={this.props.callback}>
+        {this.props.model_details.map(function(model_details: ModelDetails) {
+          let file : string = model_details.file || "unknown";
+          return (<option value={file}>{file}</option>)
+        })};
+      </select>
+    );
+  }
+}
+
+export {DropdownMelgan};

tts_frontend/src/components/model_picker/ModelPickerDropdownComponent.tsx

@@ -0,0 +1,62 @@
+import React from 'react';
+import { ModelResponse, ModelDetails } from './ModelResponse';
+import { DropdownMelgan } from './DropdownMelgan';
+import { DropdownArpabetTacotron } from './DropdownArpabetTacotron';
+
+interface Props {
+  changeArpabetTacotronCallback: (ev: React.FormEvent<HTMLSelectElement>) => void,
+  changeMelganCallback: (ev: React.FormEvent<HTMLSelectElement>) => void,
+}
+
+interface State {
+  models?: ModelResponse,
+}
+
+class ModelPickerDropdownComponent extends React.Component<Props, State> {
+
+  constructor(props: Props) {
+    super(props);
+    this.state = {
+      models: undefined,
+    };
+  }
+
+  componentDidMount() {
+    this.loadModels();
+  }
+
+  public loadModels() {
+    const url = '/models';
+    fetch(url)
+      .then(res => res.json())
+      .then(
+        (result) => {
+          const models = result;
+          this.setState({
+            models: models,
+          })
+        }
+      );
+  }
+
+  public render() {
+    let arpabet_tacotron_models : ModelDetails[] = this.state.models?.tacotron || [];
+    let melgan_models : ModelDetails[] = this.state.models?.melgan || [];
+    return (
+      <div>
+        <h3>Tacotron (Arpabet)</h3>
+        <DropdownArpabetTacotron 
+          model_details={arpabet_tacotron_models} 
+          callback={this.props.changeArpabetTacotronCallback}
+          />
+        <h3>Melgan</h3>
+        <DropdownMelgan 
+          model_details={melgan_models} 
+          callback={this.props.changeMelganCallback}
+          />
+      </div>
+    );
+  }
+}
+
+export {ModelPickerDropdownComponent};

tts_frontend/src/components/model_picker/ModelResponse.ts

@@ -0,0 +1,12 @@
+
+class ModelResponse {
+  tacotron?: ModelDetails[]
+  melgan?: ModelDetails[]
+}
+
+class ModelDetails {
+  file?: string
+  description?: string
+}
+
+export {ModelResponse, ModelDetails};

tts_service/Cargo.toml

@@ -18,10 +18,12 @@ actix-files= "0.2.1"
 actix-rt = "1.0.0"
 actix-utils = "1.0.3"
 actix-web = "2.0.0"
+anyhow = "1.0.30"
 hound = "3.4.0" # wav codec
 serde = "1.0.104"
 serde_derive = "1.0.104"
 tch = "0.1.3" # libtorch bindings (pytorch)
+toml = "0.5.6"
 
 # My libraries
 arpabet = "1.0.0"

tts_service/models.toml

@@ -0,0 +1,36 @@
+# --------------- Tacotron ---------------
+
+[[tacotron]]
+file = "/home/bt/dev/voder/tacotron_melgan/tacotron2_trump_txlearn_ljs_arpabet_ckpt30000_2020_02_06.jit"
+description = "First Trump txlearn with Arpabet"
+
+[[tacotron]]
+file = "/home/bt/dev/voder/tacotron_melgan/tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt7000.jit"
+description = "Trump txlearn with Arpabet 2020-05-14"
+
+[[tacotron]]
+file = "/home/bt/dev/voder/tacotron_melgan/tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt11500.jit"
+description = "Trump txlearn with Arpabet 2020-05-14"
+
+# ---------------- Melgan ----------------
+
+[[melgan]]
+file = "/home/bt/dev/tacotron-melgan/melgan_trump-txlearn-2020.05.05.jit"
+description = "Latest LJS transfer learned Trump."
+
+[[melgan]]
+file = "/home/bt/dev/tacotron-melgan/melgan_trump-txlearn-2020.05.05_13675.jit"
+description = "Latest LJS transfer learned Trump."
+
+[[melgan]]
+file = "/home/bt/dev/tacotron-melgan/melgan_ljs-high-pitch_container.pt"
+description = "High-pitched LJS"
+
+[[melgan]]
+file = "/home/bt/dev/tacotron-melgan/shared_melgan_container.pt"
+description = "Unknown"
+
+[[melgan]]
+file = "/home/bt/dev/tacotron-melgan/shared_melgan_container2.pt"
+description = "Unknown"
+

tts_service/src/config.rs

@@ -0,0 +1,23 @@
+use serde_derive::Deserialize;
+use std::fs;
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct ModelConfigs {
+  tacotron: Vec<ModelDetails>,
+  melgan: Vec<ModelDetails>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
+pub struct ModelDetails {
+  pub file: String,
+  pub description: String,
+}
+
+impl ModelConfigs {
+  pub fn load_from_file(filename: &str) -> Self {
+    let contents = fs::read_to_string(filename)
+        .expect("Couldn't read file");
+    toml::from_str(&contents)
+        .expect("Couldn't parse toml")
+  }
+}

tts_service/src/main.rs

@@ -1,4 +1,5 @@
 #[macro_use] extern crate actix_web;
+#[macro_use] extern crate anyhow;
 #[macro_use] extern crate serde_derive;
 
 extern crate actix_files;
@@ -7,11 +8,11 @@ extern crate hound;
 extern crate serde;
 extern crate tch;
 
+pub mod config;
 pub mod model;
+pub mod old_model;
 pub mod text;
 
-use model::TacoMelModel;
-
 use std::env;
 use std::sync::Arc;
 
@@ -21,6 +22,7 @@ use actix_cors::Cors;
 use actix_files::Files;
 use actix_web::http::{header, Method, StatusCode};
 use actix_web::middleware::Logger;
+use actix_web::web::Json;
 use actix_web::{
   App,
   HttpRequest,
@@ -31,10 +33,14 @@ use actix_web::{
   http,
   web,
 };
+
+use crate::config::ModelConfigs;
+use crate::model::model_cache::ModelCache;
+use crate::old_model::TacoMelModel;
 use crate::text::text_to_arpabet_encoding;
 
 #[get("/")]
-async fn get_root(request: HttpRequest) -> std::io::Result<HttpResponse> {
+async fn get_root(_request: HttpRequest) -> std::io::Result<HttpResponse> {
   println!("GET /");
   Ok(HttpResponse::build(StatusCode::OK)
       .content_type("text/html; charset=utf-8")
@@ -42,7 +48,7 @@ async fn get_root(request: HttpRequest) -> std::io::Result<HttpResponse> {
 }
 
 #[get("/readiness")]
-async fn get_readiness(request: HttpRequest) -> std::io::Result<HttpResponse> {
+async fn get_readiness(_request: HttpRequest) -> std::io::Result<HttpResponse> {
   println!("GET /readiness");
   Ok(HttpResponse::build(StatusCode::OK)
       .content_type("text/html; charset=utf-8")
@@ -50,18 +56,30 @@ async fn get_readiness(request: HttpRequest) -> std::io::Result<HttpResponse> {
 }
 
 #[get("/liveness")]
-async fn get_liveness(request: HttpRequest) -> std::io::Result<HttpResponse> {
+async fn get_liveness(_request: HttpRequest) -> std::io::Result<HttpResponse> {
   println!("GET /liveness");
   Ok(HttpResponse::build(StatusCode::OK)
       .content_type("text/html; charset=utf-8")
       .body("Live"))
 }
 
+#[get("/models")]
+async fn get_models(_request: HttpRequest) -> std::io::Result<Json<ModelConfigs>> {
+  println!("GET /models");
+  let model_configs = ModelConfigs::load_from_file("models.toml");
+  println!("Model Configs: {:?}", model_configs);
+
+  Ok(Json(model_configs))
+}
+
 /// For JSON payloads
 #[derive(Deserialize)]
 pub struct TtsRequest {
   text: String,
   speaker: String,
+  // The client can specify the models to use
+  arpabet_tacotron_model: Option<String>,
+  melgan_model: Option<String>,
 }
 
 /// For query strings
@@ -71,7 +89,7 @@ pub struct TtsQueryRequest {
 }
 
 //#[get("/tts")]
-async fn get_tts(request: HttpRequest,
+/*async fn get_tts(request: HttpRequest,
   query: web::Query<TtsQueryRequest>,
   model: web::Data<Arc<TacoMelModel>>)
     -> std::io::Result<HttpResponse> {
@@ -86,25 +104,41 @@ async fn get_tts(request: HttpRequest,
   Ok(HttpResponse::build(StatusCode::OK)
       .content_type("audio/wav")
       .body(wav_data))
-}
+}*/
 
 //#[post("/tts")]
 async fn post_tts(request: HttpRequest,
   query: web::Json<TtsRequest>,
-  //arpabet: web::Data<Arc<Arpabet>>,
-  model: web::Data<Arc<TacoMelModel>>)
+  model_cache: web::Data<Arc<ModelCache>>)
   -> std::io::Result<HttpResponse> {
   println!("POST /tts");
-  let text = query.text.to_string();
 
+  let tacotron_model = query.arpabet_tacotron_model
+      .as_ref()
+      .map(|s| s.clone())
+      .unwrap_or("/home/bt/dev/voder/tacotron_melgan/tacotron2_trump_txlearn_ljs_arpabet_2020-05-14_ckpt11500.jit".to_string());
+
+  let melgan_model = query.melgan_model
+      .as_ref()
+      .map(|s| s.clone())
+      .unwrap_or("/home/bt/dev/tacotron-melgan/melgan_trump-txlearn-2020.05.05_13675.jit".to_string());
+
+  let text = query.text.to_string();
+  println!("Tacotron Model: {}", tacotron_model);
+  println!("Melgan Model: {}", melgan_model);
   println!("Text: {}", text);
 
   let arpabet = Arpabet::load_cmudict();
   let encoded = text_to_arpabet_encoding(arpabet, &text);
 
-  println!("Encoded: {:?}", encoded);
+  println!("Encoded Text: {:?}", encoded);
 
-  let wav_data = model.run_tts_encoded(&encoded);
+  let mut cache = model_cache.into_inner();
+
+  let tacotron = cache.get_or_load_arbabet_tacotron(&tacotron_model).unwrap();
+  let melgan = cache.get_or_load_melgan(&melgan_model).unwrap();
+
+  let wav_data = TacoMelModel::new().run_tts_encoded(&tacotron, &melgan, &encoded);
 
   Ok(HttpResponse::build(StatusCode::OK)
       .content_type("audio/wav")
@@ -116,13 +150,18 @@ const MELGAN_MODEL : &'static str = "MELGAN_MODEL";
 const BIND_ADDRESS : &'static str = "BIND_ADDRESS";
 const ASSET_DIRECTORY : &'static str = "ASSET_DIRECTORY";
 
-pub struct AppState {
-  pub arpabet: Arpabet,
-  pub model: TacoMelModel,
-}
+//pub struct AppState {
+//  pub arpabet: Arpabet,
+//  pub model: TacoMelModel,
+//}
 
 #[actix_rt::main]
 async fn main() -> std::io::Result<()> {
+  println!("Loading configs.");
+
+  let configs = ModelConfigs::load_from_file("models.toml");
+  println!("Configs: {:?}", configs);
+
   println!("Starting service.");
 
   let bind_address = env::var(BIND_ADDRESS)
@@ -139,9 +178,10 @@ async fn main() -> std::io::Result<()> {
 
   println!("Loading models...");
 
-  let ttsEngine = TacoMelModel::create(&tacotron_filename, &melgan_filename);
+  //let ttsEngine = TacoMelModel::create(&tacotron_filename, &melgan_filename);
+  let model_cache = ModelCache::new();
 
-  let arc = web::Data::new(Arc::new(ttsEngine));
+  let arc = web::Data::new(Arc::new(model_cache));
   //let arc2 = web::Data::new(Arc::new(arpabet));
 
   println!("Starting HTTP service.");
@@ -153,13 +193,14 @@ async fn main() -> std::io::Result<()> {
       .service(Files::new("/frontend", asset_directory.clone()).show_files_listing())
       .service(
         web::resource("/tts")
-            .route(web::get().to(get_tts))
+            //.route(web::get().to(get_tts))
             .route(web::post().to(post_tts))
             .route(web::head().to(|| HttpResponse::Ok()))
       )
       .service(get_root)
       .service(get_readiness)
       .service(get_liveness)
+      .service(get_models)
       .app_data(arc.clone())
       //.app_data(arc2.clone())
     )

tts_service/src/model/arpabet_tacotron_model.rs

@@ -0,0 +1,28 @@
+use crate::model::model_container::ModelContainer;
+use anyhow::Result;
+use tch::Tensor;
+
+pub struct ArpabetTacotronModel {
+  model_container: ModelContainer,
+}
+
+impl ArpabetTacotronModel {
+  pub fn load(filename: &str) -> Result<Self> {
+    let model_container = ModelContainer::load(filename)?;
+    Ok(Self {
+      model_container,
+    })
+  }
+
+  pub fn encoded_arpabet_to_mel(&self, text_buffer: &Vec<i64>) -> Tensor {
+    let text_tensor = Tensor::of_slice(text_buffer.as_slice());
+    println!("Text tensor: {:?}", text_tensor);
+
+    let text_tensor = text_tensor.unsqueeze(0);
+    println!("Text tensor unsq: {:?}", text_tensor);
+
+    // This produces the mel, though it will require processing before
+    // sending it to MelGan.
+    self.model_container.forward(&text_tensor)
+  }
+}

tts_service/src/model/melgan_model.rs

@@ -0,0 +1,20 @@
+use crate::model::model_container::ModelContainer;
+use anyhow::Result;
+use tch::Tensor;
+
+pub struct MelganModel {
+  model_container: ModelContainer,
+}
+
+impl MelganModel {
+  pub fn load(filename: &str) -> Result<Self> {
+    let model_container = ModelContainer::load(filename)?;
+    Ok(Self {
+      model_container,
+    })
+  }
+
+  pub fn tacotron_mel_to_audio(&self, mel_tensor: &Tensor) -> Tensor {
+    self.model_container.forward(&mel_tensor)
+  }
+}

tts_service/src/model/mod.rs

@@ -0,0 +1,4 @@
+pub mod melgan_model;
+pub mod model_container;
+pub mod arpabet_tacotron_model;
+pub mod model_cache;

tts_service/src/model/model_cache.rs

@@ -0,0 +1,56 @@
+use anyhow::Error;
+use crate::model::arpabet_tacotron_model::ArpabetTacotronModel;
+use crate::model::melgan_model::MelganModel;
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::sync::Mutex;
+
+pub struct ModelCache {
+  arpabet_tacotron_models: Arc<Mutex<HashMap<String, Arc<ArpabetTacotronModel>>>>,
+  melgan_models: Arc<Mutex<HashMap<String, Arc<MelganModel>>>>,
+}
+
+impl ModelCache {
+  pub fn new() -> Self {
+    Self {
+      arpabet_tacotron_models: Arc::new(Mutex::new(HashMap::new())),
+      melgan_models: Arc::new(Mutex::new(HashMap::new())),
+    }
+  }
+
+  pub fn get_or_load_arbabet_tacotron(&self, filename: &str) -> Option<Arc<ArpabetTacotronModel>> {
+    let mut lock = self.arpabet_tacotron_models.lock()
+        .expect("should unlock");
+
+    if let Some(model) = lock.get(filename) {
+      return Some(model.clone());
+    }
+
+    match ArpabetTacotronModel::load(&filename) {
+      Ok(model) => {
+        let arc = Arc::new(model);
+        lock.insert(filename.to_string(), arc.clone());
+        Some(arc)
+      },
+      Err(_) => None,
+    }
+  }
+
+  pub fn get_or_load_melgan(&self, filename: &str) -> Option<Arc<MelganModel>> {
+    let mut lock = self.melgan_models.lock()
+        .expect("should unlock");
+
+    if let Some(model) = lock.get(filename) {
+      return Some(model.clone());
+    }
+
+    match MelganModel::load(&filename) {
+      Ok(model) => {
+        let arc = Arc::new(model);
+        lock.insert(filename.to_string(), arc.clone());
+        Some(arc)
+      },
+      Err(_) => None,
+    }
+  }
+}

tts_service/src/model/model_container.rs

@@ -0,0 +1,31 @@
+use anyhow::Result;
+use tch::CModule;
+use tch::Tensor;
+use tch::nn::Module;
+use tch::nn::ModuleT;
+
+/// Holds the loaded pytorch JIT model
+pub struct ModelContainer {
+  filename: String,
+  jit_model: CModule,
+}
+
+impl ModelContainer {
+  pub fn load(filename: &str) -> Result<Self> {
+    let model = CModule::load(filename)
+        .map_err(|e| anyhow!("jit model failed to load: {}", e))?;
+
+    Ok(Self {
+      filename: filename.to_string(),
+      jit_model: model,
+    })
+  }
+
+  pub fn get_filename(&self) -> &str {
+    &self.filename
+  }
+
+  pub fn forward(&self, tensor: &Tensor) -> Tensor {
+    self.jit_model.forward(tensor)
+  }
+}

tts_service/src/old_model.rs

@@ -10,9 +10,11 @@ use tch::nn::ModuleT;
 use std::io::Cursor;
 use std::io::BufWriter;
 
+use crate::model::arpabet_tacotron_model::ArpabetTacotronModel;
+use crate::model::model_container::ModelContainer;
+use crate::model::melgan_model::MelganModel;
+
 pub struct TacoMelModel {
-  tacotron_model: CModule,
-  melgan_model: CModule,
 }
 
 impl TacoMelModel {
@@ -21,17 +23,16 @@ impl TacoMelModel {
     CModule::load(filename).unwrap()
   }
 
-  pub fn create(tacotron_filename: &str, melgan_filename: &str) -> Self {
-    let tacotron_model = Self::load_model_file(tacotron_filename);
-    let melgan_model = Self::load_model_file(melgan_filename);
-
-    Self {
-      tacotron_model,
-      melgan_model,
-    }
+  pub fn new() -> Self {
+    Self {}
   }
 
-  pub fn run_tts(&self, text: &str) -> Vec<i16> {
+  /*pub fn run_tts_audio(&self, text: &str) -> Vec<u8> {
+    let audio_signal = self.run_tts(text);
+    Self::audio_signal_to_wav_bytes(audio_signal)
+  }
+
+  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();
@@ -49,31 +50,23 @@ impl TacoMelModel {
     println!("Text buffer: {:?}", text_buffer);
 
     self.encoded_text_to_audio_signal(&text_buffer)
+  }*/
+
+  /// Run TTS on Arpabet encoding
+  pub fn run_tts_encoded(&self, tacotron: &ArpabetTacotronModel, melgan: &MelganModel, encoded_text: &Vec<i64>) -> Vec<u8> {
+    let audio_signal = self.encoded_text_to_audio_signal(tacotron, melgan, encoded_text);
+    Self::audio_signal_to_wav_bytes(audio_signal)
   }
 
-  fn encoded_text_to_audio_signal(&self, text_buffer: &Vec<i64>) -> Vec<i16> {
-    let text_tensor = Tensor::of_slice(text_buffer.as_slice());
-    println!("Text tensor: {:?}", text_tensor);
-    let text_tensor = text_tensor.unsqueeze(0);
-    println!("Text tensor unsq: {:?}", text_tensor);
-    let mut mel_tensor = self.tacotron_model.forward(&text_tensor);
+  fn encoded_text_to_audio_signal(&self, tacotron: &ArpabetTacotronModel, melgan: &MelganModel, text_buffer: &Vec<i64>) -> Vec<i16> {
+    let mut mel_tensor = tacotron.encoded_arpabet_to_mel(&text_buffer);
     println!("\n\n>>> Mel tensor:\n{:?}\n\n", mel_tensor);
     println!("Running melgan...");
-    let audio_tensor = self.melgan_model.forward(&mel_tensor);
+    let audio_tensor = melgan.tacotron_mel_to_audio(&mel_tensor);
     println!("\n\n>>> Audio tensor:\n{:?}\n\n", audio_tensor);
     Self::audio_tensor_to_audio_signal(audio_tensor)
   }
 
-  pub fn run_tts_audio(&self, text: &str) -> Vec<u8> {
-    let audio_signal = self.run_tts(text);
-    Self::audio_signal_to_wav_bytes(audio_signal)
-  }
-
-  pub fn run_tts_encoded(&self, encoded_text: &Vec<i64>) -> Vec<u8> {
-    let audio_signal = self.encoded_text_to_audio_signal(encoded_text);
-    Self::audio_signal_to_wav_bytes(audio_signal)
-  }
-
   fn audio_tensor_to_audio_signal(mel: Tensor) -> Vec<i16> {
     let mut flat_audio_tensor = mel.squeeze();