from __future__ import annotations import json import logging import time import unicodedata import wave from collections import defaultdict from pathlib import Path from typing import Any, Callable, Optional, Sequence import numpy as np import sentencepiece as spm from scripts.axe_session import AxeSession try: import onnxruntime as ort _HAVE_ORT = True except ImportError: _HAVE_ORT = False SAMPLE_MODE_GREEDY = "greedy" SAMPLE_MODE_FIXED = "fixed" SAMPLE_MODE_FULL = "full" SENTENCE_END_PUNCTUATION = set(".!?。！？") CLAUSE_SPLIT_PUNCTUATION = set(",，、；;：:") CLOSING_PUNCTUATION = set(['"', "'", "\u2019", "\u201d", ")", "]", "}", "）", "】", "》", "」", "』"]) _MANIFEST_CANDIDATES = ("browser_poc_manifest.json",) _FULL8_IO_MAP: dict[str, tuple[list[str], list[str]]] = { "prefill": ( ["input_ids", "attention_mask"], ["global_hidden", "present_key_0", "present_value_0", "present_key_1", "present_value_1", "present_key_2", "present_value_2", "present_key_3", "present_value_3", "present_key_4", "present_value_4", "present_key_5", "present_value_5", "present_key_6", "present_value_6", "present_key_7", "present_value_7", "present_key_8", "present_value_8", "present_key_9", "present_value_9", "present_key_10", "present_value_10", "present_key_11", "present_value_11"], ), "decode": ( ["input_ids", "past_valid_lengths", "past_key_0", "past_value_0", "past_key_1", "past_value_1", "past_key_2", "past_value_2", "past_key_3", "past_value_3", "past_key_4", "past_value_4", "past_key_5", "past_value_5", "past_key_6", "past_value_6", "past_key_7", "past_value_7", "past_key_8", "past_value_8", "past_key_9", "past_value_9", "past_key_10", "past_value_10", "past_key_11", "past_value_11"], ["global_hidden", "present_key_0", "present_value_0", "present_key_1", "present_value_1", "present_key_2", "present_value_2", "present_key_3", "present_value_3", "present_key_4", "present_value_4", "present_key_5", "present_value_5", "present_key_6", "present_value_6", "present_key_7", "present_value_7", "present_key_8", "present_value_8", "present_key_9", "present_value_9", "present_key_10", "present_value_10", "present_key_11", "present_value_11"], ), "local_decoder": ( ["global_hidden", "text_token_id", "audio_prefix_token_ids"], ["text_logits", "audio_logits"], ), "local_fixed_sampled_frame": ( ["global_hidden", "repetition_seen_mask", "assistant_random_u", "audio_random_u"], ["should_continue", "frame_token_ids"], ), "codec_encode": ( ["waveform", "input_lengths"], ["audio_codes", "audio_code_lengths"], ), "codec_decode": ( ["audio_codes", "audio_code_lengths"], ["audio", "audio_lengths"], ), } _AXMODEL_FILE_MAP: dict[str, str] = { "prefill": "tts_prefill.axmodel", "decode": "tts_decode_step.axmodel", "local_decoder": "tts_local_decoder.axmodel", "local_fixed_sampled_frame": "tts_local_fixed_sampled_frame.axmodel", "codec_encode": "codec_encode.axmodel", "codec_decode": "codec_decode.axmodel", } _ONNX_FILE_MAP: dict[str, str] = { "decode": "tts_decode_step.onnx", } def _argmax(values: np.ndarray) -> int: return int(np.argmax(values)) def _softmax(values: np.ndarray) -> np.ndarray: max_v = float(np.max(values)) shifted = np.asarray(values - max_v, dtype=np.float64) exps = np.exp(shifted) return exps / np.sum(exps) def _apply_repetition_penalty( values: np.ndarray, previous_token_ids: list[int], penalty: float ) -> np.ndarray: if not previous_token_ids or penalty == 1.0: return values result = values.copy() for tid in set(int(t) for t in previous_token_ids): if 0 <= tid < result.shape[0]: result[tid] = result[tid] * penalty if result[tid] < 0 else result[tid] / penalty return result def _argmax_with_repetition_penalty( values: np.ndarray, previous_token_set: set[int], penalty: float ) -> int: best_idx, best_val = 0, float("-inf") apply_penalty = bool(previous_token_set) and penalty != 1.0 for idx, v in enumerate(values): score = float(v) if apply_penalty and idx in previous_token_set: score = score * penalty if score < 0 else score / penalty if score > best_val: best_val = score best_idx = idx return int(best_idx) def _sample_from_scores( values: np.ndarray, *, do_sample: bool, temperature: float, top_k: int, top_p: float, rng: np.random.Generator, ) -> int: if not do_sample: return _argmax(values) if temperature <= 0: raise ValueError("temperature must be positive when do_sample=True") scores = np.asarray(values, dtype=np.float32).copy() / float(temperature) if top_k > 0 and top_k < scores.shape[0]: threshold = float(np.sort(scores)[::-1][top_k - 1]) scores[scores < threshold] = float("-inf") if 0 < top_p < 1: indexed = sorted(enumerate(scores.tolist()), key=lambda x: x[1], reverse=True) sorted_scores = np.asarray([x[1] for x in indexed], dtype=np.float32) sorted_probs = _softmax(sorted_scores) remove = [False] * len(indexed) cumulative = 0.0 for i, p in enumerate(sorted_probs): cumulative += float(p) if cumulative > float(top_p): remove[i] = True for i in range(len(remove) - 1, 0, -1): remove[i] = remove[i - 1] if remove: remove[0] = False for i, should_remove in enumerate(remove): if should_remove: scores[indexed[i][0]] = float("-inf") probs = _softmax(scores) rv = float(rng.random()) for i, p in enumerate(probs): rv -= float(p) if rv <= 0: return int(i) return _argmax(scores) def _sample_assistant_text_token( text_logits: np.ndarray, manifest: dict[str, Any], generation_defaults: dict[str, Any], rng: np.random.Generator, ) -> int: candidate_ids = np.asarray( [int(manifest["tts_config"]["audio_assistant_slot_token_id"]), int(manifest["tts_config"]["audio_end_token_id"])], dtype=np.int32, ) candidate_scores = text_logits[candidate_ids] sampled_idx = _sample_from_scores( candidate_scores, do_sample=False, # deterministic: always pick the higher-scoring token temperature=float(generation_defaults["text_temperature"]), top_k=min(int(generation_defaults["text_top_k"]), int(candidate_scores.shape[0])), top_p=float(generation_defaults["text_top_p"]), rng=rng, ) return int(candidate_ids[sampled_idx]) def _sample_audio_token( audio_logits: np.ndarray, previous_token_ids: list[int], previous_token_set: set[int], generation_defaults: dict[str, Any], rng: np.random.Generator, ) -> int: penalty = float(generation_defaults["audio_repetition_penalty"]) if not bool(generation_defaults["do_sample"]): return _argmax_with_repetition_penalty(audio_logits, previous_token_set, penalty) penalized = _apply_repetition_penalty(audio_logits, previous_token_ids, penalty) return _sample_from_scores( penalized, do_sample=True, temperature=float(generation_defaults["audio_temperature"]), top_k=int(generation_defaults["audio_top_k"]), top_p=float(generation_defaults["audio_top_p"]), rng=rng, ) def _flatten3d_int32(nested: list[list[list[int]]]) -> tuple[np.ndarray, list[int]]: d0, d1, d2 = len(nested), len(nested[0]), len(nested[0][0]) data = np.zeros((d0 * d1 * d2,), dtype=np.int32) offset = 0 for i in range(d0): for j in range(d1): for k in range(d2): data[offset] = int(nested[i][j][k]) offset += 1 return data, [d0, d1, d2] def _flatten2d_int32(nested: list[list[int]]) -> tuple[np.ndarray, list[int]]: d0, d1 = len(nested), len(nested[0]) data = np.zeros((d0 * d1,), dtype=np.int32) offset = 0 for i in range(d0): for j in range(d1): data[offset] = int(nested[i][j]) offset += 1 return data, [d0, d1] def _extract_last_hidden(hidden: np.ndarray) -> np.ndarray: if hidden.ndim == 2: return hidden.astype(np.float32, copy=False) if hidden.ndim != 3 or hidden.shape[0] != 1: raise ValueError(f"Unexpected global_hidden shape: {hidden.shape}") return hidden[:, -1, :].astype(np.float32, copy=False) def _slice_channel_major_audio( audio: np.ndarray, start_sample: int = 0, end_sample: int | None = None ) -> list[np.ndarray]: if audio.ndim != 3 or audio.shape[0] != 1: raise ValueError(f"Unexpected audio tensor shape: {audio.shape}") channels = int(audio.shape[1]) total = int(audio.shape[2]) start = max(0, int(start_sample)) end = total if end_sample is None else max(start, min(int(end_sample), total)) return [audio[0, c, start:end].astype(np.float32, copy=False) for c in range(channels)] def _merge_audio_channels(channel_arrays: list[np.ndarray]) -> np.ndarray: if not channel_arrays: return np.zeros((0, 1), dtype=np.float32) if len(channel_arrays) == 1: return np.asarray(channel_arrays[0], dtype=np.float32).reshape(-1, 1) min_len = min(int(c.shape[0]) for c in channel_arrays) return np.stack([np.asarray(c[:min_len], dtype=np.float32) for c in channel_arrays], axis=1) def _concat_waveforms(waveforms: list[np.ndarray]) -> np.ndarray: if not waveforms: return np.zeros((0, 1), dtype=np.float32) non_empty = [w for w in waveforms if w.size > 0] if not non_empty: n_ch = int(waveforms[0].shape[1]) if waveforms[0].ndim == 2 else 1 return np.zeros((0, n_ch), dtype=np.float32) return np.concatenate(non_empty, axis=0) def _write_waveform_to_wav(path: str | Path, waveform: np.ndarray, sample_rate: int) -> Path: output_path = Path(path).expanduser().resolve() output_path.parent.mkdir(parents=True, exist_ok=True) audio = np.asarray(waveform, dtype=np.float32) if audio.ndim == 1: audio = audio.reshape(-1, 1) clipped = np.clip(audio, -1.0, 1.0) pcm16 = np.round(clipped * 32767.0).astype(np.int16) with wave.open(str(output_path), "wb") as wf: wf.setnchannels(int(pcm16.shape[1])) wf.setsampwidth(2) wf.setframerate(int(sample_rate)) wf.writeframes(pcm16.tobytes()) return output_path def _load_audio_numpy(path: str | Path, target_sample_rate: int, target_channels: int) -> np.ndarray: path = Path(path).expanduser().resolve() waveform: np.ndarray | None = None src_sr: int = 0 try: import soundfile as sf data, src_sr = sf.read(str(path), dtype="float32", always_2d=True) # (samples, channels) waveform = data.T # (channels, samples) except Exception: pass if waveform is None: with wave.open(str(path), "rb") as wf: src_sr = wf.getframerate() n_ch = wf.getnchannels() sw = wf.getsampwidth() raw = wf.readframes(wf.getnframes()) if sw == 2: pcm = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0 elif sw == 4: pcm = np.frombuffer(raw, dtype=np.int32).astype(np.float32) / 2147483648.0 elif sw == 1: pcm = (np.frombuffer(raw, dtype=np.uint8).astype(np.float32) - 128.0) / 128.0 else: raise ValueError(f"Unsupported sample width: {sw}") waveform = pcm.reshape(-1, n_ch).T # (channels, samples) assert waveform is not None src_channels = int(waveform.shape[0]) if src_channels != target_channels: if src_channels == 1 and target_channels > 1: waveform = np.repeat(waveform, target_channels, axis=0) elif src_channels > 1 and target_channels == 1: waveform = waveform.mean(axis=0, keepdims=True) else: raise ValueError(f"Unsupported channel conversion: {src_channels} → {target_channels}") if src_sr != target_sample_rate: try: from scipy.signal import resample_poly from math import gcd g = gcd(int(target_sample_rate), int(src_sr)) up = int(target_sample_rate) // g down = int(src_sr) // g resampled = [] for ch in range(waveform.shape[0]): resampled.append(resample_poly(waveform[ch], up, down).astype(np.float32)) waveform = np.stack(resampled, axis=0) except ImportError: logging.warning("scipy not available, using linear interpolation for audio resampling") old_len = int(waveform.shape[1]) new_len = int(round(old_len * target_sample_rate / src_sr)) old_idx = np.linspace(0, old_len - 1, new_len) resampled = [] for ch in range(waveform.shape[0]): resampled.append(np.interp(old_idx, np.arange(old_len), waveform[ch]).astype(np.float32)) waveform = np.stack(resampled, axis=0) return waveform[np.newaxis, :, :] # (1, channels, samples) def _contains_cjk(text: str) -> bool: for ch in str(text or ""): if ( "\u4e00" <= ch <= "\u9fff" or "\u3400" <= ch <= "\u4dbf" or "\u3040" <= ch <= "\u30ff" or "\uac00" <= ch <= "\ud7af" ): return True return False def _normalize_punctuation(text: str) -> str: _WORD_INTERNAL = frozenset(["'", "\u2019", "-", "\u2013"]) # don't touch these result: list[str] = [] prev_is_punct: bool = False for ch in text: is_sent_end = ch in SENTENCE_END_PUNCTUATION is_punct = is_sent_end or ( ch not in _WORD_INTERNAL and unicodedata.category(ch).startswith("P") ) if is_punct: if prev_is_punct: if is_sent_end: result[-1] = ch else: result.append(ch if is_sent_end else ",") prev_is_punct = True else: result.append(ch) prev_is_punct = False return "".join(result) def _prepare_text_for_sentence_chunking(text: str) -> str: t = str(text or "").strip() if not t: raise ValueError("Text prompt cannot be empty.") t = t.replace("\r", " ").replace("\n", " ") while " " in t: t = t.replace(" ", " ") t = _normalize_punctuation(t) while ", ," in t: t = t.replace(", ,", ",") if _contains_cjk(t): if t[-1] not in SENTENCE_END_PUNCTUATION: t += "。" return t if t[:1].islower(): t = t[:1].upper() + t[1:] if t[-1].isalnum(): t += "." if len([x for x in t.split() if x]) < 5: t = f" {t}" return t def _split_text_by_punctuation(text: str, punctuation: set[str]) -> list[str]: sentences: list[str] = [] current: list[str] = [] idx = 0 t = str(text or "") while idx < len(t): ch = t[idx] current.append(ch) if ch in punctuation: lookahead = idx + 1 while lookahead < len(t) and t[lookahead] in CLOSING_PUNCTUATION: current.append(t[lookahead]) lookahead += 1 sentence = "".join(current).strip() if sentence: sentences.append(sentence) current.clear() while lookahead < len(t) and t[lookahead].isspace(): lookahead += 1 idx = lookahead continue idx += 1 tail = "".join(current).strip() if tail: sentences.append(tail) return sentences def _join_sentence_parts(left: str, right: str) -> str: if not left: return right if not right: return left if _contains_cjk(left) or _contains_cjk(right): return left + right return f"{left} {right}" class AxTtsRuntime: def __init__( self, config_dir: str | Path, axmodel_dir: str | Path, *, onnx_dir: str | Path | None = None, use_onnx_decode: bool = False, max_new_frames: int | None = None, do_sample: bool = True, sample_mode: str | None = None, ) -> None: self.config_dir = Path(config_dir).expanduser().resolve() self.axmodel_dir = Path(axmodel_dir).expanduser().resolve() self.use_onnx_decode = bool(use_onnx_decode) if onnx_dir is not None: self.onnx_dir = Path(onnx_dir).expanduser().resolve() else: self.onnx_dir = self.axmodel_dir.parent / "onnxmodels" self.manifest_path = self._find_manifest() self.manifest_dir = self.manifest_path.parent self.manifest: dict[str, Any] = json.loads( self.manifest_path.read_text(encoding="utf-8") ) gen = self.manifest["generation_defaults"] if max_new_frames is not None: gen["max_new_frames"] = int(max_new_frames) if do_sample is not None: gen["do_sample"] = bool(do_sample) raw_mode = sample_mode if sample_mode is not None else gen.get("sample_mode", "fixed") gen["sample_mode"] = self._normalize_sample_mode(raw_mode, bool(gen["do_sample"])) gen["do_sample"] = gen["sample_mode"] != SAMPLE_MODE_GREEDY self.tts_meta: dict[str, Any] = json.loads( self._resolve_path(self.manifest["model_files"]["tts_meta"]).read_text("utf-8") ) self.codec_meta: dict[str, Any] = json.loads( self._resolve_path(self.manifest["model_files"]["codec_meta"]).read_text("utf-8") ) self.tts_static_shapes = dict(self.tts_meta.get("static_shapes", {})) self.codec_static_shapes = dict(self.codec_meta.get("static_shapes", {})) tok_rel = str(self.manifest["model_files"].get("tokenizer_model", "tokenizer.model")) self.sp_model = spm.SentencePieceProcessor(model_file=str(self._resolve_path(tok_rel))) self.rng = np.random.default_rng(1234) self._model_display_names: dict[str, str] = {} self.sessions: dict = self._create_sessions() self._model_time_stats: dict[str, float] = defaultdict(float) self._model_call_stats: dict[str, int] = defaultdict(int) logging.info("AxTtsRuntime ready: %d sessions loaded", len(self.sessions)) def _find_manifest(self) -> Path: for name in _MANIFEST_CANDIDATES: p = (self.config_dir / name).resolve() if p.is_file(): return p raise FileNotFoundError( f"browser_poc_manifest.json not found in {self.config_dir}" ) def _resolve_path(self, rel: str | Path) -> Path: p = (self.manifest_dir / Path(rel)).resolve() if p.exists(): return p fallback = (self.config_dir / Path(rel).name).resolve() if fallback.exists(): return fallback return p @staticmethod def _normalize_sample_mode(raw: str | None, do_sample: bool = True) -> str: s = str(raw or "").strip() if s in {SAMPLE_MODE_GREEDY, SAMPLE_MODE_FIXED, SAMPLE_MODE_FULL}: return s if s == "mixed3": return SAMPLE_MODE_FIXED if do_sample else SAMPLE_MODE_GREEDY return SAMPLE_MODE_GREEDY if not do_sample else SAMPLE_MODE_FIXED def _resolve_axmodel_path(self, key: str, axmodel_name: str) -> Path: flat = self.axmodel_dir / axmodel_name if key != "local_fixed_sampled_frame": return flat if flat.exists(): return flat nested = self.axmodel_dir / "build-tts_local_fixed_sampled_frame" / axmodel_name return nested def _create_sessions(self) -> dict: sessions: dict[str, Any] = {} self._session_input_names: dict[str, list[str]] = {} if self.use_onnx_decode: onnx_name = _ONNX_FILE_MAP["decode"] onnx_path = self.onnx_dir / onnx_name if not onnx_path.exists(): raise FileNotFoundError(f"ONNX not found: {onnx_path}") if not _HAVE_ORT: raise RuntimeError("onnxruntime not installed; cannot load ONNX decode_step.") sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"]) sessions["decode"] = sess self._model_display_names["decode"] = f"{onnx_path.name} [onnxruntime CPU]" self._session_input_names["decode"] = [inp.name for inp in sess.get_inputs()] logging.info("loaded decode → %s (onnxruntime CPU)", onnx_path.name) else: axmodel_name = _AXMODEL_FILE_MAP["decode"] axmodel_path = self.axmodel_dir / axmodel_name if not axmodel_path.exists(): raise FileNotFoundError(f"axmodel not found: {axmodel_path}") input_names, output_names = _FULL8_IO_MAP["decode"] sess = AxeSession(axmodel_path, input_names, output_names) sessions["decode"] = sess self._model_display_names["decode"] = f"{axmodel_path.name} [axmodel]" self._session_input_names["decode"] = [inp.name for inp in sess.get_inputs()] logging.info("loaded decode → %s (axmodel)", axmodel_path.name) for key in ("prefill", "local_decoder", "local_fixed_sampled_frame", "codec_decode", "codec_encode"): if key not in _AXMODEL_FILE_MAP: continue axmodel_name = _AXMODEL_FILE_MAP[key] axmodel_path = self._resolve_axmodel_path(key, axmodel_name) if not axmodel_path.exists(): logging.warning("axmodel not found, skipping session %s: %s", key, axmodel_path) continue input_names, output_names = _FULL8_IO_MAP[key] sess = AxeSession(axmodel_path, input_names, output_names) sessions[key] = sess self._model_display_names[key] = f"{axmodel_path.name} [axmodel]" self._session_input_names[key] = [inp.name for inp in sess.get_inputs()] logging.info("loaded %s → %s (axmodel)", key, axmodel_path.name) if "codec_encode" not in sessions: logging.info("codec_encode axmodel not found; voice clone (reference audio) is disabled.") return sessions def _reset_timing_stats(self) -> None: self._model_time_stats = defaultdict(float) self._model_call_stats = defaultdict(int) self._used_model_keys: list[str] = [] self._used_model_key_set: set[str] = set() def _mark_model_used(self, key: str) -> None: if key in self._used_model_key_set: return self._used_model_key_set.add(key) self._used_model_keys.append(key) logging.info( "[usage] first_use session=%s model=%s", key, self._model_display_names.get(key, key), ) def _log_used_model_summary(self) -> None: if not getattr(self, "_used_model_keys", None): logging.info("[usage] used_models=none") return logging.info("[usage] used_models_count=%d", len(self._used_model_keys)) for idx, key in enumerate(self._used_model_keys, start=1): logging.info( "[usage] used_model[%d]=session=%s model=%s calls=%d", idx, key, self._model_display_names.get(key, key), self._model_call_stats.get(key, 0), ) def _run_session( self, key: str, input_feed: dict[str, np.ndarray], output_names: list[str] | None = None, ) -> list[np.ndarray]: sess = self.sessions[key] actual_input_names = [inp.name for inp in sess.get_inputs()] actual_input_name_set = set(actual_input_names) filtered_input_feed = { name: value for name, value in input_feed.items() if name in actual_input_name_set } missing_inputs = [name for name in actual_input_names if name not in filtered_input_feed] if missing_inputs: raise RuntimeError( f"session={key} missing required inputs: {missing_inputs}; " f"provided={sorted(input_feed.keys())}" ) extra_inputs = sorted(name for name in input_feed if name not in actual_input_name_set) self._mark_model_used(key) start = time.perf_counter() try: outputs = sess.run(output_names, filtered_input_feed) except Exception as exc: input_summary = { name: {"shape": list(np.asarray(value).shape), "dtype": str(np.asarray(value).dtype)} for name, value in filtered_input_feed.items() } raise RuntimeError( f"session={key} model={self._model_display_names.get(key, key)} run failed; " f"used_models_so_far={self._used_model_keys}; " f"expected_inputs={actual_input_names}; extra_inputs={extra_inputs}; " f"input_summary={input_summary}" ) from exc elapsed = time.perf_counter() - start self._model_time_stats[key] += elapsed self._model_call_stats[key] += 1 return outputs def encode_text(self, text: str) -> list[int]: return [int(t) for t in self.sp_model.encode(str(text or ""), out_type=int)] def count_text_tokens(self, text: str) -> int: return len(self.encode_text(text)) def split_text_by_token_budget(self, text: str, max_tokens: int) -> list[str]: remaining = str(text or "").strip() if not remaining: return [] pieces: list[str] = [] preferred_boundary = set(CLAUSE_SPLIT_PUNCTUATION) | set(SENTENCE_END_PUNCTUATION) | {" "} while remaining: if self.count_text_tokens(remaining) <= max_tokens: pieces.append(remaining) break lo, hi, best = 1, len(remaining), 1 while lo <= hi: mid = (lo + hi) // 2 cand = remaining[:mid].strip() if not cand: lo = mid + 1 continue if self.count_text_tokens(cand) <= max_tokens: best = mid lo = mid + 1 else: hi = mid - 1 cut = best prefix = remaining[:best] preferred = -1 scan_min = max(-1, len(prefix) - 25) for si in range(len(prefix) - 1, scan_min, -1): if prefix[si] in preferred_boundary: preferred = si + 1 break if preferred > 0: cut = preferred piece = remaining[:cut].strip() if not piece: piece = remaining[:best].strip() cut = best pieces.append(piece) remaining = remaining[cut:].strip() return pieces def split_voice_clone_text(self, text: str, max_tokens: int = 75) -> list[str]: t = str(text or "").strip() if not t: return [] safe_max = max(1, int(max_tokens)) prepared = _prepare_text_for_sentence_chunking(t) sents = _split_text_by_punctuation(prepared, SENTENCE_END_PUNCTUATION) or [prepared.strip()] slices: list[tuple[int, str]] = [] for sent in sents: s = sent.strip() if not s: continue cnt = self.count_text_tokens(s) if cnt <= safe_max: slices.append((cnt, s)) continue clauses = _split_text_by_punctuation(s, CLAUSE_SPLIT_PUNCTUATION) if len(clauses) <= 1: clauses = [s] for clause in clauses: c = clause.strip() if not c: continue ccnt = self.count_text_tokens(c) if ccnt <= safe_max: slices.append((ccnt, c)) continue for piece in self.split_text_by_token_budget(c, safe_max): p = piece.strip() if p: slices.append((self.count_text_tokens(p), p)) chunks: list[str] = [] cur_chunk, cur_cnt = "", 0 for cnt, s in slices: if not cur_chunk: cur_chunk, cur_cnt = s, cnt continue if cur_cnt + cnt > safe_max: chunks.append(cur_chunk.strip()) cur_chunk, cur_cnt = s, cnt else: cur_chunk = _join_sentence_parts(cur_chunk, s) cur_cnt = self.count_text_tokens(cur_chunk) if cur_chunk: chunks.append(cur_chunk.strip()) return chunks or [t] def _split_text_by_sentence_punctuation(self, text: str) -> list[str]: t = str(text or "").strip() if not t: return [] codec_limit = self._static_decode_code_length() if codec_limit is not None: _safe_text_tokens = max(1, int(codec_limit / 3 * 0.9)) _min_tokens = max(1, int(_safe_text_tokens * 0.25)) else: _safe_text_tokens = 80 _min_tokens = 20 prepared = _prepare_text_for_sentence_chunking(t) sents = _split_text_by_punctuation(prepared, SENTENCE_END_PUNCTUATION) if not sents: sents = [prepared.strip()] result: list[str] = [] for sent in sents: s = sent.strip() if not s: continue cnt = self.count_text_tokens(s) if cnt <= _safe_text_tokens: result.append(s) continue clauses = _split_text_by_punctuation(s, CLAUSE_SPLIT_PUNCTUATION) if len(clauses) <= 1: clauses = [s] merged: list[str] = [] buf = "" buf_tokens = 0 for clause in clauses: c = clause.strip() if not c: continue ccnt = self.count_text_tokens(c) if ccnt >= _safe_text_tokens: if buf: merged.append(buf.strip()) buf = "" buf_tokens = 0 merged.extend( p.strip() for p in self.split_text_by_token_budget(c, _safe_text_tokens) if p.strip() ) elif buf_tokens > 0 and buf_tokens + ccnt > _safe_text_tokens: merged.append(buf.strip()) buf = c buf_tokens = ccnt else: buf = (buf + " " + c) if buf else c buf_tokens += ccnt if buf: merged.append(buf.strip()) for m in merged: mcnt = self.count_text_tokens(m) if mcnt <= _safe_text_tokens: result.append(m) else: result.extend( p.strip() for p in self.split_text_by_token_budget(m, _safe_text_tokens) if p.strip() ) return result or [t] def build_text_rows(self, token_ids: list[int]) -> list[list[int]]: row_width = int(self.manifest["tts_config"]["n_vq"]) + 1 rows: list[list[int]] = [] for tid in token_ids: row = [int(self.manifest["tts_config"]["audio_pad_token_id"])] * row_width row[0] = int(tid) rows.append(row) return rows def build_audio_prefix_rows( self, prompt_audio_codes: list[list[int]], slot_token_id: int | None = None ) -> list[list[int]]: row_width = int(self.manifest["tts_config"]["n_vq"]) + 1 slot = int( self.manifest["tts_config"]["audio_user_slot_token_id"] if slot_token_id is None else slot_token_id ) rows: list[list[int]] = [] for code_row in prompt_audio_codes: row = [int(self.manifest["tts_config"]["audio_pad_token_id"])] * row_width row[0] = slot for i in range(min(len(code_row), int(self.manifest["tts_config"]["n_vq"]))): row[i + 1] = int(code_row[i]) rows.append(row) return rows def build_voice_clone_request_rows( self, prompt_audio_codes: list[list[int]], text_token_ids: list[int] ) -> dict[str, list[list[int]]]: prefix_ids = [ *self.manifest["prompt_templates"]["user_prompt_prefix_token_ids"], int(self.manifest["tts_config"]["audio_start_token_id"]), ] suffix_ids = [ int(self.manifest["tts_config"]["audio_end_token_id"]), *self.manifest["prompt_templates"]["user_prompt_after_reference_token_ids"], *text_token_ids, *self.manifest["prompt_templates"]["assistant_prompt_prefix_token_ids"], int(self.manifest["tts_config"]["audio_start_token_id"]), ] rows = [ *self.build_text_rows(prefix_ids), *self.build_audio_prefix_rows(prompt_audio_codes), *self.build_text_rows(suffix_ids), ] return {"inputIds": rows, "attentionMask": [[1 for _ in rows]]} def _static_prefill_length(self) -> int | None: v = self.tts_static_shapes.get("prefill_seq") return None if v is None else int(v) def _static_decode_code_length(self) -> int | None: v = self.codec_static_shapes.get("decode_code_length") return None if v is None else int(v) def _static_encode_waveform_length(self) -> int | None: v = self.codec_static_shapes.get("waveform_length") return None if v is None else int(v) def _pad_request_rows( self, request_rows: dict[str, list[list[int]]] ) -> tuple[dict[str, list[list[int]]], int]: static_len = self._static_prefill_length() actual_len = len(request_rows["inputIds"]) if static_len is None: return request_rows, actual_len if actual_len > static_len: raise ValueError( f"static prefill_seq={static_len} < request length={actual_len}. " "Reduce prompt/text length or re-export with larger --sample-seq-len." ) if actual_len == static_len: return request_rows, actual_len row_width = int(self.manifest["tts_config"]["n_vq"]) + 1 pad_row = [int(self.manifest["tts_config"]["audio_pad_token_id"])] * row_width pad_row[0] = int(self.manifest["tts_config"]["pad_token_id"]) padded_rows = [list(r) for r in request_rows["inputIds"]] padded_rows.extend([list(pad_row) for _ in range(static_len - actual_len)]) padded_mask = [1] * actual_len + [0] * (static_len - actual_len) return {"inputIds": padded_rows, "attentionMask": [padded_mask]}, actual_len def _pad_kv( self, kv_dict: dict[str, np.ndarray], target_len: int ) -> dict[str, np.ndarray]: result: dict[str, np.ndarray] = {} for key, val in kv_dict.items(): cur = val.shape[1] if cur > target_len: padded = val[:, :target_len, :, :].astype(np.float32) elif cur < target_len: pad = np.zeros((1, target_len - cur, *val.shape[2:]), dtype=np.float32) padded = np.concatenate([val.astype(np.float32), pad], axis=1) else: padded = val.astype(np.float32) result[key] = padded return result def run_local_fixed_sampled_frame( self, global_hidden: np.ndarray, *, previous_token_sets_by_channel: list[set[int]], ) -> tuple[bool, list[int]]: n_vq = int(self.manifest["tts_config"]["n_vq"]) codebook_size = int(self.tts_meta["model_config"]["audio_codebook_sizes"][0]) rep_mask = np.zeros((1, n_vq, codebook_size), dtype=np.int32) for ci, token_set in enumerate(previous_token_sets_by_channel): for tid in token_set: if 0 <= tid < codebook_size: rep_mask[0, ci, tid] = 1 sess = self.sessions["local_fixed_sampled_frame"] session_input_names = {inp.name for inp in sess.get_inputs()} input_feed: dict = { "global_hidden": global_hidden.astype(np.float32, copy=False), "repetition_seen_mask": rep_mask, } asst_u: float | None = None if "assistant_random_u" in session_input_names: asst_u = float(self.rng.random()) input_feed["assistant_random_u"] = np.array([asst_u], dtype=np.float32) if "audio_random_u" in session_input_names: input_feed["audio_random_u"] = np.asarray( [[float(self.rng.random()) for _ in range(n_vq)]], dtype=np.float32, ) outputs = self._run_session("local_fixed_sampled_frame", input_feed) out_names = [o.name for o in sess.get_outputs()] named = dict(zip(out_names, outputs, strict=True)) frame_ids = np.asarray(named["frame_token_ids"]).reshape(-1).astype(np.int32).tolist() should_continue = bool(int(np.asarray(named["should_continue"]).reshape(-1)[0])) logging.debug( "[fixed_frame] asst_u=%.4f should_continue=%s tokens=%s", asst_u if asst_u is not None else -1.0, should_continue, frame_ids[:4], ) return should_continue, [int(x) for x in frame_ids] def _can_use_local_fixed_sampled_frame(self) -> bool: sess = self.sessions.get("local_fixed_sampled_frame") if sess is None: return False session_input_names = {inp.name for inp in sess.get_inputs()} return "audio_random_u" in session_input_names def run_local_decoder( self, global_hidden: np.ndarray, text_token_id: int, frame_prefix: list[int], ) -> tuple[np.ndarray, np.ndarray]: n_vq = int(self.manifest["tts_config"]["n_vq"]) audio_pad = int(self.manifest["tts_config"]["audio_pad_token_id"]) padded_prefix = np.full((1, n_vq - 1), audio_pad, dtype=np.int32) for i in range(min(len(frame_prefix), n_vq - 1)): padded_prefix[0, i] = int(frame_prefix[i]) outputs = self._run_session( "local_decoder", { "global_hidden": global_hidden.astype(np.float32, copy=False), "text_token_id": np.asarray([int(text_token_id)], dtype=np.int32), "audio_prefix_token_ids": padded_prefix, }, ) out_names = [o.name for o in self.sessions["local_decoder"].get_outputs()] named = dict(zip(out_names, outputs, strict=True)) return named["text_logits"].reshape(-1), named["audio_logits"] def slice_audio_channel_logits(self, audio_logits: np.ndarray, channel_index: int) -> np.ndarray: per_ch = int(audio_logits.shape[-1]) flat = audio_logits.reshape(-1) return flat[channel_index * per_ch : (channel_index + 1) * per_ch] def _infer_global_kv_len(self) -> int | None: if "decode" in self.sessions: dec_sess = self.sessions["decode"] if hasattr(dec_sess, "get_inputs"): for inp in dec_sess.get_inputs(): if inp.name == "past_key_0": dims = getattr(inp, "shape", []) or [] if len(dims) >= 2 and isinstance(dims[1], str): return None break v = self.tts_static_shapes.get("kv_seq") or self.tts_static_shapes.get("global_kv_len") if v is not None: return int(v) return 512 def generate_audio_frames( self, request_rows: dict[str, list[list[int]]], on_frame: Callable | None = None, ) -> list[list[int]]: padded_rows, actual_prefill_len = self._pad_request_rows(request_rows) return self._generate_frames_with_prefill( padded_rows, actual_prefill_length=actual_prefill_len, on_frame=on_frame, ) def _generate_frames_with_prefill( self, request_rows: dict[str, list[list[int]]], *, actual_prefill_length: int, on_frame: Callable | None = None, ) -> list[list[int]]: global_kv_len = self._infer_global_kv_len() generation_defaults = self.manifest["generation_defaults"] row_width = int(self.manifest["tts_config"]["n_vq"]) + 1 prefill_ids, prefill_dims = _flatten3d_int32([request_rows["inputIds"]]) prefill_mask, prefill_mask_dims = _flatten2d_int32(request_rows["attentionMask"]) prefill_outputs = self._run_session( "prefill", { "input_ids": prefill_ids.reshape(prefill_dims), "attention_mask": prefill_mask.reshape(prefill_mask_dims), }, ) prefill_out_names = [o.name for o in self.sessions["prefill"].get_outputs()] named_prefill = dict(zip(prefill_out_names, prefill_outputs, strict=True)) global_hidden = named_prefill["global_hidden"][:, actual_prefill_length - 1, :].astype(np.float32) past_valid_length = int(actual_prefill_length) raw_past: dict[str, np.ndarray] = { out.replace("present_", "past_"): named_prefill[out][:, :actual_prefill_length, :, :] for out in self.tts_meta["onnx"]["prefill_output_names"][1:] } use_dynamic_kv = global_kv_len is None if use_dynamic_kv: past_by_name = raw_past # shape: (1, actual_len, 12, 64)，步进增长 else: past_by_name = self._pad_kv(raw_past, global_kv_len) # shape: (1, 512, 12, 64) generated_frames: list[list[int]] = [] prev_tokens = [[] for _ in range(int(self.manifest["tts_config"]["n_vq"]))] prev_tok_sets = [set() for _ in range(int(self.manifest["tts_config"]["n_vq"]))] audio_pad_id = int(self.manifest["tts_config"]["audio_pad_token_id"]) consecutive_pad_frames = 0 _MAX_CONSECUTIVE_PAD_FRAMES = 10 # 0.8s silence → force stop (guard against quantized AXModel) for step in range(int(generation_defaults["max_new_frames"])): frame: list[int] = [] if ( generation_defaults["sample_mode"] == SAMPLE_MODE_FIXED and self._can_use_local_fixed_sampled_frame() ): should_continue, frame = self.run_local_fixed_sampled_frame( global_hidden, previous_token_sets_by_channel=prev_tok_sets ) if not should_continue: logging.info("[gen] frame %d: should_continue=False → 停止生成", step) break for ci, tok in enumerate(frame): prev_tokens[ci].append(tok) prev_tok_sets[ci].add(tok) elif "local_decoder" in self.sessions: if ( generation_defaults["sample_mode"] == SAMPLE_MODE_FIXED and "local_fixed_sampled_frame" in self.sessions and step == 0 ): logging.warning( "local_fixed_sampled_frame 缺少 audio_random_u 输入，回退到 local_decoder 随机采样路径" ) local_text_logits, _ = self.run_local_decoder(global_hidden, 0, []) next_text = _sample_assistant_text_token( local_text_logits, self.manifest, generation_defaults, self.rng ) if next_text != int(self.manifest["tts_config"]["audio_assistant_slot_token_id"]): break for ci in range(int(self.manifest["tts_config"]["n_vq"])): _, audio_logits = self.run_local_decoder(global_hidden, next_text, frame) ch_logits = self.slice_audio_channel_logits(audio_logits, ci).astype(np.float32) tok = _sample_audio_token( ch_logits, prev_tokens[ci], prev_tok_sets[ci], generation_defaults, self.rng ) frame.append(tok) prev_tokens[ci].append(tok) prev_tok_sets[ci].add(tok) else: raise RuntimeError("Neither local_fixed_sampled_frame nor local_decoder session found") if not frame: break if all(tok == audio_pad_id for tok in frame): consecutive_pad_frames += 1 if consecutive_pad_frames >= _MAX_CONSECUTIVE_PAD_FRAMES: logging.info( "[gen] %d consecutive all-pad frames → force stop (trim %d silent frames)", consecutive_pad_frames, consecutive_pad_frames, ) generated_frames = generated_frames[:-(consecutive_pad_frames - 1)] if len(generated_frames) >= consecutive_pad_frames else [] break else: consecutive_pad_frames = 0 generated_frames.append(frame) next_row = np.full((1, 1, row_width), int(self.manifest["tts_config"]["audio_pad_token_id"]), dtype=np.int32) next_row[0, 0, 0] = int(self.manifest["tts_config"]["audio_assistant_slot_token_id"]) for idx, tok in enumerate(frame): next_row[0, 0, idx + 1] = int(tok) decode_feeds: dict[str, np.ndarray] = { "input_ids": next_row, "past_valid_lengths": np.asarray([past_valid_length], dtype=np.int32), } for inp_name in self.tts_meta["onnx"]["decode_input_names"][2:]: decode_feeds[inp_name] = past_by_name[inp_name] dec_outputs = self._run_session("decode", decode_feeds) dec_out_names = [o.name for o in self.sessions["decode"].get_outputs()] named_dec = dict(zip(dec_out_names, dec_outputs, strict=True)) global_hidden = _extract_last_hidden(named_dec["global_hidden"]) past_valid_length += 1 if use_dynamic_kv: past_by_name = { out_name.replace("present_", "past_"): named_dec[out_name] for out_name in self.tts_meta["onnx"]["decode_output_names"][1:] } else: new_pos = past_valid_length - 1 # 新 token 的逻辑位置（0-indexed） new_padded: dict[str, np.ndarray] = {} for out_name in self.tts_meta["onnx"]["decode_output_names"][1:]: past_name = out_name.replace("present_", "past_") v = named_dec[out_name] buf = np.zeros((1, global_kv_len, 12, 64), dtype=np.float32) if new_pos < global_kv_len: buf[:, :new_pos, :, :] = v[:, :new_pos, :, :] buf[:, new_pos, :, :] = v[:, global_kv_len, :, :] else: buf[:, :global_kv_len - 1, :, :] = v[:, 1:global_kv_len, :, :] buf[:, global_kv_len - 1, :, :] = v[:, global_kv_len, :, :] new_padded[past_name] = buf past_by_name = new_padded if past_valid_length > global_kv_len: past_valid_length = global_kv_len if on_frame is not None: on_frame(generated_frames, step, frame) logging.info("[gen] 生成完毕: 共 %d 帧 (max_new_frames=%s)", len(generated_frames), generation_defaults["max_new_frames"]) return generated_frames def decode_full_audio(self, generated_frames: list[list[int]]) -> tuple[list[np.ndarray], int]: if not generated_frames: return [], 0 static_len = self._static_decode_code_length() n_vq = int(self.manifest["tts_config"]["n_vq"]) frame_count = len(generated_frames) if static_len is None or frame_count <= static_len: return self._decode_one_segment(generated_frames, frame_count, n_vq, static_len) all_audio: list[np.ndarray] = [] total_samples = 0 for seg_start in range(0, frame_count, static_len): seg_end = min(seg_start + static_len, frame_count) seg_frames = generated_frames[seg_start:seg_end] seg_count = seg_end - seg_start channel_arrays, seg_samples = self._decode_one_segment(seg_frames, seg_count, n_vq, static_len) all_audio.extend(channel_arrays) total_samples += seg_samples return all_audio, total_samples def _decode_one_segment( self, frames: list[list[int]], frame_count: int, n_vq: int, static_len: int | None, ) -> tuple[list[np.ndarray], int]: max_len = static_len if static_len is not None else frame_count audio_codes = np.zeros((1, max_len, n_vq), dtype=np.int32) for fi, frame in enumerate(frames): for ci in range(n_vq): audio_codes[0, fi, ci] = int(frame[ci] if ci < len(frame) else 0) outputs = self._run_session( "codec_decode", { "audio_codes": audio_codes, "audio_code_lengths": np.asarray([frame_count], dtype=np.int32), }, ) out_names = [o.name for o in self.sessions["codec_decode"].get_outputs()] named = dict(zip(out_names, outputs, strict=True)) audio_length = int(named["audio_lengths"].reshape(-1)[0]) return _slice_channel_major_audio(named["audio"], 0, audio_length), audio_length def encode_reference_audio(self, reference_audio_path: str | Path) -> list[list[int]]: codec_cfg = self.codec_meta["codec_config"] target_sr = int(codec_cfg["sample_rate"]) target_ch = int(codec_cfg["channels"]) waveform = _load_audio_numpy(reference_audio_path, target_sr, target_ch) waveform_length = int(waveform.shape[-1]) if bool(codec_cfg.get("pad_reference_audio_to_downsample_rate", False)): ds_rate = int(codec_cfg["downsample_rate"]) remainder = waveform_length % ds_rate if remainder: pad_len = ds_rate - remainder waveform = np.concatenate( [waveform, np.zeros((1, target_ch, pad_len), dtype=np.float32)], axis=2 ) waveform_length += pad_len static_wl = self._static_encode_waveform_length() num_quantizers = int(codec_cfg["num_quantizers"]) prompt_audio_codes: list[list[int]] = [] out_names = [o.name for o in self.sessions["codec_encode"].get_outputs()] if static_wl and static_wl > 0: for start in range(0, waveform_length, static_wl): chunk = waveform[..., start : start + static_wl] chunk_len = int(chunk.shape[-1]) if chunk_len <= 0: continue if chunk_len < static_wl: padded = np.zeros((1, target_ch, static_wl), dtype=np.float32) padded[..., :chunk_len] = chunk chunk = padded outputs = self._run_session( "codec_encode", { "waveform": chunk, "input_lengths": np.asarray([chunk_len], dtype=np.int32), }, ) named = dict(zip(out_names, outputs, strict=True)) codes = np.asarray(named["audio_codes"], dtype=np.int32) code_len = int(np.asarray(named["audio_code_lengths"]).reshape(-1)[0]) for fi in range(code_len): prompt_audio_codes.append([int(codes[0, fi, qi]) for qi in range(num_quantizers)]) else: outputs = self._run_session( "codec_encode", { "waveform": waveform, "input_lengths": np.asarray([waveform_length], dtype=np.int32), }, ) named = dict(zip(out_names, outputs, strict=True)) codes = np.asarray(named["audio_codes"], dtype=np.int32) code_len = int(np.asarray(named["audio_code_lengths"]).reshape(-1)[0]) for fi in range(code_len): prompt_audio_codes.append([int(codes[0, fi, qi]) for qi in range(num_quantizers)]) return prompt_audio_codes def resolve_prompt_audio_codes( self, *, voice: str | None, prompt_audio_path: str | Path | None ) -> list[list[int]]: if prompt_audio_path: return self.encode_reference_audio(prompt_audio_path) resolved_voice = str(voice or self.manifest["builtin_voices"][0]["voice"]) for v in self.manifest["builtin_voices"]: if v["voice"] == resolved_voice: return list(v["prompt_audio_codes"]) raise ValueError(f"Built-in voice not found: {resolved_voice}") def list_builtin_voices(self) -> list[dict[str, Any]]: return list(self.manifest["builtin_voices"]) def synthesize( self, *, text: str, voice: str | None = None, prompt_audio_path: str | Path | None = None, output_audio_path: str | Path | None = None, sample_mode: str | None = None, do_sample: bool = True, streaming: bool = True, max_new_frames: int | None = None, voice_clone_max_text_tokens: int = 75, seed: int | None = None, ) -> dict[str, Any]: self._reset_timing_stats() gen = self.manifest["generation_defaults"] if max_new_frames is not None: gen["max_new_frames"] = int(max_new_frames) static_codec_len = self._static_decode_code_length() if static_codec_len is not None and gen.get("max_new_frames", float("inf")) > static_codec_len: gen["max_new_frames"] = static_codec_len logging.info("capped max_new_frames=%d (static codec limit)", static_codec_len) normalized_mode = self._normalize_sample_mode(sample_mode, do_sample) gen["sample_mode"] = normalized_mode gen["do_sample"] = normalized_mode != SAMPLE_MODE_GREEDY if seed is not None: self.rng = np.random.default_rng(int(seed)) infer_start_time = time.perf_counter() try: text_chunks = self._split_text_by_sentence_punctuation(text) prompt_audio_codes = self.resolve_prompt_audio_codes( voice=voice, prompt_audio_path=prompt_audio_path ) sample_rate = int(self.codec_meta["codec_config"]["sample_rate"]) channels = int(self.codec_meta["codec_config"]["channels"]) static_decode_code_length = self._static_decode_code_length() out_path = ( Path(output_audio_path).expanduser().resolve() if output_audio_path else (Path("output.wav")).resolve() ) all_waveforms: list[np.ndarray] = [] all_generated_frames: list[list[int]] = [] final_text_chunks: list[str] = [] pending_text_chunks = list(text_chunks) wav_file = None if streaming: out_path.parent.mkdir(parents=True, exist_ok=True) wav_file = wave.open(str(out_path), "wb") wav_file.setnchannels(channels) wav_file.setsampwidth(2) wav_file.setframerate(sample_rate) try: while pending_text_chunks: chunk_idx = len(final_text_chunks) + 1 total_chunks = len(final_text_chunks) + len(pending_text_chunks) chunk_text = pending_text_chunks.pop(0) logging.info("chunk %d/%d: %r", chunk_idx, total_chunks, chunk_text[:40]) text_token_ids = self.encode_text(chunk_text) request_rows = self.build_voice_clone_request_rows(prompt_audio_codes, text_token_ids) generated_frames = self.generate_audio_frames(request_rows) channel_arrays, _ = self.decode_full_audio(generated_frames) waveform = np.asarray(_merge_audio_channels(channel_arrays), dtype=np.float32) all_waveforms.append(waveform) all_generated_frames.extend(generated_frames) final_text_chunks.append(chunk_text) if wav_file is not None and waveform.size > 0: clipped = np.clip(waveform, -1.0, 1.0) pcm16 = np.round(clipped * 32767.0).astype(np.int16) wav_file.writeframes(pcm16.tobytes()) logging.info( "stream emitted chunk %d: samples=%d duration=%.3fs", len(final_text_chunks), int(waveform.shape[0]), float(waveform.shape[0]) / float(sample_rate), ) if pending_text_chunks: pause_words = len([w for w in chunk_text.strip().split() if w]) pause_sec = 0.40 if pause_words <= 4 else 0.24 pause_samples = max(0, int(round(sample_rate * pause_sec))) if pause_samples > 0: pause_waveform = np.zeros((pause_samples, channels), dtype=np.float32) all_waveforms.append(pause_waveform) if wav_file is not None: wav_file.writeframes(np.zeros((pause_samples, channels), dtype=np.int16).tobytes()) finally: if wav_file is not None: wav_file.close() final_waveform = _concat_waveforms(all_waveforms) total_infer_time = time.perf_counter() - infer_start_time audio_duration_sec = ( float(final_waveform.shape[0]) / float(sample_rate) if sample_rate > 0 and final_waveform.size > 0 else 0.0 ) rtf = (total_infer_time / audio_duration_sec) if audio_duration_sec > 0 else float("inf") model_time_total = sum(self._model_time_stats.values()) audio_path = out_path if streaming else _write_waveform_to_wav(out_path, final_waveform, sample_rate) logging.info("已保存 %s sample_rate=%s frames=%s", audio_path, sample_rate, len(all_generated_frames)) logging.info( "[timing] total_infer_time=%.3fs audio_duration=%.3fs rtf=%.4f model_time_total=%.3fs", total_infer_time, audio_duration_sec, rtf, model_time_total, ) for key in sorted(self._model_display_names): model_name = self._model_display_names[key] logging.info( "[timing] session=%s model=%s calls=%d total=%.3fs", key, model_name, self._model_call_stats.get(key, 0), self._model_time_stats.get(key, 0.0), ) return { "audio_path": str(audio_path), "waveform": final_waveform, "sample_rate": sample_rate, "audio_token_ids": np.asarray(all_generated_frames, dtype=np.int32), "text_chunks": final_text_chunks, "sample_mode": normalized_mode, "do_sample": normalized_mode != SAMPLE_MODE_GREEDY, "streaming": bool(streaming), "timing": { "total_infer_time_sec": total_infer_time, "audio_duration_sec": audio_duration_sec, "rtf": rtf, "model_time_total_sec": model_time_total, "per_model_time_sec": dict(self._model_time_stats), "per_model_calls": dict(self._model_call_stats), "per_model_display_name": dict(self._model_display_names), "used_model_keys": list(self._used_model_keys), }, } finally: self._log_used_model_summary()