From 96a146c5e1193cd7552ef74aba9aa170c8fa1206 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 16 Jan 2020 03:17:50 -0500 Subject: [PATCH] Arpabet processing in Tacotron --- Cargo.lock | 72 +++++++++++++++++++++++++ tacotron_melgan/src/main.rs | 10 ++-- tacotron_melgan/wrap_tacotron2_model.py | 11 ++-- tts_service/Cargo.toml | 1 + tts_service/src/main.rs | 48 +++++++---------- tts_service/src/model.rs | 16 +++--- tts_service/src/text.rs | 60 +++++++++++++++++++++ 7 files changed, 172 insertions(+), 46 deletions(-) create mode 100644 tts_service/src/text.rs diff --git a/Cargo.lock b/Cargo.lock index 1761afb85e..e7189ead94 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -316,6 +316,16 @@ name = "arc-swap" version = "0.4.4" source = "registry+https://github.com/rust-lang/crates.io-index" +[[package]] +name = "arpabet" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "lazy_static 1.4.0 (registry+https://github.com/rust-lang/crates.io-index)", + "phf 0.8.0 (registry+https://github.com/rust-lang/crates.io-index)", + "regex 1.3.1 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "async-trait" version = "0.1.22" @@ -1425,6 +1435,46 @@ dependencies = [ "fixedbitset 0.1.9 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "phf" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "phf_macros 0.8.0 (registry+https://github.com/rust-lang/crates.io-index)", + "phf_shared 0.8.0 (registry+https://github.com/rust-lang/crates.io-index)", + "proc-macro-hack 0.5.11 (registry+https://github.com/rust-lang/crates.io-index)", +] + +[[package]] +name = "phf_generator" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "phf_shared 0.8.0 (registry+https://github.com/rust-lang/crates.io-index)", + "rand 0.7.2 (registry+https://github.com/rust-lang/crates.io-index)", +] + +[[package]] +name = "phf_macros" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "phf_generator 0.8.0 (registry+https://github.com/rust-lang/crates.io-index)", + "phf_shared 0.8.0 (registry+https://github.com/rust-lang/crates.io-index)", + "proc-macro-hack 0.5.11 (registry+https://github.com/rust-lang/crates.io-index)", + "proc-macro2 1.0.7 (registry+https://github.com/rust-lang/crates.io-index)", + "quote 1.0.2 (registry+https://github.com/rust-lang/crates.io-index)", + "syn 1.0.13 (registry+https://github.com/rust-lang/crates.io-index)", +] + +[[package]] +name = "phf_shared" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "siphasher 0.3.1 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "pin-project" version = "0.4.6" @@ -1607,6 +1657,7 @@ dependencies = [ "rand_chacha 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)", "rand_core 0.5.1 (registry+https://github.com/rust-lang/crates.io-index)", "rand_hc 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)", + "rand_pcg 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)", ] [[package]] @@ -1704,6 +1755,14 @@ dependencies = [ "rand_core 0.4.2 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "rand_pcg" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "rand_core 0.5.1 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "rand_xorshift" version = "0.1.1" @@ -1906,6 +1965,11 @@ dependencies = [ "libc 0.2.66 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "siphasher" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" + [[package]] name = "slab" version = "0.4.2" @@ -2215,6 +2279,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)", + "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)", "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)", @@ -2520,6 +2585,7 @@ dependencies = [ "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 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" "checksum atty 0.2.14 (registry+https://github.com/rust-lang/crates.io-index)" = "d9b39be18770d11421cdb1b9947a45dd3f37e93092cbf377614828a319d5fee8" "checksum autocfg 0.1.7 (registry+https://github.com/rust-lang/crates.io-index)" = "1d49d90015b3c36167a20fe2810c5cd875ad504b39cff3d4eae7977e6b7c1cb2" @@ -2648,6 +2714,10 @@ dependencies = [ "checksum peeking_take_while 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)" = "19b17cddbe7ec3f8bc800887bab5e717348c95ea2ca0b1bf0837fb964dc67099" "checksum percent-encoding 2.1.0 (registry+https://github.com/rust-lang/crates.io-index)" = "d4fd5641d01c8f18a23da7b6fe29298ff4b55afcccdf78973b24cf3175fee32e" "checksum petgraph 0.4.13 (registry+https://github.com/rust-lang/crates.io-index)" = "9c3659d1ee90221741f65dd128d9998311b0e40c5d3c23a62445938214abce4f" +"checksum phf 0.8.0 (registry+https://github.com/rust-lang/crates.io-index)" = "3dfb61232e34fcb633f43d12c58f83c1df82962dcdfa565a4e866ffc17dafe12" +"checksum phf_generator 0.8.0 (registry+https://github.com/rust-lang/crates.io-index)" = "17367f0cc86f2d25802b2c26ee58a7b23faeccf78a396094c13dced0d0182526" +"checksum phf_macros 0.8.0 (registry+https://github.com/rust-lang/crates.io-index)" = "7f6fde18ff429ffc8fe78e2bf7f8b7a5a5a6e2a8b58bc5a9ac69198bbda9189c" +"checksum phf_shared 0.8.0 (registry+https://github.com/rust-lang/crates.io-index)" = "c00cf8b9eafe68dde5e9eaa2cef8ee84a9336a47d566ec55ca16589633b65af7" "checksum pin-project 0.4.6 (registry+https://github.com/rust-lang/crates.io-index)" = "94b90146c7216e4cb534069fb91366de4ea0ea353105ee45ed297e2d1619e469" "checksum pin-project-internal 0.4.6 (registry+https://github.com/rust-lang/crates.io-index)" = "44ca92f893f0656d3cba8158dd0f2b99b94de256a4a54e870bd6922fcc6c8355" "checksum pin-project-lite 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)" = "e8822eb8bb72452f038ebf6048efa02c3fe22bf83f76519c9583e47fc194a422" @@ -2680,6 +2750,7 @@ dependencies = [ "checksum rand_jitter 0.1.4 (registry+https://github.com/rust-lang/crates.io-index)" = "1166d5c91dc97b88d1decc3285bb0a99ed84b05cfd0bc2341bdf2d43fc41e39b" "checksum rand_os 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)" = "7b75f676a1e053fc562eafbb47838d67c84801e38fc1ba459e8f180deabd5071" "checksum rand_pcg 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)" = "abf9b09b01790cfe0364f52bf32995ea3c39f4d2dd011eac241d2914146d0b44" +"checksum rand_pcg 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)" = "16abd0c1b639e9eb4d7c50c0b8100b0d0f849be2349829c740fe8e6eb4816429" "checksum rand_xorshift 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)" = "cbf7e9e623549b0e21f6e97cf8ecf247c1a8fd2e8a992ae265314300b2455d5c" "checksum rawpointer 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)" = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" "checksum rayon 1.3.0 (registry+https://github.com/rust-lang/crates.io-index)" = "db6ce3297f9c85e16621bb8cca38a06779ffc31bb8184e1be4bed2be4678a098" @@ -2707,6 +2778,7 @@ dependencies = [ "checksum sha1 0.6.0 (registry+https://github.com/rust-lang/crates.io-index)" = "2579985fda508104f7587689507983eadd6a6e84dd35d6d115361f530916fa0d" "checksum shlex 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)" = "7fdf1b9db47230893d76faad238fd6097fd6d6a9245cd7a4d90dbd639536bbd2" "checksum signal-hook-registry 1.2.0 (registry+https://github.com/rust-lang/crates.io-index)" = "94f478ede9f64724c5d173d7bb56099ec3e2d9fc2774aac65d34b8b890405f41" +"checksum siphasher 0.3.1 (registry+https://github.com/rust-lang/crates.io-index)" = "83da420ee8d1a89e640d0948c646c1c088758d3a3c538f943bfa97bdac17929d" "checksum slab 0.4.2 (registry+https://github.com/rust-lang/crates.io-index)" = "c111b5bd5695e56cffe5129854aa230b39c93a305372fdbb2668ca2394eea9f8" "checksum smallvec 1.1.0 (registry+https://github.com/rust-lang/crates.io-index)" = "44e59e0c9fa00817912ae6e4e6e3c4fe04455e75699d06eedc7d85917ed8e8f4" "checksum socket2 0.3.11 (registry+https://github.com/rust-lang/crates.io-index)" = "e8b74de517221a2cb01a53349cf54182acdc31a074727d3079068448c0676d85" diff --git a/tacotron_melgan/src/main.rs b/tacotron_melgan/src/main.rs index 0b5fa31c58..d4aea75c30 100644 --- a/tacotron_melgan/src/main.rs +++ b/tacotron_melgan/src/main.rs @@ -1,7 +1,7 @@ #![feature(proc_macro_hygiene, decl_macro)] -#[macro_use] extern crate rocket; -#[macro_use] extern crate rocket_contrib; +//#[macro_use] extern crate rocket; +//#[macro_use] extern crate rocket_contrib; #[macro_use] extern crate serde_derive; extern crate hound; extern crate rand; @@ -13,11 +13,11 @@ use tch::CModule; use tch::Tensor; use tch::nn::Module; use tch::nn::ModuleT; -use server::run_server; +//use server::run_server; pub mod melgan; pub mod model; -pub mod server; +//pub mod server; pub mod tacotron; pub fn main() { @@ -27,7 +27,7 @@ pub fn main() { //tacotron::run_tacotron(); //run_end_to_end(); - run_server(); + //run_server(); } fn run_end_to_end() { diff --git a/tacotron_melgan/wrap_tacotron2_model.py b/tacotron_melgan/wrap_tacotron2_model.py index af7c1e5f79..2123a05ea9 100755 --- a/tacotron_melgan/wrap_tacotron2_model.py +++ b/tacotron_melgan/wrap_tacotron2_model.py @@ -11,12 +11,17 @@ from math import sqrt from tacotron_model import Tacotron2 +tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_arpabet_ljs_checkpoint_22000' +output_filename = 'tacotron_arpabet_1.pt' + +# 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 + #torch.set_default_tensor_type('torch.DoubleTensor') #torch.set_default_tensor_type('torch.FloatTensor') print('Loading tacotron model...') #tacotron_model_file = '/home/bt/models/tacotron2-nvidia/tacotron2_statedict.pt' -tacotron_model_file = '/home/bt/models/tacotron2-nvidia/checkpoint_2000' tacotron_model = torch.load(tacotron_model_file, map_location=torch.device('cpu')) def cuda_to_cpu(model): @@ -56,7 +61,7 @@ class HParams: fp16_run = False n_mel_channels = 80 n_frames_per_step = 1 - n_symbols = 148 # NB: Determined by running inference + n_symbols = NUM_SYMBOLS_ENCODED symbols_embedding_dim = 512 encoder_kernel_size = 5 encoder_n_convolutions = 3 @@ -78,8 +83,6 @@ print('Load state dict...') module.load_state_dict(tacotron_model['state_dict']) #module.eval() # NB: Complains unless called: Did you forget call .eval() on your model? -output_filename = 'tacotron_container4.pt' - # NB: Tracing evaluates the model on input and unrolls and hardcodes branching # and loops. Scripting allows these to remain by converting the entire program # to TorchScript. Scripting seems MUCH faster, is also harder get working. diff --git a/tts_service/Cargo.toml b/tts_service/Cargo.toml index 87dd27507a..ab0a94f0c1 100644 --- a/tts_service/Cargo.toml +++ b/tts_service/Cargo.toml @@ -15,6 +15,7 @@ actix-files= "0.2.1" actix-rt = "1.0.0" actix-utils = "1.0.3" actix-web = "2.0.0" +arpabet = "1.0.0" hound = "3.4.0" # wav codec serde = "1.0.104" serde_derive = "1.0.104" diff --git a/tts_service/src/main.rs b/tts_service/src/main.rs index bbdbf6158b..1890c5f46f 100644 --- a/tts_service/src/main.rs +++ b/tts_service/src/main.rs @@ -2,17 +2,21 @@ #[macro_use] extern crate serde_derive; extern crate actix_files; +extern crate arpabet; extern crate hound; extern crate serde; extern crate tch; pub mod model; +pub mod text; use model::TacoMelModel; use std::env; use std::sync::Arc; +use arpabet::Arpabet; + use actix_cors::Cors; use actix_files::Files; use actix_web::http::{header, Method, StatusCode}; @@ -27,6 +31,7 @@ use actix_web::{ http, web, }; +use crate::text::text_to_arpabet_encoding; #[get("/")] async fn get_root(request: HttpRequest) -> std::io::Result { @@ -86,16 +91,20 @@ async fn get_tts(request: HttpRequest, //#[post("/tts")] async fn post_tts(request: HttpRequest, query: web::Json, - //query: web::Form, + //arpabet: web::Data>, model: web::Data>) -> std::io::Result { println!("POST /tts"); - let text = query.text.to_string(); println!("Text: {}", text); - let wav_data = model.run_tts_audio(&text); + let arpabet = Arpabet::load_cmudict(); + let encoded = text_to_arpabet_encoding(arpabet, &text); + + println!("Encoded: {:?}", encoded); + + let wav_data = model.run_tts_encoded(&encoded); Ok(HttpResponse::build(StatusCode::OK) .content_type("audio/wav") @@ -107,6 +116,11 @@ 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, +} + #[actix_rt::main] async fn main() -> std::io::Result<()> { println!("Starting service."); @@ -128,6 +142,7 @@ async fn main() -> std::io::Result<()> { let ttsEngine = TacoMelModel::create(&tacotron_filename, &melgan_filename); let arc = web::Data::new(Arc::new(ttsEngine)); + //let arc2 = web::Data::new(Arc::new(arpabet)); println!("Starting HTTP service."); println!("Asset directory: {}", asset_directory); @@ -135,30 +150,6 @@ async fn main() -> std::io::Result<()> { HttpServer::new(move || App::new() .wrap(Logger::default()) - // Configure CORS policy - /*.configure(|app| Cors::new() - .allowed_origin("https://localhost:12345/") - .allowed_origin("https://localhost:3000/") - .allowed_origin("https://localhost:8000/") - .allowed_methods(vec!["GET", "POST"]) - .allowed_headers(vec![http::header::AUTHORIZATION, http::header::ACCEPT]) - .allowed_header(http::header::CONTENT_TYPE) - .max_age(3600) - .resource("/tts", |r| { - r.method(http::Method::GET).f(|_| HttpResponse::Ok()); - r.method(http::Method::POST).f(|_| HttpResponse::Ok()); - }) - .register())*/ - /*.wrap( - Cors::new() - .allowed_origin("https://localhost:12345/") - .allowed_origin("https://localhost:3000/") - .allowed_origin("https://localhost:8000/") - .allowed_methods(vec!["GET", "POST"]) - .allowed_headers(vec![http::header::AUTHORIZATION, http::header::ACCEPT]) - .allowed_header(http::header::CONTENT_TYPE) - .max_age(3600) - .finish())*/ .service(Files::new("/frontend", asset_directory.clone()).show_files_listing()) .service( web::resource("/tts") @@ -169,9 +160,8 @@ async fn main() -> std::io::Result<()> { .service(get_root) .service(get_readiness) .service(get_liveness) - //.service(get_tts) - //.service(post_tts) .app_data(arc.clone()) + //.app_data(arc2.clone()) ) .bind(bind_address)? .run() diff --git a/tts_service/src/model.rs b/tts_service/src/model.rs index 6766024b5e..e3400131d8 100644 --- a/tts_service/src/model.rs +++ b/tts_service/src/model.rs @@ -48,29 +48,29 @@ impl TacoMelModel { println!("Text buffer: {:?}", text_buffer); + self.encoded_text_to_audio_signal(&text_buffer) + } + + 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); - println!("\n\n>>> Mel tensor:\n{:?}\n\n", mel_tensor); println!("Running melgan..."); - let audio_tensor = self.melgan_model.forward(&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) } diff --git a/tts_service/src/text.rs b/tts_service/src/text.rs new file mode 100644 index 0000000000..9c6f35475d --- /dev/null +++ b/tts_service/src/text.rs @@ -0,0 +1,60 @@ +use arpabet::Arpabet; +use arpabet::phoneme::{Phoneme,VowelStress,Vowel,Consonant}; +use arpabet::Polyphone; +use arpabet::extensions::Punctuation; +use arpabet::extensions::SentenceToken; + +/// Convert text to the encoding used in Nvidia/Tacotron2 natively. +pub fn text_to_standard_encoding(text: &str) -> Vec { + let copied = text.to_string().to_ascii_lowercase(); + let mut text_buffer : Vec = Vec::new(); + + for ch in copied.chars() { + // TODO: HORRIBLE EXPERIMENTAL HACK. + // Write a formal module to clean and process text + let mut v = ch as i64 - 59; + if v < 1 { + v = 11; // NB: Space + } + text_buffer.push(v); + } + + text_buffer +} + +/// Convert text to the encoding used in my Arpabet extension +pub fn text_to_arpabet_encoding(arpabet: &Arpabet, text: &str) -> Vec { + let copied = text.to_string().to_ascii_lowercase(); + let mut encoded_buffer : Vec = Vec::new(); + + let split = copied.split_whitespace(); + + let space = u8::from(Punctuation::Space) as i64; + let start= u8::from(Punctuation::StartToken) as i64; + let end = u8::from(Punctuation::EndToken) as i64; + + encoded_buffer.push(start); + + for word in split { + match arpabet.get_polyphone(word) { + None => continue, + Some(polyphone) => { + let vec : Vec = polyphone.iter() + .map(|phoneme| { + match phoneme { + Phoneme::Consonant(c) => u8::from(*c) as i64, + Phoneme::Vowel(v) => u8::from(*v) as i64, + } + }) + .collect(); + + encoded_buffer.extend(vec); + encoded_buffer.push(space); + }, + } + } + + encoded_buffer.push(end); + + encoded_buffer +}