portable_tts_runtime.py

#!/usr/bin/env python
from __future__ import annotations

import gc
import importlib.util
import json
import os
import site
import sys
import time
import wave
from dataclasses import dataclass
from pathlib import Path
from types import SimpleNamespace
from typing import Any

import numpy as np
import torch
from transformers import AutoConfig, AutoModel, AutoTokenizer, TorchAoConfig

from decoder4_features_torch import Decoder4FeatureExtractor


BUNDLE_ROOT = Path(__file__).resolve().parent


def _candidate_roots() -> list[Path]:
    roots = [
        BUNDLE_ROOT,
        BUNDLE_ROOT.parent,
        Path.cwd(),
        Path.cwd() / "moss_tts_clipper_istftnet2_release",
        BUNDLE_ROOT.parent / "moss_tts_clipper_istftnet2_release",
    ]
    seen: set[Path] = set()
    unique: list[Path] = []
    for root in roots:
        resolved = root.resolve()
        if resolved not in seen:
            seen.add(resolved)
            unique.append(resolved)
    return unique


def resolve_asset(value: str | Path) -> Path:
    path = Path(value)
    if path.exists():
        return path.resolve()
    for root in _candidate_roots():
        candidate = root / path
        if candidate.exists():
            return candidate.resolve()
    return path


def add_release_root_to_syspath(checkpoint: str | Path) -> Path:
    checkpoint_path = resolve_asset(checkpoint)
    release_root = checkpoint_path.parent
    release_root_str = str(release_root)
    if release_root_str not in sys.path:
        sys.path.insert(0, release_root_str)
    return checkpoint_path


def ensure_nvidia_library_path() -> None:
    if os.environ.get("MOSS_TTS_NVIDIA_LD_LIBRARY_PATH_READY") == "1":
        return
    lib_dirs: list[str] = []
    site_roots = list(dict.fromkeys(site.getsitepackages() + [site.getusersitepackages()]))
    for site_root in site_roots:
        nvidia_root = Path(site_root) / "nvidia"
        if not nvidia_root.exists():
            continue
        for lib_dir in nvidia_root.glob("*/lib"):
            if lib_dir.is_dir():
                lib_dirs.append(str(lib_dir))
    if not lib_dirs:
        return
    current = os.environ.get("LD_LIBRARY_PATH", "")
    current_parts = [part for part in current.split(":") if part]
    wanted = [part for part in lib_dirs if part not in current_parts]
    if not wanted:
        os.environ["MOSS_TTS_NVIDIA_LD_LIBRARY_PATH_READY"] = "1"
        return
    os.environ["LD_LIBRARY_PATH"] = ":".join(wanted + current_parts)
    os.environ["MOSS_TTS_NVIDIA_LD_LIBRARY_PATH_READY"] = "1"
    os.execv(sys.executable, [sys.executable, *sys.argv])


def choose_providers(requested: str, available: list[str]) -> list[str]:
    return [requested] if requested in available else ["CPUExecutionProvider"]


def select_dtype(device: torch.device, dtype_arg: str) -> torch.dtype:
    if device.type != "cuda":
        return torch.float32
    if dtype_arg == "fp16":
        return torch.float16
    if dtype_arg == "bf16":
        return torch.bfloat16
    if dtype_arg == "fp32":
        return torch.float32
    major, _minor = torch.cuda.get_device_capability(device)
    return torch.bfloat16 if major >= 8 else torch.float16


def save_wav_pcm16(path: Path, audio: torch.Tensor, sample_rate: int) -> None:
    pcm = (audio.clamp(-1.0, 1.0) * 32767.0).to(torch.int16).cpu().numpy()
    path.parent.mkdir(parents=True, exist_ok=True)
    with wave.open(str(path), "wb") as handle:
        handle.setnchannels(1)
        handle.setsampwidth(2)
        handle.setframerate(sample_rate)
        handle.writeframes(pcm.tobytes())


def summarize(values: list[float]) -> dict[str, float]:
    if not values:
        return {"mean": 0.0, "min": 0.0, "max": 0.0}
    return {
        "mean": sum(values) / len(values),
        "min": min(values),
        "max": max(values),
    }


def build_processor_without_audio_model(checkpoint: str | Path):
    bundled_release_root = BUNDLE_ROOT.parent
    if (bundled_release_root / "moss_tts_local_clipper_checkpoint").exists():
        bundled_release_root_str = str(bundled_release_root)
        if bundled_release_root_str not in sys.path:
            sys.path.insert(0, bundled_release_root_str)
    try:
        from moss_tts_local_clipper_checkpoint.processing_moss_tts import MossTTSDelayProcessor
    except ImportError:
        processing_file = Path(checkpoint) / "processing_moss_tts.py"
        if not processing_file.exists():
            raise
        spec = importlib.util.spec_from_file_location("moss_tts_checkpoint_processing", processing_file)
        if spec is None or spec.loader is None:
            raise ImportError(f"Could not load processor module from {processing_file}.")
        module = importlib.util.module_from_spec(spec)
        sys.modules[spec.name] = module
        spec.loader.exec_module(module)
        MossTTSDelayProcessor = module.MossTTSDelayProcessor

    config = AutoConfig.from_pretrained(checkpoint, trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
    return MossTTSDelayProcessor(tokenizer=tokenizer, model_config=config)


def load_style_extractor_class(checkpoint: str | Path):
    try:
        from moss_tts_local.style_features import BertStyleFeatureExtractor

        return BertStyleFeatureExtractor
    except ImportError:
        pass

    style_file = Path(checkpoint) / "style_features.py"
    if style_file.exists():
        spec = importlib.util.spec_from_file_location("moss_tts_checkpoint_style_features", style_file)
        if spec is None or spec.loader is None:
            raise ImportError(f"Could not load style feature module from {style_file}.")
        module = importlib.util.module_from_spec(spec)
        sys.modules[spec.name] = module
        spec.loader.exec_module(module)
        return module.BertStyleFeatureExtractor

    from moss_tts_local_clipper_checkpoint.style_features import BertStyleFeatureExtractor

    return BertStyleFeatureExtractor


class TorchDecoder4FeatureRuntime:
    def __init__(self, codec_path: str | Path, device: torch.device, dtype: torch.dtype) -> None:
        if device.type != "cuda" or dtype != torch.float16:
            raise RuntimeError("torch_fp16 decoder4 runtime requires CUDA fp16.")
        audio_tokenizer = AutoModel.from_pretrained(codec_path, trust_remote_code=True).to(
            device=device,
            dtype=dtype,
        )
        audio_tokenizer.eval()
        num_quantizers = int(audio_tokenizer.config.quantizer_kwargs.get("num_quantizers", 32))
        self.extractor = Decoder4FeatureExtractor(
            audio_tokenizer,
            num_quantizers=num_quantizers,
            output_dtype=dtype,
        ).to(device=device, dtype=dtype)
        self.extractor.eval()
        self.device = device
        self.inputs = [SimpleNamespace(name="codes"), SimpleNamespace(name="lengths")]
        del audio_tokenizer
        gc.collect()
        try:
            torch.cuda.empty_cache()
        except Exception:
            pass

    def get_inputs(self) -> list[Any]:
        return self.inputs

    def get_providers(self) -> list[str]:
        return ["torch_fp16_cuda"]

    def run(self, output_names: Any, feeds: dict[str, np.ndarray]) -> list[np.ndarray]:
        del output_names
        codes = torch.from_numpy(feeds["codes"]).to(device=self.device, dtype=torch.long)
        lengths = torch.from_numpy(feeds["lengths"]).to(device=self.device, dtype=torch.long)
        with torch.inference_mode():
            features, feature_lengths = self.extractor(codes, lengths)
        if self.device.type == "cuda":
            torch.cuda.synchronize()
        return [
            features.detach().float().cpu().numpy(),
            feature_lengths.detach().cpu().numpy(),
        ]


def run_vocoder_onnx(session: Any, features: np.ndarray, feature_lengths: np.ndarray) -> torch.Tensor:
    input_name = session.get_inputs()[0].name
    audio_np = session.run(None, {input_name: features.astype(np.float32, copy=False)})[0]
    samples = int(feature_lengths[0]) * 960
    audio = torch.from_numpy(audio_np[0, 0, :samples]).float().cpu()
    return torch.nan_to_num(audio).reshape(-1).clamp(-1.0, 1.0)


class TorchScriptVocoderRuntime:
    def __init__(
        self,
        decoder_dir: str | Path,
        device: torch.device,
        *,
        use_cudagraph: bool = False,
        bucket_frames: int = 0,
    ) -> None:
        artifact = Path(decoder_dir) / (
            "istftnet2_decoder_cuda.ts" if device.type == "cuda" else "istftnet2_decoder_cpu.ts"
        )
        if not artifact.exists():
            raise FileNotFoundError(f"Missing TorchScript vocoder artifact: {artifact}")
        self.device = device
        # Load on CPU first, then move to device via .to() so ZeroGPU's
        # torch.cuda hijack intercepts the move (torch.jit.load with
        # map_location="cuda" bypasses the hijack and fails in the main
        # process where no real GPU is attached).
        self.module = torch.jit.load(str(artifact), map_location="cpu").eval()
        if device.type == "cuda":
            self.module = self.module.to(device)
        self.use_cudagraph = bool(use_cudagraph and device.type == "cuda")
        self.bucket_frames = max(0, int(bucket_frames))
        self.graphs: dict[tuple[Any, ...], tuple[torch.cuda.CUDAGraph, torch.Tensor, torch.Tensor]] = {}

    def get_inputs(self) -> list[Any]:
        return [SimpleNamespace(name="features")]

    def get_providers(self) -> list[str]:
        suffix = "_cudagraph" if self.use_cudagraph else ""
        return [f"torchscript_{self.device.type}{suffix}"]

    def prewarm_buckets(self, frame_lengths: list[int]) -> dict[str, Any]:
        requested = [int(length) for length in frame_lengths if int(length) > 0]
        if not requested:
            return {"requested": [], "elapsed_sec": 0.0}
        if self.device.type == "cuda":
            torch.cuda.synchronize(self.device)
        start = time.perf_counter()
        warmed: list[int] = []
        with torch.inference_mode():
            for frames in requested:
                features = torch.randn(1, 768, frames, device=self.device, dtype=torch.float32)
                self.run_tensor(features)
                warmed.append(frames)
        if self.device.type == "cuda":
            torch.cuda.synchronize(self.device)
        return {"requested": requested, "warmed": warmed, "elapsed_sec": time.perf_counter() - start}

    def run_tensor(self, features_tensor: torch.Tensor) -> torch.Tensor:
        features_tensor = features_tensor.to(device=self.device, dtype=torch.float32).contiguous()
        if self.bucket_frames > 0 and features_tensor.ndim == 3:
            frames = int(features_tensor.shape[-1])
            bucketed = ((frames + self.bucket_frames - 1) // self.bucket_frames) * self.bucket_frames
            if bucketed > frames:
                padded = torch.zeros(
                    *features_tensor.shape[:-1],
                    bucketed,
                    device=features_tensor.device,
                    dtype=features_tensor.dtype,
                )
                padded[..., :frames] = features_tensor
                features_tensor = padded
        if not self.use_cudagraph:
            with torch.inference_mode():
                return self.module(features_tensor)
        key = (
            features_tensor.device.index,
            features_tensor.dtype,
            tuple(features_tensor.shape),
        )
        entry = self.graphs.get(key)
        if entry is None:
            try:
                static_features = torch.empty_like(features_tensor)
                static_features.copy_(features_tensor)
                warmup_stream = torch.cuda.Stream(device=features_tensor.device)
                warmup_stream.wait_stream(torch.cuda.current_stream(features_tensor.device))
                with torch.cuda.stream(warmup_stream), torch.inference_mode():
                    for _ in range(3):
                        static_audio = self.module(static_features)
                torch.cuda.current_stream(features_tensor.device).wait_stream(warmup_stream)

                graph = torch.cuda.CUDAGraph()
                with torch.cuda.graph(graph), torch.inference_mode():
                    static_audio = self.module(static_features)
                entry = (graph, static_features, static_audio)
                self.graphs[key] = entry
            except Exception:
                self.use_cudagraph = False
                self.graphs.clear()
                try:
                    torch.cuda.synchronize(features_tensor.device)
                except Exception:
                    pass
                with torch.inference_mode():
                    return self.module(features_tensor)

        graph, static_features, static_audio = entry
        static_features.copy_(features_tensor)
        graph.replay()
        return static_audio


def run_vocoder(
    session: Any,
    features: np.ndarray | torch.Tensor,
    feature_lengths: np.ndarray | torch.Tensor,
) -> torch.Tensor:
    if isinstance(session, TorchScriptVocoderRuntime):
        if isinstance(features, np.ndarray):
            features_tensor = torch.from_numpy(features).to(device=session.device, dtype=torch.float32)
        else:
            features_tensor = features.to(device=session.device, dtype=torch.float32)
        with torch.inference_mode():
            audio_tensor = session.run_tensor(features_tensor)
        if session.device.type == "cuda":
            torch.cuda.synchronize()
        if isinstance(feature_lengths, torch.Tensor):
            samples = int(feature_lengths[0].item()) * 960
        else:
            samples = int(feature_lengths[0]) * 960
        audio = audio_tensor[0, 0, :samples].detach().float().cpu()
        return torch.nan_to_num(audio).reshape(-1).clamp(-1.0, 1.0)

    if isinstance(features, torch.Tensor):
        features = features.detach().cpu().numpy()
    if isinstance(feature_lengths, torch.Tensor):
        feature_lengths = feature_lengths.detach().cpu().numpy()
    return run_vocoder_onnx(session, features, feature_lengths)


def parse_int_csv(value: str) -> list[int]:
    values: list[int] = []
    for part in value.split(","):
        part = part.strip()
        if not part:
            continue
        values.append(int(part))
    return values


def enable_static_cache_for_global_model(model: Any) -> None:
    type(model)._can_compile_fullgraph = True
    language_config = model.config.language_config
    for field in (
        "max_position_embeddings",
        "hidden_size",
        "num_attention_heads",
        "head_dim",
        "num_key_value_heads",
        "sliding_window",
        "layer_types",
        "num_hidden_layers",
    ):
        if hasattr(language_config, field):
            setattr(model.config, field, getattr(language_config, field))


def configure_torch_compile() -> None:
    import importlib
    import torch._inductor.config as inductor_config

    dynamo_config = importlib.import_module("torch._dynamo.config")
    dynamo_config.cache_size_limit = 256
    dynamo_config.capture_scalar_outputs = True
    inductor_config.triton.cudagraphs = False
    inductor_config.triton.cudagraph_trees = False
    inductor_config.triton.cudagraph_skip_dynamic_graphs = True


@dataclass
class OptimizedTTSConfig:
    checkpoint: str | Path = "moss_tts_local_clipper_checkpoint"
    codec_path: str | Path = "moss_audio_tokenizer"
    decoder_dir: str | Path = "istftnet2_decoder4_50hz"
    decoder4_features_onnx: str | Path = "ort_sessions/decoder4_features_fp32/model.onnx"
    decoder_runtime: str = "torchscript_cuda"
    vocoder_cudagraph: bool = False
    vocoder_bucket_frames: int = 0
    vocoder_prewarm_buckets: str = ""
    decoder4_features_runtime: str = "torch_fp16"
    decoder4_provider: str = "CUDAExecutionProvider"
    dtype: str = "fp16"
    tts_quantization: str = "none"
    torch_opt_mode: str = "static-local-cache-triton-compile-cudagraph"
    compile_mode: str = "max-autotune-no-cudagraphs"
    cache_implementation: str = "static"
    compile_global_transformer: bool = True
    global_compile_mode: str = "default"
    attn_implementation: str = "sdpa"
    fast_prepare_inputs: bool = True
    tensorrt_local: bool = False
    triton_top_p: bool = False
    triton_fused_lm_head: bool = False
    triton_qkv_cache: bool = False
    packed_local_qkv: bool = False
    packed_local_mlp: bool = False
    packed_adapter_mlp: bool = False
    packed_adapter_mlp_scope: str = "heads"
    static_packed_weights: bool = False
    triton_rmsnorm: bool = False
    tensorize_rmsnorm_eps: bool = True
    fast_control_head: bool = False
    feedback_lookup: bool = True
    local_compile_fullgraph: bool = False
    style_bert_model: str = "cirimus/modernbert-base-go-emotions"
    style_bert_layer: int = 19
    style_bert_max_length: int = 512


class OptimizedTTSRunner:
    def __init__(self, config: OptimizedTTSConfig) -> None:
        ensure_nvidia_library_path()
        torch.backends.cudnn.benchmark = False
        try:
            torch.backends.cuda.enable_cudnn_sdp(False)
            torch.backends.cuda.enable_flash_sdp(True)
            torch.backends.cuda.enable_mem_efficient_sdp(True)
            torch.backends.cuda.enable_math_sdp(True)
        except Exception:
            pass

        checkpoint = add_release_root_to_syspath(config.checkpoint)
        codec_path = resolve_asset(config.codec_path)
        decoder_dir = resolve_asset(config.decoder_dir)
        decoder4_features_onnx = resolve_asset(config.decoder4_features_onnx)

        from torch_hf_optimizations import install_torch_frame_sampler, tensorize_rmsnorm_eps

        import onnxruntime as ort

        self.config = config
        self.checkpoint = checkpoint
        self.codec_path = codec_path
        self.decoder_dir = decoder_dir
        self.decoder4_features_onnx = decoder4_features_onnx
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.dtype = select_dtype(self.device, config.dtype)

        wants_cuda = (
            config.decoder4_provider == "CUDAExecutionProvider"
            or config.decoder_runtime == "onnx_cuda"
            or config.decoder4_features_runtime == "torch_fp16"
        )
        if wants_cuda and hasattr(ort, "preload_dlls"):
            ort.preload_dlls(cuda=True, cudnn=True, directory="")
        available_providers = ort.get_available_providers()
        session_options = ort.SessionOptions()
        session_options.log_severity_level = 2
        self.available_ort_providers = available_providers

        self.processor = build_processor_without_audio_model(checkpoint)
        self.processor.model_config.sampling_rate = 48000
        if config.decoder4_features_runtime == "onnx":
            self.decoder4_session = ort.InferenceSession(
                str(decoder4_features_onnx),
                sess_options=session_options,
                providers=choose_providers(config.decoder4_provider, available_providers),
            )
        elif config.decoder4_features_runtime == "torch_fp16":
            self.decoder4_session = TorchDecoder4FeatureRuntime(codec_path, device=self.device, dtype=torch.float16)
        else:
            raise ValueError(f"Unsupported decoder4_features_runtime: {config.decoder4_features_runtime}")

        if config.decoder_runtime.startswith("torchscript_"):
            vocoder_device = torch.device(
                "cuda" if config.decoder_runtime == "torchscript_cuda" and torch.cuda.is_available() else "cpu"
            )
            self.vocoder_session = TorchScriptVocoderRuntime(
                decoder_dir,
                vocoder_device,
                use_cudagraph=config.vocoder_cudagraph,
                bucket_frames=config.vocoder_bucket_frames,
            )
        else:
            vocoder_provider = "CUDAExecutionProvider" if config.decoder_runtime == "onnx_cuda" else "CPUExecutionProvider"
            self.vocoder_session = ort.InferenceSession(
                str(decoder_dir / "istftnet2_decoder.onnx"),
                sess_options=session_options,
                providers=choose_providers(vocoder_provider, available_providers),
            )
        self.vocoder_prewarm_result: dict[str, Any] | None = None
        if isinstance(self.vocoder_session, TorchScriptVocoderRuntime):
            prewarm_buckets = parse_int_csv(config.vocoder_prewarm_buckets)
            if prewarm_buckets:
                try:
                    self.vocoder_prewarm_result = self.vocoder_session.prewarm_buckets(prewarm_buckets)
                except Exception:
                    # Prewarm can fail in ZeroGPU main process (no real GPU
                    # attached at startup). It's an optimization only.
                    self.vocoder_prewarm_result = {"requested": prewarm_buckets, "warmed": [], "elapsed_sec": 0.0, "skipped": True}
        tts_load_kwargs: dict[str, Any] = {
            "trust_remote_code": True,
            "torch_dtype": self.dtype,
            "attn_implementation": config.attn_implementation,
        }
        if config.tts_quantization != "none":
            if self.device.type != "cuda":
                raise RuntimeError("TorchAO TTS quantization requires CUDA in this runner.")
            modules_to_not_convert = None
            if config.tts_quantization in {
                "torchao_fp8_weight_only_global",
                "torchao_fp8_dynamic_activation_fp8_weight_global",
            }:
                modules_to_not_convert = [
                    "local_transformer",
                    "speech_embedding_to_local_mlp",
                    "local_to_speech_embedding_mlps",
                    "layer_norm_before_lm_heads",
                    "lm_heads",
                ]
                if config.tts_quantization == "torchao_fp8_weight_only_global":
                    quant_kind = "torchao_fp8_weight_only"
                else:
                    quant_kind = "torchao_fp8_dynamic_activation_fp8_weight"
            else:
                quant_kind = config.tts_quantization
            if quant_kind == "torchao_fp8_weight_only":
                from torchao.quantization import Float8WeightOnlyConfig

                tts_load_kwargs["quantization_config"] = TorchAoConfig(
                    Float8WeightOnlyConfig(),
                    modules_to_not_convert=modules_to_not_convert,
                )
            elif quant_kind == "torchao_fp8_dynamic_activation_fp8_weight":
                from torchao.quantization import Float8DynamicActivationFloat8WeightConfig

                tts_load_kwargs["quantization_config"] = TorchAoConfig(
                    Float8DynamicActivationFloat8WeightConfig(),
                    modules_to_not_convert=modules_to_not_convert,
                )
            else:
                raise ValueError(f"Unsupported tts_quantization: {config.tts_quantization}")
            tts_load_kwargs["device_map"] = {"": self.device.index or 0}
        self.model = AutoModel.from_pretrained(checkpoint, **tts_load_kwargs)
        if config.tts_quantization == "none":
            self.model = self.model.to(self.device)
        self.model.eval()
        self.style_feature_dim = int(getattr(self.model.config, "style_feature_dim", 0) or 0)
        self.num_emotions = int(getattr(self.model.config, "num_emotions", 0) or 0)
        self.style_extractor = None
        if self.style_feature_dim > 0 and self.style_feature_dim != 2:
            BertStyleFeatureExtractor = load_style_extractor_class(checkpoint)
            self.style_extractor = BertStyleFeatureExtractor(
                repo_id=config.style_bert_model,
                layer_index=config.style_bert_layer,
                max_length=config.style_bert_max_length,
                target_dim=self.style_feature_dim,
                device=self.device,
                dtype=self.dtype,
                attn_implementation="sdpa" if self.device.type == "cuda" else "eager",
            )
        self.tensorized_rmsnorm_eps = 0
        if config.tensorize_rmsnorm_eps:
            self.tensorized_rmsnorm_eps = tensorize_rmsnorm_eps(self.model, device=self.device)
        if config.compile_global_transformer:
            configure_torch_compile()
        if config.compile_global_transformer or config.cache_implementation == "static":
            enable_static_cache_for_global_model(self.model)
        if config.compile_global_transformer:
            if config.global_compile_mode == "default":
                global_compile_mode = None
                global_compile_options = None
                self.global_compile_mode_effective = "default"
            else:
                global_compile_mode = None
                global_compile_options = {
                    "triton.cudagraphs": False,
                    "triton.cudagraph_trees": False,
                    "triton.cudagraph_skip_dynamic_graphs": True,
                }
                self.global_compile_mode_effective = "default-no-inductor-cudagraphs"
            self.model.model.language_model = torch.compile(
                self.model.model.language_model,
                mode=global_compile_mode,
                options=global_compile_options,
                fullgraph=False,
                dynamic=True,
            )
        else:
            self.global_compile_mode_effective = "disabled"
        install_torch_frame_sampler(
            self.model,
            mode=config.torch_opt_mode,
            compile_mode=config.compile_mode,
            packed_local_qkv=config.packed_local_qkv,
            packed_local_mlp=config.packed_local_mlp,
            packed_adapter_mlp=config.packed_adapter_mlp,
            packed_adapter_mlp_scope=config.packed_adapter_mlp_scope,
            static_packed_weights=config.static_packed_weights,
            triton_rmsnorm=config.triton_rmsnorm,
            tensorrt_local=config.tensorrt_local,
            triton_top_p=config.triton_top_p,
            triton_fused_lm_head=config.triton_fused_lm_head,
            triton_qkv_cache=config.triton_qkv_cache,
            fast_prepare_inputs=config.fast_prepare_inputs,
            fast_control_head=config.fast_control_head,
            feedback_lookup=config.feedback_lookup,
            local_compile_fullgraph=config.local_compile_fullgraph,
        )
        self.sample_rate = int(self.processor.model_config.sampling_rate)

    def decode_outputs(self, outputs: list[tuple[int, torch.Tensor]]) -> tuple[torch.Tensor, float]:
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        start = time.perf_counter()
        decoded_segments: list[torch.Tensor] = []
        codes_input_name = self.decoder4_session.get_inputs()[0].name
        lengths_input_name = self.decoder4_session.get_inputs()[1].name
        for start_length, generation_ids in outputs:
            frame_tokens = generation_ids.detach().cpu()[:, 0]
            audio_codes = generation_ids.detach().cpu()[:, 1:]
            is_pad = (audio_codes == int(self.processor.model_config.audio_pad_code)).all(dim=1)
            is_eos = frame_tokens == int(self.processor.model_config.audio_end_token_id)
            non_pad = ~is_pad & ~is_eos
            if not non_pad.any():
                continue
            idx = torch.nonzero(non_pad).squeeze(1)
            breaks = torch.where(idx[1:] != idx[:-1] + 1)[0] + 1
            segments_idx = [idx] if breaks.numel() == 0 else list(torch.split(idx, breaks.tolist()))
            for segment_index, segment_idx in enumerate(segments_idx):
                segment_codes = audio_codes[segment_idx].contiguous()
                if int(segment_codes.shape[0]) <= 0:
                    continue
                codes_np = segment_codes.T.unsqueeze(1).numpy().astype(np.int64, copy=False)
                lengths_np = np.asarray([int(segment_codes.shape[0])], dtype=np.int64)
                feeds = {codes_input_name: codes_np, lengths_input_name: lengths_np}
                if hasattr(self.decoder4_session, "run_tensors"):
                    features, feature_lengths = self.decoder4_session.run_tensors(feeds)
                else:
                    features, feature_lengths = self.decoder4_session.run(None, feeds)
                segment_audio = run_vocoder(self.vocoder_session, features, feature_lengths)
                if segment_index == 0 and int(start_length) > 0:
                    trim_ratio = max(0.0, min(float(start_length) / float(segment_codes.shape[0]), 1.0))
                    if trim_ratio >= 1.0:
                        continue
                    if trim_ratio > 0.0:
                        segment_audio = segment_audio[..., int(segment_audio.shape[-1] * trim_ratio) :]
                decoded_segments.append(segment_audio)
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        elapsed = time.perf_counter() - start
        if not decoded_segments:
            raise RuntimeError("Generation did not produce decodable audio.")
        return torch.cat(decoded_segments, dim=-1).float().cpu(), elapsed

    def synthesize(
        self,
        *,
        text: str,
        language: str = "en",
        speaker_id: int = 31,
        max_new_tokens: int = 160,
        text_temperature: float = 0.0,
        text_top_p: float = 1.0,
        text_top_k: int | None = None,
        audio_temperature: float = 0.8,
        audio_top_p: float = 0.92,
        audio_top_k: int | None = None,
        audio_repetition_penalty: float = 1.0,
        n_vq_for_inference: int = 32,
        style_text: str | None = None,
        style_features: torch.Tensor | None = None,
        style_emotion_id: int | None = None,
        style_energy: float = 0.7,
    ) -> dict[str, Any]:
        conversations = [[self.processor.build_user_message(text=text, language=language)]]
        prompt_start = time.perf_counter()
        batch = self.processor(conversations, mode="generation")
        input_ids = batch["input_ids"].to(self.device)
        attention_mask = batch["attention_mask"].to(self.device)
        prompt_sec = time.perf_counter() - prompt_start

        if self.device.type == "cuda":
            torch.cuda.synchronize()
        generate_start = time.perf_counter()
        generation_kwargs = {}
        if self.config.cache_implementation != "none":
            generation_kwargs["cache_implementation"] = self.config.cache_implementation
        speaker_tensor = torch.tensor([speaker_id], device=self.device, dtype=torch.long)
        generation_kwargs["speaker_ids"] = speaker_tensor
        if self.style_feature_dim > 0:
            if style_features is None:
                if self.style_feature_dim == 2 and style_emotion_id is not None:
                    max_emotion = max(0, (self.num_emotions or 1) - 1)
                    style_features = torch.tensor(
                        [
                            [
                                float(max(0, min(max_emotion, int(style_emotion_id)))),
                                float(max(0.0, min(1.0, style_energy))),
                            ]
                        ],
                        dtype=torch.float32,
                    )
                elif self.style_extractor is None:
                    raise RuntimeError("Checkpoint expects style features, but no style extractor is available.")
                else:
                    style_features = self.style_extractor.encode((style_text or text).strip())
            generation_kwargs["style_features"] = style_features.to(device=self.device)
        with torch.inference_mode():
            outputs = self.model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                max_new_tokens=max_new_tokens,
                n_vq_for_inference=n_vq_for_inference,
                text_temperature=text_temperature,
                text_top_p=text_top_p,
                text_top_k=text_top_k,
                audio_temperature=audio_temperature,
                audio_top_p=audio_top_p,
                audio_top_k=None,
                audio_repetition_penalty=audio_repetition_penalty,
                **generation_kwargs,
            )
        if self.device.type == "cuda":
            torch.cuda.synchronize()
        generate_sec = time.perf_counter() - generate_start

        audio, decode_sec = self.decode_outputs(outputs)
        audio_sec = float(audio.numel() / self.sample_rate)
        return {
            "audio": audio,
            "audio_sec": audio_sec,
            "prompt_sec": prompt_sec,
            "generate_sec": generate_sec,
            "decode_sec": decode_sec,
            "generate_x_realtime": audio_sec / generate_sec if generate_sec else 0.0,
            "total_x_realtime": audio_sec / (generate_sec + decode_sec) if generate_sec + decode_sec else 0.0,
            "generated_tokens": int(outputs[0][1].shape[0]) if outputs else 0,
            "prompt_tokens": int(input_ids.shape[1]),
        }

    def runtime_summary(self) -> dict[str, Any]:
        return {
            "device": str(self.device),
            "dtype": str(self.dtype),
            "tts_quantization": self.config.tts_quantization,
            "processor_loads_audio_tokenizer": False,
            "audio_decode_runtime": f"{self.config.decoder4_features_runtime}_decoder4_features_plus_{self.config.decoder_runtime}_vocoder",
            "checkpoint": str(self.checkpoint),
            "codec_path": str(self.codec_path),
            "decoder_dir": str(self.decoder_dir),
            "decoder4_features_onnx": str(self.decoder4_features_onnx),
            "vocoder_onnx": str(self.decoder_dir / "istftnet2_decoder.onnx"),
            "decoder4_providers": self.decoder4_session.get_providers(),
            "vocoder_providers": self.vocoder_session.get_providers(),
            "available_ort_providers": self.available_ort_providers,
            "torch_opt_mode": self.config.torch_opt_mode,
            "compile_mode": self.config.compile_mode,
            "cache_implementation": self.config.cache_implementation,
            "compile_global_transformer": bool(self.config.compile_global_transformer),
            "global_compile_mode": self.global_compile_mode_effective,
            "fast_prepare_inputs": bool(self.config.fast_prepare_inputs),
            "tensorrt_local": bool(self.config.tensorrt_local),
            "triton_top_p": bool(self.config.triton_top_p),
            "triton_fused_lm_head": bool(self.config.triton_fused_lm_head),
            "triton_qkv_cache": bool(self.config.triton_qkv_cache),
            "packed_local_qkv": bool(self.config.packed_local_qkv),
            "packed_local_mlp": bool(self.config.packed_local_mlp),
            "packed_adapter_mlp": bool(self.config.packed_adapter_mlp),
            "packed_adapter_mlp_scope": self.config.packed_adapter_mlp_scope,
            "static_packed_weights": bool(self.config.static_packed_weights),
            "triton_rmsnorm": bool(self.config.triton_rmsnorm),
            "tensorize_rmsnorm_eps": bool(self.config.tensorize_rmsnorm_eps),
            "tensorized_rmsnorm_eps": int(self.tensorized_rmsnorm_eps),
            "fast_control_head": bool(self.config.fast_control_head),
            "feedback_lookup": bool(self.config.feedback_lookup),
            "local_compile_fullgraph": bool(self.config.local_compile_fullgraph),
            "style_feature_dim": int(self.style_feature_dim),
            "num_emotions": int(self.num_emotions),
            "vocoder_cudagraph": bool(self.config.vocoder_cudagraph),
            "vocoder_bucket_frames": int(self.config.vocoder_bucket_frames),
            "vocoder_prewarm_buckets": parse_int_csv(self.config.vocoder_prewarm_buckets),
            "vocoder_prewarm_result": self.vocoder_prewarm_result,
        }


def write_summary(path: Path, summary: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    serializable = {
        key: value
        for key, value in summary.items()
        if not isinstance(value, torch.Tensor)
    }
    path.write_text(json.dumps(serializable, indent=2))