""" 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}")