@jingyi0605/codingns4dsh 0.2.1-beta.3 → 0.2.1-beta.4

package/assets/moss/ort_cpu_runtime.py

@@ -0,0 +1,839 @@
+# 来源：OpenMOSS/MOSS-TTS-Nano，提交 8b7bcc9341b3b4ef3a3a58ba1338a7d85ff133eb。
+# 上游代码保持原样；许可 Apache-2.0，见本目录 LICENSE。
+from __future__ import annotations
+
+import json
+import math
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable
+
+import numpy as np
+import onnxruntime as ort
+
+SAMPLE_MODE_GREEDY = "greedy"
+SAMPLE_MODE_FIXED = "fixed"
+SAMPLE_MODE_FULL = "full"
+EXECUTION_PROVIDER_CPU = "cpu"
+EXECUTION_PROVIDER_CUDA = "cuda"
+
+MANIFEST_CANDIDATE_RELATIVE_PATHS = (
+    "browser_poc_manifest.json",
+    "MOSS-TTS-Nano-100M-ONNX/browser_poc_manifest.json",
+    "MOSS-TTS-Nano-ONNX-CPU/browser_poc_manifest.json",
+)
+MODEL_DIR_ALIAS_MAP = {
+    "MOSS-TTS-Nano-ONNX-CPU": "MOSS-TTS-Nano-100M-ONNX",
+    "MOSS-Audio-Tokenizer-Nano-ONNX-CPU": "MOSS-Audio-Tokenizer-Nano-ONNX",
+}
+
+
+def _argmax(values: np.ndarray) -> int:
+    return int(np.argmax(values))
+
+
+def _normalize_execution_provider(raw_execution_provider: str | None) -> str:
+    normalized = str(raw_execution_provider or EXECUTION_PROVIDER_CPU).strip().lower()
+    if normalized in {EXECUTION_PROVIDER_CPU, "CPUExecutionProvider".lower()}:
+        return EXECUTION_PROVIDER_CPU
+    if normalized in {EXECUTION_PROVIDER_CUDA, "gpu", "CUDAExecutionProvider".lower()}:
+        return EXECUTION_PROVIDER_CUDA
+    raise ValueError("execution_provider must be one of: cpu, cuda")
+
+
+def _resolve_ort_providers(execution_provider: str) -> list[Any]:
+    normalized = _normalize_execution_provider(execution_provider)
+    if normalized == EXECUTION_PROVIDER_CPU:
+        return ["CPUExecutionProvider"]
+    available_providers = set(ort.get_available_providers())
+    if "CUDAExecutionProvider" not in available_providers:
+        available = ", ".join(ort.get_available_providers()) or "none"
+        raise RuntimeError(
+            "CUDAExecutionProvider was requested, but this onnxruntime build does not expose it. "
+            "Install onnxruntime-gpu that matches your CUDA/cuDNN runtime. "
+            f"Available providers: {available}"
+        )
+    preload_dlls = getattr(ort, "preload_dlls", None)
+    if callable(preload_dlls):
+        preload_dlls()
+    return ["CUDAExecutionProvider", "CPUExecutionProvider"]
+
+
+def _flatten3d_int32(nested: list[list[list[int]]]) -> tuple[np.ndarray, list[int]]:
+    dim0 = len(nested)
+    dim1 = len(nested[0])
+    dim2 = len(nested[0][0])
+    data = np.zeros((dim0 * dim1 * dim2,), dtype=np.int32)
+    offset = 0
+    for i in range(dim0):
+        for j in range(dim1):
+            for k in range(dim2):
+                data[offset] = int(nested[i][j][k])
+                offset += 1
+    return data, [dim0, dim1, dim2]
+
+
+def _flatten2d_int32(nested: list[list[int]]) -> tuple[np.ndarray, list[int]]:
+    dim0 = len(nested)
+    dim1 = len(nested[0])
+    data = np.zeros((dim0 * dim1,), dtype=np.int32)
+    offset = 0
+    for i in range(dim0):
+        for j in range(dim1):
+            data[offset] = int(nested[i][j])
+            offset += 1
+    return data, [dim0, dim1]
+
+
+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_samples = int(audio.shape[2])
+    start = max(0, int(start_sample))
+    end = total_samples if end_sample is None else max(start, min(int(end_sample), total_samples))
+    return [audio[0, channel_index, start:end].astype(np.float32, copy=False) for channel_index in range(channels)]
+
+
+def _extract_last_hidden(hidden_states: np.ndarray) -> np.ndarray:
+    if hidden_states.ndim == 2:
+        return hidden_states.astype(np.float32, copy=False)
+    if hidden_states.ndim != 3 or hidden_states.shape[0] != 1:
+        raise ValueError(f"Unexpected global_hidden shape: {hidden_states.shape}")
+    return hidden_states[:, -1, :].astype(np.float32, copy=False)
+
+
+def _apply_repetition_penalty(values: np.ndarray, previous_token_ids: list[int], repetition_penalty: float) -> np.ndarray:
+    if not previous_token_ids or repetition_penalty == 1.0:
+        return values
+    result = values.copy()
+    for token_id in set(int(item) for item in previous_token_ids):
+        if token_id < 0 or token_id >= result.shape[0]:
+            continue
+        result[token_id] = result[token_id] * repetition_penalty if result[token_id] < 0 else result[token_id] / repetition_penalty
+    return result
+
+
+def _argmax_with_repetition_penalty(values: np.ndarray, previous_token_set: set[int], repetition_penalty: float) -> int:
+    best_index = 0
+    best_value = float("-inf")
+    apply_penalty = bool(previous_token_set) and repetition_penalty != 1.0
+    for index, value in enumerate(values):
+        score = float(value)
+        if apply_penalty and index in previous_token_set:
+            score = score * repetition_penalty if score < 0 else score / repetition_penalty
+        if score > best_value:
+            best_value = score
+            best_index = index
+    return int(best_index)
+
+
+def _softmax(values: np.ndarray) -> np.ndarray:
+    max_value = float(np.max(values))
+    shifted = np.asarray(values - max_value, dtype=np.float64)
+    exps = np.exp(shifted)
+    return exps / np.sum(exps, dtype=np.float64)
+
+
+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 not (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]:
+        sorted_desc = np.sort(scores)[::-1]
+        threshold = float(sorted_desc[top_k - 1])
+        scores[scores < threshold] = float("-inf")
+    if top_p > 0 and top_p < 1:
+        indexed = list(enumerate(scores.tolist()))
+        indexed.sort(key=lambda item: item[1], reverse=True)
+        sorted_scores = np.asarray([item[1] for item in indexed], dtype=np.float32)
+        sorted_probs = _softmax(sorted_scores)
+        remove_mask = [False] * len(indexed)
+        cumulative = 0.0
+        for index, probability in enumerate(sorted_probs):
+            cumulative += float(probability)
+            if cumulative > float(top_p):
+                remove_mask[index] = True
+        for index in range(len(remove_mask) - 1, 0, -1):
+            remove_mask[index] = remove_mask[index - 1]
+        if remove_mask:
+            remove_mask[0] = False
+        for index, should_remove in enumerate(remove_mask):
+            if should_remove:
+                scores[indexed[index][0]] = float("-inf")
+    probabilities = _softmax(scores)
+    random_value = float(rng.random())
+    for index, probability in enumerate(probabilities):
+        random_value -= float(probability)
+        if random_value <= 0:
+            return int(index)
+    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_index = _sample_from_scores(
+        candidate_scores,
+        do_sample=bool(generation_defaults["do_sample"]),
+        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_index])
+
+
+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:
+    repetition_penalty = float(generation_defaults["audio_repetition_penalty"])
+    if not bool(generation_defaults["do_sample"]):
+        return _argmax_with_repetition_penalty(audio_logits, previous_token_set, repetition_penalty)
+    penalized_scores = _apply_repetition_penalty(audio_logits, previous_token_ids, repetition_penalty)
+    return _sample_from_scores(
+        penalized_scores,
+        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 _normalize_sample_mode(raw_sample_mode: str | None, raw_do_sample: bool = True) -> str:
+    normalized = str(raw_sample_mode or "").strip()
+    if normalized in {SAMPLE_MODE_GREEDY, SAMPLE_MODE_FIXED, SAMPLE_MODE_FULL}:
+        return normalized
+    if normalized == "mixed3":
+        return SAMPLE_MODE_FIXED if raw_do_sample else SAMPLE_MODE_GREEDY
+    return SAMPLE_MODE_GREEDY if not raw_do_sample else SAMPLE_MODE_FIXED
+
+
+def _compute_stream_lead_seconds(emitted_samples_total: int, sample_rate: int, first_audio_emitted_at_seconds: float | None) -> float:
+    if not first_audio_emitted_at_seconds or sample_rate <= 0:
+        return 0.0
+    elapsed_seconds = max(0.0, time.perf_counter() - first_audio_emitted_at_seconds)
+    emitted_seconds = emitted_samples_total / float(sample_rate)
+    return emitted_seconds - elapsed_seconds
+
+
+def _resolve_stream_decode_frame_budget(
+    emitted_samples_total: int,
+    sample_rate: int,
+    first_audio_emitted_at_seconds: float | None,
+) -> int:
+    lead_seconds = _compute_stream_lead_seconds(emitted_samples_total, sample_rate, first_audio_emitted_at_seconds)
+    if not first_audio_emitted_at_seconds or lead_seconds < 0.20:
+        return 1
+    if lead_seconds < 0.55:
+        return 2
+    if lead_seconds < 1.10:
+        return 4
+    return 8
+
+
+@dataclass
+class CodecStreamingDecodeSession:
+    codec_meta: dict[str, Any]
+    session: ort.InferenceSession
+
+    def __post_init__(self) -> None:
+        self.transformer_specs = list(self.codec_meta.get("streaming_decode", {}).get("transformer_offsets", []))
+        self.attention_specs = list(self.codec_meta.get("streaming_decode", {}).get("attention_caches", []))
+        self.state_feeds: dict[str, np.ndarray] = {}
+        self.reset()
+
+    def reset(self) -> None:
+        self.state_feeds = {}
+        for spec in self.transformer_specs:
+            self.state_feeds[str(spec["input_name"])] = np.zeros(tuple(spec["shape"]), dtype=np.int32)
+        for spec in self.attention_specs:
+            self.state_feeds[str(spec["offset_input_name"])] = np.zeros(tuple(spec["offset_shape"]), dtype=np.int32)
+            self.state_feeds[str(spec["cached_keys_input_name"])] = np.zeros(tuple(spec["cache_shape"]), dtype=np.float32)
+            self.state_feeds[str(spec["cached_values_input_name"])] = np.zeros(tuple(spec["cache_shape"]), dtype=np.float32)
+            positions = np.full(tuple(spec["positions_shape"]), -1, dtype=np.int32)
+            self.state_feeds[str(spec["cached_positions_input_name"])] = positions
+
+    def run_frames(self, frame_rows: list[list[int]]) -> tuple[np.ndarray, int] | None:
+        if not frame_rows:
+            return None
+        num_quantizers = int(self.codec_meta["codec_config"]["num_quantizers"])
+        frame_count = len(frame_rows)
+        audio_codes = np.zeros((1, frame_count, num_quantizers), dtype=np.int32)
+        for frame_index, frame_row in enumerate(frame_rows):
+            for channel_index in range(num_quantizers):
+                audio_codes[0, frame_index, channel_index] = int(frame_row[channel_index] if channel_index < len(frame_row) else 0)
+        feeds: dict[str, np.ndarray] = {
+            "audio_codes": audio_codes,
+            "audio_code_lengths": np.asarray([frame_count], dtype=np.int32),
+        }
+        feeds.update(self.state_feeds)
+        outputs = self.session.run(None, feeds)
+        output_names = [output.name for output in self.session.get_outputs()]
+        named_outputs = dict(zip(output_names, outputs, strict=True))
+        for spec in self.transformer_specs:
+            self.state_feeds[str(spec["input_name"])] = named_outputs[str(spec["output_name"])]
+        for spec in self.attention_specs:
+            self.state_feeds[str(spec["offset_input_name"])] = named_outputs[str(spec["offset_output_name"])]
+            self.state_feeds[str(spec["cached_keys_input_name"])] = named_outputs[str(spec["cached_keys_output_name"])]
+            self.state_feeds[str(spec["cached_values_input_name"])] = named_outputs[str(spec["cached_values_output_name"])]
+            self.state_feeds[str(spec["cached_positions_input_name"])] = named_outputs[str(spec["cached_positions_output_name"])]
+        return (
+            named_outputs["audio"],
+            int(named_outputs["audio_lengths"].reshape(-1)[0]),
+        )
+
+
+class OrtCpuRuntime:
+    def __init__(
+        self,
+        model_dir: str | Path,
+        thread_count: int = 4,
+        max_new_frames: int | None = None,
+        do_sample: bool | None = None,
+        sample_mode: str | None = None,
+        execution_provider: str = EXECUTION_PROVIDER_CPU,
+    ) -> None:
+        self.model_dir = Path(model_dir).expanduser().resolve()
+        self.thread_count = max(1, int(thread_count))
+        self.execution_provider = _normalize_execution_provider(execution_provider)
+        self.ort_providers = _resolve_ort_providers(self.execution_provider)
+        self.manifest_path = self._resolve_manifest_path(self.model_dir)
+        self.manifest_dir = self.manifest_path.parent
+        manifest = json.loads(self.manifest_path.read_text(encoding="utf-8"))
+        self.manifest = manifest
+        if max_new_frames is not None:
+            self.manifest["generation_defaults"]["max_new_frames"] = int(max_new_frames)
+        if do_sample is not None:
+            self.manifest["generation_defaults"]["do_sample"] = bool(do_sample)
+        self.manifest["generation_defaults"]["sample_mode"] = _normalize_sample_mode(
+            sample_mode if sample_mode is not None else self.manifest["generation_defaults"].get("sample_mode"),
+            bool(self.manifest["generation_defaults"]["do_sample"]),
+        )
+        self.manifest["generation_defaults"]["do_sample"] = (
+            self.manifest["generation_defaults"]["sample_mode"] != SAMPLE_MODE_GREEDY
+        )
+        self.tts_meta_path = self.resolve_manifest_relative_path(manifest["model_files"]["tts_meta"])
+        self.codec_meta_path = self.resolve_manifest_relative_path(manifest["model_files"]["codec_meta"])
+        self.tts_meta = json.loads(self.tts_meta_path.read_text(encoding="utf-8"))
+        self.codec_meta = json.loads(self.codec_meta_path.read_text(encoding="utf-8"))
+        self.rng = np.random.default_rng(1234)
+        self.sessions = self._create_sessions()
+        self.codec_streaming_session = CodecStreamingDecodeSession(
+            codec_meta=self.codec_meta,
+            session=self.sessions["codec_decode_step"],
+        )
+
+    @staticmethod
+    def _resolve_manifest_path(model_dir: Path) -> Path:
+        tried_paths: list[Path] = []
+        for relative_path in MANIFEST_CANDIDATE_RELATIVE_PATHS:
+            candidate = (model_dir / relative_path).resolve()
+            tried_paths.append(candidate)
+            if candidate.is_file():
+                return candidate
+        joined = ", ".join(str(path_value) for path_value in tried_paths)
+        raise FileNotFoundError(f"browser_poc_manifest.json not found. tried: {joined}")
+
+    def resolve_manifest_relative_path(self, relative_path: str | Path) -> Path:
+        relative = Path(relative_path)
+        resolved = (self.manifest_dir / relative).resolve()
+        if resolved.exists():
+            return resolved
+        relative_text = str(relative).replace("\\", "/")
+        for legacy_name, canonical_name in MODEL_DIR_ALIAS_MAP.items():
+            legacy_fragment = f"/{legacy_name}/"
+            if legacy_fragment not in f"/{relative_text}/":
+                continue
+            rewritten_text = relative_text.replace(legacy_name, canonical_name)
+            rewritten = (self.manifest_dir / Path(rewritten_text)).resolve()
+            if rewritten.exists():
+                return rewritten
+        return resolved
+
+    def _session(self, path_value: Path) -> ort.InferenceSession:
+        options = ort.SessionOptions()
+        options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        options.intra_op_num_threads = self.thread_count
+        options.inter_op_num_threads = 1
+        session = ort.InferenceSession(str(path_value), sess_options=options, providers=self.ort_providers)
+        if self.execution_provider == EXECUTION_PROVIDER_CUDA and "CUDAExecutionProvider" not in session.get_providers():
+            raise RuntimeError(
+                "CUDAExecutionProvider was requested, but ONNX Runtime created a session without CUDA support "
+                f"for {path_value}. Actual providers: {session.get_providers()}"
+            )
+        return session
+
+    def _create_sessions(self) -> dict[str, ort.InferenceSession]:
+        tts_dir = self.tts_meta_path.parent
+        codec_dir = self.codec_meta_path.parent
+        return {
+            "prefill": self._session(tts_dir / self.tts_meta["files"]["prefill"]),
+            "decode": self._session(tts_dir / self.tts_meta["files"]["decode_step"]),
+            "local_decoder": self._session(tts_dir / self.tts_meta["files"]["local_decoder"]),
+            **(
+                {"local_greedy_frame": self._session(tts_dir / self.tts_meta["files"]["local_greedy_frame"])}
+                if self.tts_meta["files"].get("local_greedy_frame")
+                else {}
+            ),
+            **(
+                {"local_fixed_sampled_frame": self._session(tts_dir / self.tts_meta["files"]["local_fixed_sampled_frame"])}
+                if self.tts_meta["files"].get("local_fixed_sampled_frame")
+                else {}
+            ),
+            **(
+                {"local_cached_step": self._session(tts_dir / self.tts_meta["files"]["local_cached_step"])}
+                if self.tts_meta["files"].get("local_cached_step")
+                else {}
+            ),
+            "codec_encode": self._session(codec_dir / self.codec_meta["files"]["encode"]),
+            "codec_decode": self._session(codec_dir / self.codec_meta["files"]["decode_full"]),
+            "codec_decode_step": self._session(codec_dir / self.codec_meta["files"]["decode_step"]),
+        }
+
+    def list_builtin_voices(self) -> list[dict[str, Any]]:
+        return list(self.manifest["builtin_voices"])
+
+    def list_text_samples(self) -> list[dict[str, Any]]:
+        return list(self.manifest["text_samples"])
+
+    def warmup(self) -> None:
+        voice = self.list_builtin_voices()[0]
+        text_sample = self.list_text_samples()[0]
+        request_rows = self.build_voice_clone_request_rows(voice["prompt_audio_codes"], text_sample["text_token_ids"])
+        prefill_ids, prefill_dims = _flatten3d_int32([request_rows["inputIds"]])
+        prefill_mask, prefill_mask_dims = _flatten2d_int32(request_rows["attentionMask"])
+        outputs = self.sessions["prefill"].run(
+            None,
+            {
+                "input_ids": prefill_ids.reshape(prefill_dims),
+                "attention_mask": prefill_mask.reshape(prefill_mask_dims),
+            },
+        )
+        output_names = [output.name for output in self.sessions["prefill"].get_outputs()]
+        named_outputs = dict(zip(output_names, outputs, strict=True))
+        global_hidden = _extract_last_hidden(named_outputs["global_hidden"])
+        if "local_cached_step" in self.sessions:
+            local_past_by_name = self.create_empty_local_cached_past()
+            _text_logits, _audio_logits, _next_local_past = self.run_local_cached_step(
+                global_hidden,
+                text_token_id=0,
+                audio_token_id=0,
+                channel_index=0,
+                step_type=0,
+                past_valid_lengths=0,
+                local_past_by_name=local_past_by_name,
+            )
+        if "local_fixed_sampled_frame" in self.sessions and self.manifest["generation_defaults"]["sample_mode"] == SAMPLE_MODE_FIXED:
+            self.run_local_fixed_sampled_frame(
+                global_hidden,
+                previous_token_sets_by_channel=[set() for _ in range(int(self.manifest["tts_config"]["n_vq"]))],
+            )
+        elif "local_greedy_frame" in self.sessions and not bool(self.manifest["generation_defaults"]["do_sample"]):
+            self.run_local_greedy_frame(
+                global_hidden,
+                previous_token_sets_by_channel=[set() for _ in range(int(self.manifest["tts_config"]["n_vq"]))],
+                repetition_penalty=float(self.manifest["generation_defaults"]["audio_repetition_penalty"]),
+            )
+        else:
+            self.run_local_decoder(global_hidden, self.manifest["tts_config"]["audio_assistant_slot_token_id"], [])
+        empty_frames = [([0] * int(self.manifest["tts_config"]["n_vq"]))]
+        self.decode_full_audio(empty_frames)
+        self.codec_streaming_session.reset()
+        self.codec_streaming_session.run_frames(empty_frames)
+        self.codec_streaming_session.reset()
+
+    def build_text_rows(self, token_ids: list[int]) -> list[list[int]]:
+        rows: list[list[int]] = []
+        row_width = int(self.manifest["tts_config"]["n_vq"]) + 1
+        for token_id in token_ids:
+            row = [int(self.manifest["tts_config"]["audio_pad_token_id"])] * row_width
+            row[0] = int(token_id)
+            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]]:
+        rows: list[list[int]] = []
+        row_width = int(self.manifest["tts_config"]["n_vq"]) + 1
+        resolved_slot_token_id = int(
+            self.manifest["tts_config"]["audio_user_slot_token_id"] if slot_token_id is None else slot_token_id
+        )
+        for code_row in prompt_audio_codes:
+            row = [int(self.manifest["tts_config"]["audio_pad_token_id"])] * row_width
+            row[0] = resolved_slot_token_id
+            for index in range(min(len(code_row), int(self.manifest["tts_config"]["n_vq"]))):
+                row[index + 1] = int(code_row[index])
+            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_text_token_ids = [
+            *self.manifest["prompt_templates"]["user_prompt_prefix_token_ids"],
+            int(self.manifest["tts_config"]["audio_start_token_id"]),
+        ]
+        suffix_text_token_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_text_token_ids),
+            *self.build_audio_prefix_rows(prompt_audio_codes),
+            *self.build_text_rows(suffix_text_token_ids),
+        ]
+        return {
+            "inputIds": rows,
+            "attentionMask": [[1 for _ in rows]],
+        }
+
+    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 index in range(min(len(frame_prefix), n_vq - 1)):
+            padded_prefix[0, index] = int(frame_prefix[index])
+        outputs = self.sessions["local_decoder"].run(
+            None,
+            {
+                "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,
+            },
+        )
+        output_names = [output.name for output in self.sessions["local_decoder"].get_outputs()]
+        named_outputs = dict(zip(output_names, outputs, strict=True))
+        return named_outputs["text_logits"].reshape(-1), named_outputs["audio_logits"]
+
+    def create_empty_local_cached_past(self) -> dict[str, np.ndarray]:
+        local_layers = int(self.tts_meta["model_config"]["local_layers"])
+        local_heads = int(self.tts_meta["model_config"]["local_heads"])
+        local_head_dim = int(self.tts_meta["model_config"]["local_head_dim"])
+        return {
+            name: np.zeros((1, 0, local_heads, local_head_dim), dtype=np.float32)
+            for layer_index in range(local_layers)
+            for name in (f"local_past_key_{layer_index}", f"local_past_value_{layer_index}")
+        }
+
+    def run_local_cached_step(
+        self,
+        global_hidden: np.ndarray,
+        *,
+        text_token_id: int,
+        audio_token_id: int,
+        channel_index: int,
+        step_type: int,
+        past_valid_lengths: int,
+        local_past_by_name: dict[str, np.ndarray],
+    ) -> tuple[np.ndarray, np.ndarray, dict[str, np.ndarray]]:
+        outputs = self.sessions["local_cached_step"].run(
+            None,
+            {
+                "global_hidden": global_hidden.astype(np.float32, copy=False),
+                "text_token_id": np.asarray([int(text_token_id)], dtype=np.int32),
+                "audio_token_id": np.asarray([int(audio_token_id)], dtype=np.int32),
+                "channel_index": np.asarray([int(channel_index)], dtype=np.int32),
+                "step_type": np.asarray([int(step_type)], dtype=np.int32),
+                "past_valid_lengths": np.asarray([int(past_valid_lengths)], dtype=np.int32),
+                **local_past_by_name,
+            },
+        )
+        output_names = [output.name for output in self.sessions["local_cached_step"].get_outputs()]
+        named_outputs = dict(zip(output_names, outputs, strict=True))
+        next_local_past = {
+            output_name.replace("local_present_", "local_past_"): named_outputs[output_name]
+            for output_name in self.tts_meta["onnx"]["local_cached_output_names"][2:]
+        }
+        return named_outputs["text_logits"].reshape(-1), named_outputs["audio_logits"], next_local_past
+
+    def run_local_greedy_frame(
+        self,
+        global_hidden: np.ndarray,
+        *,
+        previous_token_sets_by_channel: list[set[int]],
+        repetition_penalty: float,
+    ) -> tuple[bool, list[int]]:
+        audio_codebook_size = int(self.tts_meta["model_config"]["audio_codebook_sizes"][0])
+        n_vq = int(self.manifest["tts_config"]["n_vq"])
+        repetition_seen_mask = np.zeros((1, n_vq, audio_codebook_size), dtype=np.int32)
+        for channel_index, token_ids in enumerate(previous_token_sets_by_channel):
+            for token_id in token_ids:
+                if 0 <= token_id < audio_codebook_size:
+                    repetition_seen_mask[0, channel_index, token_id] = 1
+        outputs = self.sessions["local_greedy_frame"].run(
+            None,
+            {
+                "global_hidden": global_hidden.astype(np.float32, copy=False),
+                "repetition_seen_mask": repetition_seen_mask,
+                "repetition_penalty": np.asarray([float(repetition_penalty)], dtype=np.float32),
+            },
+        )
+        output_names = [output.name for output in self.sessions["local_greedy_frame"].get_outputs()]
+        named_outputs = dict(zip(output_names, outputs, strict=True))
+        should_continue = bool(int(np.asarray(named_outputs["should_continue"]).reshape(-1)[0]))
+        frame_token_ids = np.asarray(named_outputs["frame_token_ids"]).reshape(-1).astype(np.int32, copy=False).tolist()
+        return should_continue, [int(item) for item in frame_token_ids]
+
+    def run_local_fixed_sampled_frame(
+        self,
+        global_hidden: np.ndarray,
+        *,
+        previous_token_sets_by_channel: list[set[int]],
+    ) -> tuple[bool, list[int]]:
+        audio_codebook_size = int(self.tts_meta["model_config"]["audio_codebook_sizes"][0])
+        n_vq = int(self.manifest["tts_config"]["n_vq"])
+        repetition_seen_mask = np.zeros((1, n_vq, audio_codebook_size), dtype=np.int32)
+        for channel_index, token_ids in enumerate(previous_token_sets_by_channel):
+            for token_id in token_ids:
+                if 0 <= token_id < audio_codebook_size:
+                    repetition_seen_mask[0, channel_index, token_id] = 1
+        assistant_random_u = np.asarray([min(0.99999994, max(0.0, float(self.rng.random())))], dtype=np.float32)
+        audio_random_u = np.asarray(
+            [[min(0.99999994, max(0.0, float(self.rng.random()))) for _ in range(n_vq)]],
+            dtype=np.float32,
+        )
+        outputs = self.sessions["local_fixed_sampled_frame"].run(
+            None,
+            {
+                "global_hidden": global_hidden.astype(np.float32, copy=False),
+                "repetition_seen_mask": repetition_seen_mask,
+                "assistant_random_u": assistant_random_u,
+                "audio_random_u": audio_random_u,
+            },
+        )
+        output_names = [output.name for output in self.sessions["local_fixed_sampled_frame"].get_outputs()]
+        named_outputs = dict(zip(output_names, outputs, strict=True))
+        frame_token_ids = np.asarray(named_outputs["frame_token_ids"]).reshape(-1).astype(np.int32, copy=False).tolist()
+        should_continue = bool(int(np.asarray(named_outputs["should_continue"]).reshape(-1)[0]))
+        return should_continue, [int(item) for item in frame_token_ids]
+
+    def slice_audio_channel_logits(self, audio_logits: np.ndarray, channel_index: int) -> np.ndarray:
+        per_channel = int(audio_logits.shape[-1])
+        flat = audio_logits.reshape(-1)
+        start = channel_index * per_channel
+        end = start + per_channel
+        return flat[start:end]
+
+    def decode_full_audio(self, generated_frames: list[list[int]]) -> tuple[list[np.ndarray], int]:
+        if not generated_frames:
+            return [], 0
+        audio_codes, dims = _flatten3d_int32([generated_frames])
+        outputs = self.sessions["codec_decode"].run(
+            None,
+            {
+                "audio_codes": audio_codes.reshape(dims),
+                "audio_code_lengths": np.asarray([len(generated_frames)], dtype=np.int32),
+            },
+        )
+        output_names = [output.name for output in self.sessions["codec_decode"].get_outputs()]
+        named_outputs = dict(zip(output_names, outputs, strict=True))
+        audio_length = int(named_outputs["audio_lengths"].reshape(-1)[0])
+        return _slice_channel_major_audio(named_outputs["audio"], 0, audio_length), audio_length
+
+    def generate_audio_frames(
+        self,
+        request_rows: dict[str, list[list[int]]],
+        on_frame: Callable[[list[list[int]], int, list[int]], None] | None = None,
+    ) -> list[list[int]]:
+        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"])
+        outputs = self.sessions["prefill"].run(
+            None,
+            {
+                "input_ids": prefill_ids.reshape(prefill_dims),
+                "attention_mask": prefill_mask.reshape(prefill_mask_dims),
+            },
+        )
+        output_names = [output.name for output in self.sessions["prefill"].get_outputs()]
+        named_outputs = dict(zip(output_names, outputs, strict=True))
+        global_hidden = _extract_last_hidden(named_outputs["global_hidden"])
+        past_valid_length = sum(int(item) for item in request_rows["attentionMask"][0])
+        past_by_name = {
+            output_name.replace("present_", "past_"): named_outputs[output_name]
+            for output_name in self.tts_meta["onnx"]["prefill_output_names"][1:]
+        }
+        generated_frames: list[list[int]] = []
+        previous_tokens_by_channel = [[] for _ in range(int(self.manifest["tts_config"]["n_vq"]))]
+        previous_token_sets_by_channel = [set() for _ in range(int(self.manifest["tts_config"]["n_vq"]))]
+
+        for step_index in range(int(generation_defaults["max_new_frames"])):
+            frame: list[int] = []
+            if "local_greedy_frame" in self.sessions and not bool(generation_defaults["do_sample"]):
+                should_continue, frame = self.run_local_greedy_frame(
+                    global_hidden,
+                    previous_token_sets_by_channel=previous_token_sets_by_channel,
+                    repetition_penalty=float(generation_defaults["audio_repetition_penalty"]),
+                )
+                if not should_continue:
+                    break
+                for channel_index, sampled_token in enumerate(frame):
+                    previous_tokens_by_channel[channel_index].append(sampled_token)
+                    previous_token_sets_by_channel[channel_index].add(sampled_token)
+            elif "local_fixed_sampled_frame" in self.sessions and generation_defaults["sample_mode"] == SAMPLE_MODE_FIXED:
+                should_continue, frame = self.run_local_fixed_sampled_frame(
+                    global_hidden,
+                    previous_token_sets_by_channel=previous_token_sets_by_channel,
+                )
+                if not should_continue:
+                    break
+                for channel_index, sampled_token in enumerate(frame):
+                    previous_tokens_by_channel[channel_index].append(sampled_token)
+                    previous_token_sets_by_channel[channel_index].add(sampled_token)
+            elif "local_cached_step" in self.sessions:
+                local_past_by_name = self.create_empty_local_cached_past()
+                local_past_valid_length = 0
+                local_text_logits, _ignored_audio_logits, local_past_by_name = self.run_local_cached_step(
+                    global_hidden,
+                    text_token_id=0,
+                    audio_token_id=0,
+                    channel_index=0,
+                    step_type=0,
+                    past_valid_lengths=local_past_valid_length,
+                    local_past_by_name=local_past_by_name,
+                )
+                local_past_valid_length += 1
+                next_text_token = _sample_assistant_text_token(
+                    local_text_logits,
+                    self.manifest,
+                    generation_defaults,
+                    self.rng,
+                )
+                if next_text_token != int(self.manifest["tts_config"]["audio_assistant_slot_token_id"]):
+                    break
+                _unused_text_logits, audio_logits, local_past_by_name = self.run_local_cached_step(
+                    global_hidden,
+                    text_token_id=next_text_token,
+                    audio_token_id=0,
+                    channel_index=0,
+                    step_type=1,
+                    past_valid_lengths=local_past_valid_length,
+                    local_past_by_name=local_past_by_name,
+                )
+                local_past_valid_length += 1
+                first_channel_logits = self.slice_audio_channel_logits(audio_logits, 0).astype(np.float32, copy=False)
+                sampled_token = _sample_audio_token(
+                    first_channel_logits,
+                    previous_tokens_by_channel[0],
+                    previous_token_sets_by_channel[0],
+                    generation_defaults,
+                    self.rng,
+                )
+                frame.append(sampled_token)
+                previous_tokens_by_channel[0].append(sampled_token)
+                previous_token_sets_by_channel[0].add(sampled_token)
+
+                previous_token = sampled_token
+                host_sampled_channel_limit = int(self.manifest["tts_config"]["n_vq"])
+                for channel_index in range(1, host_sampled_channel_limit):
+                    _unused_text_logits, audio_logits, local_past_by_name = self.run_local_cached_step(
+                        global_hidden,
+                        text_token_id=0,
+                        audio_token_id=previous_token,
+                        channel_index=channel_index - 1,
+                        step_type=2,
+                        past_valid_lengths=local_past_valid_length,
+                        local_past_by_name=local_past_by_name,
+                    )
+                    local_past_valid_length += 1
+                    channel_logits = self.slice_audio_channel_logits(audio_logits, channel_index).astype(np.float32, copy=False)
+                    sampled_token = _sample_audio_token(
+                        channel_logits,
+                        previous_tokens_by_channel[channel_index],
+                        previous_token_sets_by_channel[channel_index],
+                        generation_defaults,
+                        self.rng,
+                    )
+                    frame.append(sampled_token)
+                    previous_tokens_by_channel[channel_index].append(sampled_token)
+                    previous_token_sets_by_channel[channel_index].add(sampled_token)
+                    previous_token = sampled_token
+            else:
+                local_text_logits, _ = self.run_local_decoder(global_hidden, 0, [])
+                next_text_token = _sample_assistant_text_token(
+                    local_text_logits,
+                    self.manifest,
+                    generation_defaults,
+                    self.rng,
+                )
+                if next_text_token != int(self.manifest["tts_config"]["audio_assistant_slot_token_id"]):
+                    break
+                for channel_index in range(int(self.manifest["tts_config"]["n_vq"])):
+                    _, audio_logits = self.run_local_decoder(global_hidden, next_text_token, frame)
+                    channel_logits = self.slice_audio_channel_logits(audio_logits, channel_index).astype(np.float32, copy=False)
+                    sampled_token = _sample_audio_token(
+                        channel_logits,
+                        previous_tokens_by_channel[channel_index],
+                        previous_token_sets_by_channel[channel_index],
+                        generation_defaults,
+                        self.rng,
+                    )
+                    frame.append(sampled_token)
+                    previous_tokens_by_channel[channel_index].append(sampled_token)
+                    previous_token_sets_by_channel[channel_index].add(sampled_token)
+            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 index, token in enumerate(frame):
+                next_row[0, 0, index + 1] = int(token)
+            decode_feeds: dict[str, np.ndarray] = {
+                "input_ids": next_row,
+                "past_valid_lengths": np.asarray([past_valid_length], dtype=np.int32),
+            }
+            for input_name in self.tts_meta["onnx"]["decode_input_names"][2:]:
+                decode_feeds[input_name] = past_by_name[input_name]
+            decode_outputs = self.sessions["decode"].run(None, decode_feeds)
+            decode_output_names = [output.name for output in self.sessions["decode"].get_outputs()]
+            named_decode_outputs = dict(zip(decode_output_names, decode_outputs, strict=True))
+            global_hidden = _extract_last_hidden(named_decode_outputs["global_hidden"])
+            past_valid_length += 1
+            past_by_name = {
+                output_name.replace("present_", "past_"): named_decode_outputs[output_name]
+                for output_name in self.tts_meta["onnx"]["decode_output_names"][1:]
+            }
+            if on_frame is not None:
+                on_frame(generated_frames, step_index, frame)
+        return generated_frames
+
+__all__ = [
+    "CodecStreamingDecodeSession",
+    "EXECUTION_PROVIDER_CPU",
+    "EXECUTION_PROVIDER_CUDA",
+    "OrtCpuRuntime",
+    "SAMPLE_MODE_FIXED",
+    "SAMPLE_MODE_FULL",
+    "SAMPLE_MODE_GREEDY",
+    "_normalize_execution_provider",
+    "_normalize_sample_mode",
+    "_resolve_stream_decode_frame_budget",
+]