| 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", |
| ] |
|
|
