97 lines
3.3 KiB
Rust
97 lines
3.3 KiB
Rust
use axum::{Router, routing::post};
|
|
use std::sync::Arc;
|
|
use std::{collections::HashMap, path::PathBuf};
|
|
|
|
pub mod api;
|
|
pub mod tts;
|
|
|
|
use tts::{TtsOpts, TtsPool, load_text_to_speech};
|
|
|
|
use crate::tts::engine::load_voice_style_map;
|
|
|
|
#[tokio::main(flavor = "multi_thread")]
|
|
async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
|
dotenvy::dotenv().ok();
|
|
tracing_subscriber::fmt::init();
|
|
|
|
let model_dir =
|
|
std::env::var("SUPERTONIC_MODEL_DIR").unwrap_or_else(|_| "./assets".to_string());
|
|
let voice_style_path = std::env::var("SUPERTONIC_VOICE_STYLE")
|
|
.unwrap_or_else(|_| format!("F1={model_dir}/voice_styles/F1.json"));
|
|
let _lang = std::env::var("SUPERTONIC_LANG").unwrap_or_else(|_| "en".to_string());
|
|
let total_step: usize = std::env::var("SUPERTONIC_TOTAL_STEP")
|
|
.ok()
|
|
.and_then(|v| v.parse().ok())
|
|
.unwrap_or(8);
|
|
let speed: f32 = std::env::var("SUPERTONIC_SPEED")
|
|
.ok()
|
|
.and_then(|v| v.parse().ok())
|
|
.unwrap_or(1.05);
|
|
let silence_before_duration: f32 = std::env::var("SUPERTONIC_SILENCE_BEFORE_DUR")
|
|
.ok()
|
|
.and_then(|v| v.parse().ok())
|
|
.unwrap_or(0.0);
|
|
let silence_in_content_duration: f32 = std::env::var("SUPERTONIC_SILENCE_IN_CONTENT_DUR")
|
|
.ok()
|
|
.and_then(|v| v.parse().ok())
|
|
.unwrap_or(0.26);
|
|
let silence_after_duration: f32 = std::env::var("SUPERTONIC_SILENCE_AFTER_DUR")
|
|
.ok()
|
|
.and_then(|v| v.parse().ok())
|
|
.unwrap_or(0.4);
|
|
let workers: usize = std::env::var("SUPERTONIC_WORKERS")
|
|
.ok()
|
|
.and_then(|v| v.parse().ok())
|
|
.unwrap_or(2);
|
|
let hf_repo = std::env::var("SUPERTONIC_HF_REPO")
|
|
.unwrap_or_else(|_| "https://huggingface.co/Supertone/supertonic-3".to_string());
|
|
let sound_size_mul: f32 = std::env::var("SUPERTONIC_SIZE_MUL")
|
|
.ok()
|
|
.and_then(|v| v.parse().ok())
|
|
.unwrap_or(1.5);
|
|
|
|
let model_path = PathBuf::from(&model_dir);
|
|
tts::assets::ensure_assets(&model_path, &hf_repo)?;
|
|
|
|
let onnx_dir_for_init = model_path.join("onnx").to_string_lossy().into_owned();
|
|
let voice_style_for_init = voice_style_path
|
|
.split(",")
|
|
.filter_map(|i| {
|
|
i.split_once("=").or_else(|| {
|
|
tracing::error!("Voice style '{i}' is not valid.");
|
|
None
|
|
})
|
|
})
|
|
.map(|(k, v)| (k.to_owned(), PathBuf::from(v)))
|
|
.collect::<HashMap<_, _>>();
|
|
let voice_style_map = Arc::new(load_voice_style_map(&voice_style_for_init)?);
|
|
|
|
let pool = Arc::new(TtsPool::spawn(
|
|
workers,
|
|
move |id| {
|
|
let span = tracing::info_span!("worker", worker_id = id);
|
|
let _enter = span.enter();
|
|
let tts = load_text_to_speech(&onnx_dir_for_init)?;
|
|
|
|
Ok((tts, voice_style_map.clone()))
|
|
},
|
|
TtsOpts {
|
|
total_step,
|
|
speed,
|
|
silence_in_content_duration,
|
|
silence_before_duration,
|
|
silence_after_duration,
|
|
sound_size_mul,
|
|
},
|
|
)?);
|
|
|
|
let app = Router::new()
|
|
.route("/", post(api::handler))
|
|
.with_state(pool);
|
|
|
|
let addr = std::env::var("ADDR").unwrap_or_else(|_| "0.0.0.0:80".to_string());
|
|
let listener = tokio::net::TcpListener::bind(addr).await?;
|
|
axum::serve(listener, app).await?;
|
|
|
|
Ok(())
|
|
}
|