nano_tts_runtime.py

from __future__ import annotations

import importlib
import logging
import threading
import time
import uuid
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path
from typing import Iterator, Optional

import numpy as np
import torch
from transformers import AutoModel, AutoModelForCausalLM

MOSS_AUDIO_TOKENIZER_TYPE = "moss-audio-tokenizer-nano"

APP_DIR = Path(__file__).resolve().parent
DEFAULT_CHECKPOINT_PATH = APP_DIR / "weights" / "tts"
DEFAULT_AUDIO_TOKENIZER_PATH = APP_DIR / "weights" / "codec"
DEFAULT_PROMPT_AUDIO_DIR = APP_DIR / "asserts" / "audio"
DEFAULT_OUTPUT_DIR = APP_DIR / "generated_audio"

_DEFAULT_VOICE_FILES: dict[str, tuple[str, str]] = {
    "Junhao": ("zh_1.wav", "Chinese male voice A"),
    "Zhiming": ("zh_2.wav", "Chinese male voice B"),
    "Weiguo": ("zh_5.wav", "Chinese male voice C"),
    "Xiaoyu": ("zh_3.wav", "Chinese female voice A"),
    "Yuewen": ("zh_4.wav", "Chinese female voice B"),
    "Lingyu": ("zh_6.wav", "Chinese female voice C"),
    "Trump": ("en_1.wav", "Trump reference voice"),
    "Ava": ("en_2.wav", "English female voice A"),
    "Bella": ("en_3.wav", "English female voice B"),
    "Adam": ("en_4.wav", "English male voice A"),
    "Nathan": ("en_5.wav", "English male voice B"),
    "Yui": ("jp_2.wav", "Japanese female voice B"),
    "Aoi": ("jp_3.wav", "Japanese female voice C"),
    "Hina": ("jp_4.wav", "Japanese female voice D"),
    "Mei": ("jp_5.wav", "Japanese female voice E"),
}

DEFAULT_VOICE = "Junhao"
FLASH_ATTENTION_DTYPES = {torch.float16, torch.bfloat16}


@dataclass(frozen=True)
class VoicePreset:
    name: str
    prompt_audio_path: Path
    description: str


def build_default_voice_presets() -> dict[str, VoicePreset]:
    presets: dict[str, VoicePreset] = {}
    for voice_name, (file_name, description) in _DEFAULT_VOICE_FILES.items():
        prompt_path = (DEFAULT_PROMPT_AUDIO_DIR / file_name).resolve()
        presets[voice_name] = VoicePreset(
            name=voice_name,
            prompt_audio_path=prompt_path,
            description=description,
        )
    return presets


def resolve_device(device_arg: str) -> torch.device:
    if device_arg == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return torch.device(device_arg)


def resolve_dtype(dtype_arg: str, device: torch.device) -> torch.dtype:
    if dtype_arg == "float32":
        return torch.float32
    if dtype_arg == "float16":
        return torch.float16
    if dtype_arg == "bfloat16":
        return torch.bfloat16
    if device.type == "cuda":
        if hasattr(torch.cuda, "is_bf16_supported") and torch.cuda.is_bf16_supported():
            return torch.bfloat16
        return torch.float16
    return torch.float32


def waveform_to_numpy(waveform: torch.Tensor | np.ndarray) -> np.ndarray:
    if torch.is_tensor(waveform):
        array = waveform.detach().cpu().numpy()
    else:
        array = np.asarray(waveform)
    if array.ndim == 1:
        return array.astype(np.float32, copy=False)
    if array.ndim != 2:
        raise ValueError(f"Unsupported waveform shape: {array.shape}")
    if array.shape[0] <= 8 and array.shape[0] < array.shape[1]:
        array = array.T
    return array.astype(np.float32, copy=False)


@lru_cache(maxsize=1)
def _has_flash_attn() -> bool:
    try:
        importlib.import_module("flash_attn")
    except Exception:
        return False
    return True


class NanoTTSService:
    def __init__(
        self,
        *,
        checkpoint_path: str | Path = DEFAULT_CHECKPOINT_PATH,
        audio_tokenizer_path: str | Path = DEFAULT_AUDIO_TOKENIZER_PATH,
        device: str = "auto",
        dtype: str = "auto",
        attn_implementation: str = "auto",
        output_dir: str | Path = DEFAULT_OUTPUT_DIR,
        voice_presets: Optional[dict[str, VoicePreset]] = None,
    ) -> None:
        self.checkpoint_path = Path(checkpoint_path).expanduser().resolve()
        self.audio_tokenizer_path = Path(audio_tokenizer_path).expanduser().resolve()
        self.output_dir = Path(output_dir).expanduser().resolve()
        self.output_dir.mkdir(parents=True, exist_ok=True)

        self.voice_presets = voice_presets or build_default_voice_presets()
        self.default_voice = DEFAULT_VOICE if DEFAULT_VOICE in self.voice_presets else next(iter(self.voice_presets))

        self.device = resolve_device(device)
        self.dtype = resolve_dtype(dtype, self.device)
        self.attn_implementation = self._resolve_attn_implementation(attn_implementation)

        self._lock = threading.RLock()
        self._model = None
        self._audio_tokenizer = None
        self._checkpoint_global_attn_implementation: str | None = None
        self._checkpoint_local_attn_implementation: str | None = None
        self._configured_global_attn_implementation: str | None = None
        self._configured_local_attn_implementation: str | None = None
        self._configured_audio_tokenizer_attn_implementation: str | None = None
        self._configured_audio_tokenizer_compute_dtype: str | None = None

    def _can_use_flash_attention(self) -> bool:
        return self.device.type == "cuda" and self.dtype in FLASH_ATTENTION_DTYPES and _has_flash_attn()

    def _resolve_runtime_default_attn_implementation(self) -> str:
        return "flash_attention_2" if self._can_use_flash_attention() else "sdpa"

    def _resolve_attn_implementation(self, requested: str | None) -> str | None:
        normalized = str(requested).strip().lower() if requested is not None else "auto"
        if not normalized or normalized in {"auto", "default", "model_default"}:
            return None
        if normalized not in {"sdpa", "flash_attention_2", "eager"}:
            raise ValueError(
                "attn_implementation must be one of: model_default/auto, sdpa, flash_attention_2, eager"
            )
        if normalized == "flash_attention_2":
            if not self._can_use_flash_attention():
                logging.warning(
                    "flash_attention_2 requires CUDA, flash_attn, and fp16/bf16; falling back to sdpa "
                    "(device=%s dtype=%s flash_attn=%s)",
                    self.device,
                    self.dtype,
                    _has_flash_attn(),
                )
                return "sdpa"
        return normalized

    @staticmethod
    def _normalize_loaded_attn_implementation(attn_implementation: object) -> str:
        normalized = str(attn_implementation).strip().lower() if attn_implementation is not None else ""
        if not normalized or normalized == "none":
            return "eager"
        return normalized

    def _resolve_request_attention_implementation(
        self,
        requested: Optional[str],
    ) -> tuple[str, str, str]:
        normalized = str(requested).strip().lower() if requested is not None else ""
        resolved = self._resolve_attn_implementation(normalized)
        if resolved is not None:
            return normalized, resolved, resolved

        if self.attn_implementation is not None:
            return self.attn_implementation, self.attn_implementation, self.attn_implementation

        runtime_default = self._resolve_runtime_default_attn_implementation()
        return "auto", runtime_default, runtime_default

    @staticmethod
    def _resolve_codec_attention_implementation(tts_attn_implementation: str) -> str:
        return "flash_attention_2" if tts_attn_implementation == "flash_attention_2" else "sdpa"

    def _resolve_codec_compute_dtype(self, codec_attn_implementation: str) -> str:
        if codec_attn_implementation == "flash_attention_2":
            return "bf16" if self.dtype == torch.bfloat16 else "fp16"
        return "fp32"

    @staticmethod
    def _apply_model_attention_implementation(model, *, global_attn: str, local_attn: str) -> None:
        if hasattr(model, "_set_attention_implementation"):
            model._set_attention_implementation(global_attn, local_attn_implementation=local_attn)

    def _install_stream_decode_budget_patch(self, model) -> None:
        if self.device.type != "cuda":
            return

        model_cls = model.__class__
        if getattr(model_cls, "_nanotts_stream_decode_budget_patch_installed", False):
            return

        compute_stream_lead = getattr(model_cls, "_compute_stream_lead_seconds", None)
        resolve_stream_budget = getattr(model_cls, "_resolve_stream_decode_frame_budget", None)
        if not callable(compute_stream_lead) or not callable(resolve_stream_budget):
            return

        def _patched_resolve_stream_decode_frame_budget(
            *,
            emitted_samples_total: int,
            sample_rate: int,
            first_audio_emitted_at,
        ) -> int:
            # The upstream streaming policy starts with one decode frame
            # (about 80 ms audio), which makes CUDA realtime decode emit many
            # tiny chunks and underrun browser playback on this checkpoint.
            lead_seconds = compute_stream_lead(
                emitted_samples_total=emitted_samples_total,
                sample_rate=sample_rate,
                first_audio_emitted_at=first_audio_emitted_at,
            )
            if first_audio_emitted_at is None or lead_seconds < 0.20:
                return 4
            if lead_seconds < 0.55:
                return 6
            if lead_seconds < 1.10:
                return 8
            return 12

        model_cls._nanotts_original_resolve_stream_decode_frame_budget = resolve_stream_budget
        model_cls._resolve_stream_decode_frame_budget = staticmethod(_patched_resolve_stream_decode_frame_budget)
        model_cls._nanotts_stream_decode_budget_patch_installed = True
        logging.info("installed Nano-TTS CUDA streaming decode budget patch")

    def _discard_loaded_model_locked(self, reason: str) -> None:
        if self._model is None:
            return
        logging.warning("discarding loaded Nano-TTS model state: %s", reason)
        self._model = None
        if self.device.type == "cuda":
            torch.cuda.empty_cache()

    def _discard_loaded_audio_tokenizer_locked(self, reason: str) -> None:
        if self._audio_tokenizer is None:
            return
        logging.warning("discarding loaded Nano-TTS audio tokenizer state: %s", reason)
        self._audio_tokenizer = None
        self._configured_audio_tokenizer_attn_implementation = None
        self._configured_audio_tokenizer_compute_dtype = None
        if self.device.type == "cuda":
            torch.cuda.empty_cache()

    def _restore_model_execution_state(self, model):
        current_parameter = next(model.parameters(), None)
        if current_parameter is None or current_parameter.dtype == self.dtype:
            return model
        self._discard_loaded_model_locked(
            f"current_dtype={current_parameter.dtype} expected_dtype={self.dtype}; reloading checkpoint"
        )
        return self._load_model_locked()

    def _read_model_attention_implementation(self, model) -> tuple[str, str]:
        global_attn = self._normalize_loaded_attn_implementation(
            getattr(getattr(model, "transformer", None), "attn_implementation", None)
        )
        local_attn = self._normalize_loaded_attn_implementation(
            getattr(getattr(model, "local_transformer", None), "attn_implementation", None)
        )
        return global_attn, local_attn

    def _ensure_paths(self) -> None:
        if not self.checkpoint_path.exists():
            raise FileNotFoundError(f"Nano-TTS checkpoint not found: {self.checkpoint_path}")
        if not self.audio_tokenizer_path.exists():
            raise FileNotFoundError(f"Audio tokenizer checkpoint not found: {self.audio_tokenizer_path}")

    def _load_audio_tokenizer_locked(self, *, tts_attn_implementation: str):
        codec_attn_implementation = self._resolve_codec_attention_implementation(tts_attn_implementation)
        codec_compute_dtype = self._resolve_codec_compute_dtype(codec_attn_implementation)

        if self._audio_tokenizer is None:
            logging.info(
                "loading Nano-TTS audio tokenizer checkpoint=%s device=%s attn=%s compute_dtype=%s",
                self.audio_tokenizer_path,
                self.device,
                codec_attn_implementation,
                codec_compute_dtype,
            )
            audio_tokenizer = AutoModel.from_pretrained(
                str(self.audio_tokenizer_path),
                trust_remote_code=True,
                local_files_only=True,
            )
            if hasattr(audio_tokenizer, "eval"):
                audio_tokenizer.eval()
            self._audio_tokenizer = audio_tokenizer

        audio_tokenizer = self._audio_tokenizer
        if hasattr(audio_tokenizer, "to"):
            audio_tokenizer = audio_tokenizer.to(self.device)
        if hasattr(audio_tokenizer, "set_attention_implementation"):
            audio_tokenizer.set_attention_implementation(codec_attn_implementation)
        if hasattr(audio_tokenizer, "set_compute_dtype"):
            audio_tokenizer.set_compute_dtype(codec_compute_dtype)
        if hasattr(audio_tokenizer, "eval"):
            audio_tokenizer.eval()

        self._audio_tokenizer = audio_tokenizer
        self._configured_audio_tokenizer_attn_implementation = codec_attn_implementation
        self._configured_audio_tokenizer_compute_dtype = codec_compute_dtype
        return self._audio_tokenizer

    def _load_model_locked(self):
        if self._model is not None:
            return self._model

        self._ensure_paths()
        logging.info(
            "loading Nano-TTS checkpoint=%s audio_tokenizer=%s device=%s dtype=%s attn=%s",
            self.checkpoint_path,
            self.audio_tokenizer_path,
            self.device,
            self.dtype,
            self.attn_implementation or "model_default",
        )
        model = AutoModelForCausalLM.from_pretrained(
            str(self.checkpoint_path),
            trust_remote_code=True,
            local_files_only=True,
        )
        model.to(device=self.device, dtype=self.dtype)
        self._checkpoint_global_attn_implementation, self._checkpoint_local_attn_implementation = (
            self._read_model_attention_implementation(model)
        )
        _, default_global_attn, default_local_attn = self._resolve_request_attention_implementation(None)
        self._apply_model_attention_implementation(
            model,
            global_attn=default_global_attn,
            local_attn=default_local_attn,
        )
        self._install_stream_decode_budget_patch(model)
        model.eval()
        self._configured_global_attn_implementation, self._configured_local_attn_implementation = (
            self._read_model_attention_implementation(model)
        )
        self._model = model
        return self._model

    def get_model(self):
        with self._lock:
            return self._load_model_locked()

    def list_voice_names(self) -> list[str]:
        return list(self.voice_presets.keys())

    def get_voice_preset(self, voice_name: Optional[str]) -> VoicePreset:
        if voice_name and voice_name in self.voice_presets:
            return self.voice_presets[voice_name]
        return self.voice_presets[self.default_voice]

    def resolve_prompt_audio_path(
        self,
        *,
        voice: Optional[str] = None,
        prompt_audio_path: Optional[str | Path] = None,
    ) -> Path:
        if prompt_audio_path:
            resolved = Path(prompt_audio_path).expanduser().resolve()
            if not resolved.exists():
                raise FileNotFoundError(f"Prompt audio not found: {resolved}")
            return resolved

        preset = self.get_voice_preset(voice)
        if not preset.prompt_audio_path.exists():
            raise FileNotFoundError(f"Voice preset prompt audio not found: {preset.prompt_audio_path}")
        return preset.prompt_audio_path

    def preload(self, *, voices: Optional[list[str]] = None, load_model: bool = True) -> dict[str, object]:
        loaded_voices: list[str] = []
        if load_model:
            self.get_model()
        for voice_name in voices or [self.default_voice]:
            preset = self.get_voice_preset(voice_name)
            if preset.prompt_audio_path.exists():
                loaded_voices.append(preset.name)
        return {
            "loaded_voices": loaded_voices,
            "device": str(self.device),
            "dtype": str(self.dtype),
            "attn_implementation": self.attn_implementation or "auto",
            "checkpoint_default_attn_implementation": self._checkpoint_global_attn_implementation or "eager",
            "checkpoint_default_local_attn_implementation": self._checkpoint_local_attn_implementation or "eager",
            "configured_attn_implementation": self._configured_global_attn_implementation or "eager",
            "configured_local_attn_implementation": self._configured_local_attn_implementation or "eager",
            "configured_codec_attn_implementation": self._configured_audio_tokenizer_attn_implementation or "unknown",
            "configured_codec_compute_dtype": self._configured_audio_tokenizer_compute_dtype or "unknown",
        }

    def _build_output_path(self, prefix: str) -> Path:
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        random_suffix = uuid.uuid4().hex[:8]
        return self.output_dir / f"{prefix}_{timestamp}_{random_suffix}.wav"

    def synthesize(
        self,
        *,
        text: str,
        voice: Optional[str] = None,
        mode: str = "voice_clone",
        output_audio_path: Optional[str | Path] = None,
        prompt_audio_path: Optional[str | Path] = None,
        prompt_text: Optional[str] = None,
        max_new_frames: int = 375,
        voice_clone_max_text_tokens: int = 75,
        voice_clone_max_memory_per_sample_gb: float = 1.0,
        tts_max_batch_size: int = 0,
        codec_max_batch_size: int = 0,
        do_sample: bool = True,
        text_temperature: float = 1.0,
        text_top_p: float = 1.0,
        text_top_k: int = 50,
        audio_temperature: float = 0.8,
        audio_top_p: float = 0.95,
        audio_top_k: int = 25,
        audio_repetition_penalty: float = 1.2,
        nq: Optional[int] = None,
        seed: Optional[int] = None,
        attn_implementation: Optional[str] = None,
    ) -> dict[str, object]:
        normalized_text = str(text or "").strip()
        if not normalized_text:
            raise ValueError("text is required")

        normalized_mode = str(mode).strip().lower()
        if normalized_mode not in {"continuation", "voice_clone"}:
            raise ValueError("mode must be either 'continuation' or 'voice_clone'")

        effective_prompt_audio_path: Optional[Path] = None
        resolved_voice = self.get_voice_preset(voice).name
        if normalized_mode == "voice_clone":
            effective_prompt_audio_path = self.resolve_prompt_audio_path(
                voice=resolved_voice,
                prompt_audio_path=prompt_audio_path,
            )
        elif prompt_audio_path is not None:
            effective_prompt_audio_path = self.resolve_prompt_audio_path(prompt_audio_path=prompt_audio_path)
            if not prompt_text:
                raise ValueError("continuation mode with prompt_audio_path also requires prompt_text")

        output_path = (
            Path(output_audio_path).expanduser().resolve()
            if output_audio_path is not None
            else self._build_output_path(prefix=f"{resolved_voice}_{normalized_mode}")
        )
        output_path.parent.mkdir(parents=True, exist_ok=True)

        started_at = time.monotonic()
        with self._lock:
            model = self._load_model_locked()
            model = self._restore_model_execution_state(model)
            requested_attn_implementation, effective_global_attn_implementation, effective_local_attn_implementation = (
                self._resolve_request_attention_implementation(attn_implementation)
            )
            audio_tokenizer = self._load_audio_tokenizer_locked(
                tts_attn_implementation=effective_global_attn_implementation
            )
            self._apply_model_attention_implementation(
                model,
                global_attn=effective_global_attn_implementation,
                local_attn=effective_local_attn_implementation,
            )
            if seed is not None:
                torch.manual_seed(seed)
                if torch.cuda.is_available():
                    torch.cuda.manual_seed_all(seed)

            try:
                result = model.inference(
                    text=normalized_text,
                    output_audio_path=str(output_path),
                    mode=normalized_mode,
                    prompt_text=prompt_text,
                    prompt_audio_path=None if effective_prompt_audio_path is None else str(effective_prompt_audio_path),
                    text_tokenizer_path=str(self.checkpoint_path),
                    audio_tokenizer=audio_tokenizer,
                    device=self.device,
                    nq=nq,
                    max_new_frames=int(max_new_frames),
                    voice_clone_max_text_tokens=int(voice_clone_max_text_tokens),
                    voice_clone_max_memory_per_sample_gb=float(voice_clone_max_memory_per_sample_gb),
                    tts_max_batch_size=int(tts_max_batch_size),
                    codec_max_batch_size=int(codec_max_batch_size),
                    do_sample=bool(do_sample),
                    use_kv_cache=True,
                    text_temperature=float(text_temperature),
                    text_top_p=float(text_top_p),
                    text_top_k=int(text_top_k),
                    audio_temperature=float(audio_temperature),
                    audio_top_p=float(audio_top_p),
                    audio_top_k=int(audio_top_k),
                    audio_repetition_penalty=float(audio_repetition_penalty),
                )
            except Exception:
                self._discard_loaded_audio_tokenizer_locked(
                    "inference failed; reloading audio tokenizer on next request"
                )
                self._discard_loaded_model_locked("inference failed; reloading checkpoint on next request")
                raise
            effective_global_attn_implementation, effective_local_attn_implementation = (
                self._read_model_attention_implementation(model)
            )
            current_parameter = next(model.parameters(), None)
            if current_parameter is not None and current_parameter.dtype != self.dtype:
                self._discard_loaded_model_locked(
                    f"inference left model in dtype={current_parameter.dtype}; reloading checkpoint on next request"
                )

        waveform = result["waveform"].detach().cpu()
        waveform_numpy = waveform_to_numpy(waveform)
        return {
            "audio_path": str(output_path),
            "sample_rate": int(result["sample_rate"]),
            "waveform": waveform,
            "waveform_numpy": waveform_numpy,
            "audio_token_ids": result["audio_token_ids"],
            "reference_audio_token_ids": result["reference_audio_token_ids"],
            "elapsed_seconds": time.monotonic() - started_at,
            "voice": resolved_voice,
            "mode": normalized_mode,
            "prompt_audio_path": None if effective_prompt_audio_path is None else str(effective_prompt_audio_path),
            "requested_attn_implementation": requested_attn_implementation,
            "effective_global_attn_implementation": effective_global_attn_implementation,
            "effective_local_attn_implementation": effective_local_attn_implementation,
            "voice_clone_text_chunks": result.get("voice_clone_text_chunks"),
            "voice_clone_chunk_batch_size": result.get("voice_clone_chunk_batch_size"),
            "voice_clone_codec_batch_size": result.get("voice_clone_codec_batch_size"),
        }

    def synthesize_stream(
        self,
        *,
        text: str,
        voice: Optional[str] = None,
        mode: str = "voice_clone",
        output_audio_path: Optional[str | Path] = None,
        prompt_audio_path: Optional[str | Path] = None,
        prompt_text: Optional[str] = None,
        max_new_frames: int = 375,
        voice_clone_max_text_tokens: int = 75,
        voice_clone_max_memory_per_sample_gb: float = 1.0,
        tts_max_batch_size: int = 0,
        codec_max_batch_size: int = 0,
        do_sample: bool = True,
        text_temperature: float = 1.0,
        text_top_p: float = 1.0,
        text_top_k: int = 50,
        audio_temperature: float = 0.8,
        audio_top_p: float = 0.95,
        audio_top_k: int = 25,
        audio_repetition_penalty: float = 1.2,
        nq: Optional[int] = None,
        seed: Optional[int] = None,
        attn_implementation: Optional[str] = None,
    ) -> Iterator[dict[str, object]]:
        normalized_text = str(text or "").strip()
        if not normalized_text:
            raise ValueError("text is required")

        normalized_mode = str(mode).strip().lower()
        if normalized_mode not in {"continuation", "voice_clone"}:
            raise ValueError("mode must be either 'continuation' or 'voice_clone'")

        effective_prompt_audio_path: Optional[Path] = None
        resolved_voice = self.get_voice_preset(voice).name
        if normalized_mode == "voice_clone":
            effective_prompt_audio_path = self.resolve_prompt_audio_path(
                voice=resolved_voice,
                prompt_audio_path=prompt_audio_path,
            )
        elif prompt_audio_path is not None:
            effective_prompt_audio_path = self.resolve_prompt_audio_path(prompt_audio_path=prompt_audio_path)
            if not prompt_text:
                raise ValueError("continuation mode with prompt_audio_path also requires prompt_text")

        output_path = (
            Path(output_audio_path).expanduser().resolve()
            if output_audio_path is not None
            else self._build_output_path(prefix=f"{resolved_voice}_{normalized_mode}_stream")
        )
        output_path.parent.mkdir(parents=True, exist_ok=True)

        started_at = time.monotonic()
        final_result: dict[str, object] | None = None
        with self._lock:
            model = self._load_model_locked()
            model = self._restore_model_execution_state(model)
            requested_attn_implementation, effective_global_attn_implementation, effective_local_attn_implementation = (
                self._resolve_request_attention_implementation(attn_implementation)
            )
            audio_tokenizer = self._load_audio_tokenizer_locked(
                tts_attn_implementation=effective_global_attn_implementation
            )
            self._apply_model_attention_implementation(
                model,
                global_attn=effective_global_attn_implementation,
                local_attn=effective_local_attn_implementation,
            )
            if seed is not None:
                torch.manual_seed(seed)
                if torch.cuda.is_available():
                    torch.cuda.manual_seed_all(seed)

            try:
                for event in model.inference_stream(
                    text=normalized_text,
                    output_audio_path=str(output_path),
                    mode=normalized_mode,
                    prompt_text=prompt_text,
                    prompt_audio_path=None if effective_prompt_audio_path is None else str(effective_prompt_audio_path),
                    text_tokenizer_path=str(self.checkpoint_path),
                    audio_tokenizer=audio_tokenizer,
                    device=self.device,
                    nq=nq,
                    max_new_frames=int(max_new_frames),
                    voice_clone_max_text_tokens=int(voice_clone_max_text_tokens),
                    voice_clone_max_memory_per_sample_gb=float(voice_clone_max_memory_per_sample_gb),
                    tts_max_batch_size=int(tts_max_batch_size),
                    codec_max_batch_size=int(codec_max_batch_size),
                    do_sample=bool(do_sample),
                    use_kv_cache=True,
                    text_temperature=float(text_temperature),
                    text_top_p=float(text_top_p),
                    text_top_k=int(text_top_k),
                    audio_temperature=float(audio_temperature),
                    audio_top_p=float(audio_top_p),
                    audio_top_k=int(audio_top_k),
                    audio_repetition_penalty=float(audio_repetition_penalty),
                ):
                    event_type = str(event.get("type", ""))
                    if event_type == "audio":
                        waveform = torch.as_tensor(event["waveform"], dtype=torch.float32).cpu()
                        yield {
                            "type": "audio",
                            "waveform": waveform,
                            "waveform_numpy": waveform_to_numpy(waveform),
                            "sample_rate": int(event["sample_rate"]),
                            "chunk_index": int(event.get("chunk_index", 0)),
                            "is_pause": bool(event.get("is_pause", False)),
                            "emitted_audio_seconds": float(event.get("emitted_audio_seconds", 0.0)),
                            "lead_seconds": float(event.get("lead_seconds", 0.0)),
                        }
                        continue
                    if event_type == "result":
                        final_result = dict(event)
            except Exception:
                self._discard_loaded_audio_tokenizer_locked(
                    "streaming inference failed; reloading audio tokenizer on next request"
                )
                self._discard_loaded_model_locked("streaming inference failed; reloading checkpoint on next request")
                raise

            effective_global_attn_implementation, effective_local_attn_implementation = (
                self._read_model_attention_implementation(model)
            )
            current_parameter = next(model.parameters(), None)
            if current_parameter is not None and current_parameter.dtype != self.dtype:
                self._discard_loaded_model_locked(
                    f"streaming inference left model in dtype={current_parameter.dtype}; reloading checkpoint on next request"
                )

        if final_result is None:
            raise RuntimeError("Streaming synthesis finished without a final result.")

        waveform = torch.as_tensor(final_result["waveform"], dtype=torch.float32).cpu()
        yield {
            "type": "result",
            "audio_path": str(final_result["audio_path"]),
            "sample_rate": int(final_result["sample_rate"]),
            "waveform": waveform,
            "waveform_numpy": waveform_to_numpy(waveform),
            "audio_token_ids": final_result["audio_token_ids"],
            "reference_audio_token_ids": final_result["reference_audio_token_ids"],
            "elapsed_seconds": time.monotonic() - started_at,
            "voice": resolved_voice,
            "mode": normalized_mode,
            "prompt_audio_path": None if effective_prompt_audio_path is None else str(effective_prompt_audio_path),
            "requested_attn_implementation": requested_attn_implementation,
            "effective_global_attn_implementation": effective_global_attn_implementation,
            "effective_local_attn_implementation": effective_local_attn_implementation,
            "voice_clone_text_chunks": final_result.get("voice_clone_text_chunks"),
            "voice_clone_chunk_batch_size": final_result.get("voice_clone_chunk_batch_size"),
            "voice_clone_codec_batch_size": final_result.get("voice_clone_codec_batch_size"),
        }

    def warmup(
        self,
        *,
        text: str = "你好，欢迎使用 Nano-TTS。",
        voice: Optional[str] = None,
    ) -> dict[str, object]:
        return self.synthesize(
            text=text,
            voice=voice or self.default_voice,
            mode="voice_clone",
            output_audio_path=self.output_dir / "_warmup" / "warmup.wav",
            max_new_frames=96,
            voice_clone_max_text_tokens=75,
            do_sample=False,
            text_temperature=1.0,
            text_top_p=1.0,
            text_top_k=50,
            audio_temperature=0.8,
            audio_top_p=0.95,
            audio_top_k=25,
            audio_repetition_penalty=1.0,
        )