From a8b2bf11d2f2e9282e2182c2aad932d58777e3dc Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Fri, 3 Jul 2026 16:00:30 +0800 Subject: [PATCH] [Feature] Support MOSS-TTS-Local-v1.5 Streaming (#4804) Signed-off-by: gcanlin --- tests/worker/test_omni_gpu_model_runner.py | 6 +- .../benchmarks/data_modules/seed_tts_eval.py | 70 +- vllm_omni/benchmarks/patch/patch.py | 34 +- vllm_omni/deploy/moss_tts_local.yaml | 6 +- .../entrypoints/openai/serving_speech.py | 15 +- .../models/moss_tts/audio_tokenizer_v2.py | 1959 +++++++++++++++++ .../configuration_moss_audio_tokenizer_v2.py | 359 +++ .../moss_tts/modeling_moss_tts_codec.py | 471 +++- .../moss_tts/modeling_moss_tts_local.py | 3 +- .../moss_tts/modeling_moss_tts_local_depth.py | 21 +- .../moss_tts/modeling_moss_tts_talker.py | 285 ++- .../models/moss_tts/pipeline.py | 11 +- .../stage_input_processors/moss_tts.py | 149 +- .../worker/gpu_generation_model_runner.py | 8 + vllm_omni/worker/gpu_model_runner.py | 14 +- 15 files changed, 3222 insertions(+), 189 deletions(-) create mode 100644 vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py create mode 100644 vllm_omni/model_executor/models/moss_tts/configuration_moss_audio_tokenizer_v2.py diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index f562d18ce..e68dd6c43 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/tests/worker/test_omni_gpu_model_runner.py @@ -275,7 +275,7 @@ def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker runner = _make_runner(req_ids=("r1",), hidden_size=4) runner.requests["r1"].sampling_params = SimpleNamespace( seed=42, - extra_args={"qwen3_tts_request_seed": 42}, + extra_args={"tts_local_seed": 42}, ) runner.talker_mtp = CaptureTalkerMTP() runner.vllm_config = SimpleNamespace( @@ -319,11 +319,11 @@ def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch): runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4) runner.requests["r1"].sampling_params = SimpleNamespace( seed=11, - extra_args={"qwen3_tts_request_seed": 11}, + extra_args={"tts_local_seed": 11}, ) runner.requests["r2"].sampling_params = SimpleNamespace( seed=22, - extra_args={"qwen3_tts_request_seed": 22}, + extra_args={"tts_local_seed": 22}, ) runner.talker_mtp = CaptureTalkerMTP() runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={})) diff --git a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py index eb06dde70..efa312dec 100644 --- a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py +++ b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py @@ -36,7 +36,9 @@ Enable with ``SEED_TTS_WER_EVAL=1`` or ``--seed-tts-wer-eval``. Install optional pip install 'vllm-omni[dev]' Env: ``SEED_TTS_EVAL_DEVICE`` (e.g. ``cuda:0``, ``cpu``); ``SEED_TTS_HF_WHISPER_MODEL`` -defaults to ``openai/whisper-large-v3`` (override for debugging only). +defaults to ``openai/whisper-large-v3`` (override for debugging only). Set +``SEED_TTS_WER_SAVE_AUDIO_DIR`` to save the captured 24 kHz mono WAV used by +WER evaluation for each synthesized utterance. """ from __future__ import annotations @@ -50,6 +52,7 @@ import string import tempfile import threading import wave +from pathlib import Path from typing import Any import numpy as np @@ -87,6 +90,30 @@ def pcm_s16le_mono_to_wav_bytes(pcm: bytes, *, sample_rate: int = 24000) -> byte return buf.getvalue() +def _safe_filename_part(value: Any, *, default: str = "item", max_len: int = 96) -> str: + text = str(value or "").strip() or default + safe = "".join(ch if ch.isalnum() or ch in ("-", "_", ".") else "_" for ch in text) + safe = safe.strip("._") or default + return safe[:max_len] + + +def _save_seed_tts_eval_audio( + pcm: bytes, + *, + output_dir: Path | None, + index: int, + utterance_id: Any, + locale: str, +) -> str | None: + if output_dir is None: + return None + output_dir.mkdir(parents=True, exist_ok=True) + stem = f"{index:05d}_{_safe_filename_part(utterance_id)}_{_safe_filename_part(locale)}" + path = output_dir / f"{stem}.wav" + path.write_bytes(pcm_s16le_mono_to_wav_bytes(pcm, sample_rate=24000)) + return str(path) + + def _get_eval_device() -> str: explicit = os.environ.get("SEED_TTS_EVAL_DEVICE", "").strip() if explicit: @@ -512,13 +539,18 @@ def compute_seed_tts_wer_metrics( sim_skipped_no_ref = 0 utmos_failed = 0 utmos_on = _eval_submetric_enabled("SEED_TTS_UTMOS_EVAL", default=False) + save_audio_raw = os.environ.get("SEED_TTS_WER_SAVE_AUDIO_DIR", "").strip() + save_audio_dir = Path(save_audio_raw).expanduser() if save_audio_raw else None + saved_audio = 0 + save_audio_failed = 0 - for req, out in zip(input_requests, outputs, strict=True): + for index, (req, out) in enumerate(zip(input_requests, outputs, strict=True)): assert isinstance(req, SeedTTSSampleRequest) ref = req.prompt locale = req.seed_tts_locale or "en" row_lang = "zh" if locale.lower().startswith("zh") else "en" utmos_v: float | None = None + audio_path: str | None = None if not out.success: request_failed += 1 @@ -545,6 +577,23 @@ def compute_seed_tts_wer_metrics( } ) continue + try: + audio_path = _save_seed_tts_eval_audio( + pcm, + output_dir=save_audio_dir, + index=index, + utterance_id=req.seed_tts_utterance_id, + locale=locale, + ) + if audio_path: + saved_audio += 1 + except OSError as e: + save_audio_failed += 1 + logger.warning( + "Seed-TTS WER audio save failed for utterance=%s: %s", + req.seed_tts_utterance_id, + e, + ) wav_16k = _pcm_s16le_to_f32_16k(pcm) if len(wav_16k) == 0: @@ -555,6 +604,7 @@ def compute_seed_tts_wer_metrics( "utterance_id": req.seed_tts_utterance_id, "locale": locale, "error": "empty_audio", + "audio_path": audio_path, } ) continue @@ -609,6 +659,7 @@ def compute_seed_tts_wer_metrics( "locale": locale, "error": "asr_exception", "detail": str(e)[:500], + "audio_path": audio_path, } ) continue @@ -621,6 +672,7 @@ def compute_seed_tts_wer_metrics( "utterance_id": req.seed_tts_utterance_id, "locale": locale, "error": "empty_asr", + "audio_path": audio_path, } ) continue @@ -637,6 +689,7 @@ def compute_seed_tts_wer_metrics( "locale": locale, "error": "wer_compute_failed", "detail": str(e)[:500], + "audio_path": audio_path, } ) continue @@ -673,6 +726,8 @@ def compute_seed_tts_wer_metrics( "reference_raw": raw_truth, "asr_raw": raw_hypo, } + if audio_path: + row["audio_path"] = audio_path if sim_v is not None: row["sim"] = sim_v if utmos_v is not None: @@ -697,7 +752,11 @@ def compute_seed_tts_wer_metrics( "seed_tts_utmos_mean": statistics.fmean(utmos_values) if utmos_values else None, "seed_tts_utmos_median": statistics.median(utmos_values) if utmos_values else None, "seed_tts_utmos_failed": utmos_failed, + "seed_tts_saved_audio": saved_audio, + "seed_tts_save_audio_failed": save_audio_failed, } + if save_audio_dir is not None: + result["seed_tts_save_audio_dir"] = str(save_audio_dir) if include_per_item: result["seed_tts_wer_eval_items"] = items return result @@ -729,6 +788,13 @@ def print_seed_tts_wer_summary(metrics: dict[str, Any]) -> None: print("{:<40} {:<10}".format("Request failed:", metrics.get("seed_tts_request_failed", 0))) print("{:<40} {:<10}".format("No PCM captured:", metrics.get("seed_tts_no_pcm", 0))) print("{:<40} {:<10}".format("ASR / WER failed:", metrics.get("seed_tts_asr_failed", 0))) + save_dir = metrics.get("seed_tts_save_audio_dir") + if save_dir: + print("{:<40} {:<10}".format("Saved eval WAVs:", metrics.get("seed_tts_saved_audio", 0))) + print("{:<40} {}".format("Saved eval WAV dir:", save_dir)) + failed = int(metrics.get("seed_tts_save_audio_failed", 0) or 0) + if failed: + print("{:<40} {:<10}".format("Save eval WAV errors:", failed)) if sim_ev or metrics.get("seed_tts_sim_skipped_no_ref") or metrics.get("seed_tts_sim_failed"): print("{:<40} {:<10}".format("SIM evaluated (higher ~ closer):", sim_ev)) sm = metrics.get("seed_tts_sim_mean") diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index a756c97f0..d6e28319c 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -137,6 +137,30 @@ def _audio_pcm_format() -> tuple[int, int]: return sample_rate, channels +def _pcm_s16le_to_seed_tts_wer_bytes( + pcm_bytes: bytes, + *, + sample_rate: int, + channels: int, +) -> bytes: + """Normalize streamed raw PCM to the 24 kHz mono PCM used by Seed-TTS WER.""" + if not pcm_bytes: + return b"" + channels = max(1, int(channels)) + pcm = np.frombuffer(pcm_bytes, dtype=np.int16) + if channels > 1: + usable = (pcm.size // channels) * channels + pcm = pcm[:usable].reshape(-1, channels).mean(axis=1) + pcm_f32 = pcm.astype(np.float32) / 32767.0 + if int(sample_rate) != 24000 and pcm_f32.size: + from vllm.multimodal.audio import AudioResampler + + resampler = AudioResampler(target_sr=24000) + pcm_f32 = resampler.resample(pcm_f32, orig_sr=int(sample_rate)) + pcm_f32 = np.clip(pcm_f32, -1.0, 1.0) + return (pcm_f32 * 32767).astype(np.int16).tobytes() + + get_samples_old = datasets.get_samples _DEFAULT_DAILY_OMNI_REPO = "liarliar/Daily-Omni" @@ -1216,7 +1240,15 @@ async def async_request_openai_audio_speech( output.audio_continuity_ok = continuity.is_continuous output.audio_underrun_event_count = continuity.underrun_event_count if pcm_capture is not None and pcm_capture: - output.tts_output_pcm_bytes = bytes(pcm_capture) + try: + output.tts_output_pcm_bytes = _pcm_s16le_to_seed_tts_wer_bytes( + bytes(pcm_capture), + sample_rate=sample_rate, + channels=channels, + ) + except Exception as ex: + logger.warning("Seed-TTS WER PCM normalization failed: %s", ex) + output.tts_output_pcm_bytes = bytes(pcm_capture) elif capture_wer_pcm: ct = response.headers.get("Content-Type", "") logger.warning( diff --git a/vllm_omni/deploy/moss_tts_local.yaml b/vllm_omni/deploy/moss_tts_local.yaml index 7e474caa6..1455ccc09 100644 --- a/vllm_omni/deploy/moss_tts_local.yaml +++ b/vllm_omni/deploy/moss_tts_local.yaml @@ -14,8 +14,10 @@ connectors: connector_get_sleep_s: 0.005 connector_get_max_wait_first_chunk: 1000 connector_get_max_wait: 300 - codec_chunk_frames: 15 # 15 frames * 80 ms/frame = 1.2 s chunks + initial_codec_chunk_frames: 1 # 1 frame * 80 ms/frame = 80 ms first audio chunk + codec_chunk_frames: 15 # 15 frames * 80 ms/frame = 1.2 s steady chunks codec_left_context_frames: 0 + codec_stream_slots: 8 stages: - stage_id: 0 @@ -37,7 +39,7 @@ stages: repetition_penalty: 1.1 - stage_id: 1 - max_num_seqs: 1 + max_num_seqs: 8 gpu_memory_utilization: 0.15 enforce_eager: true enable_prefix_caching: false diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index ac37a63a2..78bb0ab2c 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -3705,15 +3705,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): sampling_params_list[0].max_tokens = int(request.max_new_tokens) + 1 if request.seed is not None and sampling_params_list: - if sampling_params_list is self.engine_client.default_sampling_params_list: - import copy + import copy - sampling_params_list = copy.deepcopy(sampling_params_list) - sampling_params_list[0].seed = request.seed - if self._tts_model_type == "qwen3_tts": - if sampling_params_list[0].extra_args is None: - sampling_params_list[0].extra_args = {} - sampling_params_list[0].extra_args["qwen3_tts_request_seed"] = request.seed + sampling_params_list = copy.deepcopy(sampling_params_list) + stage0_params = sampling_params_list[0] + stage0_params.seed = request.seed + if stage0_params.extra_args is None: + stage0_params.extra_args = {} + stage0_params.extra_args["tts_local_seed"] = request.seed generator = self.engine_client.generate( prompt=prompt, diff --git a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py new file mode 100644 index 000000000..f3ba18c36 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py @@ -0,0 +1,1959 @@ +# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch MossAudioTokenizer model.""" + +from __future__ import annotations + +import copy +import math +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from typing import cast + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.modeling_utils import PreTrainedAudioTokenizerBase +from transformers.utils import ModelOutput, auto_docstring, logging + +from .configuration_moss_audio_tokenizer_v2 import MossAudioTokenizerConfig + +logger = logging.get_logger(__name__) + + +# ============================================================================= +# Output Classes +# ============================================================================= + + +@dataclass +@auto_docstring +class MossAudioTokenizerEncoderOutput(ModelOutput): + r""" + audio_codes (`torch.LongTensor` of shape `(num_quantizers, batch_size, sequence_length)`, *optional*): + Discrete audio codes computed using the encoder and quantizer. + audio_codes_lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Valid lengths for each sample's audio codes. + encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, hidden_size, sequence_length)`, *optional*): + Hidden states from the encoder before quantization. + """ + + audio_codes: torch.Tensor | None = None + audio_codes_lengths: torch.Tensor | None = None + encoder_hidden_states: torch.Tensor | None = None + + +@dataclass +@auto_docstring +class MossAudioTokenizerDecoderOutput(ModelOutput): + r""" + audio (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`, *optional*): + Decoded audio waveform. + audio_lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Valid lengths for each sample's audio. + """ + + audio: torch.Tensor | None = None + audio_lengths: torch.Tensor | None = None + + +@dataclass +@auto_docstring +class MossAudioTokenizerOutput(ModelOutput): + r""" + audio (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`, *optional*): + Decoded audio waveform. + audio_lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Valid lengths for each sample's audio. + audio_codes (`torch.LongTensor` of shape `(num_quantizers, batch_size, sequence_length)`, *optional*): + Discrete audio codes computed using the encoder and quantizer. + audio_codes_lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Valid lengths for each sample's audio codes. + """ + + audio: torch.Tensor | None = None + audio_lengths: torch.Tensor | None = None + audio_codes: torch.Tensor | None = None + audio_codes_lengths: torch.Tensor | None = None + + +# ============================================================================= +# Streaming Module Base Classes +# ============================================================================= + + +@dataclass +class StreamingState: + """Base state for streaming modules.""" + + batch_size: int + device: torch.device + + def __post_init__(self): + self.exec_mask = torch.ones(self.batch_size, dtype=torch.bool, device=self.device) + + def set_exec_mask(self, exec_mask: torch.Tensor): + self.exec_mask[:] = exec_mask + + def reset(self, reset_mask: torch.Tensor) -> None: + self.exec_mask[:] = torch.where(reset_mask, torch.ones_like(self.exec_mask), self.exec_mask) + + def __enter__(self): + # ExitStack expects a context manager; returning self is conventional and useful for debugging. + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + pass + + +class StreamingModule(nn.Module): + """Base class for streaming components.""" + + def __init__(self) -> None: + super().__init__() + self._streaming_state: StreamingState | None = None + self._streaming_detached: bool = False + self._cached_children: list[tuple[str, StreamingModule]] | None = None + + @property + def is_streaming(self): + return self._streaming_state is not None + + def _apply_named_streaming(self, fn): + def _handle_module(prefix: str, module: nn.Module): + if isinstance(module, StreamingModule): + if module._streaming_detached and prefix != "": + return + if self._cached_children is None: + raise RuntimeError("Internal error: _cached_children should be initialized before traversal.") + self._cached_children.append((prefix, module)) + for name, child in module.named_children(): + new_prefix = f"{prefix}.{name}" if prefix else name + _handle_module(new_prefix, child) + + if self._cached_children is None: + self._cached_children = [] + _handle_module("", self) + for name, child in self._cached_children: + fn(name, child) + + def _start_streaming(self, batch_size: int, exit_stack: ExitStack): + def _start_streaming_fn(name: str, module: StreamingModule): + if module._streaming_state is not None: + raise RuntimeError(f"{name} is already streaming!") + state = module._init_streaming_state(batch_size) + exit_stack.enter_context(state) + module._streaming_state = state + + self._apply_named_streaming(_start_streaming_fn) + + def _stop_streaming(self) -> None: + def _stop_streaming_fn(name: str, module: StreamingModule): + module._streaming_state = None + + self._apply_named_streaming(_stop_streaming_fn) + + def _init_streaming_state(self, batch_size: int) -> StreamingState: + device = next(iter(self.parameters())).device + return StreamingState(batch_size, device) + + def streaming(self, batch_size: int) -> ExitStack: + """Context manager to enter streaming mode.""" + exit_stack = ExitStack() + self._start_streaming(batch_size, exit_stack) + exit_stack.callback(self._stop_streaming) + return exit_stack + + def _set_streaming_exec_mask(self, exec_mask: torch.Tensor) -> None: + """Set the active slot mask on every live streaming module.""" + + def _set_exec_mask(_: str, module: StreamingModule) -> None: + state = module._streaming_state + if state is not None: + state.set_exec_mask(exec_mask.to(state.device)) + + self._apply_named_streaming(_set_exec_mask) + + +class StreamingContainer(StreamingModule): + """Container for streaming modules.""" + + pass + + +# ============================================================================= +# Normalization Layers +# ============================================================================= + + +class MossAudioTokenizerRMSNorm(nn.Module): + """Root Mean Square Layer Normalization.""" + + def __init__( + self, + dim: int, + eps: float = 1e-5, + dtype: torch.dtype | None = None, + device=None, + ): + super().__init__() + self.eps = eps + self.dtype = dtype + self.alpha = nn.Parameter(torch.full((1, 1, dim), 1.0, requires_grad=True, device=device, dtype=dtype)) + + def forward(self, x: torch.Tensor): + x_dtype = x.dtype + if self.dtype is not None: + x = x.to(self.dtype) + var = self.eps + torch.mean(x**2, dim=2, keepdim=True) + y = (x * (self.alpha.to(var) * torch.rsqrt(var))).to(x_dtype) + return y + + +class MossAudioTokenizerLayerScale(nn.Module): + """Layer scale from Touvron et al. 2021.""" + + def __init__( + self, + channels: int, + init: float = 1e-4, + channel_last: bool = True, + device=None, + dtype=None, + ): + super().__init__() + self.channel_last = channel_last + self.scale = nn.Parameter(torch.full((channels,), init, requires_grad=True, device=device, dtype=dtype)) + + def forward(self, x: torch.Tensor): + if self.channel_last: + return self.scale * x + else: + return self.scale[:, None] * x + + +def create_norm_fn(norm_type: str, dim: int, **kwargs) -> nn.Module: + """Create normalization module.""" + if norm_type == "layer_norm": + return nn.LayerNorm(dim, eps=1e-5, **kwargs) + elif norm_type in {"rms_norm"}: + return MossAudioTokenizerRMSNorm(dim, eps=1e-5, **kwargs) + elif norm_type in {"rms_norm_f32"}: + kwargs.pop("dtype", None) + return MossAudioTokenizerRMSNorm(dim, eps=1e-8, dtype=torch.float, **kwargs) + else: + raise ValueError(f"Unknown norm type: {norm_type}") + + +# ============================================================================= +# Rotary Position Embedding +# ============================================================================= + + +def apply_rope( + q: torch.Tensor, + k: torch.Tensor, + offset: torch.Tensor, + max_period: float = 10_000, + time_before_heads: bool = False, +): + """Apply rotary position embedding.""" + if time_before_heads: + B, T, H, D = q.shape + else: + B, H, T, D = q.shape + if k.shape != q.shape: + raise ValueError(f"Expected k.shape == q.shape, got k={tuple(k.shape)} q={tuple(q.shape)}") + if D <= 0 or (D % 2) != 0: + raise ValueError(f"RoPE requires an even last dimension, got D={D}") + + ds = torch.arange(D // 2, device=q.device, dtype=torch.float32) + freqs = torch.exp(ds * (-math.log(max_period) * 2 / D)) + ts = offset.float().view(-1, 1) + torch.arange(T, device=q.device, dtype=torch.float32) + + if time_before_heads: + ts = ts.view(B, -1, 1, 1) + else: + ts = ts.view(B, 1, -1, 1) + + dims = q.shape[:-1] + q = q.view(*dims, D // 2, 2) + k = k.view(*dims, D // 2, 2) + + qr, qi = q[..., 0].float(), q[..., 1].float() + kr, ki = k[..., 0].float(), k[..., 1].float() + + rotr = torch.cos(freqs * ts) + roti = torch.sin(freqs * ts) + + qor = qr * rotr - qi * roti + qoi = qr * roti + qi * rotr + kor = kr * rotr - ki * roti + koi = kr * roti + ki * rotr + + dtype = q.dtype + qo = torch.stack([qor.to(dtype), qoi.to(dtype)], dim=-1) + ko = torch.stack([kor.to(dtype), koi.to(dtype)], dim=-1) + + return qo.view(*dims, D), ko.view(*dims, D) + + +class MossAudioTokenizerRotaryEmbedding(nn.Module): + """Rotary positional embedding (RoPE).""" + + def __init__(self, max_period: float = 10000.0): + super().__init__() + self.max_period = max_period + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + offset: torch.Tensor, + time_before_heads: bool = False, + ): + return apply_rope(q, k, offset, self.max_period, time_before_heads) + + +# ============================================================================= +# Gating Modules +# ============================================================================= + + +class MossAudioTokenizerActivationGating(nn.Module): + """Gating FFN layer with activation.""" + + def __init__(self, dim: int, dim_feedforward: int, activation, **factory_kwargs): + super().__init__() + if dim_feedforward == 4 * dim: + hidden = (21 * dim) // 8 + else: + hidden = (2 * dim_feedforward) // 3 + + self.linear_in = nn.Linear(dim, 2 * hidden, bias=False, **factory_kwargs) + self.linear_out = nn.Linear(hidden, dim, bias=False, **factory_kwargs) + self.activation = activation + + def forward(self, x: torch.Tensor): + x = self.linear_in(x) + B, T, _ = x.shape + x = x.view(B, T, 2, -1) + x = self.activation(x[..., 0, :]) * x[..., 1, :] + x = self.linear_out(x) + return x + + +def _get_activation(name: str): + if name in ["sigmoid", "tanh", "relu"]: + return getattr(torch, name) + elif name in ["leaky_relu", "elu", "gelu", "silu", "mish", "softsign"]: + return getattr(F, name) + elif name == "identity": + return nn.Identity() + else: + raise ValueError(f"Unknown activation {name}") + + +def make_gating(name: str, dim: int, dim_feedforward: int, **factory_kwargs) -> nn.Module: + return MossAudioTokenizerActivationGating(dim, dim_feedforward, _get_activation(name), **factory_kwargs) + + +# ============================================================================= +# Positional Embeddings +# ============================================================================= + + +def create_sin_embedding( + positions: torch.Tensor, + dim: int, + max_period: float = 10000, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Create sinusoidal positional embedding with shape [B, T, C].""" + if dim % 2 != 0: + raise ValueError(f"Sinusoidal embedding requires even dim, got dim={dim}") + half_dim = dim // 2 + if half_dim <= 1: + raise ValueError(f"Sinusoidal embedding requires dim >= 4, got dim={dim}") + positions = positions.to(dtype) + adim = torch.arange(half_dim, device=positions.device, dtype=dtype).view(1, 1, -1) + max_period_tensor = torch.full([], max_period, device=positions.device, dtype=dtype) + phase = positions / (max_period_tensor ** (adim / (half_dim - 1))) + return torch.cat([torch.cos(phase), torch.sin(phase)], dim=-1) + + +# ============================================================================= +# KV Cache for Attention +# ============================================================================= + + +class KVCacheResult: + """Container for KV cache results that supports tuple unpacking.""" + + __slots__ = ("keys", "values", "positions") + + def __init__(self, keys: torch.Tensor, values: torch.Tensor, positions: torch.Tensor): + self.keys = keys + self.values = values + self.positions = positions + + def __iter__(self): + """Allow unpacking as (keys, values, positions).""" + return iter((self.keys, self.values, self.positions)) + + @staticmethod + def from_kv(keys: torch.Tensor, values: torch.Tensor) -> KVCacheResult: + B, H, T, D = keys.shape + positions = torch.arange(T, device=keys.device, dtype=torch.long) + return KVCacheResult(keys, values, positions.expand(B, -1)) + + +class RingKVCache: + """Efficient streaming KVCache compatible with CUDA Graph.""" + + def __init__( + self, + batch_size: int, + num_heads: int, + dim_per_head: int, + capacity: int, + respect_exec_mask: bool = True, + device: torch.device = torch.device("cuda"), + dtype: torch.dtype = torch.bfloat16, + ): + self.capacity = capacity + self.cache = torch.zeros( + (2, batch_size, num_heads, capacity, dim_per_head), + device=device, + dtype=dtype, + ) + self.respect_exec_mask = respect_exec_mask + if self.respect_exec_mask: + self.end_offset = torch.zeros(batch_size, device=device, dtype=torch.long) + else: + self.end_offset = torch.zeros(1, device=device, dtype=torch.long) + + def reset(self, reset_mask: torch.Tensor) -> None: + self.end_offset[:] = torch.where(reset_mask, torch.zeros_like(self.end_offset), self.end_offset) + + def complete(self, k: torch.Tensor, v: torch.Tensor, exec_mask: torch.Tensor) -> KVCacheResult: + B, H, T, D = k.shape + if T <= 0: + raise ValueError(f"Expected T > 0, got T={T}") + + indexes = torch.arange(T, device=self.end_offset.device, dtype=self.end_offset.dtype) + indexes = indexes + self.end_offset.view(-1, 1) + indexes = indexes % self.capacity + + if self.respect_exec_mask: + this_indexes = indexes.view(B, 1, T, 1).expand(-1, H, T, D) + self.cache[0].scatter_(2, this_indexes, k) + self.cache[1].scatter_(2, this_indexes, v) + else: + self.cache[0].index_copy_(2, indexes[0], k) + self.cache[1].index_copy_(2, indexes[0], v) + + keys = self.cache[0] + values = self.cache[1] + + indexes = torch.arange(self.capacity, device=self.end_offset.device, dtype=torch.long) + last_offset = self.end_offset.view(-1, 1) + T - 1 + end_index = last_offset % self.capacity + delta = indexes - end_index + + positions = torch.where( + delta <= 0, + last_offset + delta, + last_offset + delta - self.capacity, + ) + + if self.respect_exec_mask: + self.end_offset[:] = torch.where(exec_mask, self.end_offset + T, self.end_offset) + else: + self.end_offset.add_(T) + + invalid = indexes >= self.end_offset.view(-1, 1) + positions = torch.where(invalid, torch.full_like(positions, -1), positions) + + return KVCacheResult(keys, values, positions) + + +# ============================================================================= +# Multi-Head Attention +# ============================================================================= + + +@dataclass +class MHAState(StreamingState): + kv_cache: RingKVCache | None + offset: torch.Tensor + offset_cpu: int + + def reset(self, reset_mask: torch.Tensor): + super().reset(reset_mask) + self.offset[:] = torch.where(reset_mask, torch.zeros_like(self.offset), self.offset) + if self.kv_cache is not None: + self.kv_cache.reset(reset_mask) + self.offset_cpu = 0 + + +def apply_weights_per_step( + modules: nn.ModuleList, + schedule: list[int] | None, + x: torch.Tensor, + offset: int | None, +) -> torch.Tensor: + """Apply different weights for each time step.""" + if len(modules) == 1: + return modules[0](x) + + if offset is None: + raise ValueError("offset must be provided when using per-step weights (len(modules) > 1).") + ys = [] + B, T, C = x.shape + for t in range(T): + module_index = t + offset + if schedule is not None: + if module_index >= len(schedule) or module_index < 0: + raise ValueError( + f"weights_per_step_schedule is too short for module_index={module_index} (len={len(schedule)})." + ) + module_index = schedule[module_index] + if module_index >= len(modules) or module_index < 0: + raise ValueError(f"module_index={module_index} out of range for len(modules)={len(modules)}.") + y = modules[module_index](x[:, t : t + 1]) + ys.append(y) + return torch.cat(ys, 1) + + +class MossAudioTokenizerMultiheadAttention(StreamingModule): + """Multi-head attention with streaming support.""" + + def __init__( + self, + embed_dim: int, + num_heads: int, + causal: bool = False, + context: int | None = None, + rope: MossAudioTokenizerRotaryEmbedding | None = None, + weights_per_step: int = 0, + weights_per_step_schedule: list[int] | None = None, + device=None, + dtype=None, + ): + super().__init__() + factory_kwargs = {"device": device, "dtype": dtype} + + self.embed_dim = embed_dim + self.causal = causal + self.context = context + self.rope = rope + self.num_heads = num_heads + self.weights_per_step = weights_per_step + self.weights_per_step_schedule = weights_per_step_schedule + + out_dim = 3 * embed_dim + mult = 1 + if weights_per_step: + mult = max(weights_per_step_schedule) + 1 if weights_per_step_schedule else weights_per_step + self.mult = mult + + self.out_projs = nn.ModuleList( + [nn.Linear(embed_dim, embed_dim, bias=False, **factory_kwargs) for _ in range(mult)] + ) + self.in_projs = nn.ModuleList( + [nn.Linear(embed_dim, out_dim, bias=False, **factory_kwargs) for _ in range(mult)] + ) + + self._register_load_state_dict_pre_hook(self._load_hook, with_module=True) + + @staticmethod + def _load_hook(module, state_dict, prefix, *_): + mappings = { + "in_proj_weight": "in_projs.{i}.weight", + "in_proj.weight": "in_projs.{i}.weight", + "out_proj.weight": "out_projs.{i}.weight", + } + mult = module.mult + for suffix in ["", "_scb"]: + for source, target in mappings.items(): + this_source = prefix + source + suffix + if this_source in state_dict: + weight = state_dict[this_source] + _, *OD = weight.shape + weight = weight.view(mult, -1, *OD) + for i in range(mult): + state_dict[prefix + target.format(i=i) + suffix] = weight[i] + state_dict.pop(this_source) + + def _init_streaming_state(self, batch_size: int) -> MHAState: + in_proj = cast(nn.Linear, self.in_projs[0]) + device = cast(torch.device, in_proj.weight.device) + dtype = cast(torch.dtype, in_proj.weight.dtype) + + dim_per_head = self.embed_dim // self.num_heads + if self.context is None: + capacity = self.weights_per_step if self.weights_per_step else 1024 + else: + capacity = self.context + + kv_cache = RingKVCache( + batch_size, + self.num_heads, + dim_per_head, + capacity, + respect_exec_mask=not self.weights_per_step, + device=cast(torch.device, device), + dtype=cast(torch.dtype, dtype), + ) + return MHAState( + batch_size, + cast(torch.device, device), + kv_cache, + offset=torch.zeros(batch_size, device=cast(torch.device, device), dtype=torch.long), + offset_cpu=0, + ) + + def _complete_kv(self, k, v) -> KVCacheResult: + state = cast(MHAState | None, self._streaming_state) + if state is None: + return KVCacheResult.from_kv(k, v) + if state.kv_cache is None: + return KVCacheResult.from_kv(k, v) + return state.kv_cache.complete(k, v, state.exec_mask) + + def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor): + state = cast(MHAState | None, self._streaming_state) + B, T = query.shape[:2] + + if state is None: + offset = torch.zeros(B, device=query.device, dtype=torch.long) + offset_cpu = 0 + else: + offset = state.offset + offset_cpu = state.offset_cpu + + projected = apply_weights_per_step(self.in_projs, self.weights_per_step_schedule, query, offset_cpu) + dim_per_head = self.embed_dim // self.num_heads + projected = projected.reshape(B, T, 3, self.num_heads, dim_per_head).permute(2, 0, 3, 1, 4) + q, k, v = projected[0], projected[1], projected[2] + + if self.rope: + q, k = self.rope(q, k, offset, time_before_heads=False) + + k, v, pos_k = self._complete_kv(k, v) + pos_k = pos_k[:, None] + + if self.causal: + pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(-1, 1) + delta = pos_q - pos_k + attn_bias = (pos_k >= 0) & (delta >= 0) + if self.context is not None: + attn_bias = attn_bias & (delta < self.context) + attn_bias = attn_bias[:, None] + else: + attn_bias = None + + x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0) + x = x.transpose(1, 2).reshape(B, T, self.embed_dim) + x = apply_weights_per_step(self.out_projs, self.weights_per_step_schedule, x, offset_cpu) + + if state is not None: + state.offset[:] = torch.where(state.exec_mask, state.offset + T, state.offset) + state.offset_cpu += T + return x + + +# ============================================================================= +# Transformer Layer +# ============================================================================= + + +@dataclass +class LayerState(StreamingState): + offset_cpu: int = 0 + + def reset(self, reset_mask: torch.Tensor): + super().reset(reset_mask) + self.offset_cpu = 0 + + +class MossAudioTokenizerTransformerLayer(StreamingModule): + """Transformer layer with streaming support.""" + + def __init__( + self, + d_model: int, + num_heads: int, + dim_feedforward: int = 2048, + causal: bool = False, + context: int | None = None, + rope: MossAudioTokenizerRotaryEmbedding | None = None, + norm: str = "layer_norm", + layer_scale: float | None = None, + gating: str = "none", + weights_per_step: int = 0, + weights_per_step_schedule: list[int] | None = None, + activation=F.gelu, + device=None, + dtype=None, + ): + super().__init__() + factory_kwargs = {"device": device, "dtype": dtype} + + self.self_attn = MossAudioTokenizerMultiheadAttention( + embed_dim=d_model, + num_heads=num_heads, + causal=causal, + context=context, + rope=rope, + weights_per_step=weights_per_step, + weights_per_step_schedule=weights_per_step_schedule, + **factory_kwargs, + ) + self.norm1 = create_norm_fn(norm, d_model, **factory_kwargs) + self.norm2 = create_norm_fn(norm, d_model, **factory_kwargs) + + self.weights_per_step = weights_per_step + self.weights_per_step_schedule = weights_per_step_schedule + self.gating: nn.Module | nn.ModuleList | None = None + self.linear1: nn.Module | None = None + self.linear2: nn.Module | None = None + self.activation = activation + + num_weights = 1 + if weights_per_step: + num_weights = max(weights_per_step_schedule) + 1 if weights_per_step_schedule else weights_per_step + + if gating == "none": + self.linear1 = nn.Linear(d_model, dim_feedforward, bias=False, **factory_kwargs) + self.linear2 = nn.Linear(dim_feedforward, d_model, bias=False, **factory_kwargs) + else: + if weights_per_step: + dim_ff_list = [dim_feedforward] * num_weights if isinstance(dim_feedforward, int) else dim_feedforward + self.gating = nn.ModuleList( + [make_gating(gating, d_model, dim, **factory_kwargs) for dim in dim_ff_list] + ) + else: + self.gating = make_gating(gating, d_model, dim_feedforward, **factory_kwargs) + + if layer_scale is None: + self.layer_scale_1 = nn.Identity() + self.layer_scale_2 = nn.Identity() + else: + self.layer_scale_1 = MossAudioTokenizerLayerScale( + channels=d_model, init=layer_scale, channel_last=True, **cast(dict[str, object], factory_kwargs) + ) + self.layer_scale_2 = MossAudioTokenizerLayerScale( + channels=d_model, init=layer_scale, channel_last=True, **cast(dict[str, object], factory_kwargs) + ) + + def _init_streaming_state(self, batch_size: int) -> LayerState: + device = next(iter(self.parameters())).device + return LayerState(batch_size, device, offset_cpu=0) + + def _ff_block(self, x: torch.Tensor) -> torch.Tensor: + state = self._streaming_state + offset = state.offset_cpu if isinstance(state, LayerState) else 0 + + x_orig = x + x = self.norm2(x) + + if self.gating is None: + assert self.linear1 is not None + assert self.linear2 is not None + update = self.linear2(self.activation(self.linear1(x))) + else: + if self.weights_per_step: + assert isinstance(self.gating, nn.ModuleList) + update = apply_weights_per_step(self.gating, self.weights_per_step_schedule, x, offset) + else: + update = self.gating(x) + return x_orig.to(update) + self.layer_scale_2(update) + + def _sa_block(self, x: torch.Tensor): + x_orig = x + x = self.norm1(x) + update = self.self_attn(x, x, x) + return x_orig.to(update) + self.layer_scale_1(update) + + def forward(self, x: torch.Tensor): + x = self._sa_block(x) + x = self._ff_block(x) + state = self._streaming_state + if state is not None: + assert isinstance(state, LayerState) + state.offset_cpu += x.shape[1] + return x + + +# ============================================================================= +# Streaming Transformer +# ============================================================================= + + +@dataclass +class TransformerState(StreamingState): + offsets: torch.Tensor + + def reset(self, reset_mask: torch.Tensor): + super().reset(reset_mask) + self.offsets[:] = torch.where(reset_mask, torch.zeros_like(self.offsets), self.offsets) + + +class MossAudioTokenizerTransformer(StreamingModule): + """Transformer with streaming/causal support.""" + + def __init__( + self, + d_model: int, + num_heads: int, + num_layers: int, + dim_feedforward: int = 2048, + causal: bool = False, + context: int | None = None, + positional_embedding: str = "sin", + max_period: float = 10_000, + positional_scale: float = 1.0, + device=None, + dtype=None, + **kwargs, + ): + super().__init__() + if d_model % num_heads != 0: + raise ValueError(f"d_model must be divisible by num_heads, got d_model={d_model}, num_heads={num_heads}") + + self.positional_embedding = positional_embedding + self.max_period = max_period + self.positional_scale = positional_scale + + self.rope: MossAudioTokenizerRotaryEmbedding | None = None + if positional_embedding in {"rope", "sin_rope"}: + self.rope = MossAudioTokenizerRotaryEmbedding(max_period=max_period) + + self.layers = nn.ModuleList() + for _ in range(num_layers): + self.layers.append( + MossAudioTokenizerTransformerLayer( + d_model=d_model, + num_heads=num_heads, + dim_feedforward=dim_feedforward, + causal=causal, + context=context, + rope=self.rope, + device=device, + dtype=dtype, + **kwargs, + ) + ) + + def _init_streaming_state(self, batch_size: int) -> TransformerState: + device = next(self.parameters()).device + return TransformerState( + batch_size, + device, + offsets=torch.zeros(batch_size, device=device, dtype=torch.long), + ) + + def forward(self, x: torch.Tensor, *args, **kwargs): + B, T, C = x.shape + state = self._streaming_state + offsets = ( + torch.zeros(1, dtype=torch.long, device=x.device) + if state is None + else ( + state.offsets + if isinstance(state, TransformerState) + else torch.zeros(1, dtype=torch.long, device=x.device) + ) + ) + + if self.positional_embedding in {"sin", "sin_rope"}: + positions = torch.arange(T, device=x.device).view(1, -1, 1) + positions = positions + offsets.view(-1, 1, 1) + pos_emb = create_sin_embedding(positions, C, max_period=self.max_period, dtype=x.dtype) + x = x + self.positional_scale * pos_emb + + for layer in self.layers: + x = layer(x, *args, **kwargs) + + if state is not None: + assert isinstance(state, TransformerState) + state.offsets[:] = torch.where(state.exec_mask, state.offsets + T, state.offsets) + return x + + +class MossAudioTokenizerProjectedTransformer(StreamingContainer): + """Transformer with input/output projections.""" + + def __init__( + self, + input_dimension: int, + output_dimension: int, + d_model: int, + *, + conv_layout: bool = False, + module_type: str, + **kwargs, + ): + super().__init__() + self.module_type = module_type + self.downsample_ratio: int = 1 + self.input_dimension = input_dimension + self.output_dimension = output_dimension + + self.input_proj = nn.Linear(input_dimension, d_model, bias=False) + self.transformer = MossAudioTokenizerTransformer(d_model=d_model, **kwargs) + self.conv_layout = conv_layout + self.output_proj = nn.Linear(d_model, output_dimension, bias=False) + + def forward(self, x, input_lengths, *args, **kwargs): + x = self.input_proj(x.transpose(1, 2)) # (B, D, T) -> (B, T, D) + x = self.transformer(x, *args, **kwargs) + x = self.output_proj(x).transpose(1, 2) # (B, T, D) -> (B, D, T) + return x, input_lengths + + +# ============================================================================= +# Patched Pretransform Module +# ============================================================================= + + +class MossAudioTokenizerPatchedPretransform(nn.Module): + """Patching module for downsampling/upsampling.""" + + def __init__(self, patch_size: int, is_downsample: bool, module_type: str, **kwargs): + super().__init__() + self.patch_size = patch_size + self.downsample_ratio: int = patch_size + self.is_downsample = is_downsample + self.module_type = module_type + + def encode(self, x, input_lengths): + b, d, _ = x.shape + h = self.patch_size + x = x.reshape(b, d, -1, h).permute(0, 1, 3, 2).reshape(b, d * h, -1) + # We pad the input waveform to a multiple of `downsample_rate` before applying the encoder. + # Use a ceil division to match that padding and avoid dropping the last (partially padded) frame. + output_lengths = input_lengths // self.patch_size + return x, output_lengths + + def decode(self, x, input_lengths): + b, dh, length = x.shape + h = self.patch_size + d = dh // h + x = x.reshape(b, d, h, length).permute(0, 1, 3, 2).reshape(b, d, length * h) + output_lengths = input_lengths * self.patch_size + return x, output_lengths + + def forward(self, x, input_lengths): + if self.is_downsample: + return self.encode(x, input_lengths) + else: + return self.decode(x, input_lengths) + + +# ============================================================================= +# Vector Quantization +# ============================================================================= + + +def WNConv1d(*args, **kwargs): + """Weight-normalized Conv1d.""" + return nn.utils.parametrizations.weight_norm(nn.Conv1d(*args, **kwargs)) + + +class MossAudioTokenizerVectorQuantize(nn.Module): + """Single codebook vector quantization (inference only).""" + + def __init__( + self, + input_dim: int, + codebook_size: int, + codebook_dim: int, + **kwargs, + ): + super().__init__() + self.input_dim = input_dim + self.codebook_size = codebook_size + self.codebook_dim = codebook_dim + + if input_dim != codebook_dim: + self.in_proj = WNConv1d(input_dim, codebook_dim, kernel_size=1) + self.out_proj = WNConv1d(codebook_dim, input_dim, kernel_size=1) + else: + self.in_proj = nn.Identity() + self.out_proj = nn.Identity() + + self.codebook = nn.Embedding(codebook_size, codebook_dim) + + @torch.no_grad() + def forward(self, z: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Args: + z: Input tensor of shape (B, D, T) + Returns: + z_q: Quantized tensor of shape (B, D, T) + indices: Code indices of shape (B, T) + z_e: Encoded tensor before quantization + """ + z = z.float() + z_e = self.in_proj(z).float() + + encodings = z_e.transpose(1, 2).reshape(-1, z_e.shape[1]) + + codebook_weight = self.codebook.weight + dist = ( + encodings.pow(2).sum(1, keepdim=True) + - 2 * encodings @ codebook_weight.float().t() + + codebook_weight.float().pow(2).sum(1, keepdim=True).t() + ) + + indices = (-dist).max(1)[1] + indices = indices.reshape(z.size(0), -1) + + z_q = self.decode_code(indices) + z_q = self.out_proj(z_q).float() + + return z_q, indices, z_e + + def decode_code(self, embed_id: torch.Tensor) -> torch.Tensor: + """Decode code indices to embeddings.""" + return self.codebook(embed_id).transpose(1, 2).float() + + +class MossAudioTokenizerLFQ(nn.Module): + """LFQ (inference-only) used by ResidualLFQ.""" + + def __init__( + self, + input_dim: int, + codebook_size: int, + codebook_dim: int, + **kwargs, + ): + super().__init__() + self.input_dim = input_dim + self.codebook_size = codebook_size + self.codebook_dim = codebook_dim + + if self.input_dim != self.codebook_dim: + self.in_proj = WNConv1d(self.input_dim, self.codebook_dim, kernel_size=1) + self.out_proj = WNConv1d(self.codebook_dim, self.input_dim, kernel_size=1) + else: + self.in_proj = nn.Identity() + self.out_proj = nn.Identity() + + self.codebook = nn.Embedding(codebook_size, codebook_dim) + + @torch.no_grad() + def forward(self, z: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize z into codebook vectors.""" + z = z.float() + z_e = self.in_proj(z).float() + z_q, indices = self.decode_latents(z_e) + z_q = (z_e + (z_q - z_e).detach()).float() + z_q = self.out_proj(z_q).float() + return z_q, indices, z_e + + def embed_code(self, embed_id: torch.Tensor) -> torch.Tensor: + return F.embedding(embed_id, self.codebook.weight) + + def decode_code_wo_out_proj(self, embed_id: torch.Tensor) -> torch.Tensor: + return self.embed_code(embed_id).transpose(1, 2) + + def decode_code(self, embed_id: torch.Tensor) -> torch.Tensor: + z_q = self.decode_code_wo_out_proj(embed_id).float() + z_q = self.out_proj(z_q).float() + return z_q + + def decode_latents(self, latents: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Match training LFQ: L2-normalize then argmin squared distance.""" + encodings = latents.transpose(1, 2).reshape(-1, latents.shape[1]).float() + codebook = self.codebook.weight.float() + + encodings = F.normalize(encodings) + codebook = F.normalize(codebook) + + dist = ( + encodings.pow(2).sum(1, keepdim=True) + - 2 * encodings @ codebook.t() + + codebook.pow(2).sum(1, keepdim=True).t() + ) + indices = (-dist).max(1)[1] + indices = indices.reshape(latents.size(0), -1) + z_q = self.decode_code_wo_out_proj(indices).float() + return z_q, indices + + +class MossAudioTokenizerResidualVQ(nn.Module): + """Residual Vector Quantization (inference only).""" + + def __init__( + self, + input_dim: int = 1024, + rvq_dim: int | None = None, + output_dim: int | None = None, + num_quantizers: int = 32, + codebook_size: int = 1024, + codebook_dim: int = 8, + **kwargs, + ): + super().__init__() + self.input_dim = input_dim + self.rvq_dim = rvq_dim or input_dim + self.output_dim = output_dim or input_dim + self.num_quantizers = num_quantizers + self.codebook_size = codebook_size + self.codebook_dim = codebook_dim + + self.input_proj = ( + WNConv1d(input_dim, self.rvq_dim, kernel_size=1) if input_dim != self.rvq_dim else nn.Identity() + ) + self.output_proj = ( + WNConv1d(self.rvq_dim, self.output_dim, kernel_size=1) if self.rvq_dim != self.output_dim else nn.Identity() + ) + + self.quantizers = nn.ModuleList( + [ + MossAudioTokenizerVectorQuantize( + input_dim=self.rvq_dim, + codebook_size=codebook_size, + codebook_dim=codebook_dim, + **kwargs, + ) + for _ in range(num_quantizers) + ] + ) + + @torch.no_grad() + def forward( + self, + z: torch.Tensor, + input_length: torch.Tensor, + n_quantizers: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Args: + z: Input tensor of shape (B, D, T) + input_length: Valid lengths for each sample (B,) + n_quantizers: Number of quantizers to use + Returns: + quantized_out: Quantized output (B, D, T) + all_indices: All code indices (N, B, T) + output_length: Output lengths (B,) + """ + z = self.input_proj(z) + + batch_size, _, max_time = z.shape + mask = torch.arange(max_time, device=z.device).expand(batch_size, max_time) < input_length.unsqueeze(1) + + quantized_out = torch.zeros_like(z, dtype=torch.float32) + residual = z.clone().float() + all_indices = [] + + n_quantizers = n_quantizers or self.num_quantizers + + for i, quantizer in enumerate(self.quantizers): + if i >= n_quantizers: + break + + masked_residual = residual * mask.unsqueeze(1) + z_q_i, indices_i, _ = quantizer(masked_residual) + + update_mask = mask.unsqueeze(1) + quantized_out = quantized_out + z_q_i * update_mask + residual = residual - z_q_i * update_mask + all_indices.append(indices_i) + + all_indices = torch.stack(all_indices) # (N, B, T) + quantized_out = self.output_proj(quantized_out) + + return quantized_out, all_indices, input_length + + def decode_codes(self, codes: torch.Tensor) -> torch.Tensor: + """Decode codes from multiple quantizers to embeddings.""" + nq, B, T = codes.shape + emb = torch.zeros(B, self.rvq_dim, T, device=codes.device, dtype=torch.float32) + + for i, quantizer in enumerate(self.quantizers[:nq]): + quantizer = cast(MossAudioTokenizerVectorQuantize, quantizer) + quantized_i = quantizer.decode_code(codes[i]) + emb += quantized_i + + emb = self.output_proj(emb) + return emb + + +class MossAudioTokenizerResidualLFQ(nn.Module): + """Residual LFQ (inference only).""" + + def __init__( + self, + input_dim: int = 1024, + rvq_dim: int | None = None, + output_dim: int | None = None, + num_quantizers: int = 32, + codebook_size: int = 1024, + codebook_dim: int = 8, + **kwargs, + ): + super().__init__() + self.input_dim = input_dim + self.rvq_dim = rvq_dim or input_dim + self.output_dim = output_dim or input_dim + self.num_quantizers = num_quantizers + self.codebook_size = codebook_size + self.codebook_dim = codebook_dim + + self.input_proj = ( + WNConv1d(input_dim, self.rvq_dim, kernel_size=1) if input_dim != self.rvq_dim else nn.Identity() + ) + self.output_proj = ( + WNConv1d(self.rvq_dim, self.output_dim, kernel_size=1) if self.rvq_dim != self.output_dim else nn.Identity() + ) + + self.quantizers = nn.ModuleList( + [ + MossAudioTokenizerLFQ( + input_dim=self.rvq_dim, + codebook_size=codebook_size, + codebook_dim=codebook_dim, + **kwargs, + ) + for _ in range(num_quantizers) + ] + ) + + @torch.no_grad() + def forward( + self, + z: torch.Tensor, + input_length: torch.Tensor, + n_quantizers: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Inference quantization.""" + z = self.input_proj(z).float() + + batch_size, _, max_time = z.shape + mask = torch.arange(max_time, device=z.device).expand(batch_size, max_time) < input_length.unsqueeze(1) + + quantized_out = torch.zeros_like(z, dtype=torch.float32) + residual = z.clone().float() + all_indices = [] + + n_quantizers = n_quantizers or self.num_quantizers + for i, quantizer in enumerate(self.quantizers): + if i >= n_quantizers: + break + + masked_residual = residual * mask.unsqueeze(1) + z_q_i, indices_i, _ = quantizer(masked_residual) + + update_mask = mask.unsqueeze(1) + quantized_out = quantized_out + z_q_i * update_mask + residual = residual - z_q_i * update_mask + all_indices.append(indices_i) + + all_indices = ( + torch.stack(all_indices) + if all_indices + else torch.empty(0, batch_size, max_time, device=z.device, dtype=torch.long) + ) + quantized_out = self.output_proj(quantized_out) + return quantized_out, all_indices, input_length + + def decode_codes(self, codes: torch.Tensor) -> torch.Tensor: + nq, B, T = codes.shape + emb = torch.zeros(B, self.rvq_dim, T, device=codes.device, dtype=torch.float32) + for i, quantizer in enumerate(self.quantizers[:nq]): + quantizer = cast(MossAudioTokenizerLFQ, quantizer) + emb += quantizer.decode_code(codes[i]).float() + emb = self.output_proj(emb) + return emb + + +# ============================================================================= +# Main Model Classes +# ============================================================================= + + +@auto_docstring +class MossAudioTokenizerPreTrainedModel(PreTrainedAudioTokenizerBase): + """Base class for MossAudioTokenizer models.""" + + config_class = MossAudioTokenizerConfig + base_model_prefix = "" + main_input_name = "input_values" + input_modalities = "audio" + supports_gradient_checkpointing = False + _no_split_modules = [ + "MossAudioTokenizerTransformerLayer", + "MossAudioTokenizerResidualVQ", + "MossAudioTokenizerResidualLFQ", + ] + + +@auto_docstring( + custom_intro=""" + The MossAudioTokenizer neural audio codec model for audio tokenization and synthesis. + """ +) +class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): + """ + MossAudioTokenizer model for audio tokenization and synthesis. + + This model can encode audio waveforms into discrete tokens and decode + tokens back into audio waveforms. + """ + + def __init__(self, config: MossAudioTokenizerConfig): + super().__init__(config) + + self.config = config + _ = config.version + self.sampling_rate = config.sampling_rate + self.downsample_rate = config.downsample_rate + self.number_channels = getattr(config, "number_channels", getattr(config, "num_channels", 1)) + self.enable_channel_interleave = getattr(config, "enable_channel_interleave", True) + self.causal_transformer_context_duration = config.causal_transformer_context_duration + + # Build encoder + channel_interleave_factor = ( + self.number_channels if self.enable_channel_interleave and self.number_channels > 1 else 1 + ) + current_frame_rate: float = float(self.sampling_rate * channel_interleave_factor) + self.encoder = nn.ModuleList() + + for encoder_kwargs_i in config.encoder_kwargs: + encoder_kwargs_i = dict(encoder_kwargs_i) # Make a copy + if encoder_kwargs_i["module_type"] == "PatchedPretransform": + self.encoder.append(MossAudioTokenizerPatchedPretransform(**encoder_kwargs_i, is_downsample=True)) + elif encoder_kwargs_i["module_type"] == "Transformer": + context_duration = float( + encoder_kwargs_i.pop("context_duration", self.causal_transformer_context_duration) + ) + self.encoder.append( + MossAudioTokenizerProjectedTransformer( + **encoder_kwargs_i, + context=int(round(current_frame_rate * context_duration)), + ) + ) + current_frame_rate /= self.encoder[-1].downsample_ratio + + # Build quantizer + quantizer_kwargs = dict(config.quantizer_kwargs) + quantizer_type = quantizer_kwargs.get("quantizer_type", getattr(config, "quantizer_type", "rvq")) + if quantizer_type in {"rvq", "spec_rvq"}: + self.quantizer = MossAudioTokenizerResidualVQ(**quantizer_kwargs) + elif quantizer_type in {"rlfq", "random_prefix_rlfq"}: + self.quantizer = MossAudioTokenizerResidualLFQ(**quantizer_kwargs) + else: + raise ValueError(f"Unsupported quantizer_type: {quantizer_type}") + + # Build decoder + decoder_kwargs_list = copy.deepcopy(config.decoder_kwargs) + self.decoder = nn.ModuleList() + + for decoder_kwargs_i in decoder_kwargs_list: + decoder_kwargs_i = dict(decoder_kwargs_i) + if decoder_kwargs_i["module_type"] == "PatchedPretransform": + self.decoder.append(MossAudioTokenizerPatchedPretransform(**decoder_kwargs_i, is_downsample=False)) + elif decoder_kwargs_i["module_type"] == "Transformer": + context_duration = float( + decoder_kwargs_i.pop("context_duration", self.causal_transformer_context_duration) + ) + self.decoder.append( + MossAudioTokenizerProjectedTransformer( + **decoder_kwargs_i, + context=int(round(current_frame_rate * context_duration)), + ) + ) + current_frame_rate *= self.decoder[-1].downsample_ratio + + expected_output_frame_rate = float(self.sampling_rate * channel_interleave_factor) + if int(round(current_frame_rate)) != int(round(expected_output_frame_rate)): + raise ValueError( + "Decoder stack does not invert the encoder frame rate correctly: " + f"got current_frame_rate={current_frame_rate}, expected={expected_output_frame_rate}." + ) + + self.post_init() + + def _start_streaming(self, batch_size: int): + """Start streaming mode for all modules.""" + + def _start(module): + if isinstance(module, StreamingModule): + module._streaming_state = module._init_streaming_state(batch_size) + + self.apply(_start) + + def _stop_streaming(self): + """Stop streaming mode for all modules.""" + + def _stop(module): + if isinstance(module, StreamingModule): + module._streaming_state = None + + self.apply(_stop) + + def _set_streaming_exec_mask(self, exec_mask: torch.Tensor) -> None: + """Set the active slot mask on every live streaming module.""" + + def _set(module): + if isinstance(module, StreamingModule) and module._streaming_state is not None: + module._streaming_state.set_exec_mask(exec_mask.to(module._streaming_state.device)) + + self.apply(_set) + + @contextmanager + def streaming(self, batch_size: int = 1): + """Context manager for streaming mode.""" + self._start_streaming(batch_size) + try: + yield + finally: + self._stop_streaming() + + def _prepare_waveform_batch(self, wav_list: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + if len(wav_list) == 0: + raise ValueError("`wav_list` must contain at least one waveform.") + + device = wav_list[0].device + dtype = wav_list[0].dtype + batch_size = len(wav_list) + input_lengths = torch.zeros(batch_size, device=device, dtype=torch.long) + + normalized_wavs: list[torch.Tensor] = [] + for i, wav in enumerate(wav_list): + if self.number_channels == 1: + if wav.dim() == 1: + wav_i = wav.unsqueeze(0) + elif wav.dim() == 2 and wav.shape[0] == 1: + wav_i = wav + else: + raise ValueError( + f"Expected wav_list[{i}] to have shape `(T,)` or `(1, T)` for a mono model, " + f"got {tuple(wav.shape)}." + ) + else: + if wav.dim() != 2 or wav.shape[0] != self.number_channels: + raise ValueError( + f"Expected wav_list[{i}] to have shape `({self.number_channels}, T)`, got {tuple(wav.shape)}." + ) + wav_i = wav + + normalized_wavs.append(wav_i) + input_lengths[i] = wav_i.shape[-1] + + max_length = int(input_lengths.max().item()) if batch_size > 0 else 0 + input_values = torch.zeros( + batch_size, + self.number_channels, + max_length, + device=device, + dtype=dtype, + ) + for i, wav_i in enumerate(normalized_wavs): + input_values[i, :, : wav_i.shape[-1]] = wav_i + return input_values, input_lengths + + def _prepare_codes_batch( + self, + codes_list: list[torch.Tensor], + num_quantizers: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, int]: + if len(codes_list) == 0: + raise ValueError("`codes_list` must contain at least one code tensor.") + + batch_size = len(codes_list) + device = codes_list[0].device + nqs = [codes.shape[0] for codes in codes_list] + if num_quantizers is None: + num_quantizers = nqs[0] + if any(nq != num_quantizers for nq in nqs): + raise ValueError( + "All elements in `codes_list` must have the same number of quantizers when " + "`num_quantizers` is None. Pass `num_quantizers=...` to decode a common prefix." + ) + else: + min_nq = min(nqs) + if min_nq < num_quantizers: + raise ValueError( + "`num_quantizers` must be <= the number of quantizers for every element in `codes_list`. " + f"Got num_quantizers={num_quantizers}, min(codes.shape[0])={min_nq}." + ) + + max_length = max(codes.shape[-1] for codes in codes_list) + audio_codes = torch.zeros(num_quantizers, batch_size, max_length, device=device, dtype=torch.long) + audio_codes_lengths = torch.zeros(batch_size, device=device, dtype=torch.long) + + for i, codes in enumerate(codes_list): + codes_i = codes[:num_quantizers] + audio_codes[:, i, : codes_i.shape[-1]] = codes_i + audio_codes_lengths[i] = codes_i.shape[-1] + return audio_codes, audio_codes_lengths, num_quantizers + + def _flatten_channels_for_codec( + self, + input_values: torch.Tensor, + input_lengths: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if input_values.dim() != 3: + raise ValueError(f"Expected `input_values` with shape `(B, C, T)`, got {tuple(input_values.shape)}.") + if input_values.shape[1] != self.number_channels: + raise ValueError( + f"Expected `input_values.shape[1] == {self.number_channels}`, got {input_values.shape[1]}." + ) + + if input_values.shape[-1] % self.downsample_rate != 0: + pad_length = self.downsample_rate - (input_values.shape[-1] % self.downsample_rate) + input_values = F.pad(input_values, (0, pad_length)) + + if self.number_channels > 1 and self.enable_channel_interleave: + input_values = input_values.transpose(1, 2).contiguous().view(input_values.shape[0], 1, -1) + input_lengths = input_lengths * self.number_channels + return input_values, input_lengths + + def _restore_channels_from_codec( + self, + output_values: torch.Tensor, + output_lengths: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.number_channels == 1 or not self.enable_channel_interleave: + return output_values.float(), output_lengths + + output_values = ( + output_values.squeeze(1) + .contiguous() + .view(output_values.shape[0], -1, self.number_channels) + .transpose(1, 2) + .contiguous() + .float() + ) + output_lengths = torch.div(output_lengths, self.number_channels, rounding_mode="floor") + return output_values, output_lengths + + @torch.no_grad() + def batch_encode( + self, wav_list: list[torch.Tensor], num_quantizers: int | None = None + ) -> MossAudioTokenizerEncoderOutput: + """Batch encode a list of audio waveforms. + + Args: + wav_list: List of audio tensors, each of shape `(num_samples,)`. + num_quantizers: Number of quantizers to use. By default, all quantizers are used. + + Returns: + [`MossAudioTokenizerEncoderOutput`] with `audio_codes` and `audio_codes_lengths`. + """ + input_values, input_lengths = self._prepare_waveform_batch(wav_list) + return self._encode_frame(input_values, input_lengths, n_quantizers=num_quantizers) + + @torch.no_grad() + def batch_decode( + self, codes_list: list[torch.Tensor], num_quantizers: int | None = None + ) -> MossAudioTokenizerDecoderOutput: + """Batch decode a list of audio codes. + + Args: + codes_list: List of audio code tensors, each of shape `(num_quantizers, codes_length)`. + num_quantizers: If provided, decode only the first `num_quantizers` quantizers from each element in + `codes_list`. If omitted, all elements in `codes_list` must have the same number of quantizers. + + Returns: + [`MossAudioTokenizerDecoderOutput`] with `audio` and `audio_lengths`. + """ + audio_codes, audio_codes_lengths, _ = self._prepare_codes_batch( + codes_list, + num_quantizers=num_quantizers, + ) + return self._decode_frame(audio_codes, audio_codes_lengths) + + @torch.no_grad() + def _encode_frame( + self, + input_values: torch.Tensor, + input_lengths: torch.Tensor | None = None, + n_quantizers: int | None = None, + ) -> MossAudioTokenizerEncoderOutput: + """Tokenize audio waveform into discrete tokens.""" + # Handle input shape + if input_values.dim() == 1: + input_values = input_values.view(1, 1, -1) + elif input_values.dim() == 2: + if self.number_channels == 1: + input_values = input_values.unsqueeze(1) + else: + input_values = input_values.unsqueeze(0) + + B, _, T = input_values.shape + device = input_values.device + + if input_lengths is None: + input_lengths = torch.full((B,), T, device=device, dtype=torch.long) + + input_values, input_lengths = self._flatten_channels_for_codec(input_values, input_lengths) + + # Encode + e, e_lengths = input_values, input_lengths + for encoder_module in self.encoder: + e, e_lengths = encoder_module(e, e_lengths) + + # Quantize + quantizer = cast(MossAudioTokenizerResidualVQ | MossAudioTokenizerResidualLFQ, self.quantizer) + zq, audio_codes, audio_codes_lengths = quantizer(e, e_lengths, n_quantizers) + + return MossAudioTokenizerEncoderOutput( + audio_codes=audio_codes, audio_codes_lengths=audio_codes_lengths, encoder_hidden_states=e + ) + + @torch.no_grad() + def _decode_frame( + self, + codes: torch.Tensor, + codes_lengths: torch.Tensor | None = None, + ) -> MossAudioTokenizerDecoderOutput: + """Detokenize discrete tokens into audio waveform.""" + nq, B, T = codes.shape + device = codes.device + + if codes_lengths is None: + codes_lengths = torch.full((B,), T, device=device, dtype=torch.long) + + # Decode from codes + quantizer = cast(MossAudioTokenizerResidualVQ | MossAudioTokenizerResidualLFQ, self.quantizer) + zq = quantizer.decode_codes(codes) + + d, d_lengths = zq, codes_lengths + for decoder_module in self.decoder: + d, d_lengths = decoder_module(d, d_lengths) + + d, d_lengths = self._restore_channels_from_codec(d, d_lengths) + return MossAudioTokenizerDecoderOutput(audio=d, audio_lengths=d_lengths) + + def encode( # type: ignore[override] + self, + input_values: torch.Tensor, + padding_mask: torch.Tensor | None = None, + num_quantizers: int | None = None, + return_dict: bool | None = None, + chunk_duration: float | None = None, + ) -> MossAudioTokenizerEncoderOutput | tuple[torch.Tensor, torch.Tensor]: + """ + Encodes the input audio waveform into discrete codes. + + Args: + input_values (`torch.Tensor` of shape `(batch_size, channels, sequence_length)`): + Float values of the input audio waveform. + padding_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to indicate valid audio samples. + num_quantizers (`int`, *optional*): + Number of quantizers to use. By default, all quantizers are used. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + chunk_duration (`float`, *optional*): + If provided, encode the input waveform in successive chunks of `chunk_duration` seconds while keeping a + streaming KV cache for the causal transformers. + + `chunk_duration` must be <= `config.causal_transformer_context_duration`, and + `chunk_duration * config.sampling_rate` must be divisible by `config.downsample_rate`. + + Returns: + `MossAudioTokenizerEncoderOutput` or tuple containing audio codes and lengths. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + # Handle input shape + if input_values.dim() == 1: + input_values = input_values.view(1, 1, -1) + elif input_values.dim() == 2: + if self.number_channels == 1: + input_values = input_values.unsqueeze(1) + else: + input_values = input_values.unsqueeze(0) + + B, _, T = input_values.shape + device = input_values.device + + if padding_mask is not None: + input_lengths = padding_mask.sum(dim=-1).long() + else: + input_lengths = torch.full((B,), T, device=device, dtype=torch.long) + + if chunk_duration is None: + encoder_output = self._encode_frame(input_values, input_lengths, num_quantizers) + else: + if chunk_duration <= 0: + raise ValueError("`chunk_duration` must be > 0 when provided.") + if chunk_duration > self.causal_transformer_context_duration: + raise ValueError( + "`chunk_duration` must be <= `config.causal_transformer_context_duration` " + f"({self.causal_transformer_context_duration}), got {chunk_duration}." + ) + if B != 1: + raise ValueError("Streaming encode via `chunk_duration` currently only supports batch_size=1.") + + chunk_length = int(round(chunk_duration * self.sampling_rate)) + if chunk_length <= 0: + raise ValueError("`chunk_duration` is too small and results in chunk_length <= 0.") + if chunk_length % self.downsample_rate != 0: + raise ValueError( + "`chunk_duration * config.sampling_rate` must be divisible by `config.downsample_rate`. " + f"Got chunk_length={chunk_length}, downsample_rate={self.downsample_rate}." + ) + + input_length = int(input_lengths[0].item()) + if input_length <= chunk_length: + encoder_output = self._encode_frame(input_values[..., :input_length], input_lengths, num_quantizers) + else: + codes_chunks: list[torch.Tensor] = [] + hidden_chunks: list[torch.Tensor] = [] + + with ExitStack() as exit_stack: + for encoder_module in self.encoder: + if isinstance(encoder_module, StreamingModule): + exit_stack.enter_context(encoder_module.streaming(batch_size=B)) + + for start_idx in range(0, input_length, chunk_length): + input_length_i = min(chunk_length, input_length - start_idx) + if input_length_i <= 0: + break + + input_lengths_i = torch.tensor([input_length_i], device=device, dtype=torch.long) + input_values_i = input_values[..., start_idx : start_idx + input_length_i] + result_i = self._encode_frame(input_values_i, input_lengths_i, num_quantizers) + + if result_i.audio_codes is None or result_i.audio_codes_lengths is None: + raise RuntimeError("Internal error: `_encode_frame` returned empty audio codes.") + if result_i.encoder_hidden_states is None: + raise RuntimeError("Internal error: `_encode_frame` returned empty encoder hidden states.") + + codes_length_i = result_i.audio_codes_lengths + codes_chunks.append(result_i.audio_codes[:, :, : codes_length_i[0]]) + hidden_chunks.append(result_i.encoder_hidden_states[:, :, : codes_length_i[0]]) + + audio_codes = torch.cat(codes_chunks, dim=-1) + encoder_hidden_states = torch.cat(hidden_chunks, dim=-1) + audio_codes_lengths = torch.tensor([audio_codes.shape[-1]], device=device, dtype=torch.long) + encoder_output = MossAudioTokenizerEncoderOutput( + audio_codes=audio_codes, + audio_codes_lengths=audio_codes_lengths, + encoder_hidden_states=encoder_hidden_states, + ) + + if not return_dict: + assert encoder_output.audio_codes is not None + assert encoder_output.audio_codes_lengths is not None + return ( + cast(torch.Tensor, encoder_output.audio_codes), + cast(torch.Tensor, encoder_output.audio_codes_lengths), + ) + return encoder_output + + def decode( # type: ignore[override] + self, + audio_codes: torch.Tensor, + padding_mask: torch.Tensor | None = None, + return_dict: bool | None = None, + chunk_duration: float | None = None, + num_quantizers: int | None = None, + ) -> MossAudioTokenizerDecoderOutput | tuple[torch.Tensor]: + """ + Decodes the given codes into an output audio waveform. + + Args: + audio_codes (`torch.LongTensor` of shape `(num_quantizers, batch_size, sequence_length)`): + Discrete code embeddings computed using `model.encode`. + padding_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to indicate valid code positions. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + chunk_duration (`float`, *optional*): + If provided, decode the input codes in successive chunks of `chunk_duration` seconds while keeping a + streaming KV cache for the causal transformers. + + num_quantizers (`int`, *optional*): + Number of quantizers to use. By default, all quantizers in `audio_codes` are used. + + `chunk_duration` must be <= `config.causal_transformer_context_duration`, and + `chunk_duration * config.sampling_rate` must be divisible by `config.downsample_rate`. + + Returns: + `MossAudioTokenizerDecoderOutput` or tuple containing decoded audio. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if audio_codes.dim() == 2: + audio_codes = audio_codes.unsqueeze(1) # nq, T -> nq, B=1, T + + if num_quantizers is not None: + if num_quantizers > audio_codes.shape[0]: + raise ValueError( + f"`num_quantizers` ({num_quantizers}) must be <= audio_codes.shape[0] ({audio_codes.shape[0]})." + ) + audio_codes = audio_codes[:num_quantizers] + + _, B, T = audio_codes.shape + device = audio_codes.device + + if padding_mask is not None: + codes_lengths = padding_mask.sum(dim=-1).long() + else: + codes_lengths = torch.full((B,), T, device=device, dtype=torch.long) + + if chunk_duration is None: + decoder_output = self._decode_frame(audio_codes, codes_lengths) + else: + if chunk_duration <= 0: + raise ValueError("`chunk_duration` must be > 0 when provided.") + if chunk_duration > self.causal_transformer_context_duration: + raise ValueError( + "`chunk_duration` must be <= `config.causal_transformer_context_duration` " + f"({self.causal_transformer_context_duration}), got {chunk_duration}." + ) + if B != 1: + raise ValueError("Streaming decode via `chunk_duration` currently only supports batch_size=1.") + + chunk_length = int(round(chunk_duration * self.sampling_rate)) + if chunk_length <= 0: + raise ValueError("`chunk_duration` is too small and results in chunk_length <= 0.") + if chunk_length % self.downsample_rate != 0: + raise ValueError( + "`chunk_duration * config.sampling_rate` must be divisible by `config.downsample_rate`. " + f"Got chunk_length={chunk_length}, downsample_rate={self.downsample_rate}." + ) + + chunk_frame_length = chunk_length // self.downsample_rate + codes_length = int(codes_lengths[0].item()) + if codes_length <= chunk_frame_length: + decoder_output = self._decode_frame(audio_codes[..., :codes_length], codes_lengths) + else: + wav_chunks: list[torch.Tensor] = [] + with ExitStack() as exit_stack: + for decoder_module in self.decoder: + if isinstance(decoder_module, StreamingModule): + exit_stack.enter_context(decoder_module.streaming(batch_size=B)) + + for start_idx in range(0, codes_length, chunk_frame_length): + codes_length_i = min(chunk_frame_length, codes_length - start_idx) + if codes_length_i <= 0: + break + + codes_lengths_i = torch.tensor([codes_length_i], device=device, dtype=torch.long) + codes_i = audio_codes[:, :, start_idx : start_idx + codes_length_i] + result_i = self._decode_frame(codes_i, codes_lengths_i) + + if result_i.audio is None or result_i.audio_lengths is None: + raise RuntimeError("Internal error: `_decode_frame` returned empty audio.") + + wav_chunks.append(result_i.audio[:, :, : result_i.audio_lengths[0]]) + + wav = torch.cat(wav_chunks, dim=-1) + audio_lengths = torch.tensor([wav.shape[-1]], device=device, dtype=torch.long) + decoder_output = MossAudioTokenizerDecoderOutput(audio=wav, audio_lengths=audio_lengths) + + if not return_dict: + assert decoder_output.audio is not None + return (cast(torch.Tensor, decoder_output.audio),) + return decoder_output + + @auto_docstring + def forward( + self, + input_values: torch.FloatTensor | None = None, + padding_mask: torch.BoolTensor | None = None, + audio_codes: torch.Tensor | None = None, + num_quantizers: int | None = None, + return_dict: bool | None = None, + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None] | MossAudioTokenizerOutput: # type: ignore[override] + r""" + input_values (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`, *optional*): + Raw audio input converted to Float. + padding_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid computing on padding token indices. Mask values selected in `[0, 1]`: + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + audio_codes (`torch.LongTensor` of shape `(num_quantizers, batch_size, sequence_length)`, *optional*): + Discrete code embeddings computed using `model.encode`. + num_quantizers (`int`, *optional*): + Number of quantizers (codebooks) to use. By default, all quantizers are used. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + + Examples: + + ```python + >>> import torch + >>> from transformers import MossAudioTokenizerModel + + >>> model = MossAudioTokenizerModel.from_pretrained("moss_audio_tokenizer-model") + + >>> # Create dummy audio input + >>> audio = torch.randn(1, 1, 24000) # 1 second of audio at 24kHz + + >>> outputs = model(input_values=audio) + >>> audio_codes = outputs.audio_codes + >>> audio_values = outputs.audio + ``` + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + output_audio_codes: torch.Tensor | None = None + output_audio_codes_lengths: torch.Tensor | None = None + output_audio: torch.Tensor | None = None + output_audio_lengths: torch.Tensor | None = None + decoded_from_encoded_codes = False + + # Encode if input_values provided + if input_values is not None: + encoder_output = self.encode(input_values, padding_mask, num_quantizers, return_dict=True) + encoder_output = cast(MossAudioTokenizerEncoderOutput, encoder_output) + output_audio_codes = encoder_output.audio_codes + output_audio_codes_lengths = encoder_output.audio_codes_lengths + + # If codes not provided separately, use encoded codes for decoding + if audio_codes is None: + audio_codes = output_audio_codes + decoded_from_encoded_codes = True + + # Decode if codes available + if audio_codes is not None: + # If we're decoding the codes we just produced, use the computed lengths so we don't decode padded garbage. + if decoded_from_encoded_codes and output_audio_codes_lengths is not None: + decoder_output = self._decode_frame(audio_codes, output_audio_codes_lengths) + else: + decoder_output = self.decode( + audio_codes, + padding_mask=padding_mask, + return_dict=True, + num_quantizers=num_quantizers, + ) + decoder_output = cast(MossAudioTokenizerDecoderOutput, decoder_output) + output_audio = decoder_output.audio + output_audio_lengths = decoder_output.audio_lengths + + if not return_dict: + return (output_audio_codes, output_audio, output_audio_lengths) + + return MossAudioTokenizerOutput( + audio=output_audio, + audio_lengths=output_audio_lengths, + audio_codes=output_audio_codes, + audio_codes_lengths=output_audio_codes_lengths, + ) + + +__all__ = ["MossAudioTokenizerModel", "MossAudioTokenizerPreTrainedModel"] diff --git a/vllm_omni/model_executor/models/moss_tts/configuration_moss_audio_tokenizer_v2.py b/vllm_omni/model_executor/models/moss_tts/configuration_moss_audio_tokenizer_v2.py new file mode 100644 index 000000000..c518742a2 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/configuration_moss_audio_tokenizer_v2.py @@ -0,0 +1,359 @@ +# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MossAudioTokenizer model configuration""" + +from typing import Any + +from transformers.configuration_utils import PreTrainedConfig +from transformers.utils import logging + +logger = logging.get_logger(__name__) + + +class MossAudioTokenizerConfig(PreTrainedConfig): + r""" + This is the configuration class to store the configuration of a [`MossAudioTokenizerModel`]. It is used to + instantiate a MossAudioTokenizer model according to the specified arguments, defining the model architecture. + + Instantiating a configuration with the defaults will yield a similar configuration to that of the + [VoiceAgentGroup/moss_audio_tokenizer](https://huggingface.co/VoiceAgentGroup/moss_audio_tokenizer) architecture. + + Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PreTrainedConfig`] for more information. + + Args: + sampling_rate (`int`, *optional*, defaults to 24000): + The sampling rate at which the audio waveform should be digitalized expressed in hertz (Hz). + downsample_rate (`int`, *optional*, defaults to 1920): + Total downsampling rate from waveform to tokens. + causal_transformer_context_duration (`float`, *optional*, defaults to 10.0): + Context duration in seconds for causal transformer. + encoder_kwargs (`list[dict]`, *optional*): + List of encoder module configurations. Each dict specifies a module type and its parameters. + decoder_kwargs (`list[dict]`, *optional*): + List of decoder module configurations in execution order. + quantizer_type (`str`, *optional*, defaults to `"rvq"`): + Quantizer type. Options include `"rvq"`, `"spec_rvq"`, `"rlfq"`, `"random_prefix_rlfq"`. + quantizer_kwargs (`dict`, *optional*): + Configuration for the quantizer including `input_dim`, `rvq_dim`, `output_dim`, `num_quantizers`, + `codebook_size`, and `codebook_dim`. + + Example: + + ```python + >>> from transformers import MossAudioTokenizerModel, MossAudioTokenizerConfig + + >>> # Initializing a MossAudioTokenizer style configuration + >>> configuration = MossAudioTokenizerConfig() + + >>> # Initializing a model (with random weights) from the configuration + >>> model = MossAudioTokenizerModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "moss-audio-tokenizer" + + # Backward-compatible alias used by some checkpoints. + attribute_map = {"sample_rate": "sampling_rate"} + + sampling_rate: int + downsample_rate: int + causal_transformer_context_duration: float + encoder_kwargs: list[dict[str, Any]] + decoder_kwargs: list[dict[str, Any]] + number_channels: int + enable_channel_interleave: bool + attention_implementation: str + compute_dtype: str + codec_weight_dtype: str + quantizer_type: str + quantizer_kwargs: dict[str, Any] + + def __init__( + self, + version: str | None = None, + sampling_rate: int = 48000, + downsample_rate: int = 3840, + causal_transformer_context_duration: float = 10.0, + encoder_kwargs: list[dict[str, Any]] | None = None, + decoder_kwargs: list[dict[str, Any]] | None = None, + number_channels: int = 2, + enable_channel_interleave: bool = True, + attention_implementation: str = "sdpa", + compute_dtype: str = "fp32", + codec_weight_dtype: str = "fp32", + quantizer_type: str = "rlfq", + quantizer_kwargs: dict[str, Any] | None = None, + **kwargs, + ): + # Some checkpoints might include an incorrect/legacy `model_type` (e.g. "speech_tokenizer"). + # We drop it to avoid overriding the class-level `model_type`. + kwargs.pop("model_type", None) + if "channels_numbers" in kwargs: + number_channels = kwargs.pop("channels_numbers") + if "enable_channel_interleave" in kwargs: + enable_channel_interleave = kwargs.pop("enable_channel_interleave") + if "attention_backend" in kwargs and attention_implementation == "sdpa": + attention_implementation = kwargs.pop("attention_backend") + if "codec_compute_dtype" in kwargs and compute_dtype == "fp32": + compute_dtype = kwargs.pop("codec_compute_dtype") + if "codec_load_dtype" in kwargs and codec_weight_dtype == "fp32": + codec_weight_dtype = kwargs.pop("codec_load_dtype") + + # `version` is accepted for compatibility but not used in modeling. + self.version = version + self.sampling_rate = sampling_rate + self.downsample_rate = downsample_rate + self.causal_transformer_context_duration = causal_transformer_context_duration + self.number_channels = number_channels + self.enable_channel_interleave = enable_channel_interleave + self.attention_implementation = attention_implementation + self.compute_dtype = compute_dtype + self.codec_weight_dtype = codec_weight_dtype + # Default encoder configuration + if encoder_kwargs is None: + encoder_kwargs = [ + { + "module_type": "PatchedPretransform", + "patch_size": 240, + }, + { + "module_type": "Transformer", + "input_dimension": 240, + "output_dimension": 384, + "d_model": 768, + "num_heads": 12, + "num_layers": 12, + "dim_feedforward": 3072, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + }, + { + "module_type": "PatchedPretransform", + "patch_size": 2, + }, + { + "module_type": "Transformer", + "input_dimension": 768, + "output_dimension": 384, + "d_model": 768, + "num_heads": 12, + "num_layers": 12, + "dim_feedforward": 3072, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + }, + { + "module_type": "PatchedPretransform", + "patch_size": 2, + }, + { + "module_type": "Transformer", + "input_dimension": 768, + "output_dimension": 640, + "d_model": 768, + "num_heads": 12, + "num_layers": 12, + "dim_feedforward": 3072, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + }, + { + "module_type": "PatchedPretransform", + "patch_size": 2, + }, + { + "module_type": "Transformer", + "input_dimension": 1280, + "output_dimension": 768, + "d_model": 1280, + "num_heads": 20, + "num_layers": 32, + "dim_feedforward": 5120, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + }, + ] + self.encoder_kwargs = encoder_kwargs + + # Default decoder configuration (execution order) + if decoder_kwargs is None: + decoder_kwargs = [ + { + "module_type": "Transformer", + "input_dimension": 768, + "output_dimension": 1280, + "d_model": 1280, + "num_heads": 20, + "num_layers": 32, + "dim_feedforward": 5120, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + }, + { + "module_type": "PatchedPretransform", + "patch_size": 2, + }, + { + "module_type": "Transformer", + "input_dimension": 640, + "output_dimension": 768, + "d_model": 768, + "num_heads": 12, + "num_layers": 12, + "dim_feedforward": 3072, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + }, + { + "module_type": "PatchedPretransform", + "patch_size": 2, + }, + { + "module_type": "Transformer", + "input_dimension": 384, + "output_dimension": 768, + "d_model": 768, + "num_heads": 12, + "num_layers": 12, + "dim_feedforward": 3072, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + }, + { + "module_type": "PatchedPretransform", + "patch_size": 2, + }, + { + "module_type": "Transformer", + "input_dimension": 384, + "output_dimension": 768, + "d_model": 768, + "num_heads": 12, + "num_layers": 12, + "dim_feedforward": 3072, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + }, + { + "module_type": "PatchedPretransform", + "patch_size": 2, + }, + { + "module_type": "Transformer", + "input_dimension": 384, + "output_dimension": 240, + "d_model": 768, + "num_heads": 12, + "num_layers": 12, + "dim_feedforward": 3072, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + }, + { + "module_type": "PatchedPretransform", + "patch_size": 240, + }, + ] + self.decoder_kwargs = decoder_kwargs + + # Default quantizer configuration + if quantizer_kwargs is None: + quantizer_kwargs = { + "input_dim": 768, + "rvq_dim": 512, + "output_dim": 768, + "num_quantizers": 32, + "codebook_size": 1024, + "codebook_dim": 8, + "quantizer_type": "rlfq", + } + + # Handle quantizer_type from kwargs or config + kw_qtype = quantizer_kwargs.get("quantizer_type", None) + if kw_qtype is not None: + self.quantizer_type = kw_qtype + else: + self.quantizer_type = quantizer_type + quantizer_kwargs["quantizer_type"] = quantizer_type + + self.quantizer_kwargs = quantizer_kwargs + + super().__init__(**kwargs) + + @property + def num_quantizers(self) -> int: + """Return the number of quantizers from quantizer_kwargs.""" + return self.quantizer_kwargs.get("num_quantizers", 32) + + @property + def codebook_size(self) -> int: + """Return the codebook size from quantizer_kwargs.""" + return self.quantizer_kwargs.get("codebook_size", 4096) + + @property + def frame_rate(self) -> float: + """Return the frame rate (tokens per second).""" + return self.sampling_rate / self.downsample_rate + + +__all__ = ["MossAudioTokenizerConfig"] diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py index 7c5d47653..60f455240 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py @@ -5,6 +5,7 @@ from __future__ import annotations +import contextlib from collections.abc import Iterable from typing import Any @@ -19,6 +20,12 @@ from vllm_omni.model_executor.models.moss_tts.audio_tokenizer import ( MossAudioTokenizerConfig, MossAudioTokenizerModel, ) +from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import ( + MossAudioTokenizerModel as MossAudioTokenizerV2Model, +) +from vllm_omni.model_executor.models.moss_tts.configuration_moss_audio_tokenizer_v2 import ( + MossAudioTokenizerConfig as MossAudioTokenizerV2Config, +) from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import ( MossTTSCUDAGraphCodecWrapper, ) @@ -27,6 +34,96 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput logger = init_logger(__name__) +class _MossCodecStreamSession: + """Persistent streaming decode session for vendored MOSS-Audio-Tokenizer-v2.""" + + def __init__( + self, + codec: nn.Module, + *, + stream_slots: int, + n_vq: int, + ) -> None: + self._codec = codec + self._stream_slots = int(stream_slots) + self._batch_size = self._stream_slots + self._n_vq = int(n_vq) + self._device = next(codec.parameters()).device + self._free_stream_slots = list(range(self._stream_slots)) + self._exit_stack = contextlib.ExitStack() + self._closed = False + with torch.no_grad(): + self._exit_stack.enter_context(codec.streaming(self._batch_size)) + + def acquire(self) -> int | None: + if not self._free_stream_slots: + return None + return self._free_stream_slots.pop() + + def release(self, slot: int) -> None: + if self._closed: + return + self.reset_slots([slot]) + self._free_stream_slots.append(slot) + + def reset_slots(self, slots: list[int]) -> None: + if not slots: + return + reset_mask = torch.zeros(self._batch_size, dtype=torch.bool, device=self._device) + reset_mask[slots] = True + + def _reset(module: nn.Module) -> None: + state = getattr(module, "_streaming_state", None) + if state is not None: + state.reset(reset_mask.to(state.device)) + + with torch.no_grad(): + self._codec.apply(_reset) + + def close(self) -> None: + if self._closed: + return + with torch.no_grad(): + self._exit_stack.close() + self._closed = True + + @torch.no_grad() + def step(self, slot_codes: dict[int, torch.Tensor]) -> dict[int, torch.Tensor]: + if not slot_codes: + return {} + step_lengths = {int(codes.shape[1]) for codes in slot_codes.values()} + if len(step_lengths) != 1: + raise ValueError(f"MOSS codec streaming step needs uniform T, got {sorted(step_lengths)}") + (step_t,) = step_lengths + codes_step = torch.zeros( + self._n_vq, + self._batch_size, + step_t, + dtype=torch.long, + device=self._device, + ) + codes_lengths = torch.zeros(self._batch_size, dtype=torch.long, device=self._device) + exec_mask = torch.zeros(self._batch_size, dtype=torch.bool, device=self._device) + for slot, codes in slot_codes.items(): + codes_step[:, slot, :] = codes.to(device=self._device, dtype=torch.long) + codes_lengths[slot] = int(codes.shape[1]) + exec_mask[slot] = True + + self._codec._set_streaming_exec_mask(exec_mask) + result = self._codec._decode_frame(codes_step, codes_lengths) + if result.audio is None: + return {} + audio = result.audio.detach().to("cpu", torch.float32) + lengths = result.audio_lengths.detach().to("cpu") if result.audio_lengths is not None else None + out: dict[int, torch.Tensor] = {} + for slot in slot_codes: + wav = audio[slot] + if lengths is not None: + wav = wav[..., : int(lengths[slot].item())] + out[slot] = wav.contiguous() + return out + + class MossTTSCodecDecoder(nn.Module): """Stage-1 decoder for all MOSS-TTS variants. @@ -75,6 +172,12 @@ class MossTTSCodecDecoder(nn.Module): self._cuda_graph_wrapper: MossTTSCUDAGraphCodecWrapper | None = None self._n_channels: int = 1 self._sr_tensor = torch.tensor(self._OUTPUT_SAMPLE_RATE, dtype=torch.int32) + self._stream_session: _MossCodecStreamSession | None = None + self._stream_slots: int = self._connector_int("codec_stream_slots", default=0) + self._stream_max_step_frames: int = self._connector_int("codec_max_step_frames", default=100) + self._stream_req_slots: dict[str, int] = {} + self._stream_pending_codes: dict[str, list[torch.Tensor]] = {} + self._stream_starved_reqs: set[str] = set() # ------------------------------------------------------------------ # vLLM-Omni stubs (codec has no AR loop) @@ -112,9 +215,9 @@ class MossTTSCodecDecoder(nn.Module): chunk transfer adapter assigns those to ``request.prompt_token_ids``, which arrives here as ``input_ids`` concatenated across all requests. Per-request slice boundaries are computed from - ``kwargs["num_scheduled_tokens"]`` (a list of token counts, one per - request). ``runtime_additional_information`` carries per-request - metadata such as ``left_context_size``. + ``kwargs["seq_token_counts"]`` (token counts, one per request). + ``runtime_additional_information`` carries per-request metadata such as + ``left_context_size``. Returns ------- @@ -123,8 +226,8 @@ class MossTTSCodecDecoder(nn.Module): multimodal_outputs["sr"] — list of scalar int32 tensors """ sr_tensor = self._sr_tensor - empty = torch.zeros((0,), dtype=torch.float32) - info_list: list[dict[str, Any]] = runtime_additional_information or [{}] + empty = self._empty_audio() + info_list: list[dict[str, Any]] = list(runtime_additional_information or [{}]) num_req = max(len(info_list), 1) if self._codec is None: @@ -140,24 +243,47 @@ class MossTTSCodecDecoder(nn.Module): audios: list[torch.Tensor] = [empty] * num_req srs: list[torch.Tensor] = [sr_tensor] * num_req device = next(self._codec.parameters()).device + streaming_work: list[tuple[int, str, torch.Tensor, bool]] = [] if input_ids is None or input_ids.numel() == 0: + for i, wav in self._finish_empty_streaming_requests(info_list).items(): + audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav return OmniOutput( text_hidden_states=None, multimodal_outputs={"model_outputs": audios, "sr": srs}, ) - # ``input_ids`` is concatenated across all requests; split by - # query_start_loc-style offsets from ``num_scheduled_tokens`` in - # kwargs if available, else assume one request. + # ``input_ids`` is concatenated across all requests. vLLM-Omni runners + # pass the per-request lengths via the shared code2wav contract + # ``seq_token_counts``. ids_flat = input_ids.reshape(-1).to(dtype=torch.long) - num_scheduled_tokens = kwargs.get("num_scheduled_tokens") - if isinstance(num_scheduled_tokens, (list, tuple)) and len(num_scheduled_tokens) == num_req: - offsets = [0] - for n in num_scheduled_tokens: - offsets.append(offsets[-1] + int(n)) - else: - offsets = [0, int(ids_flat.shape[0])] + token_counts = self._normalize_seq_token_counts(kwargs.get("seq_token_counts")) + if token_counts is None: + raise RuntimeError( + "MossTTS codec requires seq_token_counts; otherwise concatenated " + "codec tokens cannot be split per request." + ) + if sum(token_counts) != int(ids_flat.shape[0]): + raise RuntimeError( + "MossTTS codec seq_token_counts mismatch: " + f"counts={token_counts}, sum={sum(token_counts)}, input_tokens={int(ids_flat.shape[0])}." + ) + + num_req = len(token_counts) + if len(info_list) < num_req: + info_list.extend({} for _ in range(num_req - len(info_list))) + elif len(info_list) > num_req: + info_list = info_list[:num_req] + if len(audios) < num_req: + audios.extend(empty for _ in range(num_req - len(audios))) + srs.extend(sr_tensor for _ in range(num_req - len(srs))) + elif len(audios) > num_req: + audios = audios[:num_req] + srs = srs[:num_req] + + offsets = [0] + for n in token_counts: + offsets.append(offsets[-1] + int(n)) for i, info in enumerate(info_list): if i + 1 >= len(offsets): @@ -165,6 +291,14 @@ class MossTTSCodecDecoder(nn.Module): seg = ids_flat[offsets[i] : offsets[i + 1]] if seg.numel() == 0: continue + meta = (info.get("meta", {}) if isinstance(info, dict) else {}) or {} + finished = bool(meta.get("stream_finished", meta.get("finished", False))) + streaming_enabled = bool(meta.get("codec_streaming", False)) + code_flat_numel = meta.get("code_flat_numel") + if streaming_enabled and finished and code_flat_numel is not None and int(code_flat_numel) == 0: + for _, wav in self._finish_empty_streaming_requests([info]).items(): + audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav + continue if seg.numel() % self._n_vq != 0: logger.warning( "MossTTS codec input length %d not divisible by n_vq %d; skipping.", @@ -181,14 +315,19 @@ class MossTTSCodecDecoder(nn.Module): codebook_size = self._codec.config.codebook_size codes_nq_t = codes_nq_t.clamp_(0, int(codebook_size) - 1) - meta = (info.get("meta", {}) if isinstance(info, dict) else {}) or {} - left_ctx = self._runtime_value(info, meta, "left_context_size", 0) + left_ctx = meta.get("left_context_size", 0) if isinstance(left_ctx, (list, tuple)): left_ctx = int(left_ctx[0]) if left_ctx else 0 elif isinstance(left_ctx, torch.Tensor): left_ctx = int(left_ctx.reshape(-1)[0].item()) if left_ctx.numel() else 0 left_ctx = int(left_ctx) + req_key = self._runtime_request_key(info, meta, i) + + if streaming_enabled: + streaming_work.append((i, req_key, codes_nq_t, finished)) + continue + if self._cuda_graph_wrapper is not None: out = self._cuda_graph_wrapper.decode(codes_nq_t) else: @@ -216,19 +355,246 @@ class MossTTSCodecDecoder(nn.Module): ) wav = wav[..., trim:] - audios[i] = wav.reshape(-1) if self._n_channels == 1 else wav + audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav + + if streaming_work: + for i, wav in self._decode_streaming_batch(streaming_work).items(): + audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav return OmniOutput( text_hidden_states=None, multimodal_outputs={"model_outputs": audios, "sr": srs}, ) + def _finish_empty_streaming_requests(self, info_list: list[dict[str, Any]]) -> dict[int, torch.Tensor]: + """Release codec stream state for empty finish sentinels. + + Stage-0 can finish on a step that emits no new audio frame. The stage + input processor forwards that as an empty payload with finished=true. + If a request never acquired a stream slot, do not offline-decode its + buffered codes here: streaming requests must only emit client deltas. + """ + session = self._stream_session + outputs: dict[int, torch.Tensor] = {} + if session is None: + return outputs + for i, info in enumerate(info_list): + if not isinstance(info, dict): + continue + meta = (info.get("meta", {}) or {}) if isinstance(info.get("meta", {}), dict) else {} + if not bool(meta.get("codec_streaming", False)): + continue + finished = bool(meta.get("stream_finished", meta.get("finished", False))) + if not finished: + continue + req_key = self._runtime_request_key(info, meta, i) + slot = self._stream_req_slots.get(req_key) + pending = req_key in self._stream_pending_codes + if slot is not None or pending: + if pending and slot is None: + logger.warning( + "MOSS codec stream request %s finished before a codec stream slot became available; " + "dropping buffered codes instead of offline-decoding a non-delta waveform.", + req_key, + ) + self._finish_stream_request(req_key, session, slot) + return outputs + @staticmethod - def _runtime_value(info: Any, meta: dict[str, Any], name: str, default: Any = None) -> Any: - if name in meta: - return meta[name] - if isinstance(info, dict) and name in info: - return info[name] + def _normalize_seq_token_counts(value: Any) -> list[int] | None: + if value is None: + return None + if not isinstance(value, (list, tuple)): + raise TypeError( + "MossTTS codec expects seq_token_counts to be a list/tuple of per-request token counts, " + f"got {type(value).__name__}." + ) + counts = [int(item) for item in value] + if not counts: + return None + for count in counts: + if count < 0: + raise ValueError(f"MossTTS codec seq_token_counts must be non-negative, got {counts}.") + return counts + + def _runtime_request_key(self, info: Any, meta: dict[str, Any], index: int) -> str: + for value in ( + meta.get("req_id"), + info.get("request_id") if isinstance(info, dict) else None, + ): + if isinstance(value, (list, tuple)): + value = value[0] if value else None + if value is not None: + return str(value) + return f"moss-codec-stream-{index}" + + def _empty_audio(self) -> torch.Tensor: + if self._n_channels > 1: + return torch.zeros((self._n_channels, 0), dtype=torch.float32) + return torch.zeros((0,), dtype=torch.float32) + + def _ensure_stream_session(self) -> _MossCodecStreamSession | None: + if self._codec is None: + return None + if self._stream_session is not None: + return self._stream_session + slots = self._stream_slots + if slots <= 0: + scheduler_cfg = getattr(self.vllm_config, "scheduler_config", None) + slots = int(getattr(scheduler_cfg, "max_num_seqs", 1) or 1) + self._stream_session = _MossCodecStreamSession( + self._codec, + stream_slots=max(1, slots), + n_vq=self._n_vq, + ) + return self._stream_session + + def _decode_streaming_batch( + self, + items: list[tuple[int, str, torch.Tensor, bool]], + ) -> dict[int, torch.Tensor]: + session = self._ensure_stream_session() + if session is None: + return {} + + outputs: dict[int, torch.Tensor] = {} + grouped: dict[int, list[tuple[int, str, int, torch.Tensor, bool]]] = {} + max_step_frames = max(1, int(self._stream_max_step_frames)) + + for output_index, request_id, codes_nq_t, finished in items: + pending = self._stream_pending_codes.get(request_id) + slot = self._stream_req_slots.get(request_id) + if slot is None: + slot = session.acquire() + if slot is None: + self._append_stream_pending(request_id, codes_nq_t) + if request_id not in self._stream_starved_reqs: + logger.warning( + "MOSS codec streaming slots exhausted; buffering %s until a stream slot is available.", + request_id, + ) + self._stream_starved_reqs.add(request_id) + if finished: + logger.warning( + "MOSS codec stream request %s finished before a codec stream slot became available; " + "dropping buffered codes instead of offline-decoding a non-delta waveform.", + request_id, + ) + self._finish_stream_request(request_id, session, None) + continue + self._stream_req_slots[request_id] = slot + + if pending: + self._append_stream_pending(request_id, codes_nq_t) + replay_codes = self._pop_stream_pending(request_id) + wav = self._decode_stream_slot_sequence(session, slot, replay_codes) + if wav is not None: + outputs[output_index] = wav + if finished: + self._finish_stream_request(request_id, session, slot) + continue + + if int(codes_nq_t.shape[1]) > max_step_frames: + wav = self._decode_stream_slot_sequence(session, slot, codes_nq_t) + if wav is not None: + outputs[output_index] = wav + if finished: + self._finish_stream_request(request_id, session, slot) + continue + + grouped.setdefault(int(codes_nq_t.shape[1]), []).append( + (output_index, request_id, slot, codes_nq_t, finished) + ) + + for group in grouped.values(): + plan = {slot: codes_nq_t for _, _, slot, codes_nq_t, _ in group} + decoded = session.step(plan) + for output_index, request_id, slot, _, finished in group: + wav = decoded.get(slot) + if wav is not None: + outputs[output_index] = wav + if finished: + self._finish_stream_request(request_id, session, slot) + + return outputs + + def _append_stream_pending(self, request_id: str, codes_nq_t: torch.Tensor) -> None: + self._stream_pending_codes.setdefault(request_id, []).append( + codes_nq_t.detach().to("cpu", torch.long).contiguous() + ) + + def _pop_stream_pending(self, request_id: str) -> torch.Tensor: + pending = self._stream_pending_codes.pop(request_id, []) + if not pending: + return torch.empty((self._n_vq, 0), dtype=torch.long) + return torch.cat(pending, dim=1).contiguous() + + def _decode_stream_slot_sequence( + self, + session: _MossCodecStreamSession, + slot: int, + codes_nq_t: torch.Tensor, + ) -> torch.Tensor | None: + if codes_nq_t.numel() == 0: + return None + max_step_frames = max(1, int(self._stream_max_step_frames)) + parts: list[torch.Tensor] = [] + for start in range(0, int(codes_nq_t.shape[1]), max_step_frames): + chunk = codes_nq_t[:, start : start + max_step_frames] + decoded = session.step({slot: chunk}) + wav = decoded.get(slot) + if wav is not None: + parts.append(wav) + if not parts: + return None + return torch.cat(parts, dim=-1) if len(parts) > 1 else parts[0] + + def _finish_stream_request( + self, + request_id: str, + session: _MossCodecStreamSession, + slot: int | None, + ) -> None: + if slot is not None: + session.release(slot) + self._stream_req_slots.pop(request_id, None) + self._stream_pending_codes.pop(request_id, None) + self._stream_starved_reqs.discard(request_id) + + def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: + """Release codec streaming slots when requests finish outside payload flow. + + Normal streaming completion releases slots from ``_decode_streaming_batch`` + when the Stage-0 payload carries ``finished=True``. Client disconnects + and engine-side aborts can finish a request without delivering that + terminal payload, so the runner calls this hook from its finished-request + path to avoid leaking stream slots and buffered codes. + """ + session = self._stream_session + for req_id in finished_req_ids: + request_id = str(req_id) + slot = self._stream_req_slots.get(request_id) + has_state = ( + slot is not None or request_id in self._stream_pending_codes or request_id in self._stream_starved_reqs + ) + if not has_state: + continue + if session is not None: + self._finish_stream_request(request_id, session, slot) + else: + self._stream_req_slots.pop(request_id, None) + self._stream_pending_codes.pop(request_id, None) + self._stream_starved_reqs.discard(request_id) + + def _connector_int(self, name: str, default: int = 0) -> int: + model_cfg = getattr(self.vllm_config, "model_config", None) + connector_cfg = getattr(model_cfg, "stage_connector_config", None) + if isinstance(connector_cfg, dict): + extra_cfg: dict | None = connector_cfg.get("extra", connector_cfg) + else: + extra_cfg = getattr(connector_cfg, "extra", None) + if isinstance(extra_cfg, dict) and name in extra_cfg: + return int(extra_cfg[name]) return default # ------------------------------------------------------------------ @@ -272,14 +638,12 @@ class MossTTSCodecDecoder(nn.Module): (".self_attn.out_projs.0.", ".attn.out_proj."), (".linear1.", ".ff1."), (".linear2.", ".ff2."), - # v2 (MOSS-Audio-Tokenizer-v2) uses different submodule names for - # the same attention/FFN sublayers — no trailing "s"/index on - # in_proj/out_proj, and an `ffn.{0,2}` Sequential instead of - # separate linear1/linear2 attributes. - (".self_attn.in_proj.", ".attn.in_proj."), - (".self_attn.out_proj.", ".attn.out_proj."), - (".ffn.0.", ".ff1."), - (".ffn.2.", ".ff2."), + # v2 checkpoint names use singular in_proj/out_proj and ffn.{0,2}; + # the vendored module keeps the original MOSS layer names. + (".self_attn.in_proj.", ".self_attn.in_projs.0."), + (".self_attn.out_proj.", ".self_attn.out_projs.0."), + (".ffn.0.", ".linear1."), + (".ffn.2.", ".linear2."), (".layer_scale_1.", ".ls1."), (".layer_scale_2.", ".ls2."), (".input_proj.", ".in_proj."), @@ -292,21 +656,39 @@ class MossTTSCodecDecoder(nn.Module): return name.replace(src, dst) return name - loaded_total = 0 + loaded_names: set[str] = set() skipped: list[str] = [] + shape_mismatches: list[tuple[str, str, tuple[int, ...], tuple[int, ...]]] = [] for name, tensor in codec_weights: # Try direct name first (e.g. ``quantizer.input_proj.*`` exists # under the same name in both layouts), then the remap (transformer # submodules need ``.linear1.``→``.ff1.`` etc.). tgt = name if name in params_dict else _remap(name) if tgt in params_dict: + expected_shape = tuple(params_dict[tgt].shape) + actual_shape = tuple(tensor.shape) + if expected_shape != actual_shape: + shape_mismatches.append((name, tgt, actual_shape, expected_shape)) + continue default_weight_loader(params_dict[tgt], tensor) - loaded_total += 1 + loaded_names.add(tgt) else: skipped.append(name) + + missing = sorted(set(params_dict) - loaded_names) + if missing or skipped or shape_mismatches: + raise RuntimeError( + "MOSS Audio Tokenizer weights were not fully loaded: " + f"loaded={len(loaded_names)}/{len(params_dict)} " + f"missing={len(missing)} skipped={len(skipped)} " + f"shape_mismatches={len(shape_mismatches)}; " + f"first_missing={missing[:5]} " + f"first_skipped={skipped[:5]} " + f"first_shape_mismatches={shape_mismatches[:3]}" + ) logger.info( "MOSS Audio Tokenizer weights: loaded=%d/%d skipped=%d (first skipped: %s)", - loaded_total, + len(loaded_names), len(params_dict), len(skipped), skipped[:3] if skipped else "none", @@ -316,7 +698,15 @@ class MossTTSCodecDecoder(nn.Module): codec.to(device=device, dtype=torch.float32) codec.eval() self._codec = codec - self._n_channels = int(getattr(codec_cfg, "number_channels", 1) or 1) + inferred_channels = 2 if "v2" in codec_path.lower() else 1 + self._n_channels = int( + getattr( + codec_cfg, + "number_channels", + getattr(codec_cfg, "num_channels", inferred_channels), + ) + or inferred_channels + ) self._sr_tensor = torch.tensor(int(codec_cfg.sampling_rate), dtype=torch.int32) logger.info( @@ -336,16 +726,13 @@ class MossTTSCodecDecoder(nn.Module): def _build_codec(self, codec_path: str) -> tuple[Any, nn.Module]: try: - from transformers import AutoConfig, AutoModel - - codec_cfg = AutoConfig.from_pretrained(codec_path, trust_remote_code=True) - codec = AutoModel.from_config(codec_cfg, trust_remote_code=True) - logger.info("Using MOSS Audio Tokenizer remote-code classes from %s", codec_path) + codec_cfg = MossAudioTokenizerV2Config.from_pretrained(codec_path) + codec = MossAudioTokenizerV2Model(codec_cfg) + logger.info("Using vendored MOSS Audio Tokenizer v2 classes from %s", codec_path) return codec_cfg, codec except Exception: logger.exception( - "Failed to instantiate official MOSS Audio Tokenizer via HF remote code; " - "falling back to vendored codec." + "Failed to instantiate vendored MOSS Audio Tokenizer v2; falling back to legacy vendored codec." ) codec_cfg = MossAudioTokenizerConfig.from_pretrained(codec_path) diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py index bf01a83bf..2fd0beefd 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py @@ -126,6 +126,7 @@ def _sample_token( top_k: int, top_p: float, do_sample: bool, + generator: torch.Generator | None = None, ) -> torch.Tensor: """Top-k + top-p sampling (matches upstream's ``sample_token`` for the inference branch). @@ -152,7 +153,7 @@ def _sample_token( probs = F.softmax(logits, dim=-1) flat = probs.reshape(-1, probs.shape[-1]) - sampled = torch.multinomial(flat, num_samples=1).reshape(probs.shape[:-1]) + sampled = torch.multinomial(flat, num_samples=1, generator=generator).reshape(probs.shape[:-1]) return sampled diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py index cb2e95551..e0d871b46 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py @@ -160,6 +160,7 @@ class MossTTSLocalDepthTransformer(nn.Module): text_top_p: float = 1.0, repetition_penalty: float = 1.0, history_per_codebook: list[list[int]] | None = None, + generator: torch.Generator | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Generate one audio frame for batch B. @@ -183,7 +184,16 @@ class MossTTSLocalDepthTransformer(nn.Module): local_hidden = hidden[:, 0, :] binary_logits = local_text_lm_head(local_hidden).float() - binary_choice = _sample_token(binary_logits, text_temperature, text_top_k, text_top_p, do_sample) + # This is a binary continue/stop gate. The checkpoint expects sampling + # here; greedy argmax is biased toward "continue" and may never stop. + binary_choice = _sample_token( + binary_logits, + text_temperature, + text_top_k, + text_top_p, + do_sample, + generator=generator, + ) should_continue = binary_choice.eq(0) import os as _os @@ -209,7 +219,14 @@ class MossTTSLocalDepthTransformer(nn.Module): pos = sel > 0 sel = torch.where(pos, sel / repetition_penalty, sel * repetition_penalty) channel_logits.index_copy_(-1, hist_t, sel) - channel_token = _sample_token(channel_logits, temperature, top_k, top_p, do_sample) + channel_token = _sample_token( + channel_logits, + temperature, + top_k, + top_p, + do_sample, + generator=generator, + ) codes[:, channel_index] = channel_token if channel_index + 1 < n_vq: diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py index 9f4037e4b..9758c922e 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py @@ -1139,26 +1139,19 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): info["audio_codes"] = {"current": new_codes, "accumulated": updated_acc} state["step"] = int(state.get("step", 0)) + 1 - per_req_codes[i] = updated_acc + per_req_codes[i] = new_codes.unsqueeze(0) have_codes = True if not have_codes: return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) - # Emit one accumulated-codes tensor per request as a list; the splitter - # routes ``element[idx]`` to request idx, so each Stage-1 request decodes - # its own codes instead of request 0's. - # The realtime variant emits raw codes (no delay pattern), so we signal - # that to the chunk processor via a 1-D bool tensor in ``meta.finished`` - # (already a tensor field on the schema). The processor reads its truth - # value as "skip de-delay" only for realtime; the delay path doesn't - # populate ``meta`` at all. - device = hidden.device + # Emit one current-frame tensor per request as a list; the splitter + # routes ``element[idx]`` to request idx, so each Stage-1 request receives + # only its own newly generated raw codec row. return OmniOutput( text_hidden_states=hidden, multimodal_outputs={ "codes": {"audio": per_req_codes}, - "meta": {"finished": torch.tensor([True], dtype=torch.bool, device=device)}, }, ) @@ -1266,7 +1259,7 @@ class MossTTSLocalTalkerForGeneration(nn.Module): * Per-step audio generation runs ``MossTTSLocalDepthTransformer`` (1-layer GPT2-style local transformer, interleaved RoPE, ``n_vq=12`` codebooks -- structurally different from Realtime's Qwen3-style - ``CodePredictorBaseModel``) inside ``make_omni_output``. Stop + ``CodePredictorBaseModel``) in the runner-owned ``talker_mtp`` path. Stop condition: the local transformer's binary head selects the "audio_end" candidate (logits index 1). """ @@ -1327,6 +1320,8 @@ class MossTTSLocalTalkerForGeneration(nn.Module): # Stacked embedding cache built after load_weights() for vectorised # additive-fusion gathers (avoids an n_vq-iteration Python loop). self._stacked_audio_emb_w: torch.Tensor | None = None + self.mtp_hidden_size = hidden_size + self.talker_mtp_accepts_req_infos = True self.gpu_resident_buffer_keys: set[tuple[str, str]] = { ("audio_codes", "current"), @@ -1384,6 +1379,15 @@ class MossTTSLocalTalkerForGeneration(nn.Module): num_rows, ): if state is not None and state.get("is_stopping"): + forced_im_end_count = int(state.get("_forced_im_end_count", 0)) + if forced_im_end_count == 1: + logger.warning( + "MOSS-TTS Local request is still running after " + "im_end_token_id=%d was forced once; verify stage-0 " + "stop_token_ids/eos includes this id.", + self.im_end_token_id, + ) + state["_forced_im_end_count"] = forced_im_end_count + 1 logits[row_start:row_end, self.im_end_token_id] = 0.0 else: logits[row_start:row_end, self.audio_assistant_slot_token_id] = 0.0 @@ -1455,38 +1459,159 @@ class MossTTSLocalTalkerForGeneration(nn.Module): else torch.full((self.n_vq,), self.audio_pad_token_id, dtype=torch.long, device=device) ) - max_new_frames_req = info_dict.get("max_new_frames") - if isinstance(max_new_frames_req, (list, tuple)) and max_new_frames_req: - max_new_frames_req = max_new_frames_req[0] - try: - max_new_frames = int(max_new_frames_req) if max_new_frames_req is not None else -1 - except (TypeError, ValueError): - max_new_frames = -1 + max_new_frames = info_dict.get("max_new_frames", [-1])[0] info_update: dict[str, Any] = { - "audio_state": {"is_stopping": False, "step": 0, "max_new_frames": max_new_frames}, - "audio_codes": {"current": current_codes}, + "audio_state": { + "is_stopping": False, + "step": 0, + "max_new_frames": max_new_frames, + }, + "audio_codes": {"current": current_codes, "emit": False}, "ref_offset": ref_offset + span_len, } return input_ids, embeds, info_update - # Decode step: input_ids is the forced text token from compute_logits. + # Decode step: return the plain forced text-token embedding. The local + # depth transformer runs in OmniGPUModelRunner.talker_mtp before the + # backbone forward; it adds the freshly sampled audio-code embedding + # and stores the generated code frame in model_intermediate_buffer. text_embed = self.model.embed_tokens(input_ids.reshape(1)) # (1, H) - prev_codes = (info_dict.get("audio_codes", {}) or {}).get("current") - if isinstance(prev_codes, torch.Tensor) and prev_codes.numel() == self.n_vq: - codes = prev_codes.to(device=device, dtype=torch.long).unsqueeze(0) # (1, n_vq) - combined = text_embed + self._audio_embed(codes) + hs = info_dict.get("hidden_states", {}) or {} + last_hidden = hs.get("last") + if isinstance(last_hidden, torch.Tensor) and last_hidden.numel() > 0: + mtp_hidden = last_hidden.to(device=device, dtype=text_embed.dtype).reshape(1, -1) else: - combined = text_embed - return input_ids, combined, {} + mtp_hidden = torch.zeros((1, self.hidden_size), device=device, dtype=text_embed.dtype) + return input_ids, text_embed, {"mtp_inputs": (mtp_hidden, torch.zeros_like(mtp_hidden))} def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]: if hidden_states.numel() == 0: return {} return {"hidden_states": {"last": hidden_states[-1].detach()}} + @torch.inference_mode() + def talker_mtp( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor, + last_talker_hidden: torch.Tensor, + text_step: torch.Tensor, + do_sample: bool | None = None, + temperature: float | None = None, + top_k: int | None = None, + top_p: float | None = None, + generator: torch.Generator | None = None, + req_infos: list[dict[str, Any]] | None = None, + **_: Any, + ) -> tuple[torch.Tensor, None]: + bsz = int(input_embeds.shape[0]) + dev = input_embeds.device + input_embeds_out = input_embeds.reshape(bsz, -1).clone() + req_infos = req_infos or [{} for _ in range(bsz)] + + def write_update(info: dict[str, Any], update: dict[str, Any]) -> None: + for key, value in update.items(): + if isinstance(value, dict): + target = info.setdefault(key, {}) + target.update(value) + else: + info[key] = value + + # Match MOSS-TTS Local v1.5's model-card/SGLang defaults. The generic + # stage SamplingParams are for the backbone token loop and should not + # silently retune the local depth transformer's audio-channel sampler. + audio_temperature = 1.7 + audio_top_k = 25 + audio_top_p = 0.8 + do_sample_val = True if do_sample is None else bool(do_sample) + + for i in range(bsz): + info = req_infos[i] if i < len(req_infos) and isinstance(req_infos[i], dict) else {} + state = dict(info.get("audio_state", {}) or {}) + step = int(state.get("step", 0)) + max_new_frames = int(state.get("max_new_frames", -1)) + acc = (info.get("audio_codes", {}) or {}).get("accumulated") + + if state.get("is_stopping"): + write_update( + info, + { + "audio_state": state, + "audio_codes": { + "current": input_embeds.new_empty((0, self.n_vq), dtype=torch.long), + "emit": False, + }, + }, + ) + continue + + rep_window = 50 + if isinstance(acc, torch.Tensor) and acc.numel() > 0: + tail = acc[-rep_window:].long().cpu().tolist() + hist_per_cb = [[row[cb] for row in tail] for cb in range(self.n_vq)] + else: + hist_per_cb = [[] for _ in range(self.n_vq)] + + should_continue_t, new_codes_b = self.local_transformer.generate_frame( + last_talker_hidden[i : i + 1], + self.audio_lm_heads, + self.audio_embeddings, + self.local_text_lm_head, + n_vq=self.n_vq, + do_sample=do_sample_val, + temperature=audio_temperature, + top_k=audio_top_k, + top_p=audio_top_p, + repetition_penalty=1.0, + history_per_codebook=hist_per_cb, + generator=generator, + ) + should_continue = bool(should_continue_t[0].item()) + new_codes = new_codes_b.squeeze(0).to(device=dev, dtype=torch.long) + + if not should_continue or (0 <= max_new_frames <= step): + state["is_stopping"] = True + state["step"] = step + 1 + write_update( + info, + { + "audio_state": state, + "audio_codes": { + "current": input_embeds.new_empty((0, self.n_vq), dtype=torch.long), + "accumulated": acc, + "emit": False, + }, + }, + ) + continue + + if isinstance(acc, torch.Tensor) and acc.numel() > 0: + updated_acc = torch.cat([acc.to(dev), new_codes.unsqueeze(0)], dim=0) + else: + updated_acc = new_codes.unsqueeze(0) + + state["step"] = step + 1 + current = new_codes.unsqueeze(0) + input_embeds_out[i : i + 1] = input_embeds_out[i : i + 1] + self._audio_embed(current).to( + dtype=input_embeds_out.dtype + ) + write_update( + info, + { + "audio_state": state, + "audio_codes": { + "current": current, + "accumulated": updated_acc, + "emit": True, + }, + }, + ) + + return input_embeds_out, None + # ------------------------------------------------------------------ - # Per-step audio generation via local transformer + # Package runner-generated audio frames # ------------------------------------------------------------------ def make_omni_output( @@ -1511,91 +1636,41 @@ class MossTTSLocalTalkerForGeneration(nn.Module): self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts] self._batch_state_spans = kwargs.get("request_token_spans") - spans = kwargs.get("request_token_spans") - per_req_codes: list[torch.Tensor] = [hidden.new_empty((0, self.n_vq), dtype=torch.long) for _ in info_dicts] + per_req_codes: list[torch.Tensor] = [] have_codes = False - - if hidden.numel() > 0 and info_dicts: - num_rows = hidden.shape[0] - if spans is None or len(spans) != len(info_dicts): - raise RuntimeError( - "MOSS-TTS make_omni_output requires request_token_spans " - f"for {len(info_dicts)} request infos, got {0 if spans is None else len(spans)}" - ) - for i, info in enumerate(info_dicts): - if not isinstance(info, dict): - continue - row_start, row_end = spans[i] - row_end = min(int(row_end), num_rows) - if row_start >= row_end: - continue - - state = info.get("audio_state", {}) or {} - if state.get("is_stopping"): - continue # already stopped -- no more audio frames - - last_h = hidden[row_end - 1].unsqueeze(0) # (1, H) - - # Sampling parameters mirror Realtime's defaults; repetition - # penalty applies over a 50-frame rolling window per codebook. - rep_window = 50 - acc_for_hist = (info.get("audio_codes", {}) or {}).get("accumulated") - if isinstance(acc_for_hist, torch.Tensor) and acc_for_hist.numel() > 0: - tail = acc_for_hist[-rep_window:].long().cpu().tolist() - hist_per_cb = [[row[cb] for row in tail] for cb in range(self.n_vq)] - else: - hist_per_cb = [[] for _ in range(self.n_vq)] - - should_continue_t, new_codes = self.local_transformer.generate_frame( - last_h, - self.audio_lm_heads, - self.audio_embeddings, - self.local_text_lm_head, - n_vq=self.n_vq, - do_sample=True, - temperature=0.8, - top_k=30, - top_p=0.6, - repetition_penalty=1.1, - history_per_codebook=hist_per_cb, - ) - should_continue = bool(should_continue_t[0].item()) - new_codes = new_codes.squeeze(0) # (n_vq,) - - step = int(state.get("step", 0)) - max_new_frames = int(state.get("max_new_frames", -1)) - if not should_continue or (0 <= max_new_frames <= step): - state["is_stopping"] = True - state["step"] = step + 1 - info["audio_codes"] = { - "current": new_codes, - "accumulated": (info.get("audio_codes", {}) or {}).get("accumulated"), - } - continue # don't append the stop frame to accumulated - - acc = (info.get("audio_codes", {}) or {}).get("accumulated") - if isinstance(acc, torch.Tensor) and acc.numel() > 0: - updated_acc = torch.cat([acc.to(new_codes.device), new_codes.unsqueeze(0)], dim=0) - else: - updated_acc = new_codes.unsqueeze(0) - - info["audio_codes"] = {"current": new_codes, "accumulated": updated_acc} - state["step"] = step + 1 - per_req_codes[i] = updated_acc + for info in info_dicts: + current = (info.get("audio_codes", {}) or {}).get("current") if isinstance(info, dict) else None + emit = bool((info.get("audio_codes", {}) or {}).get("emit")) if isinstance(info, dict) else False + if emit and isinstance(current, torch.Tensor) and current.numel() > 0: + codes = current.to(device=hidden.device, dtype=torch.long) + if codes.dim() == 1: + codes = codes.unsqueeze(0) + per_req_codes.append(codes) have_codes = True + else: + per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long)) if not have_codes: return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) - # Local-v1.5 emits raw codes (no delay pattern), so signal that to the - # chunk processor the same way Realtime does, via ``meta.finished`` - # (read as "skip de-delay" by ``talker2codec_async_chunk``). - device = hidden.device + # Local-v1.5 emits the current raw code frame (no delay pattern); the + # pipeline routes those rows through ``talker2codec_raw_async_chunk``. + per_req_ref_codes: list[torch.Tensor] = [] + for info in info_dicts: + ref_codes = (info.get("codes", {}) or {}).get("ref") if isinstance(info, dict) else None + if isinstance(ref_codes, torch.Tensor) and ref_codes.numel() > 0: + ref_codes = ref_codes.detach() + if ref_codes.dim() == 1 and ref_codes.numel() % self.n_vq == 0: + ref_codes = ref_codes.view(-1, self.n_vq) + if ref_codes.dim() == 2 and ref_codes.shape[1] == self.n_vq: + per_req_ref_codes.append(ref_codes.to(device=hidden.device, dtype=torch.long)) + continue + per_req_ref_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long)) + return OmniOutput( text_hidden_states=hidden, multimodal_outputs={ - "codes": {"audio": per_req_codes}, - "meta": {"finished": torch.tensor([True], dtype=torch.bool, device=device)}, + "codes": {"audio": per_req_codes, "ref": per_req_ref_codes}, }, ) diff --git a/vllm_omni/model_executor/models/moss_tts/pipeline.py b/vllm_omni/model_executor/models/moss_tts/pipeline.py index d0cec08e5..c6cf5c86c 100644 --- a/vllm_omni/model_executor/models/moss_tts/pipeline.py +++ b/vllm_omni/model_executor/models/moss_tts/pipeline.py @@ -31,7 +31,7 @@ MOSS_TTS_PIPELINE = PipelineConfig( input_sources=(), owns_tokenizer=True, engine_output_type="latent", - async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_async_chunk"), + async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_delay_async_chunk"), sampling_constraints={ "detokenize": False, }, @@ -62,7 +62,7 @@ MOSS_TTS_REALTIME_PIPELINE = PipelineConfig( input_sources=(), owns_tokenizer=True, engine_output_type="latent", - async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_async_chunk"), + async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_raw_async_chunk"), sampling_constraints={"detokenize": False}, ), StagePipelineConfig( @@ -91,8 +91,11 @@ MOSS_TTS_LOCAL_PIPELINE = PipelineConfig( input_sources=(), owns_tokenizer=True, engine_output_type="latent", - async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_async_chunk"), - sampling_constraints={"detokenize": False}, + async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_raw_async_chunk"), + sampling_constraints={ + "detokenize": False, + "stop_token_ids": [151645], + }, ), StagePipelineConfig( stage_id=1, diff --git a/vllm_omni/model_executor/stage_input_processors/moss_tts.py b/vllm_omni/model_executor/stage_input_processors/moss_tts.py index 48c0f98cb..3dfe2a603 100644 --- a/vllm_omni/model_executor/stage_input_processors/moss_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/moss_tts.py @@ -5,6 +5,8 @@ from __future__ import annotations +from collections import defaultdict +from collections.abc import Mapping from typing import Any import torch @@ -15,6 +17,8 @@ from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct logger = init_logger(__name__) +_MOSS_AUDIO_PAD_CODE = 1024 + # --------------------------------------------------------------------------- # Helpers @@ -90,13 +94,13 @@ def talker2codec( # --------------------------------------------------------------------------- -def talker2codec_async_chunk( +def talker2codec_delay_async_chunk( transfer_manager: Any, multimodal_output: dict[str, Any] | None, request: Any, is_finished: bool = False, ) -> OmniPayloadStruct | None: - """Emit accumulated audio codes to Stage 1 as they arrive from Stage 0. + """Emit delay-patterned accumulated audio codes to Stage 1. State is maintained in ``transfer_manager`` keyed by request ID. A chunk is forwarded to Stage 1 when either: @@ -138,14 +142,6 @@ def talker2codec_async_chunk( req_state["accumulated"] = new_rows else: req_state["accumulated"] = torch.cat([req_state["accumulated"], new_rows], dim=0) - # Realtime variant emits raw (un-delay-patterned) codes; record that - # so the emit step below skips the apply_de_delay_pattern transform. - # Realtime sets ``meta.finished`` to a 1-D bool tensor; the delay - # variant leaves ``meta`` empty. - meta_in = pooling_output.get("meta", {}) or {} - flag = meta_in.get("finished") - if isinstance(flag, torch.Tensor) and flag.numel() >= 1 and bool(flag.reshape(-1)[0].item()): - req_state["skip_dedelay"] = True acc = req_state["accumulated"] if acc is None or acc.numel() == 0: @@ -181,15 +177,12 @@ def talker2codec_async_chunk( # 1. de-delay ``(T+nq-1, nq)`` → ``(T, nq)`` # 2. drop rows that are entirely pad (separators between audio segments, # and the leading text-mode rows that precede the first audio_start). - # The realtime talker emits raw codes (no delay), so steps 1 is skipped. chunk_codes_long = chunk_codes.to(torch.long).cpu().contiguous() # (T_chunk, NQ) nq = int(chunk_codes_long.shape[1]) t_chunk = int(chunk_codes_long.shape[0]) audio_pad_code = 1024 # MOSS-TTS audio_pad_code; same value across variants. - if req_state.get("skip_dedelay"): - de_delayed = chunk_codes_long - elif t_chunk > nq: + if t_chunk > nq: de_delayed = chunk_codes_long.new_zeros((t_chunk - nq + 1, nq)) for i in range(nq): de_delayed[:, i] = chunk_codes_long[i : i + de_delayed.shape[0], i] @@ -224,9 +217,133 @@ def talker2codec_async_chunk( codes=CodesStruct(audio=codec_flat), meta=MetaStruct( left_context_size=left_context, - finished=torch.tensor(is_finished, dtype=torch.bool), + finished=torch.tensor(bool(is_finished), dtype=torch.bool), ), ) -__all__ = ["talker2codec", "talker2codec_async_chunk"] +def talker2codec_raw_async_chunk( + transfer_manager: Any, + multimodal_output: dict[str, Any] | None, + request: Any, + is_finished: bool = False, +) -> OmniPayloadStruct | None: + """Async processor for MOSS-TTS Local/Realtime raw codec rows. + + Stage 0 emits newly generated raw codec rows shaped ``[T, n_vq]`` (normally + ``[1, n_vq]`` per decode step). This processor buffers those new rows until + a codec chunk is ready, then forwards the chunk to Stage 1. No delay-pattern + de-delay is applied on this path. + """ + external_req_id = getattr(request, "external_req_id", None) + req_id = str(external_req_id if external_req_id is not None else getattr(request, "request_id", id(request))) + + if not hasattr(transfer_manager, "code_prompt_token_ids"): + transfer_manager.code_prompt_token_ids = defaultdict(list) + if not hasattr(transfer_manager, "request_payload"): + transfer_manager.request_payload = {} + if not hasattr(transfer_manager, "put_req_chunk"): + transfer_manager.put_req_chunk = defaultdict(int) + + pending_frames = transfer_manager.code_prompt_token_ids[req_id] + + if isinstance(multimodal_output, Mapping): + codes_dict = multimodal_output.get("codes", {}) or {} + new_frames = codes_dict.get("audio") + if isinstance(new_frames, torch.Tensor) and new_frames.numel() > 0: + frames_cpu = new_frames.detach().to("cpu", torch.long).contiguous() + if frames_cpu.ndim == 1: + frames_cpu = frames_cpu.reshape(1, -1) + if frames_cpu.ndim != 2: + raise ValueError(f"MOSS raw codec frames must be 2-D, got {tuple(frames_cpu.shape)}") + valid_rows = frames_cpu.ne(_MOSS_AUDIO_PAD_CODE).any(dim=1) + for frame in frames_cpu[valid_rows]: + pending_frames.append(frame.clone()) + # Raw/local streaming should mirror the non-streaming path: the codec + # decodes only generated audio rows. Reference audio conditions the + # talker, but feeding its codes into the codec streaming state adds a + # long first-packet prime step and changes the decoder state relative + # to non-streaming output. + + connector = getattr(transfer_manager, "connector", None) + raw_cfg = getattr(connector, "config", {}) or {} + cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {} + cfg = cfg if isinstance(cfg, dict) else {} + chunk_frames = int(cfg.get("codec_chunk_frames", 15) or 15) + initial_chunk_frames = int(cfg.get("initial_codec_chunk_frames") or 0) + if chunk_frames <= 0: + raise ValueError(f"codec_chunk_frames must be positive for MOSS raw streaming, got {chunk_frames}") + if initial_chunk_frames < 0: + raise ValueError(f"initial_codec_chunk_frames must be non-negative, got {initial_chunk_frames}") + if initial_chunk_frames > chunk_frames: + logger.warning( + "initial_codec_chunk_frames=%d > codec_chunk_frames=%d, clamping.", + initial_chunk_frames, + chunk_frames, + ) + initial_chunk_frames = chunk_frames + + pending = len(pending_frames) + emitted_any = int(transfer_manager.put_req_chunk.get(req_id, 0)) > 0 + threshold = initial_chunk_frames if initial_chunk_frames > 0 and not emitted_any else chunk_frames + if pending <= 0: + if is_finished: + transfer_manager.code_prompt_token_ids.pop(req_id, None) + transfer_manager.request_payload.pop(req_id, None) + return OmniPayloadStruct( + # A non-empty sentinel is required so Stage-1 is scheduled. + # ``code_flat_numel=0`` tells the codec this is a control-only + # finish packet, not an audio code. + codes=CodesStruct(audio=torch.tensor([0], dtype=torch.long)), + meta=MetaStruct( + req_id=[req_id], + left_context_size=0, + codec_streaming=True, + codec_chunk_frames=0, + codec_left_context_frames=0, + code_flat_numel=0, + stream_finished=torch.tensor(True, dtype=torch.bool), + finished=torch.tensor(True, dtype=torch.bool), + ), + request_id=req_id, + ) + return None + if not is_finished and pending < threshold: + return None + + emit_frames = pending if is_finished else threshold + chunk_rows = pending_frames[:emit_frames] + del pending_frames[:emit_frames] + chunk_codes = torch.stack( + [row.to(torch.long).cpu() for row in chunk_rows], + dim=0, + ).contiguous() + finished = bool(is_finished and len(pending_frames) == 0) + + codec_flat = chunk_codes.transpose(0, 1).contiguous().reshape(-1).to(torch.long) + + if finished: + transfer_manager.code_prompt_token_ids.pop(req_id, None) + transfer_manager.request_payload.pop(req_id, None) + + return OmniPayloadStruct( + codes=CodesStruct(audio=codec_flat), + meta=MetaStruct( + req_id=[req_id], + left_context_size=0, + codec_streaming=True, + codec_chunk_frames=int(chunk_codes.shape[0]), + codec_left_context_frames=0, + code_flat_numel=int(codec_flat.numel()), + stream_finished=torch.tensor(finished, dtype=torch.bool), + finished=torch.tensor(finished, dtype=torch.bool), + ), + request_id=req_id, + ) + + +__all__ = [ + "talker2codec", + "talker2codec_delay_async_chunk", + "talker2codec_raw_async_chunk", +] diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 172121c14..c8fe92eba 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -76,6 +76,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin def _update_request_states(self, scheduler_output: SchedulerOutput): # remove requests + # Some stateful vocoder model may need to clean the state + # to avoid the leak of slots when the requests have been aborted. + if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"): + self.model.on_requests_finished(scheduler_output.finished_req_ids) for req_id in scheduler_output.finished_req_ids: self.input_batch.remove_request(req_id) scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys() @@ -789,6 +793,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin if hasattr(self.model, "get_dummy_runtime_additional_information"): runtime_addi = self.model.get_dummy_runtime_additional_information(num_reqs) model_kwargs["runtime_additional_information"] = runtime_addi + # Generation-stage code2wav models split flat input_ids using this + # shared contract. Dummy/profile runs use padded synthetic inputs, + # so expose one synthetic segment covering the actual tensor length. + model_kwargs["seq_token_counts"] = [int(num_tokens_padded)] if self.uses_mrope: positions = self.mrope_positions.gpu[:, :num_tokens_padded] diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index a283039f3..e40c220b8 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1821,7 +1821,9 @@ class OmniGPUModelRunner(GPUModelRunner): def _explicit_talker_seed(req_id: str) -> int | None: sampling_params = getattr(self.requests[req_id], "sampling_params", None) extra_args = getattr(sampling_params, "extra_args", None) if sampling_params is not None else None - seed = extra_args.get("qwen3_tts_request_seed") if isinstance(extra_args, dict) else None + seed = None + if isinstance(extra_args, dict): + seed = extra_args.get("tts_local_seed") return int(seed) if seed is not None else None if decode_batch_size > 1 and any(_explicit_talker_seed(req_id) is not None for req_id in decode_req_ids): @@ -1869,6 +1871,11 @@ class OmniGPUModelRunner(GPUModelRunner): } if generator is not None: talker_kwargs["generator"] = generator + if getattr(self.model, "talker_mtp_accepts_req_infos", False): + talker_kwargs["req_ids"] = decode_req_ids + talker_kwargs["req_infos"] = [ + self.model_intermediate_buffer.setdefault(req_id, {}) for req_id in decode_req_ids + ] with current_omni_platform.set_forward_context( None, self.vllm_config, cudagraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc ): @@ -1888,8 +1895,9 @@ class OmniGPUModelRunner(GPUModelRunner): start_offsets = [int(self.query_start_loc.cpu[id_to_index[req_id]]) for req_id in decode_req_ids] for idx, (req_id, start_offset) in enumerate(zip(decode_req_ids, start_offsets, strict=True)): inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1] - update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}} - self._merge_additional_information_update(req_id, update_dict) + if code_predictor_codes is not None: + update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}} + self._merge_additional_information_update(req_id, update_dict) def _model_forward( self,