First Version
This commit is contained in:
@@ -0,0 +1,13 @@
|
||||
[package]
|
||||
name = "plugin-tts-orpheus-3b"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
core-api = { path = "../core-api" }
|
||||
anyhow = "1"
|
||||
async-trait = "0.1"
|
||||
serde_json = "1"
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
tracing = "0.1"
|
||||
reqwest = { version = "0.13", default-features = false, features = ["rustls-no-provider", "charset", "http2", "system-proxy", "json"] }
|
||||
@@ -0,0 +1,314 @@
|
||||
//! Orpheus TTS 3B plugin.
|
||||
//!
|
||||
//! On start, writes the embedded `orpheus_server.py` (bundled via
|
||||
//! `include_str!`) to `models/orpheus-3b/`, spawns it as a subprocess, reads
|
||||
//! the bound port from its stdout, then registers itself as a [`TextToSpeech`]
|
||||
//! provider with the TTS manager.
|
||||
//!
|
||||
//! The subprocess loads the Orpheus 3B model from HuggingFace (auto-download
|
||||
//! on first run, cached in `models/orpheus-3b/`) and exposes a minimal HTTP
|
||||
//! server on a random OS-assigned port.
|
||||
//!
|
||||
//! # Required secret
|
||||
//!
|
||||
//! Set before enabling the plugin:
|
||||
//! ```
|
||||
//! set_secret("HUGGINGFACE_TOKEN", "hf_...")
|
||||
//! ```
|
||||
//! Get a token at <https://huggingface.co/settings/tokens>.
|
||||
//!
|
||||
//! # Config (stored in `plugins` SQLite table)
|
||||
//!
|
||||
//! ```json
|
||||
//! {
|
||||
//! "quantization": "int8",
|
||||
//! "voice": "tara"
|
||||
//! }
|
||||
//! ```
|
||||
//!
|
||||
//! | Field | Values | Default |
|
||||
//! |-------|--------|---------|
|
||||
//! | `quantization` | `"none"` \| `"int8"` \| `"int4"` | `"int8"` |
|
||||
//! | `voice` | `"tara"` \| `"dan"` \| `"leah"` \| `"zac"` \| `"zoe"` \| `"mia"` \| `"julia"` \| `"leo"` | `"tara"` |
|
||||
|
||||
use std::sync::{Arc, atomic::{AtomicBool, Ordering}};
|
||||
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{Value, json};
|
||||
use tokio::io::{AsyncBufReadExt, BufReader};
|
||||
use tokio::process::{Child, Command};
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::{info, warn};
|
||||
|
||||
const ORPHEUS_SERVER_PY: &str = include_str!("orpheus_server.py");
|
||||
|
||||
use core_api::plugin::{Plugin, PluginContext};
|
||||
use core_api::secrets;
|
||||
use core_api::tts::TextToSpeech;
|
||||
|
||||
const PLUGIN_ID: &str = "orpheus_tts_3b";
|
||||
const MODEL_DIR: &str = "models/orpheus-3b";
|
||||
const PROVIDER_ID: &str = "orpheus_tts_3b";
|
||||
const SERVER_PY_NAME: &str = "orpheus_server.py";
|
||||
|
||||
// ── Config ────────────────────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Clone, PartialEq, Debug)]
|
||||
struct OrpheusTtsConfig {
|
||||
quantization: String,
|
||||
voice: String,
|
||||
}
|
||||
|
||||
impl OrpheusTtsConfig {
|
||||
fn from_value(v: &Value) -> Self {
|
||||
Self {
|
||||
quantization: v["quantization"].as_str().unwrap_or("int8").to_string(),
|
||||
voice: v["voice"].as_str().unwrap_or("tara").to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── OrpheusSynthesiser ────────────────────────────────────────────────────────
|
||||
|
||||
/// Calls the local Orpheus Python server to synthesise audio.
|
||||
struct OrpheusSynthesiser {
|
||||
port: u16,
|
||||
default_voice: String,
|
||||
http: reqwest::Client,
|
||||
}
|
||||
|
||||
impl OrpheusSynthesiser {
|
||||
fn new(port: u16, default_voice: impl Into<String>) -> Self {
|
||||
Self {
|
||||
port,
|
||||
default_voice: default_voice.into(),
|
||||
http: reqwest::Client::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl TextToSpeech for OrpheusSynthesiser {
|
||||
fn id(&self) -> &str { PROVIDER_ID }
|
||||
fn name(&self) -> &str { "Orpheus TTS 3B" }
|
||||
fn description(&self) -> Option<&str> {
|
||||
Some("Local Orpheus TTS 3B — high-quality expressive speech, runs on-device.")
|
||||
}
|
||||
fn instructions(&self) -> Option<&str> {
|
||||
Some("\
|
||||
Orpheus TTS supports inline emotion tags. Insert them directly in the text where the effect should occur.\n\
|
||||
Supported tags: <laugh>, <chuckle>, <sigh>, <cough>, <sniffle>, <groan>, <yawn>, <gasp>\n\
|
||||
Example: \"I told him the meeting was at nine, not eleven. <sigh> He showed up at noon. <chuckle> Classic.\"\
|
||||
")
|
||||
}
|
||||
|
||||
async fn synthesize(&self, text: &str, instructions: Option<&str>) -> Result<Vec<u8>> {
|
||||
let voice = instructions
|
||||
.and_then(|s| s.split_whitespace().next()) // first word as voice override
|
||||
.unwrap_or(&self.default_voice);
|
||||
|
||||
let url = format!("http://127.0.0.1:{}/synthesize", self.port);
|
||||
|
||||
let resp = self.http
|
||||
.post(&url)
|
||||
.json(&json!({
|
||||
"text": text,
|
||||
"voice": voice,
|
||||
"instructions": instructions,
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| anyhow!("orpheus_tts: request failed: {e}"))?;
|
||||
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let msg = resp.text().await.unwrap_or_default();
|
||||
anyhow::bail!("orpheus_tts: server error {status}: {msg}");
|
||||
}
|
||||
|
||||
Ok(resp.bytes().await.map(|b| b.to_vec())
|
||||
.map_err(|e| anyhow!("orpheus_tts: failed to read bytes: {e}"))?)
|
||||
}
|
||||
}
|
||||
|
||||
// ── Plugin inner state ────────────────────────────────────────────────────────
|
||||
|
||||
struct Inner {
|
||||
child: Child,
|
||||
port: u16,
|
||||
config: OrpheusTtsConfig,
|
||||
script_path: std::path::PathBuf,
|
||||
}
|
||||
|
||||
// ── OrpheusTtsPlugin ──────────────────────────────────────────────────────────
|
||||
|
||||
pub struct OrpheusTtsPlugin {
|
||||
running: AtomicBool,
|
||||
inner: Mutex<Option<Inner>>,
|
||||
}
|
||||
|
||||
impl OrpheusTtsPlugin {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
running: AtomicBool::new(false),
|
||||
inner: Mutex::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
async fn do_start(&self, config: &OrpheusTtsConfig, ctx: &PluginContext) -> Result<()> {
|
||||
std::fs::create_dir_all(MODEL_DIR)
|
||||
.context("orpheus_tts: failed to create model dir")?;
|
||||
|
||||
// Write the embedded script to the model dir so it can be executed.
|
||||
let script_path = std::path::Path::new(MODEL_DIR).join(SERVER_PY_NAME);
|
||||
std::fs::write(&script_path, ORPHEUS_SERVER_PY)
|
||||
.context("orpheus_tts: failed to write embedded server script")?;
|
||||
|
||||
// HuggingFace token — required for gated repos. Passed as env var so
|
||||
// transformers/huggingface_hub pick it up automatically.
|
||||
let hf_token = secrets::require(&ctx.secrets, "HUGGINGFACE_TOKEN").await?;
|
||||
|
||||
let mut child = Command::new("python3")
|
||||
.args([
|
||||
script_path.to_str().unwrap(),
|
||||
"--model-dir", MODEL_DIR,
|
||||
"--quantization", &config.quantization,
|
||||
"--default-voice", &config.voice,
|
||||
])
|
||||
.env("HF_TOKEN", &hf_token)
|
||||
.stdout(std::process::Stdio::piped())
|
||||
.stderr(std::process::Stdio::inherit())
|
||||
.spawn()
|
||||
.context("orpheus_tts: failed to spawn python3")?;
|
||||
|
||||
// Read stdout until we see "PORT:<n>" — the server prints it once bound.
|
||||
let stdout = child.stdout.take()
|
||||
.ok_or_else(|| anyhow!("orpheus_tts: no stdout from subprocess"))?;
|
||||
let mut lines = BufReader::new(stdout).lines();
|
||||
let port = loop {
|
||||
match lines.next_line().await? {
|
||||
None => anyhow::bail!("orpheus_tts: subprocess exited before printing port"),
|
||||
Some(line) => {
|
||||
if let Some(p) = line.strip_prefix("PORT:") {
|
||||
break p.trim().parse::<u16>()
|
||||
.context("orpheus_tts: invalid port from subprocess")?;
|
||||
}
|
||||
// Forward other startup lines as info.
|
||||
info!("orpheus_tts(py): {line}");
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
info!(port, "orpheus_tts: python server ready");
|
||||
|
||||
let synthesiser = Arc::new(OrpheusSynthesiser::new(port, &config.voice));
|
||||
ctx.tts_registry.register(Arc::clone(&synthesiser) as _).await;
|
||||
|
||||
self.running.store(true, Ordering::Relaxed);
|
||||
*self.inner.lock().await = Some(Inner {
|
||||
child,
|
||||
port,
|
||||
config: config.clone(),
|
||||
script_path,
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn do_stop(&self, ctx: &PluginContext) {
|
||||
ctx.tts_registry.unregister(PROVIDER_ID).await;
|
||||
if let Some(mut inner) = self.inner.lock().await.take() {
|
||||
let _ = inner.child.kill().await;
|
||||
let _ = std::fs::remove_file(&inner.script_path);
|
||||
}
|
||||
self.running.store(false, Ordering::Relaxed);
|
||||
info!("orpheus_tts: stopped");
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Plugin for OrpheusTtsPlugin {
|
||||
fn id(&self) -> &str { PLUGIN_ID }
|
||||
fn name(&self) -> &str { "Orpheus TTS 3B" }
|
||||
fn description(&self) -> &str {
|
||||
"Local text-to-speech using Orpheus 3B. Expressive, high-quality, runs fully on-device. \
|
||||
Requires ~7 GB VRAM (fp16), ~4 GB (int8), or ~2.5 GB (int4). \
|
||||
Requires secret: HUGGINGFACE_TOKEN (HuggingFace access token — \
|
||||
get one at https://huggingface.co/settings/tokens, then call \
|
||||
set_secret(\"HUGGINGFACE_TOKEN\", \"hf_...\")). \
|
||||
See docs/tts-providers.md for full setup instructions."
|
||||
}
|
||||
fn is_running(&self) -> bool { self.running.load(Ordering::Relaxed) }
|
||||
|
||||
fn config_schema(&self) -> Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"quantization": {
|
||||
"type": "string",
|
||||
"enum": ["none", "int8", "int4"],
|
||||
"default": "int8",
|
||||
"description": "bitsandbytes precision: none=fp16 (~7 GB VRAM), int8 (~4 GB), int4 (~2.5 GB)"
|
||||
},
|
||||
"voice": {
|
||||
"type": "string",
|
||||
"enum": ["tara", "dan", "leah", "zac", "zoe", "mia", "julia", "leo"],
|
||||
"default": "tara",
|
||||
"description": "Default voice. Can be overridden per synthesis call via instructions."
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn reload(&self, enabled: bool, config: Value, ctx: PluginContext) -> Result<()> {
|
||||
let new_cfg = OrpheusTtsConfig::from_value(&config);
|
||||
let is_running = self.is_running();
|
||||
|
||||
let config_changed = self.inner.lock().await
|
||||
.as_ref()
|
||||
.map(|i| i.config != new_cfg)
|
||||
.unwrap_or(false);
|
||||
|
||||
match (enabled, is_running) {
|
||||
(true, false) => self.do_start(&new_cfg, &ctx).await?,
|
||||
(false, true) => self.do_stop(&ctx).await,
|
||||
(true, true) if config_changed => {
|
||||
info!("orpheus_tts: config changed — restarting");
|
||||
self.do_stop(&ctx).await;
|
||||
self.do_start(&new_cfg, &ctx).await?;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn start(&self, ctx: PluginContext) -> Result<()> {
|
||||
// start() is called by the plugin manager; reload() handles the normal path.
|
||||
// This is a no-op here — reload() does the real work.
|
||||
let _ = ctx;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn stop(&self) -> Result<()> {
|
||||
warn!("orpheus_tts: stop() called without ctx — cannot unregister from TtsManager");
|
||||
if let Some(mut inner) = self.inner.lock().await.take() {
|
||||
let _ = inner.child.kill().await;
|
||||
}
|
||||
self.running.store(false, Ordering::Relaxed);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn runtime_status(&self) -> Option<Value> {
|
||||
let inner = self.inner.try_lock().ok()?;
|
||||
let inner = inner.as_ref()?;
|
||||
Some(json!({
|
||||
"port": inner.port,
|
||||
"quantization": inner.config.quantization,
|
||||
"voice": inner.config.voice,
|
||||
}))
|
||||
}
|
||||
|
||||
fn as_any(&self) -> &dyn std::any::Any { self }
|
||||
fn as_arc_any(self: Arc<Self>) -> Arc<dyn std::any::Any + Send + Sync> { self }
|
||||
}
|
||||
@@ -0,0 +1,228 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Orpheus TTS 3B inference server.
|
||||
|
||||
Started by the plugin-tts-orpheus-3b Rust plugin. Prints "PORT:<n>" to stdout
|
||||
once the HTTP server is bound so the plugin knows which port to connect to.
|
||||
|
||||
The model is downloaded from HuggingFace on first run and cached in --model-dir.
|
||||
|
||||
Endpoints
|
||||
---------
|
||||
POST /synthesize
|
||||
Body: {"text": "...", "voice": "tara", "instructions": "..."}
|
||||
Returns: audio/wav bytes
|
||||
|
||||
GET /health
|
||||
Returns: {"status": "ok"}
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import socket
|
||||
import sys
|
||||
import threading
|
||||
|
||||
import numpy as np
|
||||
import scipy.io.wavfile as wavfile
|
||||
import torch
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.responses import Response
|
||||
from huggingface_hub import snapshot_download
|
||||
from pydantic import BaseModel
|
||||
from snac import SNAC
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
import uvicorn
|
||||
|
||||
# ── Model IDs ────────────────────────────────────────────────────────────────
|
||||
|
||||
ORPHEUS_MODEL_ID = "canopylabs/orpheus-3b-0.1-ft"
|
||||
SNAC_MODEL_ID = "hubertsiuzdak/snac_24khz"
|
||||
SAMPLE_RATE = 24000
|
||||
|
||||
VALID_VOICES = {"tara", "dan", "leah", "zac", "zoe", "mia", "julia", "leo"}
|
||||
|
||||
# ── Globals set at startup ────────────────────────────────────────────────────
|
||||
|
||||
model = None
|
||||
tokenizer = None
|
||||
snac_model = None
|
||||
default_voice = "tara"
|
||||
if torch.cuda.is_available():
|
||||
device = "cuda"
|
||||
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
device = "mps"
|
||||
# Some transformer ops are not yet implemented on MPS; fall back to CPU.
|
||||
os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
|
||||
else:
|
||||
device = "cpu"
|
||||
|
||||
# ── Model loading ─────────────────────────────────────────────────────────────
|
||||
|
||||
def load_model(model_dir: str, quantization: str) -> None:
|
||||
global model, tokenizer, snac_model
|
||||
|
||||
print(f"[orpheus] loading model (quantization={quantization}, device={device})", flush=True)
|
||||
|
||||
hf_cache = os.path.join(model_dir, "hf_cache")
|
||||
os.makedirs(hf_cache, exist_ok=True)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
ORPHEUS_MODEL_ID,
|
||||
cache_dir=hf_cache,
|
||||
)
|
||||
|
||||
load_kwargs: dict = {
|
||||
"cache_dir": hf_cache,
|
||||
"torch_dtype": torch.float16 if device in ("cuda", "mps") else torch.float32,
|
||||
"device_map": "auto" if device == "cuda" else None,
|
||||
"low_cpu_mem_usage": True,
|
||||
}
|
||||
|
||||
# bitsandbytes quantization is CUDA-only — skip on MPS and CPU.
|
||||
if device == "cuda":
|
||||
if quantization == "int8":
|
||||
load_kwargs["load_in_8bit"] = True
|
||||
elif quantization == "int4":
|
||||
load_kwargs["load_in_4bit"] = True
|
||||
elif quantization != "none":
|
||||
print(f"[orpheus] quantization '{quantization}' not supported on {device}, running fp16", flush=True)
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(ORPHEUS_MODEL_ID, **load_kwargs)
|
||||
if device != "cuda": # for cuda, device_map="auto" already handles placement
|
||||
model = model.to(device)
|
||||
model.eval()
|
||||
|
||||
snac_model = SNAC.from_pretrained(SNAC_MODEL_ID, cache_dir=hf_cache).to(device)
|
||||
snac_model.eval()
|
||||
|
||||
print("[orpheus] model loaded", flush=True)
|
||||
|
||||
|
||||
# ── Inference ─────────────────────────────────────────────────────────────────
|
||||
|
||||
def _tokens_to_audio(token_ids: list[int]) -> np.ndarray:
|
||||
"""Decode Orpheus audio token stream via SNAC to a float32 waveform."""
|
||||
# Orpheus uses a 7-level SNAC codec; tokens are interleaved in groups of 7.
|
||||
# Filter to valid audio token range (typically 128266–129290 for 24 kHz SNAC).
|
||||
audio_token_start = 128266
|
||||
audio_tokens = [t - audio_token_start for t in token_ids if t >= audio_token_start]
|
||||
|
||||
if len(audio_tokens) < 7:
|
||||
return np.zeros(0, dtype=np.float32)
|
||||
|
||||
# Trim to multiple of 7.
|
||||
n = (len(audio_tokens) // 7) * 7
|
||||
audio_tokens = audio_tokens[:n]
|
||||
|
||||
layers = [[] for _ in range(7)]
|
||||
for i, tok in enumerate(audio_tokens):
|
||||
layers[i % 7].append(tok)
|
||||
|
||||
with torch.no_grad():
|
||||
codes = [
|
||||
torch.tensor(layer, dtype=torch.long, device=device).unsqueeze(0)
|
||||
for layer in layers
|
||||
]
|
||||
audio = snac_model.decode(codes)
|
||||
|
||||
return audio.squeeze().cpu().float().numpy()
|
||||
|
||||
|
||||
def synthesize_text(text: str, voice: str, instructions: str | None) -> bytes:
|
||||
voice = voice if voice in VALID_VOICES else default_voice
|
||||
|
||||
# Build prompt in Orpheus format.
|
||||
prompt = f"<|audio|>{voice}: {text}<|eot_id|>"
|
||||
if instructions:
|
||||
prompt = f"<|audio|>{voice}: {text} [style: {instructions}]<|eot_id|>"
|
||||
|
||||
inputs = tokenizer(prompt, return_tensors="pt").to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
output_ids = model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=4096,
|
||||
do_sample=True,
|
||||
temperature=0.7,
|
||||
repetition_penalty=1.1,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
)
|
||||
|
||||
# Strip the prompt tokens; keep only newly generated tokens.
|
||||
new_tokens = output_ids[0][inputs["input_ids"].shape[1]:].tolist()
|
||||
waveform = _tokens_to_audio(new_tokens)
|
||||
|
||||
if waveform.size == 0:
|
||||
raise RuntimeError("orpheus: decoding produced no audio samples")
|
||||
|
||||
# Encode to 16-bit WAV in memory.
|
||||
pcm = (waveform * 32767).astype(np.int16)
|
||||
buf = io.BytesIO()
|
||||
wavfile.write(buf, SAMPLE_RATE, pcm)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
# ── FastAPI app ───────────────────────────────────────────────────────────────
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
|
||||
class SynthesizeRequest(BaseModel):
|
||||
text: str
|
||||
voice: str | None = None
|
||||
instructions: str | None = None
|
||||
|
||||
|
||||
@app.post("/synthesize")
|
||||
def synthesize(req: SynthesizeRequest):
|
||||
if not req.text.strip():
|
||||
raise HTTPException(status_code=400, detail="text is empty")
|
||||
try:
|
||||
audio = synthesize_text(
|
||||
req.text,
|
||||
req.voice or default_voice,
|
||||
req.instructions,
|
||||
)
|
||||
return Response(content=audio, media_type="audio/wav")
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health():
|
||||
return {"status": "ok"}
|
||||
|
||||
|
||||
# ── Entry point ───────────────────────────────────────────────────────────────
|
||||
|
||||
def main() -> None:
|
||||
global default_voice
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model-dir", default="models/orpheus-3b")
|
||||
parser.add_argument("--quantization", default="int8", choices=["none", "int8", "int4"])
|
||||
parser.add_argument("--default-voice", default="tara")
|
||||
args = parser.parse_args()
|
||||
|
||||
default_voice = args.default_voice
|
||||
|
||||
load_model(args.model_dir, args.quantization)
|
||||
|
||||
# Bind on port 0 — OS assigns a free port.
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
sock.bind(("127.0.0.1", 0))
|
||||
port = sock.getsockname()[1]
|
||||
sock.close()
|
||||
|
||||
# Print port for the Rust plugin to read.
|
||||
print(f"PORT:{port}", flush=True)
|
||||
|
||||
config = uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning")
|
||||
server = uvicorn.Server(config)
|
||||
server.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user