#!/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))