first
This commit is contained in:
@@ -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}")
|
||||
Reference in New Issue
Block a user