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();