From 9772cabd9935152daa2bd9e3f83e80ed88bd7d38 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 15 May 2020 01:20:52 -0400 Subject: [PATCH] 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 | 17 ++++ tacotron_melgan/wrap_tacotron2_model.py | 75 ++---------------- tts_frontend/src/TtsService.ts | 4 + tts_frontend/src/components/MainComponent.tsx | 2 +- tts_frontend/src/components/TextInput.tsx | 32 ++++++-- .../model_picker/DropdownArpabetTacotron.tsx | 30 +++++++ .../model_picker/DropdownMelgan.tsx | 30 +++++++ .../ModelPickerDropdownComponent.tsx | 62 +++++++++++++++ .../components/model_picker/ModelResponse.ts | 12 +++ tts_service/Cargo.toml | 2 + tts_service/models.toml | 36 +++++++++ tts_service/src/config.rs | 23 ++++++ tts_service/src/main.rs | 79 ++++++++++++++----- .../src/model/arpabet_tacotron_model.rs | 28 +++++++ tts_service/src/model/melgan_model.rs | 20 +++++ tts_service/src/model/mod.rs | 4 + tts_service/src/model/model_cache.rs | 56 +++++++++++++ tts_service/src/model/model_container.rs | 31 ++++++++ tts_service/src/{model.rs => old_model.rs} | 49 +++++------- 19 files changed, 469 insertions(+), 123 deletions(-) create mode 100644 tts_frontend/src/components/model_picker/DropdownArpabetTacotron.tsx create mode 100644 tts_frontend/src/components/model_picker/DropdownMelgan.tsx create mode 100644 tts_frontend/src/components/model_picker/ModelPickerDropdownComponent.tsx create mode 100644 tts_frontend/src/components/model_picker/ModelResponse.ts create mode 100644 tts_service/models.toml create mode 100644 tts_service/src/config.rs create mode 100644 tts_service/src/model/arpabet_tacotron_model.rs create mode 100644 tts_service/src/model/melgan_model.rs create mode 100644 tts_service/src/model/mod.rs create mode 100644 tts_service/src/model/model_cache.rs create mode 100644 tts_service/src/model/model_container.rs rename tts_service/src/{model.rs => old_model.rs} (72%) diff --git a/Cargo.lock b/Cargo.lock index 1e87fbb59a..955c0cb1b5 100644 --- a/Cargo.lock +++ b/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" diff --git a/tacotron_melgan/wrap_tacotron2_model.py b/tacotron_melgan/wrap_tacotron2_model.py index 7261ea1aad..5243ef9880 100755 --- a/tacotron_melgan/wrap_tacotron2_model.py +++ b/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 diff --git a/tts_frontend/src/TtsService.ts b/tts_frontend/src/TtsService.ts index cfd86d578f..e2b0ec6679 100644 --- a/tts_frontend/src/TtsService.ts +++ b/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. diff --git a/tts_frontend/src/components/MainComponent.tsx b/tts_frontend/src/components/MainComponent.tsx index cd8527f3c2..3f486bfc6d 100644 --- a/tts_frontend/src/components/MainComponent.tsx +++ b/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 { constructor(props: Props) { super(props); this.state = { utterances: new Array() }; - //this.setState({utterances: new Array()}); } /** Add a new audio track. */ diff --git a/tts_frontend/src/components/TextInput.tsx b/tts_frontend/src/components/TextInput.tsx index edb3345970..05458024e9 100644 --- a/tts_frontend/src/components/TextInput.tsx +++ b/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 { this.setState({text: text}); } + handleArpabetTacotronModelChange = (ev: React.FormEvent) => { + console.log('callback A'); + const model = (ev.target as HTMLInputElement).value; + this.setState({arpabet_tacotron_model: model}); + } + + handleMelganModelChange = (ev: React.FormEvent) => { + console.log('callback B'); + const model = (ev.target as HTMLInputElement).value; + this.setState({melgan_model: model}); + } + makeRequest = (ev: React.FormEvent) => { console.log("Form Submit"); @@ -50,8 +67,10 @@ class TextInput extends React.Component { } 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 { .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 { 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 { return (
+ ); diff --git a/tts_frontend/src/components/model_picker/DropdownArpabetTacotron.tsx b/tts_frontend/src/components/model_picker/DropdownArpabetTacotron.tsx new file mode 100644 index 0000000000..ad2f7d58e8 --- /dev/null +++ b/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) => void, + model_details: ModelDetails[] +} + +interface State { +} + +class DropdownArpabetTacotron extends React.Component { + + constructor(props: Props) { + super(props); + } + + public render() { + return ( + + ); + } +} + +export {DropdownArpabetTacotron}; diff --git a/tts_frontend/src/components/model_picker/DropdownMelgan.tsx b/tts_frontend/src/components/model_picker/DropdownMelgan.tsx new file mode 100644 index 0000000000..a29edabdb1 --- /dev/null +++ b/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) => void, + model_details: ModelDetails[] +} + +interface State { +} + +class DropdownMelgan extends React.Component { + + constructor(props: Props) { + super(props); + } + + public render() { + return ( + + ); + } +} + +export {DropdownMelgan}; diff --git a/tts_frontend/src/components/model_picker/ModelPickerDropdownComponent.tsx b/tts_frontend/src/components/model_picker/ModelPickerDropdownComponent.tsx new file mode 100644 index 0000000000..2d7c08ab31 --- /dev/null +++ b/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) => void, + changeMelganCallback: (ev: React.FormEvent) => void, +} + +interface State { + models?: ModelResponse, +} + +class ModelPickerDropdownComponent extends React.Component { + + 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 ( +
+

Tacotron (Arpabet)

+ +

Melgan

+ +
+ ); + } +} + +export {ModelPickerDropdownComponent}; diff --git a/tts_frontend/src/components/model_picker/ModelResponse.ts b/tts_frontend/src/components/model_picker/ModelResponse.ts new file mode 100644 index 0000000000..24741bca15 --- /dev/null +++ b/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}; diff --git a/tts_service/Cargo.toml b/tts_service/Cargo.toml index c1e1303619..3b246798c5 100644 --- a/tts_service/Cargo.toml +++ b/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" diff --git a/tts_service/models.toml b/tts_service/models.toml new file mode 100644 index 0000000000..ff8881d609 --- /dev/null +++ b/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" + diff --git a/tts_service/src/config.rs b/tts_service/src/config.rs new file mode 100644 index 0000000000..c0979c7f2b --- /dev/null +++ b/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, + melgan: Vec, +} + +#[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") + } +} diff --git a/tts_service/src/main.rs b/tts_service/src/main.rs index 1890c5f46f..bcbdc3fa97 100644 --- a/tts_service/src/main.rs +++ b/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 { +async fn get_root(_request: HttpRequest) -> std::io::Result { 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 { } #[get("/readiness")] -async fn get_readiness(request: HttpRequest) -> std::io::Result { +async fn get_readiness(_request: HttpRequest) -> std::io::Result { 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 { } #[get("/liveness")] -async fn get_liveness(request: HttpRequest) -> std::io::Result { +async fn get_liveness(_request: HttpRequest) -> std::io::Result { 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> { + 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, + melgan_model: Option, } /// 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, model: web::Data>) -> std::io::Result { @@ -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, - //arpabet: web::Data>, - model: web::Data>) + model_cache: web::Data>) -> std::io::Result { 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()) ) diff --git a/tts_service/src/model/arpabet_tacotron_model.rs b/tts_service/src/model/arpabet_tacotron_model.rs new file mode 100644 index 0000000000..52dbeacd68 --- /dev/null +++ b/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 { + let model_container = ModelContainer::load(filename)?; + Ok(Self { + model_container, + }) + } + + pub fn encoded_arpabet_to_mel(&self, text_buffer: &Vec) -> 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) + } +} diff --git a/tts_service/src/model/melgan_model.rs b/tts_service/src/model/melgan_model.rs new file mode 100644 index 0000000000..3047bf4622 --- /dev/null +++ b/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 { + 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) + } +} diff --git a/tts_service/src/model/mod.rs b/tts_service/src/model/mod.rs new file mode 100644 index 0000000000..a0ff78a3e0 --- /dev/null +++ b/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; diff --git a/tts_service/src/model/model_cache.rs b/tts_service/src/model/model_cache.rs new file mode 100644 index 0000000000..a4fe4f2dce --- /dev/null +++ b/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>>>, + melgan_models: Arc>>>, +} + +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> { + 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> { + 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, + } + } +} diff --git a/tts_service/src/model/model_container.rs b/tts_service/src/model/model_container.rs new file mode 100644 index 0000000000..54fc6ccdae --- /dev/null +++ b/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 { + 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) + } +} diff --git a/tts_service/src/model.rs b/tts_service/src/old_model.rs similarity index 72% rename from tts_service/src/model.rs rename to tts_service/src/old_model.rs index e3400131d8..a4347bfaca 100644 --- a/tts_service/src/model.rs +++ b/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 { + /*pub fn run_tts_audio(&self, text: &str) -> Vec { + let audio_signal = self.run_tts(text); + Self::audio_signal_to_wav_bytes(audio_signal) + } + + fn run_tts(&self, text: &str) -> Vec { println!("Text : {:?}", text); let copied = text.to_string().to_ascii_lowercase(); let mut text_buffer : Vec = 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) -> Vec { + 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) -> Vec { - 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) -> Vec { + 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 { - let audio_signal = self.run_tts(text); - Self::audio_signal_to_wav_bytes(audio_signal) - } - - pub fn run_tts_encoded(&self, encoded_text: &Vec) -> Vec { - 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 { let mut flat_audio_tensor = mel.squeeze();