valenti/vocero-client
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

first

Browse Source
This commit is contained in:
valenti
2026-08-26 17:42:00 -04:00
commit ad95d530ad
134 changed files with 31750 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
speech_to_speech/TTS/README.md
+155
View File
@@ -0,0 +1,155 @@
# TTS Summary
## Available TTS modes (`--tts`)
Runtime-supported values in `s2s_pipeline.py`:
- `chatTTS` → `chatTTS_handler.py`
- `facebookMMS` → `facebookmms_handler.py`
- `pocket` → `pocket_tts_handler.py`
- `kokoro` → `kokoro_handler.py`
- `qwen3` → `qwen3_tts_handler.py`
Deprecated TTS implementations, including MeloTTS, live in [`../../../archive/TTS`](../../../archive/TTS) and are no longer wired into `s2s_pipeline.py`.
## Usage
### 1) ChatTTS (`--tts chatTTS`)
Primary args prefix: `--chat_tts_*`
```bash
python s2s_pipeline.py \
--tts chatTTS \
--chat_tts_device cuda \
--chat_tts_stream true \
--chat_tts_chunk_size 512
```
### 2) Facebook MMS (`--tts facebookMMS`)
Primary args prefix: `--facebook_mms_*` plus `--tts_language`
```bash
python s2s_pipeline.py \
--tts facebookMMS \
--facebook_mms_device cuda \
--tts_language en
```
This handler maps STT language codes (e.g. `en`, `fr`, `es`) to MMS model suffixes (e.g. `eng`, `fra`, `spa`) and reloads the model on language changes.
### 3) Pocket TTS (`--tts pocket`)
Primary args prefix: `--pocket_tts_*`
```bash
python s2s_pipeline.py \
--tts pocket \
--pocket_tts_voice jean \
--pocket_tts_device cpu \
--pocket_tts_sample_rate 16000
```
Available preset voices include:
`alba`, `marius`, `javert`, `jean`, `fantine`, `cosette`, `eponine`, `azelma`.
### 4) Kokoro (`--tts kokoro`)
Primary args prefix: `--kokoro_*`
```bash
python s2s_pipeline.py \
--tts kokoro \
--kokoro_device auto \
--kokoro_voice bm_fable \
--kokoro_lang_code b
```
Behavior:
- Uses MLX backend on Apple Silicon (`mlx-community/Kokoro-82M-bf16`)
- Uses native kokoro pipeline otherwise (`hexgrad/Kokoro-82M`)
- Can auto-switch voice/language based on STT language code mapping
### 5) Qwen3-TTS (`--tts qwen3`)
Primary args prefix: `--qwen3_tts_*`
```bash
python s2s_pipeline.py \
--tts qwen3 \
--qwen3_tts_model_name Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
--qwen3_tts_device cuda \
--qwen3_tts_backend ggml \
--qwen3_tts_speaker Aiden \
--qwen3_tts_language auto \
--qwen3_tts_non_streaming_mode True
```
Behavior:
- Uses `faster-qwen3-tts` on non-macOS platforms, defaulting to the GGML backend. Pass `--qwen3_tts_backend torch` to use the CUDA-graphs backend instead.
- Uses `mlx-audio` on Apple Silicon and auto-maps `Qwen/...` model IDs to `mlx-community/...`, defaulting to the `6bit` MLX variant unless the model name already pins a suffix.
- Supports MLX quantization overrides on Apple Silicon via `--qwen3_tts_mlx_quantization bf16|4bit|6bit|8bit`.
- Keeps the existing voice-clone/custom-voice/voice-design handler flow intact.
- Defaults to the CustomVoice model with speaker `Aiden`, so no reference audio is required. Voice-clone/base models can still use `--qwen3_tts_ref_audio`.
Install notes for Linux GGML:
- The default PyPI `qwentts-cpp-python` wheel targets CUDA 12.8.
- If that wheel does not match your CUDA runtime, install one of the Hugging Face wheelhouse builds before installing `speech-to-speech`.
```bash
pip install "qwentts-cpp-python==0.3.1+cu130" \
-f https://huggingface.co/datasets/andito/qwentts-cpp-python-wheels/tree/main/whl/cu130
pip install speech-to-speech
```
Available wheelhouse directories include `cu124`, `cu128`, `cu130`, and `cpu`.
Example for Apple Silicon using the default 6-bit MLX variant:
```bash
python s2s_pipeline.py \
--tts qwen3 \
--qwen3_tts_model_name Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
--qwen3_tts_speaker Aiden
```
You can override the default and select `bf16`, `4bit`, or `8bit` explicitly:
```bash
python s2s_pipeline.py \
--tts qwen3 \
--qwen3_tts_model_name Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
--qwen3_tts_mlx_quantization 4bit \
--qwen3_tts_speaker Aiden
```
To benchmark the Apple Silicon MLX variants side by side:
```bash
.venv/bin/python benchmark_tts.py \
--handlers qwen3 \
--iterations 3 \
--qwen3_mlx_quantizations bf16 4bit 6bit 8bit
```
This will run separate benchmark entries for `qwen3[bf16]`, `qwen3[4bit]`, `qwen3[6bit]`, and `qwen3[8bit]`.
## Setup
### Low-latency GPU setup
```bash
python s2s_pipeline.py \
--stt whisper \
--tts pocket
```
### Apple Silicon setup
```bash
python s2s_pipeline.py \
--local_mac_optimal_settings
```
`--tts pocket` and `--tts kokoro` are also valid options on macOS.
speech_to_speech/TTS/__init__.py
+1
View File
@@ -0,0 +1 @@
# TTS module
speech_to_speech/TTS/chatTTS_handler.py
+119
View File
@@ -0,0 +1,119 @@
from __future__ import annotations
import logging
from threading import Event
from typing import Any, Iterator
import ChatTTS
import librosa
import numpy as np
import torch
from rich.console import Console
from speech_to_speech.baseHandler import BaseHandler
from speech_to_speech.pipeline.cancel_scope import CancelScope
from speech_to_speech.pipeline.handler_types import TTSIn, TTSOut
from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, EndOfResponse
from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker
logging.basicConfig(
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)
console = Console()
class ChatTTSHandler(BaseHandler[TTSIn, TTSOut]):
def setup(
self,
should_listen: Event,
device: str = "cuda",
gen_kwargs: dict[str, Any] = {}, # Unused
stream: bool = True,
chunk_size: int = 512,
cancel_scope: CancelScope | None = None,
speculative_turns: SpeculativeTurnTracker | None = None,
) -> None:
self.should_listen = should_listen
self.cancel_scope = cancel_scope
self.speculative_turns = speculative_turns
self.device = device
self.model = ChatTTS.Chat()
self.model.load(compile=False) # Doesn't work for me with True
self.chunk_size = chunk_size
self.stream = stream
rnd_spk_emb = self.model.sample_random_speaker()
self.params_infer_code = ChatTTS.Chat.InferCodeParams(
spk_emb=rnd_spk_emb,
)
self.warmup()
def warmup(self) -> None:
logger.info(f"Warming up {self.__class__.__name__}")
_ = self.model.infer("text")
def process(self, tts_input: TTSIn) -> Iterator[TTSOut]:
speculative_turns = getattr(self, "speculative_turns", None)
if isinstance(tts_input, EndOfResponse):
if speculative_turns and not speculative_turns.is_latest_after_reopen_grace(
tts_input.turn_id,
tts_input.turn_revision,
):
return
yield AUDIO_RESPONSE_DONE
return
if speculative_turns and not speculative_turns.is_latest_after_reopen_grace(
tts_input.turn_id,
tts_input.turn_revision,
):
logger.debug("Dropping stale TTS input for turn=%s rev=%s", tts_input.turn_id, tts_input.turn_revision)
return
if speculative_turns:
speculative_turns.commit(tts_input.turn_id, tts_input.turn_revision)
text = tts_input.text
_cancel_gen = self.cancel_scope.generation if self.cancel_scope else None
console.print(f"[green]ASSISTANT: {text}")
if self.device == "mps":
import time
start = time.time()
torch.mps.synchronize() # Waits for all kernels in all streams on the MPS device to complete.
torch.mps.empty_cache() # Frees all memory allocated by the MPS device.
_ = time.time() - start # Removing this line makes it fail more often. I'm looking into it.
wavs_gen = self.model.infer(text, params_infer_code=self.params_infer_code, stream=self.stream)
if self.stream:
wavs = [np.array([])]
for gen in wavs_gen:
if (
_cancel_gen is not None
and self.cancel_scope is not None
and self.cancel_scope.is_stale(_cancel_gen)
):
logger.info("TTS generation cancelled (interruption)")
return
if gen[0] is None or len(gen[0]) == 0:
return
audio_chunk = librosa.resample(gen[0], orig_sr=24000, target_sr=16000)
audio_chunk = (audio_chunk * 32768).astype(np.int16)[0]
while len(audio_chunk) > self.chunk_size:
yield audio_chunk[: self.chunk_size] # Return the first chunk_size samples of the audio data
audio_chunk = audio_chunk[self.chunk_size :] # Remove the samples that have already been returned
yield np.pad(audio_chunk, (0, self.chunk_size - len(audio_chunk)))
else:
wavs = wavs_gen
if len(wavs[0]) == 0:
return
audio_chunk = librosa.resample(wavs[0], orig_sr=24000, target_sr=16000)
audio_chunk = (audio_chunk * 32768).astype(np.int16)
for i in range(0, len(audio_chunk), self.chunk_size):
yield np.pad(
audio_chunk[i : i + self.chunk_size],
(0, self.chunk_size - len(audio_chunk[i : i + self.chunk_size])),
)
speech_to_speech/TTS/facebookmms_handler.py
+212
View File
@@ -0,0 +1,212 @@
from __future__ import annotations
import logging
from threading import Event
from typing import Any, Iterator
import librosa
import numpy as np
import torch
from rich.console import Console
from transformers import AutoTokenizer, VitsModel
from speech_to_speech.baseHandler import BaseHandler
from speech_to_speech.pipeline.cancel_scope import CancelScope
from speech_to_speech.pipeline.handler_types import TTSIn, TTSOut
from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, EndOfResponse
from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker
logging.basicConfig(format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.DEBUG)
logger = logging.getLogger(__name__)
console = Console()
WHISPER_LANGUAGE_TO_FACEBOOK_LANGUAGE = {
"en": "eng", # English
"fr": "fra", # French
"es": "spa", # Spanish
"ko": "kor", # Korean
"hi": "hin", # Hindi
"ar": "ara", # Arabic
"hy": "hyw", # Armenian
"az": "azb", # Azerbaijani
"bu": "bul", # Bulgarian
"ca": "cat", # Catalan
"nl": "nld", # Dutch
"fi": "fin", # Finnish
"de": "deu", # German
"el": "ell", # Greek
"he": "heb", # Hebrew
"hu": "hun", # Hungarian
"is": "isl", # Icelandic
"id": "ind", # Indonesian
"ka": "kan", # Kannada
"kk": "kaz", # Kazakh
"lv": "lav", # Latvian
"zl": "zlm", # Malay
"ma": "mar", # Marathi
"fa": "fas", # Persian
"po": "pol", # Polish
"pt": "por", # Portuguese
"ro": "ron", # Romanian
"ru": "rus", # Russian
"sw": "swh", # Swahili
"sv": "swe", # Swedish
"tg": "tgl", # Tagalog
"ta": "tam", # Tamil
"th": "tha", # Thai
"tu": "tur", # Turkish
"uk": "ukr", # Ukrainian
"ur": "urd", # Urdu
"vi": "vie", # Vietnamese
"cy": "cym", # Welsh
}
class FacebookMMSTTSHandler(BaseHandler[TTSIn, TTSOut]):
def setup(
self,
should_listen: Event,
device: str = "cuda",
torch_dtype: str = "float32",
language: str = "en",
stream: bool = True,
chunk_size: int = 512,
cancel_scope: CancelScope | None = None,
speculative_turns: SpeculativeTurnTracker | None = None,
**kwargs: Any,
) -> None:
self.should_listen = should_listen
self.cancel_scope = cancel_scope
self.speculative_turns = speculative_turns
self.device = device
self.torch_dtype = getattr(torch, torch_dtype)
self.stream = stream
self.chunk_size = chunk_size
self.language = language
self._initial_language = self.language
self.load_model(self.language)
self.warmup()
def load_model(self, language_code: str) -> None:
try:
model_name = f"facebook/mms-tts-{WHISPER_LANGUAGE_TO_FACEBOOK_LANGUAGE[language_code]}"
logger.info(f"Loading model: {model_name}")
self.model = VitsModel.from_pretrained(model_name).to(self.device) # type: ignore[arg-type]
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.language = language_code
except KeyError:
logger.warning(f"Unsupported language: {language_code}. Falling back to English.")
self.load_model("en")
def warmup(self) -> None:
logger.info(f"Warming up {self.__class__.__name__}")
self.generate_audio("Hello, this is a test")
def generate_audio(self, text: str) -> torch.Tensor | None:
if not text:
logger.warning("Received empty text input")
return None
try:
logger.debug(f"Tokenizing text: {text}")
logger.debug(f"Current language: {self.language}")
logger.debug(f"Tokenizer: {self.tokenizer}")
inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
input_ids = inputs.input_ids.to(self.device).long()
attention_mask = inputs.attention_mask.to(self.device)
logger.debug(f"Input IDs shape: {input_ids.shape}, dtype: {input_ids.dtype}")
logger.debug(f"Input IDs: {input_ids}")
if input_ids.numel() == 0:
logger.error("Input IDs tensor is empty")
return None
with torch.no_grad():
output = self.model(input_ids=input_ids, attention_mask=attention_mask)
logger.debug(f"Output waveform shape: {output.waveform.shape}")
return output.waveform
except Exception as e:
logger.error(f"Error in generate_audio: {str(e)}")
logger.exception("Full traceback:")
return None
def process(self, tts_input: TTSIn) -> Iterator[TTSOut]:
speculative_turns = getattr(self, "speculative_turns", None)
if isinstance(tts_input, EndOfResponse):
if speculative_turns and not speculative_turns.is_latest_after_reopen_grace(
tts_input.turn_id,
tts_input.turn_revision,
):
return
yield AUDIO_RESPONSE_DONE
return
if speculative_turns and not speculative_turns.is_latest_after_reopen_grace(
tts_input.turn_id,
tts_input.turn_revision,
):
logger.debug("Dropping stale TTS input for turn=%s rev=%s", tts_input.turn_id, tts_input.turn_revision)
return
if speculative_turns:
speculative_turns.commit(tts_input.turn_id, tts_input.turn_revision)
gen = self.cancel_scope.generation if self.cancel_scope else None
language_code = tts_input.language_code
text = tts_input.text
console.print(f"[green]ASSISTANT: {text}")
logger.debug(f"Processing text: {text}")
logger.debug(f"Language code: {language_code}")
if language_code is not None and self.language != language_code:
try:
logger.info(f"Switching language from {self.language} to {language_code}")
self.load_model(language_code)
except KeyError:
console.print(
f"[red]Language {language_code} not supported by Facebook MMS. Using {self.language} instead."
)
logger.warning(f"Unsupported language: {language_code}")
audio_output = self.generate_audio(text)
if audio_output is None or audio_output.numel() == 0:
logger.warning("No audio output generated")
return
audio_numpy = audio_output.cpu().numpy().squeeze()
logger.debug(f"Raw audio shape: {audio_numpy.shape}, dtype: {audio_numpy.dtype}")
audio_resampled = librosa.resample(audio_numpy, orig_sr=self.model.config.sampling_rate, target_sr=16000)
logger.debug(f"Resampled audio shape: {audio_resampled.shape}, dtype: {audio_resampled.dtype}")
audio_int16 = (audio_resampled * 32768).astype(np.int16)
logger.debug(f"Final audio shape: {audio_int16.shape}, dtype: {audio_int16.dtype}")
if self.stream:
for i in range(0, len(audio_int16), self.chunk_size):
if gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(gen):
logger.info("TTS generation cancelled (interruption)")
return
chunk = audio_int16[i : i + self.chunk_size]
yield np.pad(chunk, (0, self.chunk_size - len(chunk)))
else:
for i in range(0, len(audio_int16), self.chunk_size):
if gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(gen):
logger.info("TTS generation cancelled (interruption)")
return
yield np.pad(
audio_int16[i : i + self.chunk_size],
(0, self.chunk_size - len(audio_int16[i : i + self.chunk_size])),
)
def on_session_end(self) -> None:
if self.language != self._initial_language:
self.load_model(self._initial_language)
logger.debug("Facebook MMS TTS session state reset")
speech_to_speech/TTS/kokoro_handler.py
+418
View File
@@ -0,0 +1,418 @@
"""
Kokoro TTS Handler
Supports NVIDIA Kokoro TTS model for high-quality multilingual speech synthesis.
- On Apple Silicon (MPS): Uses mlx-audio with mlx-community/Kokoro-82M-bf16
- On CUDA/CPU: Uses native kokoro library with hexgrad/Kokoro-82M
Model supports 8 languages with multiple voices per language.
"""
from __future__ import annotations
import logging
from sys import platform
from threading import Event
from typing import Any, Iterator, Optional
import numpy as np
from rich.console import Console
from speech_to_speech.baseHandler import BaseHandler
from speech_to_speech.pipeline.cancel_scope import CancelScope
from speech_to_speech.pipeline.handler_types import TTSIn, TTSOut
from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, EndOfResponse
from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker
from speech_to_speech.utils.mlx_lock import MLXLockContext
logger = logging.getLogger(__name__)
console = Console()
# Language code mapping from Whisper/langdetect language codes to Kokoro lang codes
WHISPER_LANGUAGE_TO_KOKORO_LANG = {
"en": "b", # British English
"ja": "j", # Japanese
"zh": "z", # Chinese
"fr": "f", # French
"es": "e", # Spanish
"it": "i", # Italian
"pt": "p", # Portuguese
"hi": "h", # Hindi
# Additional European languages from Parakeet - map to closest Kokoro language
"de": "b", # German -> British English (no German voice available)
"nl": "b", # Dutch -> British English
"pl": "b", # Polish -> British English
"ru": "b", # Russian -> British English
"uk": "b", # Ukrainian -> British English
}
# Default voices for each Kokoro language code
# These are native voices that sound natural for each language
# Voice naming: first letter = language, second letter = gender (f=female, m=male)
# Available voices per language:
# a (American): af_alloy, af_aoede, af_bella, af_heart, af_jessica, af_kore, af_nicole, af_nova, af_river, af_sarah, af_sky
# am_adam, am_echo, am_eric, am_fenrir, am_liam, am_michael, am_onyx, am_puck, am_santa
# b (British): bf_alice, bf_emma, bf_isabella, bf_lily, bm_daniel, bm_fable, bm_george, bm_lewis
# e (Spanish): ef_dora, em_alex, em_santa
# f (French): ff_siwis
# h (Hindi): hf_alpha, hf_beta, hm_omega, hm_psi
# i (Italian): if_sara, im_nicola
# j (Japanese): jf_alpha, jf_gongitsune, jf_nezumi, jf_tebukuro, jm_kumo
# p (Portuguese): pf_dora, pm_alex, pm_santa
# z (Chinese): zf_xiaobei, zf_xiaoni, zf_xiaoxiao, zf_xiaoyi, zm_yunjian, zm_yunxi, zm_yunxia, zm_yunyang
KOKORO_LANG_DEFAULT_VOICES = {
"a": "af_heart", # American English female
"b": "bm_fable", # British English male
"e": "ef_dora", # Spanish female
"f": "ff_siwis", # French female
"h": "hf_alpha", # Hindi female
"i": "if_sara", # Italian female
"j": "jf_alpha", # Japanese female
"p": "pf_dora", # Portuguese female
"z": "zf_xiaobei", # Chinese female
}
class KokoroTTSHandler(BaseHandler[TTSIn, TTSOut]):
"""
Handles Text-to-Speech using Kokoro TTS model.
On Apple Silicon (MPS): Uses mlx-audio with the MLX-converted model.
On CUDA/CPU: Uses native kokoro library for optimal performance.
Kokoro 82M is a 82M parameter multilingual TTS model
supporting 8 languages with multiple voices per language.
"""
def setup(
self,
should_listen: Event,
model_name: Optional[str] = None,
device: str = "auto",
voice: str = "bm_fable",
lang_code: str = "b",
speed: float = 1.0,
blocksize: int = 512,
gen_kwargs: dict[str, Any] | None = None,
cancel_scope: CancelScope | None = None,
speculative_turns: SpeculativeTurnTracker | None = None,
) -> None:
"""
Initialize the Kokoro TTS model.
Args:
model_name: Model identifier. Defaults are:
- MPS: "mlx-community/Kokoro-82M-bf16"
- CUDA/CPU: "hexgrad/Kokoro-82M"
device: Device to use ("auto", "cuda", "mps", "cpu")
voice: Voice identifier (e.g. "bm_fable")
lang_code: Language code (e.g. "b" for British English)
speed: Speech speed multiplier
blocksize: Audio chunk size for streaming
gen_kwargs: Unused, for pipeline compatibility
"""
self.should_listen = should_listen
self.voice = voice
self.lang_code = lang_code
self.speed = speed
self.blocksize = blocksize
self.cancel_scope = cancel_scope
self.speculative_turns = speculative_turns
# Determine device
if device == "auto":
if platform == "darwin":
self.device = "mps"
else:
import torch
self.device = "cuda" if torch.cuda.is_available() else "cpu"
else:
self.device = device
# Set default model based on device
if model_name is None:
if self.device == "mps":
model_name = "mlx-community/Kokoro-82M-bf16"
else:
model_name = "hexgrad/Kokoro-82M"
self.model_name = model_name
logger.info(f"Loading Kokoro model: {model_name} on {self.device}")
if self.device == "mps":
self._setup_mlx(model_name)
else:
self._setup_kokoro(model_name)
self._initial_voice = self.voice
self._initial_lang_code = self.lang_code
self.warmup()
def _setup_mlx(self, model_name: str) -> None:
"""Setup for Apple Silicon using mlx-audio."""
try:
from mlx_audio.tts.utils import load_model
self.backend = "mlx"
self.model = load_model(model_name)
logger.info("MLX Audio Kokoro model loaded successfully")
# Get or create the pipeline for our language and preload the voice
# This avoids the voice being reloaded on every generate() call
self._pipeline = self.model._get_pipeline(self.lang_code)
self._voice_tensor = self._pipeline.load_voice(self.voice)
logger.info(f"Preloaded voice: {self.voice}")
# Preload voices for common languages to avoid download delays during inference
self._preload_multilingual_voices()
except ImportError as e:
message = str(e)
if "misaki" in message or "espeakng_loader" in message:
raise ImportError(
"Kokoro TTS on Apple Silicon requires additional mlx-audio TTS dependencies. "
f"Missing dependency: {message}. "
"Install with: pip install misaki espeakng-loader"
) from e
raise ImportError(
"mlx-audio is required for Kokoro TTS on Apple Silicon. Install with: pip install mlx-audio"
) from e
def _setup_kokoro(self, model_name: str) -> None:
"""Setup for CUDA/CPU using native kokoro library."""
try:
from kokoro import KPipeline
self.backend = "kokoro"
self.pipeline = KPipeline(lang_code=self.lang_code)
logger.info("Native Kokoro pipeline loaded successfully")
except ImportError as e:
raise ImportError(
"kokoro is required for Kokoro TTS on CUDA/CPU. "
"Install with: pip install kokoro>=0.9.2 soundfile\n"
"Also ensure espeak-ng is installed: apt-get install espeak-ng (Linux) or brew install espeak-ng (macOS)"
) from e
def _preload_multilingual_voices(self) -> None:
"""Preload voices for common languages to avoid download delays during inference."""
# Only preload a few commonly used language voices to avoid excessive startup time
# Users speaking other languages will experience a one-time download delay
# preload_langs = ["a", "b", "e", "f", "h", "i", "j", "p", "z"] # English, Spanish, French, Hindi, Italian, Japanese, Portuguese
preload_langs = ["a", "e", "f"] # English, Spanish, French
for lang_code in preload_langs:
voice = KOKORO_LANG_DEFAULT_VOICES.get(lang_code)
if voice:
try:
logger.info(f"Preloading voice for language '{lang_code}': {voice}")
pipeline = self.model._get_pipeline(lang_code)
pipeline.load_voice(voice)
except Exception as e:
logger.warning(f"Failed to preload voice {voice}: {e}")
def warmup(self) -> None:
"""Warm up the model with a dummy inference."""
logger.info(f"Warming up {self.__class__.__name__}")
# Run a short dummy inference to warm up the model
if self.backend == "mlx":
for _ in self.model.generate(
text="Hello",
voice=self.voice,
speed=self.speed,
lang_code=self.lang_code,
):
pass
else:
for _ in self.pipeline("Hello", voice=self.voice, speed=self.speed):
pass
logger.info(f"{self.__class__.__name__} warmed up")
def process(self, tts_input: TTSIn) -> Iterator[TTSOut]:
"""
Process text input and generate audio output.
Yields:
Audio chunks as numpy int16 arrays
"""
speculative_turns = getattr(self, "speculative_turns", None)
if isinstance(tts_input, EndOfResponse):
if speculative_turns and not speculative_turns.is_latest_after_reopen_grace(
tts_input.turn_id,
tts_input.turn_revision,
):
return
yield AUDIO_RESPONSE_DONE
return
if speculative_turns and not speculative_turns.is_latest_after_reopen_grace(
tts_input.turn_id,
tts_input.turn_revision,
):
logger.debug("Dropping stale TTS input for turn=%s rev=%s", tts_input.turn_id, tts_input.turn_revision)
return
if speculative_turns:
speculative_turns.commit(tts_input.turn_id, tts_input.turn_revision)
runtime_config = tts_input.runtime_config
response = tts_input.response
language_code = tts_input.language_code
text = tts_input.text
voice: Optional[str] = None
if response and response.audio and response.audio.output and response.audio.output.voice:
voice = str(response.audio.output.voice)
if not voice and runtime_config:
audio_cfg = runtime_config.session.audio
audio_output = audio_cfg.output if audio_cfg is not None else None
voice = str(audio_output.voice) if audio_output is not None and audio_output.voice else None
if voice:
self.voice = voice
if self.backend == "mlx":
yield from self._process_mlx(text, language_code)
else:
yield from self._process_kokoro(text, language_code)
def _process_mlx(self, llm_sentence: str, language_code: Optional[str] = None) -> Iterator[np.ndarray]:
"""Process using MLX backend with Apple Silicon optimizations."""
from scipy.signal import resample_poly
gen = self.cancel_scope.generation if self.cancel_scope else None
with MLXLockContext(handler_name="KokoroTTS", timeout=10.0):
if language_code is not None:
new_lang_code = WHISPER_LANGUAGE_TO_KOKORO_LANG.get(language_code, self.lang_code)
if new_lang_code != self.lang_code:
new_voice = KOKORO_LANG_DEFAULT_VOICES.get(new_lang_code, self.voice)
logger.info(
f"Language change detected: {self.lang_code} -> {new_lang_code}, voice: {self.voice} -> {new_voice}"
)
try:
new_pipeline = self.model._get_pipeline(new_lang_code)
new_voice_tensor = new_pipeline.load_voice(new_voice)
self.lang_code = new_lang_code
self.voice = new_voice
self._pipeline = new_pipeline
self._voice_tensor = new_voice_tensor
except Exception as e:
logger.warning(
f"Failed to switch language/voice: {e}. Keeping current language: {self.lang_code}"
)
console.print(f"[green]ASSISTANT: {llm_sentence}")
# Generate audio using the preloaded pipeline directly
# This avoids the voice reload that happens in model.generate()
for result in self._pipeline(
text=llm_sentence,
voice=self.voice,
speed=self.speed,
):
if result.audio is None:
continue
# result.audio is an mx.array with shape (1, samples), convert to numpy and squeeze
audio = np.array(result.audio, dtype=np.float32).squeeze(0)
# Trim silence from start and end of audio
# Kokoro generates ~250ms of silence at the start and variable silence at the end
threshold = 0.01
abs_audio = np.abs(audio)
# Find first and last samples above threshold
above_threshold = abs_audio > threshold
if np.any(above_threshold):
start_idx = int(np.argmax(above_threshold))
end_idx = len(audio) - int(np.argmax(above_threshold[::-1]))
# Add small padding (5ms = 120 samples at 24kHz) to avoid cutting speech
padding = 120
start_idx = max(0, start_idx - padding)
end_idx = min(len(audio), end_idx + padding)
audio = audio[start_idx:end_idx]
# Kokoro outputs at 24kHz, resample to 16kHz for the pipeline
# Using scipy's polyphase resampling (fast and high quality)
# 16000/24000 = 2/3, so up=2, down=3
audio = resample_poly(audio, up=2, down=3)
# Convert to int16 format
audio = (audio * 32768).astype(np.int16)
# Yield audio in fixed-size chunks
for i in range(0, len(audio), self.blocksize):
if gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(gen):
logger.info("TTS generation cancelled (interruption)")
return
chunk = audio[i : i + self.blocksize]
# Pad the last chunk if necessary
if len(chunk) < self.blocksize:
chunk = np.pad(chunk, (0, self.blocksize - len(chunk)))
logger.debug(f"TTS yielding audio chunk: {len(chunk)} samples")
yield chunk
def _process_kokoro(self, llm_sentence: str, language_code: Optional[str] = None) -> Iterator[np.ndarray]:
"""Process using native kokoro library."""
from scipy.signal import resample_poly
gen = self.cancel_scope.generation if self.cancel_scope else None
if language_code is not None:
new_lang_code = WHISPER_LANGUAGE_TO_KOKORO_LANG.get(language_code, self.lang_code)
if new_lang_code != self.lang_code:
new_voice = KOKORO_LANG_DEFAULT_VOICES.get(new_lang_code, self.voice)
logger.info(
f"Language change detected: {self.lang_code} -> {new_lang_code}, voice: {self.voice} -> {new_voice}"
)
self.lang_code = new_lang_code
self.voice = new_voice
from kokoro import KPipeline
self.pipeline = KPipeline(lang_code=self.lang_code)
console.print(f"[green]ASSISTANT: {llm_sentence}")
# Generate audio using Kokoro
# The pipeline yields tuples of (graphemes, phonemes, audio)
for gs, ps, audio in self.pipeline(llm_sentence, voice=self.voice, speed=self.speed):
if audio is None:
continue
# Ensure audio is numpy array
if not isinstance(audio, np.ndarray):
audio = np.array(audio, dtype=np.float32)
else:
audio = audio.astype(np.float32)
# Kokoro outputs at 24kHz, resample to 16kHz for the pipeline
# Using scipy's polyphase resampling (fast and high quality)
# 16000/24000 = 2/3, so up=2, down=3
audio = resample_poly(audio, up=2, down=3)
# Convert to int16 format
audio = (audio * 32768).astype(np.int16)
# Yield audio in fixed-size chunks
for i in range(0, len(audio), self.blocksize):
if gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(gen):
logger.info("TTS generation cancelled (interruption)")
return
chunk = audio[i : i + self.blocksize]
# Pad the last chunk if necessary
if len(chunk) < self.blocksize:
chunk = np.pad(chunk, (0, self.blocksize - len(chunk)))
yield chunk
def on_session_end(self) -> None:
self.voice = self._initial_voice
self.lang_code = self._initial_lang_code
if self.backend == "mlx":
try:
self._pipeline = self.model._get_pipeline(self.lang_code)
self._voice_tensor = self._pipeline.load_voice(self.voice)
except Exception as e:
logger.warning(f"Failed to restore initial voice/language on session end: {e}")
else:
from kokoro import KPipeline
self.pipeline = KPipeline(lang_code=self.lang_code)
logger.debug("Kokoro TTS session state reset")
speech_to_speech/TTS/pocket_tts_handler.py
+227
View File
@@ -0,0 +1,227 @@
from __future__ import annotations
import logging
from threading import Event
from time import perf_counter
from typing import Any, Iterator
import numpy as np
from rich.console import Console
from speech_to_speech.baseHandler import BaseHandler
from speech_to_speech.pipeline.cancel_scope import CancelScope
from speech_to_speech.pipeline.handler_types import TTSIn, TTSOut
from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, EndOfResponse
from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker
logger = logging.getLogger(__name__)
console = Console()
class PocketTTSHandler(BaseHandler[TTSIn, TTSOut]):
"""
Handler for Pocket TTS model from Kyutai Labs.
Supports streaming audio generation with voice cloning.
"""
def setup(
self,
should_listen: Event,
device: str = "cpu",
voice: str = "alba", # Default voice from catalog
sample_rate: int = 16000, # Match the pipeline's audio output (LocalAudioStreamer uses 16kHz)
blocksize: int = 512,
max_tokens: int = 50,
gen_kwargs: dict[str, Any] | None = None, # For compatibility with pipeline, not used
cancel_scope: CancelScope | None = None,
speculative_turns: SpeculativeTurnTracker | None = None,
) -> None:
"""
Initialize Pocket TTS handler.
Args:
should_listen: Event to control when to start listening again
device: Device to run model on ('cpu', 'cuda', 'mps')
voice: Voice to use. Can be:
- A preset name: 'alba', 'marius', 'javert', 'jean', 'fantine', 'cosette', 'eponine', 'azelma'
- A local audio file path
- A Hugging Face path like "hf://kyutai/tts-voices/..."
sample_rate: Output sample rate (pocket-tts generates at 24kHz and will be resampled to this rate). Default 16kHz matches the pipeline's audio output.
blocksize: Size of audio blocks to yield
max_tokens: Maximum tokens to generate
"""
self.should_listen = should_listen
self.cancel_scope = cancel_scope
self.speculative_turns = speculative_turns
self.device = device
self.voice = voice
self.sample_rate = sample_rate
self.blocksize = blocksize
self.max_tokens = max_tokens
# Suppress verbose logging from pocket_tts library
logging.getLogger("pocket_tts").setLevel(logging.WARNING)
logging.getLogger("pocket_tts.models.tts_model").setLevel(logging.WARNING)
logging.getLogger("pocket_tts.utils.utils").setLevel(logging.WARNING)
# Import and load model
from pocket_tts import TTSModel
logger.info("Loading Pocket TTS model")
self.model = TTSModel.load_model()
# Move model to specified device
if device == "cuda":
self.model = self.model.cuda()
elif device == "mps":
self.model = self.model.to("mps")
elif device != "cpu":
self.model = self.model.to(device)
logger.info(f"Pocket TTS model moved to {device}")
# Load voice state
logger.info(f"Loading voice: {voice}")
self.voice_state = self.model.get_state_for_audio_prompt(voice)
logger.info(f"Pocket TTS model sample rate: {self.model.sample_rate}")
@property
def min_time_to_debug(self) -> float:
"""
Override to suppress logging for individual audio chunks.
Pocket TTS yields many small chunks (~10-20ms each), which would flood logs.
Only log if a chunk takes unusually long (>100ms).
"""
return 0.1 # 100ms threshold
def process(self, tts_input: TTSIn) -> Iterator[TTSOut]:
speculative_turns = getattr(self, "speculative_turns", None)
if isinstance(tts_input, EndOfResponse):
if speculative_turns and not speculative_turns.is_latest_after_reopen_grace(
tts_input.turn_id,
tts_input.turn_revision,
):
return
yield AUDIO_RESPONSE_DONE
return
if speculative_turns and not speculative_turns.is_latest_after_reopen_grace(
tts_input.turn_id,
tts_input.turn_revision,
):
logger.debug("Dropping stale TTS input for turn=%s rev=%s", tts_input.turn_id, tts_input.turn_revision)
return
if speculative_turns:
speculative_turns.commit(tts_input.turn_id, tts_input.turn_revision)
gen = self.cancel_scope.generation if self.cancel_scope else None
language_code = tts_input.language_code
text = tts_input.text
logger.debug(f"Received language code: {language_code}")
console.print(f"[green]ASSISTANT: {text}")
# Generate audio stream
logger.debug(f"Generating audio for: {text[:50]}...")
pipeline_start = perf_counter()
first_chunk = True
# Calculate target chunk size in original sample rate
# We need enough samples so that after resampling we get blocksize samples
needs_resampling = self.model.sample_rate != self.sample_rate
if needs_resampling:
from scipy.signal import resample_poly
# Target chunk size in original sample rate (before resampling)
resample_ratio = self.model.sample_rate / self.sample_rate
target_chunk_size = int(self.blocksize * resample_ratio)
# Calculate up/down factors for polyphase resampling
# For 24kHz -> 16kHz: 16000/24000 = 2/3
from math import gcd
g = gcd(self.sample_rate, self.model.sample_rate)
self._resample_up = self.sample_rate // g
self._resample_down = self.model.sample_rate // g
else:
target_chunk_size = self.blocksize
# Buffer to accumulate audio until we have enough for one block
audio_buffer = []
buffer_size = 0
for audio_chunk in self.model.generate_audio_stream(
self.voice_state,
text,
max_tokens=self.max_tokens,
copy_state=True, # Don't modify the original voice state
):
if gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(gen):
logger.info("TTS generation cancelled (interruption)")
return
if first_chunk:
logger.debug(f"Time to first audio: {perf_counter() - pipeline_start:.3f}s")
first_chunk = False
# Convert from torch tensor to numpy and add to buffer
audio_np = audio_chunk.cpu().numpy()
audio_buffer.append(audio_np)
buffer_size += len(audio_np)
# Process buffer when we have enough samples for at least one block
while buffer_size >= target_chunk_size:
# Concatenate buffer
concatenated = np.concatenate(audio_buffer)
# Extract enough samples for one block
chunk_to_process = concatenated[:target_chunk_size]
remainder = concatenated[target_chunk_size:]
# Resample this chunk if needed
if needs_resampling:
chunk_resampled = resample_poly(
chunk_to_process,
up=self._resample_up,
down=self._resample_down,
)
else:
chunk_resampled = chunk_to_process
# Convert to int16 format expected by audio output
audio_int16 = (chunk_resampled * 32768).astype(np.int16)
# Ensure exact blocksize (pad or trim if resampling caused slight size differences)
if len(audio_int16) < self.blocksize:
audio_int16 = np.pad(audio_int16, (0, self.blocksize - len(audio_int16)))
elif len(audio_int16) > self.blocksize:
audio_int16 = audio_int16[: self.blocksize]
yield audio_int16
# Update buffer with remainder
audio_buffer = [remainder] if len(remainder) > 0 else []
buffer_size = len(remainder)
# Process any remaining audio in buffer
if audio_buffer and buffer_size > 0:
concatenated = np.concatenate(audio_buffer)
# Resample the remainder
if needs_resampling:
concatenated = resample_poly(
concatenated,
up=self._resample_up,
down=self._resample_down,
)
# Convert to int16
audio_int16 = (concatenated * 32768).astype(np.int16)
# Yield in blocks
for i in range(0, len(audio_int16), self.blocksize):
chunk = audio_int16[i : i + self.blocksize]
# Pad last chunk if needed
if len(chunk) < self.blocksize:
chunk = np.pad(chunk, (0, self.blocksize - len(chunk)))
yield chunk
speech_to_speech/TTS/qwen3_tts_handler.py
+903
View File
@@ -0,0 +1,903 @@
"""
Qwen3 TTS Handler
- On Apple Silicon: Uses mlx-audio with MLX-converted Qwen3-TTS models.
- On CUDA/CPU: Uses faster-qwen3-tts for low-latency streaming.
"""
from __future__ import annotations
import logging
import math
import re
import tempfile
import unicodedata
from collections.abc import Callable
from pathlib import Path
from sys import platform
from threading import Event
from time import perf_counter
from typing import Any, Iterator, Optional
import numpy as np
import torch
from openai.types.realtime.realtime_response_create_params import RealtimeResponseCreateParams
from rich.console import Console
from speech_to_speech.api.openai_realtime.runtime_config import RuntimeConfig
from speech_to_speech.baseHandler import BaseHandler
from speech_to_speech.pipeline.cancel_scope import CancelScope
from speech_to_speech.pipeline.control import SESSION_END, is_control_message
from speech_to_speech.pipeline.handler_types import TTSIn, TTSOut
from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, PIPELINE_END, EndOfResponse, TTSInput
from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker
from speech_to_speech.utils.mlx_lock import MLXLockContext
logger = logging.getLogger(__name__)
console = Console()
DEFAULT_MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
DEFAULT_MLX_MODEL = "mlx-community/Qwen3-TTS-12Hz-1.7B-CustomVoice-6bit"
DEFAULT_REF_TEXT = "I'm confused why some people have super short timelines, yet at the same time are bullish on scaling up reinforcement learning atop LLMs. If we're actually close to a human-like learner, then this whole approach of training on verifiable outcomes."
DEFAULT_FASTER_STREAMING_CHUNK_SIZE = 8
DEFAULT_MLX_STREAMING_CHUNK_SIZE = 4
DEFAULT_QWEN3_TTS_MAX_NEW_TOKENS = 1536
MIN_QWEN3_TTS_UTTERANCE_TOKENS = 360
VALID_MLX_QUANTIZATION_SUFFIXES = ("bf16", "4bit", "6bit", "8bit")
VALID_FASTER_BACKENDS = ("ggml", "torch")
MLX_STREAMING_TOKENS_PER_SECOND = 12.5
PIPELINE_SR = 16000
ESTIMATED_QWEN3_WORDS_PER_SECOND = 2.6
ESTIMATED_QWEN3_CHARS_PER_SECOND = 14.0
QWEN3_TOKEN_SAFETY_MARGIN = 1.35
QWEN3_BASE_PROMPT_SECONDS = 1.0
QWEN3_PUNCTUATION_PAUSE_SECONDS = 0.5
QWEN3_LANGUAGE_ALIASES = {
"zh": "chinese",
"zh-cn": "chinese",
"zh-hans": "chinese",
"zh-hant": "chinese",
"cmn": "chinese",
"en": "english",
"en-us": "english",
"en-gb": "english",
"ja": "japanese",
"jp": "japanese",
"ko": "korean",
"kr": "korean",
"de": "german",
"fr": "french",
"ru": "russian",
"pt": "portuguese",
"pt-br": "portuguese",
"pt-pt": "portuguese",
"es": "spanish",
"it": "italian",
}
class Qwen3TTSHandler(BaseHandler[TTSIn, TTSOut]):
"""
Handles Text-to-Speech using Qwen3-TTS.
Backend selection:
- Apple Silicon (Darwin): mlx-audio
- Other platforms: faster-qwen3-tts
Supports three generation modes depending on the loaded model:
- Voice cloning (ref_audio + ref_text)
- Custom voice (preset speakers)
- Voice design (instruct prompt)
"""
def setup(
self,
should_listen: Event,
model_name: str = DEFAULT_MODEL,
device: str = "cuda",
dtype: str | torch.dtype = "auto",
attn_implementation: str = "eager",
backend: str = "ggml",
ref_audio: str | Path | None = None,
ref_text: str = DEFAULT_REF_TEXT,
language: str = "auto",
speaker: Optional[str] = "Aiden",
instruct: Optional[str] = None,
xvec_only: bool = False,
parity_mode: bool = False,
non_streaming_mode: bool | None = True,
mlx_quantization: Optional[str] = None,
streaming_chunk_size: int | None = None,
max_new_tokens: int = DEFAULT_QWEN3_TTS_MAX_NEW_TOKENS,
blocksize: int = 512,
gen_kwargs: dict[str, Any] | None = None,
cancel_scope: CancelScope | None = None,
speculative_turns: SpeculativeTurnTracker | None = None,
) -> None:
self.cancel_scope = cancel_scope
self.speculative_turns = speculative_turns
self.should_listen = should_listen
self.requested_device = device
self.ref_audio = ref_audio
self.ref_text = ref_text
self.language = self._normalize_language(language)
self.speaker = speaker
self.instruct = instruct
self.xvec_only = xvec_only
self.parity_mode = parity_mode
self.non_streaming_mode = non_streaming_mode
self.faster_backend = self._normalize_faster_backend(backend)
self.mlx_quantization = self._normalize_mlx_quantization(mlx_quantization)
self.max_new_tokens = max_new_tokens
self.blocksize = blocksize
self.dtype: torch.dtype | None | str = None
self.gen_kwargs = gen_kwargs or {}
self._mlx_ref_audio_cache: dict[str, Any] = {}
self._mlx_temp_ref_audio_files: set[str] = set()
self.backend = "mlx" if platform == "darwin" else "faster_qwen3_tts"
self.streaming_chunk_size = self._resolve_streaming_chunk_size(streaming_chunk_size)
if self.backend == "mlx":
self.device = "mps"
self.model_name = self._resolve_mlx_model_name(model_name)
logger.info(f"Loading Qwen3-TTS model: {self.model_name} via mlx-audio on Apple Silicon")
if self.non_streaming_mode is not None:
logger.debug(
"qwen3_tts_non_streaming_mode=%s is ignored on Apple Silicon because "
"mlx-audio does not expose non_streaming_mode yet.",
self.non_streaming_mode,
)
model_quantization = self._model_name_quantization_suffix(self.model_name)
if model_quantization and model_quantization != "bf16":
logger.info(
"Using MLX quantized Qwen3-TTS variant: %s",
model_quantization,
)
self._setup_mlx(self.model_name)
else:
self.device = device
self.model_name = model_name
logger.info(
"Loading Qwen3-TTS model: %s via faster-qwen3-tts (%s backend)",
self.model_name,
self.faster_backend,
)
self._setup_faster(
model_name=self.model_name,
dtype=dtype,
attn_implementation=attn_implementation,
backend=self.faster_backend,
)
logger.info(
"Using Qwen3-TTS streaming chunk size %d (~%.0fms audio per chunk) on %s",
self.streaming_chunk_size,
self.streaming_chunk_size / MLX_STREAMING_TOKENS_PER_SECOND * 1000,
self.backend,
)
self._initial_speaker = self.speaker
self._initial_ref_audio = self.ref_audio
self.warmup()
def _setup_faster(self, model_name: str, dtype: Any, attn_implementation: str, backend: str) -> None:
try:
import torch
except ImportError as e:
raise ImportError("torch is required. Install with: pip install torch") from e
if dtype == "auto":
self.dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
elif isinstance(dtype, str):
self.dtype = getattr(torch, dtype)
else:
self.dtype = dtype
try:
from faster_qwen3_tts import FasterQwen3TTS
except ImportError as e:
raise ImportError(
"faster-qwen3-tts is required for Qwen3 TTS on non-macOS platforms. "
"Install with: pip install 'faster-qwen3-tts[ggml]'"
) from e
self.model = FasterQwen3TTS.from_pretrained(
model_name,
device=self.device,
dtype=self.dtype,
attn_implementation=attn_implementation,
backend=backend,
)
logger.info("Qwen3-TTS model loaded")
def _setup_mlx(self, model_name: str) -> None:
try:
from mlx_audio.tts.utils import load_model
self.model = load_model(model_name)
except ImportError as e:
message = str(e)
if any(
dep in message
for dep in (
"misaki",
"spacy",
"phonemizer",
"espeakng_loader",
)
):
raise ImportError(
"Qwen3-TTS on Apple Silicon requires mlx-audio and its TTS dependencies. "
f"Missing dependency: {message}. "
"Install with: pip install mlx-audio misaki spacy phonemizer-fork espeakng-loader"
) from e
raise ImportError(
"mlx-audio is required for Qwen3 TTS on Apple Silicon. Install with: pip install mlx-audio"
) from e
logger.info("MLX Audio Qwen3-TTS model loaded")
def _normalize_faster_backend(self, backend: Any) -> str:
value = str(backend or "ggml").strip().lower()
if value not in VALID_FASTER_BACKENDS:
raise ValueError(
f"Unsupported qwen3_tts_backend value {backend!r}. Supported values: {', '.join(VALID_FASTER_BACKENDS)}"
)
return value
def _normalize_mlx_quantization(self, mlx_quantization: Any) -> Optional[str]:
if mlx_quantization is None:
return None
value = str(mlx_quantization).strip().lower()
if value in ("", "none", "default"):
return None
if value not in VALID_MLX_QUANTIZATION_SUFFIXES:
raise ValueError(
"Unsupported qwen3_tts_mlx_quantization value "
f"{mlx_quantization!r}. Supported values: {', '.join(VALID_MLX_QUANTIZATION_SUFFIXES)}"
)
return value
def _apply_mlx_quantization_suffix(self, model_name: str) -> str:
if self.mlx_quantization is None:
return model_name
desired_suffix = f"-{self.mlx_quantization}"
for suffix in VALID_MLX_QUANTIZATION_SUFFIXES:
current_suffix = f"-{suffix}"
if model_name.endswith(current_suffix):
return model_name[: -len(current_suffix)] + desired_suffix
return f"{model_name}{desired_suffix}"
def _model_name_quantization_suffix(self, model_name: str) -> Optional[str]:
if not model_name:
return None
for suffix in VALID_MLX_QUANTIZATION_SUFFIXES:
if model_name.endswith(f"-{suffix}"):
return suffix
return None
def _resolve_mlx_model_name(self, model_name: str) -> str:
if not model_name:
return self._apply_mlx_quantization_suffix(DEFAULT_MLX_MODEL)
if model_name.startswith("mlx-community/"):
if self.mlx_quantization is None and self._model_name_quantization_suffix(model_name) is None:
return f"{model_name}-6bit"
return self._apply_mlx_quantization_suffix(model_name)
if model_name.startswith("Qwen/"):
mapped = model_name.replace("Qwen/", "mlx-community/", 1)
if self._model_name_quantization_suffix(mapped) is None:
if self.mlx_quantization is None:
return f"{mapped}-6bit"
mapped = f"{mapped}-bf16"
return self._apply_mlx_quantization_suffix(mapped)
return model_name
def _resolve_streaming_chunk_size(self, streaming_chunk_size: int | None) -> int:
if streaming_chunk_size is not None:
return max(1, int(streaming_chunk_size))
if self.backend == "mlx":
return DEFAULT_MLX_STREAMING_CHUNK_SIZE
return DEFAULT_FASTER_STREAMING_CHUNK_SIZE
def _normalize_language(self, language: str | None) -> str:
if language is None:
return "auto"
normalized = str(language).strip().replace("_", "-").lower()
if not normalized:
return "auto"
return QWEN3_LANGUAGE_ALIASES.get(normalized, normalized)
def _infer_model_type_from_name(self) -> str:
name = (self.model_name or "").lower()
if "voicedesign" in name:
return "voice_design"
if "customvoice" in name:
return "custom_voice"
return "base"
def _resolve_audio_path(self, audio: Any) -> Path | None:
if not isinstance(audio, (str, Path)) or not audio:
return None
candidate = Path(audio).expanduser()
repo_root = Path(__file__).resolve().parents[1]
search_paths = []
if candidate.is_absolute():
search_paths.append(candidate)
else:
search_paths.append(Path.cwd() / candidate)
search_paths.append(repo_root / candidate)
seen = set()
for path in search_paths:
normalized = str(path)
if normalized in seen:
continue
seen.add(normalized)
if path.exists():
return path.resolve()
return None
def _prepare_mlx_ref_audio(self, ref_audio: Any) -> Any:
if self.backend != "mlx" or ref_audio is None:
return ref_audio
if not isinstance(ref_audio, (str, Path)):
return ref_audio
resolved_path = self._resolve_audio_path(ref_audio)
if resolved_path is None:
raise FileNotFoundError(
"Qwen3-TTS on Apple Silicon requires qwen3_tts_ref_audio to point to "
f"a readable audio file. Got: {ref_audio!r}"
)
cache_key = str(resolved_path)
cached_path = self._mlx_ref_audio_cache.get(cache_key)
if cached_path and Path(cached_path).exists():
return cached_path
try:
import soundfile as sf
from scipy.signal import resample_poly
waveform, sample_rate = sf.read(str(resolved_path), always_2d=False, dtype="float32")
waveform = np.asarray(waveform, dtype=np.float32)
if waveform.ndim > 1:
waveform = waveform.mean(axis=1)
target_sample_rate = getattr(self.model, "sample_rate", 24000)
if sample_rate != target_sample_rate:
gcd = np.gcd(int(sample_rate), int(target_sample_rate))
waveform = resample_poly(
waveform,
up=int(target_sample_rate) // gcd,
down=int(sample_rate) // gcd,
)
sample_rate = target_sample_rate
with tempfile.NamedTemporaryFile(
prefix="qwen3_ref_",
suffix=".wav",
delete=False,
) as temp_file:
normalized_path = temp_file.name
sf.write(
normalized_path,
waveform,
sample_rate,
format="WAV",
subtype="PCM_16",
)
except Exception as e:
raise RuntimeError(f"Failed to normalize Qwen3-TTS reference audio {resolved_path}: {e}") from e
self._mlx_ref_audio_cache[cache_key] = normalized_path
self._mlx_temp_ref_audio_files.add(normalized_path)
return normalized_path
def _apply_session_voice_override(
self,
model_type: str,
runtime_config: RuntimeConfig | None = None,
response: RealtimeResponseCreateParams | None = None,
) -> None:
session_voice: Optional[str] = None
if response and response.audio and response.audio.output:
resp_voice = response.audio.output.voice
session_voice = str(resp_voice) if resp_voice else None
if not session_voice and runtime_config is not None:
audio = runtime_config.session.audio
output = audio.output if audio is not None else None
sess_voice = output.voice if output is not None else None
session_voice = str(sess_voice) if sess_voice else None
if not session_voice:
return
if model_type == "custom_voice":
supported_speakers = self._supported_speakers()
if supported_speakers is not None:
speakers_by_lower = {speaker.lower(): speaker for speaker in supported_speakers}
speaker = speakers_by_lower.get(session_voice.lower())
if speaker is None:
supported_text = ", ".join(sorted(supported_speakers, key=str.lower)) or "unknown"
logger.warning(
"Ignoring Qwen3-TTS session voice override %r because it is not a supported "
"CustomVoice speaker. Supported speakers: %s",
session_voice,
supported_text,
)
return
session_voice = speaker
self.speaker = session_voice
self.ref_audio = None
return
if self._resolve_audio_path(session_voice) is not None:
self.ref_audio = session_voice
return
logger.warning(
"Ignoring Qwen3-TTS session voice override because it is not an audio file path: %r",
session_voice,
)
def warmup(self) -> None:
logger.info(f"Warming up {self.__class__.__name__}")
if self.backend == "faster_qwen3_tts":
if self.parity_mode:
logger.info("Qwen3-TTS parity mode enabled: skipping backend warmup")
else:
try:
self.model.warmup(prefill_len=100)
except Exception as e:
logger.warning("Qwen3-TTS backend warmup failed: %s", e)
try:
for _ in self._warmup_process("Hello, this is a warmup."):
pass
logger.info(f"{self.__class__.__name__} warmed up")
except Exception as e:
logger.warning(f"Warmup generation failed: {e}")
def _model_type(self) -> str:
if self.backend == "mlx":
config = getattr(self.model, "config", None)
return getattr(config, "tts_model_type", None) or self._infer_model_type_from_name()
inner = getattr(getattr(self.model, "model", None), "model", None)
return getattr(inner, "tts_model_type", None) or self._infer_model_type_from_name()
def _supported_speakers(self) -> list[str] | None:
for candidate in (
self.model,
getattr(self.model, "model", None),
getattr(getattr(self.model, "model", None), "model", None),
):
get_speakers = getattr(candidate, "get_supported_speakers", None)
if callable(get_speakers):
speakers = get_speakers()
if speakers is None:
return None
return [str(speaker) for speaker in speakers if speaker]
return None
def _resolve_speaker(self) -> Optional[str]:
if self.speaker:
return self.speaker
for candidate in (
self.model,
getattr(getattr(self.model, "model", None), "model", None),
):
get_speakers = getattr(candidate, "get_supported_speakers", None)
if callable(get_speakers):
speakers = list(get_speakers() or [])
if speakers:
return speakers[0]
return None
def _to_int16(self, audio: np.ndarray) -> np.ndarray:
return np.clip(audio * 32768, -32768, 32767).astype(np.int16)
def _estimate_max_new_tokens(self, text: Optional[str]) -> int:
text = (text or "").strip()
chunk_size = max(1, int(getattr(self, "streaming_chunk_size", 1)))
configured_cap = max(1, int(getattr(self, "max_new_tokens", DEFAULT_QWEN3_TTS_MAX_NEW_TOKENS)))
if not text:
return min(configured_cap, MIN_QWEN3_TTS_UTTERANCE_TOKENS)
word_count = len(re.findall(r"\w+", text, flags=re.UNICODE))
char_count = len(re.sub(r"\s+", "", text))
word_seconds = word_count / ESTIMATED_QWEN3_WORDS_PER_SECOND if word_count else 0.0
char_seconds = char_count / ESTIMATED_QWEN3_CHARS_PER_SECOND if char_count else 0.0
punctuation_count = sum(unicodedata.category(ch).startswith("P") for ch in text)
punctuation_seconds = punctuation_count * QWEN3_PUNCTUATION_PAUSE_SECONDS
estimated_seconds = max(word_seconds, char_seconds) + punctuation_seconds + QWEN3_BASE_PROMPT_SECONDS
estimated_tokens = math.ceil(estimated_seconds * MLX_STREAMING_TOKENS_PER_SECOND * QWEN3_TOKEN_SAFETY_MARGIN)
aligned_tokens = max(
chunk_size,
math.ceil(estimated_tokens / chunk_size) * chunk_size,
)
requested_tokens = max(MIN_QWEN3_TTS_UTTERANCE_TOKENS, aligned_tokens)
resolved_tokens = min(configured_cap, requested_tokens)
if resolved_tokens < requested_tokens:
logger.warning(
"Qwen3-TTS estimated %d codec tokens for a %d-character utterance, "
"but max_new_tokens is capped at %d; output may still truncate.",
requested_tokens,
len(text),
configured_cap,
)
logger.debug(
"Qwen3-TTS using max_new_tokens=%d for utterance with %d words and %d chars",
resolved_tokens,
word_count,
char_count,
)
return resolved_tokens
def _warmup_process(self, llm_sentence: str) -> Iterator[bytes | np.ndarray]:
model_type = self._model_type()
if self.ref_audio:
yield from self._process_voice_clone(llm_sentence)
elif model_type == "custom_voice":
yield from self._process_custom_voice(llm_sentence)
elif model_type == "voice_design":
yield from self._process_voice_design(llm_sentence)
else:
raise ValueError(
"Qwen3-TTS Base model requires ref_audio for voice cloning. "
"Provide qwen3_tts_ref_audio or use a CustomVoice/VoiceDesign model."
)
def _resample_to_pipeline_sr(self, audio: np.ndarray, sr: int) -> np.ndarray:
if sr == PIPELINE_SR:
return audio
from scipy.signal import resample_poly
gcd = np.gcd(PIPELINE_SR, sr)
return resample_poly(audio, up=PIPELINE_SR // gcd, down=sr // gcd)
def _prepare_audio_chunk(self, item: Any) -> tuple[np.ndarray | None, int | None]:
if isinstance(item, tuple):
audio_chunk, sr, _timing = item
return np.asarray(audio_chunk, dtype=np.float32), sr
audio = getattr(item, "audio", None)
if audio is None:
return None, None
audio_chunk = np.asarray(audio, dtype=np.float32).squeeze()
sr = getattr(item, "sample_rate", None) or PIPELINE_SR
return audio_chunk, sr
def _stream(self, gen: Any, label: str) -> Iterator[bytes | np.ndarray]:
"""Common streaming loop: log TTFA and RTF, yield int16 chunks."""
cancel_gen = self.cancel_scope.generation if self.cancel_scope else None
start = perf_counter()
total_samples = 0
first_chunk = True
found_speech = False
leftover = np.array([], dtype=np.int16)
for item in gen:
if cancel_gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(cancel_gen):
logger.info("TTS generation cancelled (interruption)")
return
audio_chunk, sr = self._prepare_audio_chunk(item)
if audio_chunk is None or sr is None or audio_chunk.size == 0:
continue
if first_chunk:
logger.info(f"Qwen3-TTS TTFA: {perf_counter() - start:.2f}s ({label})")
first_chunk = False
audio_chunk = self._resample_to_pipeline_sr(audio_chunk, sr)
audio_chunk = self._to_int16(audio_chunk)
# Trim the initial silent ramp-up, but keep enough preroll to avoid
# shaving soft initial phonemes at the start of the utterance.
if not found_speech:
threshold = int(32768 * 0.01)
above = np.abs(audio_chunk) > threshold
if not np.any(above):
continue
start_idx = max(0, int(np.argmax(above)) - int(PIPELINE_SR * 0.040))
audio_chunk = audio_chunk[start_idx:]
found_speech = True
audio_chunk = np.concatenate([leftover, audio_chunk])
n = (len(audio_chunk) // self.blocksize) * self.blocksize
for i in range(0, n, self.blocksize):
yield audio_chunk[i : i + self.blocksize]
total_samples += self.blocksize
leftover = audio_chunk[n:]
if len(leftover) > 0:
chunk = np.pad(leftover, (0, self.blocksize - len(leftover)))
yield chunk
total_samples += len(leftover)
generation_time = perf_counter() - start
audio_duration = total_samples / PIPELINE_SR
rtf = audio_duration / generation_time if generation_time > 0 else 0
logger.info(
f"Qwen3-TTS generated {audio_duration:.2f}s audio in {generation_time:.2f}s (RTF: {rtf:.2f}, {label})"
)
def _coalesce_pending_tts_input(self, current_input: TTSInput) -> tuple[str, Optional[str], bool]:
"""Combine already-queued text chunks before the next TTS synthesis call."""
if not hasattr(self.queue_in, "mutex") or not hasattr(self.queue_in, "queue"):
return current_input.text, current_input.language_code, False
text = current_input.text
language_code = current_input.language_code
parts = [text.strip()] if text and text.strip() else []
saw_end_of_response = False
with self.queue_in.mutex:
while self.queue_in.queue:
next_item = self.queue_in.queue[0]
if is_control_message(next_item, SESSION_END.kind):
break
if isinstance(next_item, bytes) and next_item == PIPELINE_END:
break
if isinstance(next_item, EndOfResponse):
saw_end_of_response = True
break
if not isinstance(next_item, TTSInput):
break
if current_input.turn_id != next_item.turn_id or current_input.turn_revision != next_item.turn_revision:
break
if (
language_code is not None
and next_item.language_code is not None
and next_item.language_code != language_code
):
break
self.queue_in.queue.popleft()
if next_item.text.strip():
parts.append(next_item.text.strip())
if language_code is None:
language_code = next_item.language_code
combined_text = " ".join(parts).strip()
return combined_text, language_code, saw_end_of_response
def process(self, tts_input: TTSIn) -> Iterator[TTSOut]:
speculative_turns = getattr(self, "speculative_turns", None)
if isinstance(tts_input, EndOfResponse):
if speculative_turns and not speculative_turns.is_latest_after_reopen_grace(
tts_input.turn_id,
tts_input.turn_revision,
):
return
yield AUDIO_RESPONSE_DONE
return
if speculative_turns and not speculative_turns.is_latest_after_reopen_grace(
tts_input.turn_id,
tts_input.turn_revision,
):
logger.debug("Dropping stale TTS input for turn=%s rev=%s", tts_input.turn_id, tts_input.turn_revision)
return
if speculative_turns:
speculative_turns.commit(tts_input.turn_id, tts_input.turn_revision)
runtime_config = tts_input.runtime_config
response = tts_input.response
coalesced_text, _language_code, _saw_end_of_response = self._coalesce_pending_tts_input(tts_input)
text = coalesced_text or "Hello."
model_type = self._model_type()
self._apply_session_voice_override(model_type, runtime_config, response)
console.print(f"[green]ASSISTANT: {text}")
try:
if self.ref_audio:
audio_iter = self._process_voice_clone(text)
elif model_type == "custom_voice":
audio_iter = self._process_custom_voice(text)
elif model_type == "voice_design":
audio_iter = self._process_voice_design(text)
else:
raise ValueError(
"Qwen3-TTS Base model requires ref_audio for voice cloning. "
"Provide qwen3_tts_ref_audio or use a CustomVoice/VoiceDesign model."
)
first_audio = True
for audio_chunk in audio_iter:
if first_audio:
self._log_first_audio_latency(tts_input)
first_audio = False
yield audio_chunk
except Exception as e:
logger.error(f"Error during Qwen3-TTS generation: {e}", exc_info=True)
def _log_first_audio_latency(self, tts_input: TTSInput) -> None:
if tts_input.speech_stopped_at_s is None:
return
latency_s = perf_counter() - tts_input.speech_stopped_at_s
if latency_s < 0:
return
logger.info(
"Last speech detected to first speech out: %.3fs (turn=%s rev=%s)",
latency_s,
tts_input.turn_id,
tts_input.turn_revision,
)
def _mlx_streaming_interval(self) -> float:
return max(1, self.streaming_chunk_size) / MLX_STREAMING_TOKENS_PER_SECOND
def _mlx_stream_kwargs(self, max_tokens: int) -> dict[str, Any]:
return {
"max_tokens": max_tokens,
"verbose": False,
"stream": True,
"streaming_interval": self._mlx_streaming_interval(),
**self.gen_kwargs,
}
def _stream_mlx_generation(
self,
generation_fn: Callable,
label: str,
max_tokens: int,
**generation_kwargs: Any,
) -> Iterator[bytes | np.ndarray]:
with MLXLockContext(handler_name="Qwen3TTS", timeout=10.0) as acquired:
if not acquired:
raise TimeoutError("Timed out waiting for MLX lock")
yield from self._stream(
generation_fn(
**self._mlx_stream_kwargs(max_tokens=max_tokens),
**generation_kwargs,
),
label=label,
)
def _process_voice_clone(self, text: str) -> Iterator[bytes | np.ndarray]:
utterance_max_new_tokens = self._estimate_max_new_tokens(text)
if self.backend == "mlx":
if self.xvec_only:
logger.warning("mlx-audio Qwen3-TTS does not support xvec_only; ignoring it")
if self.parity_mode:
logger.info("Qwen3-TTS parity mode is CUDA-specific and is ignored on mlx-audio")
yield from self._stream_mlx_generation(
self.model.generate,
label="voice_clone_mlx",
max_tokens=utterance_max_new_tokens,
text=text,
ref_audio=self._prepare_mlx_ref_audio(self.ref_audio),
ref_text=self.ref_text,
lang_code=self.language,
)
return
yield from self._stream(
self.model.generate_voice_clone_streaming(
text=text,
language=self.language,
ref_audio=self.ref_audio,
ref_text=self.ref_text,
xvec_only=self.xvec_only,
chunk_size=self.streaming_chunk_size,
max_new_tokens=utterance_max_new_tokens,
parity_mode=self.parity_mode,
non_streaming_mode=self.non_streaming_mode,
),
label="voice_clone_parity" if self.parity_mode else "voice_clone",
)
def _process_custom_voice(self, text: str) -> Iterator[bytes | np.ndarray]:
utterance_max_new_tokens = self._estimate_max_new_tokens(text)
speaker = self._resolve_speaker()
if not speaker:
raise ValueError(
"CustomVoice generation requires a speaker. "
"Set qwen3_tts_speaker or use a voice-clone model with ref_audio."
)
if self.backend == "mlx":
yield from self._stream_mlx_generation(
self.model.generate_custom_voice,
label="custom_voice_mlx",
max_tokens=utterance_max_new_tokens,
text=text,
speaker=speaker,
language=self.language,
instruct=self.instruct,
)
return
yield from self._stream(
self.model.generate_custom_voice_streaming(
text=text,
speaker=speaker,
language=self.language,
instruct=self.instruct,
chunk_size=self.streaming_chunk_size,
max_new_tokens=utterance_max_new_tokens,
non_streaming_mode=self.non_streaming_mode,
),
label="custom_voice",
)
def _process_voice_design(self, text: str) -> Iterator[bytes | np.ndarray]:
utterance_max_new_tokens = self._estimate_max_new_tokens(text)
if self.backend == "mlx":
yield from self._stream_mlx_generation(
self.model.generate_voice_design,
label="voice_design_mlx",
max_tokens=utterance_max_new_tokens,
text=text,
instruct=self.instruct,
language=self.language,
)
return
yield from self._stream(
self.model.generate_voice_design_streaming(
text=text,
instruct=self.instruct,
language=self.language,
chunk_size=self.streaming_chunk_size,
max_new_tokens=utterance_max_new_tokens,
non_streaming_mode=self.non_streaming_mode,
),
label="voice_design",
)
def on_session_end(self) -> None:
self.speaker = self._initial_speaker
self.ref_audio = self._initial_ref_audio
logger.debug("Qwen3-TTS session state reset")
def cleanup(self) -> None:
try:
del self.model
for path in list(getattr(self, "_mlx_temp_ref_audio_files", set())):
try:
Path(path).unlink(missing_ok=True)
except Exception:
pass
if self.backend == "mlx":
try:
import mlx.core as mx
mx.clear_cache()
except Exception:
pass
else:
import torch
if torch.cuda.is_available():
torch.cuda.empty_cache()
logger.info("Qwen3-TTS handler cleaned up")
except Exception as e:
logger.warning(f"Cleanup error: {e}")
speech_to_speech/TTS/ref_audio.wav
BIN
View File
Binary file not shown.