diff --git a/benchmarks/tts/model_configs.yaml b/benchmarks/tts/model_configs.yaml index 9acf4e47e3e..2ec7addca35 100644 --- a/benchmarks/tts/model_configs.yaml +++ b/benchmarks/tts/model_configs.yaml @@ -80,6 +80,17 @@ models: task_extra_body: voice_clone: {} + # MOSS-TTS-Local-Transformer-v1.5: Qwen3 backbone + 1-layer GPT2-style local + # depth transformer (n_vq=12), paired with the 48 kHz stereo + # MOSS-Audio-Tokenizer-v2 codec. Same voice-clone request shape as the + # other MOSS-TTS variants. + OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5: + supported_tasks: [voice_clone] + backend: openai-audio-speech + endpoint: /v1/audio/speech + task_extra_body: + voice_clone: {} + # MOSS-TTSD synthesizes [S1]/[S2] dialogue. The seed-tts voice_clone dataset # has single-speaker rows; bench piggybacks on it (both speakers share the # same ref_audio). For real dialogue evaluation use the ttsd dataset. diff --git a/examples/offline_inference/text_to_speech/moss_tts/end2end.py b/examples/offline_inference/text_to_speech/moss_tts/end2end.py index fb47ad8c9d3..51bf52168b3 100644 --- a/examples/offline_inference/text_to_speech/moss_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/moss_tts/end2end.py @@ -1,11 +1,12 @@ """Offline inference example for MOSS-TTS variants. -Covers all five models: - - OpenMOSS-Team/MOSS-TTS (8B, general TTS) - - OpenMOSS-Team/MOSS-TTSD-v1.0 (8B, dialogue TTS) - - OpenMOSS-Team/MOSS-SoundEffect (8B, sound effect synthesis) - - OpenMOSS-Team/MOSS-TTS-Realtime (1.7B, low-latency streaming) - - OpenMOSS-Team/MOSS-VoiceGenerator(1.7B, voice design) +Covers six models: + - OpenMOSS-Team/MOSS-TTS (8B, general TTS) + - OpenMOSS-Team/MOSS-TTSD-v1.0 (8B, dialogue TTS) + - OpenMOSS-Team/MOSS-SoundEffect (8B, sound effect synthesis) + - OpenMOSS-Team/MOSS-TTS-Realtime (1.7B, low-latency streaming) + - OpenMOSS-Team/MOSS-VoiceGenerator (1.7B, voice design) + - OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 (n_vq=12, 48 kHz stereo) Usage ----- @@ -210,12 +211,18 @@ def run_tts(args: argparse.Namespace) -> None: "OpenMOSS-Team/MOSS-SoundEffect": "moss_sound_effect", "OpenMOSS-Team/MOSS-TTS-Realtime": "moss_tts_realtime", "OpenMOSS-Team/MOSS-VoiceGenerator": "moss_voice_generator", + "OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5": "moss_tts_local", } deploy_config = deploy_map.get(args.model, "moss_tts") is_sound_effect = "SoundEffect" in args.model is_voice_gen = "VoiceGenerator" in args.model is_realtime = "Realtime" in args.model + # Local-v1.5 ships its own AutoProcessor (processor_config.json + + # processing_moss_tts.py), so it reuses the same _build_unified_codes + # path as the Delay variants below — just a different n_vq/codec. + is_local = "Local" in args.model + output_sr = 48000 if is_local else 24000 ref_audio = None builder_kwargs: dict = {} @@ -314,27 +321,33 @@ def run_tts(args: argparse.Namespace) -> None: wav = mm.get("model_outputs") if wav is None: continue + # Local-v1.5's codec emits stereo chunks shaped (C, T); keep the + # channel axis intact (don't flatten) so concatenation below stays + # along the time axis. Mono variants stay (T,) as before. if isinstance(wav, list): - non_empty = [w.reshape(-1) for w in wav if isinstance(w, torch.Tensor) and w.numel() > 0] + non_empty = [w for w in wav if isinstance(w, torch.Tensor) and w.numel() > 0] if non_empty and t_first_audio is None: t_first_audio = _time.monotonic() chunks.extend(non_empty) elif isinstance(wav, torch.Tensor) and wav.numel() > 0: if t_first_audio is None: t_first_audio = _time.monotonic() - chunks.append(wav.reshape(-1)) + chunks.append(wav) t_gen_end = _time.monotonic() if not chunks: print("No audio in output — check model loading logs.") return - wav_t = torch.cat(chunks).float().cpu() - sf.write(args.output, wav_t.numpy(), 24000) - print(f"Saved {len(wav_t) / 24000:.2f}s of audio to {args.output}") + cat_dim = -1 if chunks[0].dim() == 2 else 0 + wav_t = torch.cat(chunks, dim=cat_dim).float().cpu() + n_samples = wav_t.shape[-1] + # soundfile wants (T,) mono or (T, C) multi-channel. + sf.write(args.output, (wav_t.transpose(0, 1) if wav_t.dim() == 2 else wav_t).numpy(), output_sr) + print(f"Saved {n_samples / output_sr:.2f}s of audio to {args.output}") ttfa_ms = (t_first_audio - t_gen_start) * 1000.0 if t_first_audio is not None else float("nan") gen_total_ms = (t_gen_end - t_gen_start) * 1000.0 - print(f"TTFA_MS={ttfa_ms:.1f} GEN_TOTAL_MS={gen_total_ms:.1f} AUDIO_S={len(wav_t) / 24000:.2f}") + print(f"TTFA_MS={ttfa_ms:.1f} GEN_TOTAL_MS={gen_total_ms:.1f} AUDIO_S={n_samples / output_sr:.2f}") def main() -> None: diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 0ed06fa4b2a..091240690e5 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -52,6 +52,8 @@ from vllm_omni.metrics import definitions as defs _AUDIO_CONTINUITY_THRESHOLD_ENV = "VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S" +_AUDIO_SAMPLE_RATE_ENV = "VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE" +_AUDIO_CHANNELS_ENV = "VLLM_OMNI_BENCH_AUDIO_CHANNELS" RETURN_STAGE_METRICS_FIELD = "return_stage_metrics" _IMAGE_STAGE_METRICS_BACKENDS = frozenset({"openai-image-edits-omni"}) _PRINT_STAGE = False @@ -109,6 +111,32 @@ def _audio_continuity_threshold_s() -> float: return max(value, 0.0) +def _audio_pcm_format() -> tuple[int, int]: + """Return ``(sample_rate, channels)`` of the streamed PCM response. + + Defaults to 24 kHz mono (Qwen3-TTS / VoxCPM2 / most MOSS-TTS variants). + Override for models with a different codec output format (e.g. + MOSS-TTS-Local-Transformer-v1.5, which streams 48 kHz stereo PCM) via + ``VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE`` / ``VLLM_OMNI_BENCH_AUDIO_CHANNELS`` - + getting this wrong silently skews audio_duration/audio_rtf/underrun stats. + """ + sample_rate = 24000 + channels = 1 + raw_sr = os.environ.get(_AUDIO_SAMPLE_RATE_ENV) + if raw_sr: + try: + sample_rate = int(raw_sr) + except ValueError: + logger.warning("Invalid %s=%r; using default %d", _AUDIO_SAMPLE_RATE_ENV, raw_sr, sample_rate) + raw_ch = os.environ.get(_AUDIO_CHANNELS_ENV) + if raw_ch: + try: + channels = int(raw_ch) + except ValueError: + logger.warning("Invalid %s=%r; using default %d", _AUDIO_CHANNELS_ENV, raw_ch, channels) + return sample_rate, channels + + get_samples_old = datasets.get_samples _DEFAULT_DAILY_OMNI_REPO = "liarliar/Daily-Omni" @@ -1133,10 +1161,10 @@ async def async_request_openai_audio_speech( output = MixRequestFuncOutput() output.prompt_len = request_func_input.prompt_len - # PCM format: 16-bit signed, 24 kHz, mono - sample_rate = 24000 + # PCM format: 16-bit signed; sample_rate/channels are model-dependent + # (see _audio_pcm_format docstring). + sample_rate, channels = _audio_pcm_format() sample_width = 2 # 16-bit = 2 bytes - channels = 1 st = time.perf_counter() output.start_time = st diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index dc5843a2697..69c00489032 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -65,7 +65,11 @@ ) from vllm_omni.model_executor.models.ming_tts.pipeline import MING_TTS_PIPELINE from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO_4_5_PIPELINE -from vllm_omni.model_executor.models.moss_tts.pipeline import MOSS_TTS_PIPELINE, MOSS_TTS_REALTIME_PIPELINE +from vllm_omni.model_executor.models.moss_tts.pipeline import ( + MOSS_TTS_LOCAL_PIPELINE, + MOSS_TTS_PIPELINE, + MOSS_TTS_REALTIME_PIPELINE, +) from vllm_omni.model_executor.models.moss_tts_nano.pipeline import MOSS_TTS_NANO_PIPELINE from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import ( QWEN2_5_OMNI_PIPELINE, @@ -112,6 +116,7 @@ "moss_tts_nano": MOSS_TTS_NANO_PIPELINE, "moss_tts_delay": MOSS_TTS_PIPELINE, "moss_tts_realtime": MOSS_TTS_REALTIME_PIPELINE, + "moss_tts_local": MOSS_TTS_LOCAL_PIPELINE, "minicpmo_4_5": MINICPMO_4_5_PIPELINE, "higgs_audio_v2": HIGGS_AUDIO_V2_PIPELINE, "higgs_multimodal_qwen3": HIGGS_AUDIO_V3_PIPELINE, diff --git a/vllm_omni/deploy/moss_tts_local.yaml b/vllm_omni/deploy/moss_tts_local.yaml new file mode 100644 index 00000000000..7e474caa63e --- /dev/null +++ b/vllm_omni/deploy/moss_tts_local.yaml @@ -0,0 +1,56 @@ +# MOSS-TTS-Local-Transformer-v1.5 (MossTTSLocalModel, n_vq=12, 48 kHz stereo) +# HF: OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 +# Codec: OpenMOSS-Team/MOSS-Audio-Tokenizer-v2 (48 kHz, stereo, channel-interleaved RVQ) +# Uses a 1-layer GPT2-style local depth transformer (no delay warm-up), reset every +# frame — same per-frame timing (80 ms/frame, 12.5 fps) as MOSS-TTS-Realtime. +async_chunk: true + +connectors: + shm: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + codec_streaming: true + 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 + codec_left_context_frames: 0 + +stages: + - stage_id: 0 + max_num_seqs: 8 + gpu_memory_utilization: 0.60 + enable_prefix_caching: false + async_scheduling: true + max_num_batched_tokens: 4096 + max_model_len: 4096 + devices: "0" + output_connectors: + to_stage_1: shm + default_sampling_params: + temperature: 1.0 + top_p: 0.95 + top_k: 50 + max_tokens: 4096 + seed: 42 + repetition_penalty: 1.1 + + - stage_id: 1 + max_num_seqs: 1 + gpu_memory_utilization: 0.15 + enforce_eager: true + enable_prefix_caching: false + async_scheduling: true + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "0" + input_connectors: + from_stage_0: shm + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + repetition_penalty: 1.0 diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 83bb5b1c161..d5e21d419cb 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -41,6 +41,10 @@ def _register_omni_hf_configs() -> None: IndexTTS2Config, ) from vllm_omni.model_executor.models.ming_tts.config_ming_tts import MingDenseConfig + from vllm_omni.model_executor.models.moss_tts.configuration_moss_tts import ( + MossTTSLocalConfig, + MossTTSRealtimeConfig, + ) from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( Qwen3TTSConfig, ) @@ -63,6 +67,8 @@ def _register_omni_hf_configs() -> None: for model_type, config_cls in [ ("dense", MingDenseConfig), ("indextts2", IndexTTS2Config), + ("moss_tts_local", MossTTSLocalConfig), + ("moss_tts_realtime", MossTTSRealtimeConfig), ("qwen3_tts", Qwen3TTSConfig), ("cosyvoice3", CosyVoice3Config), ("glm_tts", GLMTTSConfig), diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index b5b98a3597b..da1418a098b 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -81,6 +81,7 @@ _MING_TTS_MODEL_STAGES = {"ming_tts"} _MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"} _MOSS_TTS_FULL_MODEL_STAGES = {"moss_tts", "moss_tts_codec"} +_MOSS_TTS_LOCAL_MODEL_STAGES = {"moss_tts_local", "moss_tts_local_codec"} _HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"} _HIGGS_V3_TTS_MODEL_STAGES = {"higgs_audio_v3"} _GLM_TTS_MODEL_STAGES = {"glm_tts"} @@ -99,6 +100,7 @@ | _MING_TTS_MODEL_STAGES | _MOSS_TTS_MODEL_STAGES | _MOSS_TTS_FULL_MODEL_STAGES + | _MOSS_TTS_LOCAL_MODEL_STAGES | _GLM_TTS_MODEL_STAGES | _STEP_AUDIO2_TTS_MODEL_STAGES | _INDEXTTS2_TTS_MODEL_STAGES @@ -182,6 +184,18 @@ def _create_wav_header(sample_rate: int, num_channels: int = 1, bits_per_sample: return header +def _infer_audio_num_channels(audio: np.ndarray) -> int: + """Infer channel count before streaming PCM bytes are wrapped as WAV.""" + if audio.ndim == 3 and audio.shape[0] == 1: + audio = audio[0] + if audio.ndim == 2: + if audio.shape[0] in (1, 2): + return int(audio.shape[0]) + if audio.shape[1] in (1, 2): + return int(audio.shape[1]) + return 1 + + def _sanitize_filename(filename: str) -> str: """Sanitize filename to prevent path traversal attacks. @@ -688,6 +702,8 @@ def _detect_tts_model_type(self) -> str | None: return "moss_tts_nano" if model_stage in _MOSS_TTS_FULL_MODEL_STAGES: return "moss_tts" + if model_stage in _MOSS_TTS_LOCAL_MODEL_STAGES: + return "moss_tts" if model_stage in _HIGGS_AUDIO_V2_TTS_MODEL_STAGES: return "higgs_audio_v2" if model_stage in _HIGGS_V3_TTS_MODEL_STAGES: @@ -1675,6 +1691,8 @@ def _detect_moss_variant(self) -> str: name = "" if "realtime" in name: return "realtime" + if "local" in name: + return "local" if "ttsd" in name: return "ttsd" if "soundeffect" in name: @@ -1701,9 +1719,17 @@ def _validate_moss_tts_request(self, request: OpenAICreateSpeechRequest) -> str return "Input text cannot be empty" v = self._moss_variant - if v in (None, "tts", "realtime"): + if v in (None, "tts", "realtime", "local"): if request.ref_audio is None: - label = "MOSS-TTS-Nano" if v is None else ("MOSS-TTS-Realtime" if v == "realtime" else "MOSS-TTS") + label = ( + "MOSS-TTS-Nano" + if v is None + else ( + "MOSS-TTS-Realtime" + if v == "realtime" + else ("MOSS-TTS-Local-Transformer" if v == "local" else "MOSS-TTS") + ) + ) return f"{label} requires 'ref_audio' (reference audio for voice cloning)." return self._validate_ref_audio_format(request.ref_audio) @@ -1806,11 +1832,26 @@ async def _build_moss_tts_params(self, request: OpenAICreateSpeechRequest) -> di params["prompt_audio_array"] = [[wav_list, sr]] return params - # ---- MossTTSDelay family (tts/ttsd/sound_effect/voice_generator): - # call the upstream processor server-side to produce unified codes ---- + # ---- MossTTSDelay family (tts/ttsd/sound_effect/voice_generator) + # and MOSS-TTS-Local-Transformer-v1.5: call the upstream processor + # server-side to produce unified codes. Local-v1.5 ships its own + # AutoProcessor (processor_config.json + processing_moss_tts.py) and + # reuses this exact build_user_message/encode_audios_from_wav path in + # the offline example (examples/.../moss_tts/end2end.py: + # _build_unified_codes) -- it is NOT in the same boat as Realtime + # (no processor_config.json there), so it must not fall back to the + # prompt_audio_array path above (which the talker's preprocess() + # never reads -- info_dict["codes"]["ref"] is the only thing it + # consumes, so skipping this path silently drops all voice-clone + # conditioning and produces unconditioned/garbage audio online). ---- proc = self._get_moss_processor() n_vq = int(getattr(proc.model_config, "n_vq", 32)) - sr_target = int(getattr(proc.model_config, "sampling_rate", 24000)) + # Local-v1.5 encodes reference audio at a fixed 24 kHz working rate + # regardless of its 48 kHz stereo *output* codec -- mirrors the + # offline example's hardcoded encode_audios_from_wav(sampling_rate=24000) + # for this variant; proc.model_config.sampling_rate there is the + # output rate (48000), the wrong value to resample the reference into. + sr_target = 24000 if v == "local" else int(getattr(proc.model_config, "sampling_rate", 24000)) # Reference-audio encoding + speaker caching lives in the model package # (moss_tts.reference_encoder), mirroring Fish Speech / CosyVoice3 / @@ -1837,7 +1878,7 @@ async def _encode_ref(ref_str: str) -> torch.Tensor: ) user_kwargs: dict[str, Any] = {"text": request.input or ""} - if v in ("tts", "realtime"): + if v in ("tts", "local"): user_kwargs["reference"] = [await _encode_ref(request.ref_audio)] elif v == "ttsd": refs = [await _encode_ref(request.ref_audio)] @@ -2654,7 +2695,12 @@ async def _generate_audio_chunks( assert sr_raw is not None, ( "First audio chunk must include sample rate metadata for WAV streaming" ) - wav_header = _create_wav_header(sample_rate=sample_rate_val, num_channels=1, bits_per_sample=16) + num_channels = _infer_audio_num_channels(np.asarray(chunk_np)) + wav_header = _create_wav_header( + sample_rate=sample_rate_val, + num_channels=num_channels, + bits_per_sample=16, + ) yield wav_header first_chunk = False diff --git a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py index 62d0c6568fb..36d98c2d94b 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py @@ -54,5 +54,5 @@ class MossTTSNanoAdapter(_MossTTSAdapterBase): @register_tts_adapter class MossTTSAdapter(_MossTTSAdapterBase): - stage_keys = frozenset({"moss_tts", "moss_tts_codec"}) + stage_keys = frozenset({"moss_tts", "moss_tts_codec", "moss_tts_local", "moss_tts_local_codec"}) name = "moss_tts" diff --git a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py index 89c52485a39..b34b5b94431 100644 --- a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py +++ b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py @@ -45,6 +45,8 @@ def __init__( 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 = 1, + enable_channel_interleave: bool = True, quantizer_type: str = "rlfq", quantizer_kwargs: dict[str, Any] | None = None, **kwargs: Any, @@ -54,6 +56,8 @@ def __init__( 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.encoder_kwargs = encoder_kwargs or _default_encoder_kwargs() self.decoder_kwargs = decoder_kwargs or _default_decoder_kwargs() if quantizer_kwargs is None: @@ -201,13 +205,21 @@ class _Attention(nn.Module): """Causal multi-head self-attention with RoPE, no streaming KV cache.""" def __init__( - self, embed_dim: int, num_heads: int, causal: bool, max_period: float, device=None, dtype=None + self, + embed_dim: int, + num_heads: int, + causal: bool, + max_period: float, + context: int | None = None, + device=None, + dtype=None, ) -> None: super().__init__() self.embed_dim = embed_dim self.num_heads = num_heads self.causal = causal self.max_period = max_period + self.context = context kw = {"device": device, "dtype": dtype} self.in_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=False, **kw) self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False, **kw) @@ -234,7 +246,38 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: qkv = self.in_proj(x).reshape(B, T, 3, H, D).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] # each (B, H, T, D) q, k = _apply_rope(q, k, self.max_period) - out = F.scaled_dot_product_attention(q, k, v, is_causal=self.causal) + if self.context is not None and self.context < T: + # Local-windowed causal attention: query i may only see keys in + # [i - context + 1, i]. Matches upstream's per-stage receptive + # field — wider context windows are not numerically equivalent. + # + # A dense (T, T) mask is O(T^2) memory (e.g. 4 GiB at T=65536) — + # for long utterances decoded in a single non-streaming call this + # OOMs. Process queries in blocks against only their reachable + # key range so peak mask memory is O(block * (block + context)) + # regardless of T. + block = 4096 + if T <= block: + positions = torch.arange(T, device=x.device) + delta = positions.view(-1, 1) - positions.view(1, -1) + mask = (delta >= 0) & (delta < self.context) + out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) + else: + outs = [] + for start in range(0, T, block): + end = min(start + block, T) + k_start = max(0, start - self.context + 1) + q_pos = torch.arange(start, end, device=x.device) + k_pos = torch.arange(k_start, end, device=x.device) + delta = q_pos.view(-1, 1) - k_pos.view(1, -1) + blk_mask = (delta >= 0) & (delta < self.context) + out_blk = F.scaled_dot_product_attention( + q[:, :, start:end], k[:, :, k_start:end], v[:, :, k_start:end], attn_mask=blk_mask + ) + outs.append(out_blk) + out = torch.cat(outs, dim=2) + else: + out = F.scaled_dot_product_attention(q, k, v, is_causal=self.causal) out = out.transpose(1, 2).reshape(B, T, self.embed_dim) return self.out_proj(out) @@ -248,6 +291,7 @@ def __init__( causal: bool, max_period: float, layer_scale: float | None, + context: int | None = None, device=None, dtype=None, ) -> None: @@ -255,7 +299,7 @@ def __init__( kw = {"device": device, "dtype": dtype} self.norm1 = nn.LayerNorm(d_model, eps=1e-5, **kw) self.norm2 = nn.LayerNorm(d_model, eps=1e-5, **kw) - self.attn = _Attention(d_model, num_heads, causal, max_period, device, dtype) + self.attn = _Attention(d_model, num_heads, causal, max_period, context, device, dtype) self.ff1 = nn.Linear(d_model, dim_feedforward, bias=False, **kw) self.ff2 = nn.Linear(dim_feedforward, d_model, bias=False, **kw) ls_kw = cast(dict[str, object], kw) @@ -282,13 +326,16 @@ def __init__( causal: bool, max_period: float, layer_scale: float | None, + context: int | None = None, device=None, dtype=None, ) -> None: super().__init__() self.layers = nn.ModuleList( [ - _TransformerLayer(d_model, num_heads, dim_feedforward, causal, max_period, layer_scale, device, dtype) + _TransformerLayer( + d_model, num_heads, dim_feedforward, causal, max_period, layer_scale, context, device, dtype + ) for _ in range(num_layers) ] ) @@ -307,11 +354,12 @@ def __init__( ) -> None: super().__init__() self.downsample_ratio: int = 1 - self.in_proj = nn.Linear(input_dimension, d_model, bias=False) if d_model != input_dimension else nn.Identity() + # Upstream always materializes a learned Linear here regardless of + # whether input_dimension == d_model — substituting Identity when + # the dims happen to match silently drops a trained projection. + self.in_proj = nn.Linear(input_dimension, d_model, bias=False) self.transformer = _Transformer(d_model=d_model, **kwargs) - self.out_proj = ( - nn.Linear(d_model, output_dimension, bias=False) if d_model != output_dimension else nn.Identity() - ) + self.out_proj = nn.Linear(d_model, output_dimension, bias=False) def forward(self, x: torch.Tensor, lengths: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: x = self.in_proj(x.transpose(1, 2)) # (B, D, T) → (B, T, d_model) @@ -456,10 +504,21 @@ def decode_codes(self, codes: torch.Tensor) -> torch.Tensor: def _build_modules( - specs: list[dict[str, Any]], is_downsample: bool, context_duration: float, sampling_rate: float -) -> nn.ModuleList: + specs: list[dict[str, Any]], is_downsample: bool, default_context_duration: float, frame_rate: float +) -> tuple[nn.ModuleList, float]: + """Build encoder/decoder stage modules. + + ``frame_rate`` tracks the running frames-per-second at the current stage + (updated after each module by its downsample/upsample ratio, mirroring + upstream's ``current_frame_rate``). Each ``Transformer`` module's local + attention window (``context``, in frames) is derived from its own + ``context_duration`` (seconds) at the rate *at that point in the stack* — + not a global constant — so it must be computed here, not by the caller. + Returns the final rate too, since the decoder's starting rate is the + encoder's bottleneck rate, not the raw sampling rate. + """ modules: list[nn.Module] = [] - rate = float(sampling_rate) + rate = float(frame_rate) for spec in specs: spec = dict(spec) if spec["module_type"] == "PatchedPretransform": @@ -471,16 +530,18 @@ def _build_modules( spec.pop("positional_embedding", None) spec.pop("norm", None) spec.pop("causal", None) + context_duration = float(spec.pop("context_duration", default_context_duration)) m = _ProjectedTransformer( module_type="Transformer", causal=True, max_period=spec.pop("max_period", 10000), layer_scale=spec.pop("layer_scale", None), + context=int(round(rate * context_duration)), **spec, ) modules.append(m) rate = rate / m.downsample_ratio if is_downsample else rate * m.downsample_ratio - return nn.ModuleList(modules) + return nn.ModuleList(modules), rate class MossAudioTokenizerModel(PreTrainedModel): @@ -494,12 +555,27 @@ def __init__(self, config: MossAudioTokenizerConfig) -> None: super().__init__(config) self.sampling_rate = config.sampling_rate self.downsample_rate = config.downsample_rate + # Real v1 checkpoints store these as an explicit JSON `null` rather + # than omitting the key, so `getattr(..., default)` doesn't apply — + # normalize None to the same defaults here. + self.number_channels = getattr(config, "number_channels", 1) or 1 + self.enable_channel_interleave = getattr(config, "enable_channel_interleave", True) + if self.enable_channel_interleave is None: + self.enable_channel_interleave = True ctx = config.causal_transformer_context_duration - sr = float(config.sampling_rate) - self.encoder = _build_modules(config.encoder_kwargs, is_downsample=True, context_duration=ctx, sampling_rate=sr) - self.decoder = _build_modules( - copy.deepcopy(config.decoder_kwargs), is_downsample=False, context_duration=ctx, sampling_rate=sr + channel_interleave_factor = ( + self.number_channels if (self.enable_channel_interleave and self.number_channels > 1) else 1 + ) + frame_rate = float(config.sampling_rate) * channel_interleave_factor + self.encoder, bottleneck_rate = _build_modules( + config.encoder_kwargs, is_downsample=True, default_context_duration=ctx, frame_rate=frame_rate + ) + self.decoder, _ = _build_modules( + copy.deepcopy(config.decoder_kwargs), + is_downsample=False, + default_context_duration=ctx, + frame_rate=bottleneck_rate, ) kw = dict(config.quantizer_kwargs) @@ -512,15 +588,29 @@ def batch_encode( wav_list: list[torch.Tensor], num_quantizers: int | None = None, ) -> MossAudioTokenizerEncoderOutput: - """Encode a list of 1-D waveform tensors → RVQ codes (NQ, B, T).""" + """Encode a list of waveform tensors → RVQ codes (NQ, B, T). + + Each element of ``wav_list`` is ``(T,)``/``(1, T)`` for a mono codec + (``number_channels == 1``) or ``(number_channels, T)`` otherwise. + """ device = wav_list[0].device B = len(wav_list) - max_len = max(w.shape[-1] for w in wav_list) - x = torch.zeros(B, 1, max_len, device=device) + C = self.number_channels + normalized: list[torch.Tensor] = [] lengths = torch.zeros(B, device=device, dtype=torch.long) for i, w in enumerate(wav_list): - x[i, 0, : w.shape[-1]] = w - lengths[i] = w.shape[-1] + if C == 1: + w_i = w.unsqueeze(0) if w.dim() == 1 else w + else: + if w.dim() != 2 or w.shape[0] != C: + raise ValueError(f"Expected wav_list[{i}] to have shape ({C}, T), got {tuple(w.shape)}.") + w_i = w + normalized.append(w_i) + lengths[i] = w_i.shape[-1] + max_len = int(lengths.max().item()) + x = torch.zeros(B, C, max_len, device=device, dtype=wav_list[0].dtype) + for i, w_i in enumerate(normalized): + x[i, :, : w_i.shape[-1]] = w_i return self._encode(x, lengths, num_quantizers) @torch.no_grad() @@ -541,6 +631,45 @@ def batch_decode( lengths[i] = c.shape[-1] return self._decode(codes, lengths) + def _flatten_channels_for_codec( + self, + input_values: torch.Tensor, + input_lengths: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Pad to a ``downsample_rate`` multiple, then interleave channels into one stream. + + Mirrors upstream's ``_flatten_channels_for_codec``: for a stereo + codec, ``(B, C, T)`` becomes ``(B, 1, T*C)`` with samples interleaved + frame-by-frame across channels, so the mono encoder/decoder stack can + process a multi-channel waveform unchanged. + """ + 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]: + """Inverse of ``_flatten_channels_for_codec``: de-interleave back to ``(B, C, T)``.""" + 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 + def _encode( self, x: torch.Tensor, @@ -549,9 +678,7 @@ def _encode( ) -> MossAudioTokenizerEncoderOutput: if x.dim() == 2: x = x.unsqueeze(1) - T = x.shape[-1] - if T % self.downsample_rate: - x = F.pad(x, (0, self.downsample_rate - T % self.downsample_rate)) + x, lengths = self._flatten_channels_for_codec(x, lengths) e, e_len = x, lengths for m in self.encoder: e, e_len = m(e, e_len) @@ -567,6 +694,7 @@ def _decode( d, d_len = z, lengths for m in self.decoder: d, d_len = m(d, d_len) + d, d_len = self._restore_channels_from_codec(d, d_len) return MossAudioTokenizerDecoderOutput(audio=d, audio_lengths=d_len) diff --git a/vllm_omni/model_executor/models/moss_tts/configuration_moss_tts.py b/vllm_omni/model_executor/models/moss_tts/configuration_moss_tts.py index 46147f27c97..42f103f4221 100644 --- a/vllm_omni/model_executor/models/moss_tts/configuration_moss_tts.py +++ b/vllm_omni/model_executor/models/moss_tts/configuration_moss_tts.py @@ -17,6 +17,7 @@ from transformers import AutoConfig from transformers.configuration_utils import PretrainedConfig +from transformers.models.gpt2.configuration_gpt2 import GPT2Config from transformers.models.qwen3 import Qwen3Config @@ -235,7 +236,80 @@ def get_text_config(self, **_: object) -> Qwen3Config: return self.language_config +class MossTTSLocalConfig(PretrainedConfig): + """Config for MossTTSLocalModel (MOSS-TTS-Local-Transformer-v1.5). + + Like ``MossTTSDelayConfig``, the HF checkpoint stores a nested Qwen3 config + (``qwen3_config``) which we unwrap via ``get_text_config()``. It also stores + a nested GPT2-style config (``gpt2_config``) describing the single-layer + local depth transformer that decodes the ``n_vq`` audio codebooks per frame + (see ``modeling_moss_tts_local_depth.py``). + """ + + model_type = "moss_tts_local" + + def __init__( + self, + qwen3_config: dict | None = None, + gpt2_config: dict | None = None, + n_vq: int = 12, + audio_vocab_size: int = 1024, + audio_pad_token_id: int = 1024, + pad_token_id: int = 151643, + im_start_token_id: int = 151644, + im_end_token_id: int = 151645, + audio_start_token_id: int = 151669, + audio_end_token_id: int = 151670, + audio_user_slot_token_id: int = 151654, + audio_assistant_slot_token_id: int = 151656, + sampling_rate: int = 48000, + audio_tokenizer_name_or_path: str = "OpenMOSS-Team/MOSS-Audio-Tokenizer-v2", + local_text_head_mode: str = "binary", + **kwargs: object, + ) -> None: + if qwen3_config is None: + qwen3_config = {} + if isinstance(qwen3_config, dict): + qwen3_config = dict(qwen3_config) + qwen3_config.pop("model_type", None) + self.qwen3_config = Qwen3Config(**qwen3_config) + else: + self.qwen3_config = qwen3_config + + if isinstance(gpt2_config, dict): + gpt2_config = dict(gpt2_config) + gpt2_config.pop("model_type", None) + self.gpt2_config = GPT2Config(**gpt2_config) + else: + self.gpt2_config = gpt2_config + + super().__init__(pad_token_id=pad_token_id, **kwargs) + + self.n_vq = n_vq + self.audio_vocab_size = audio_vocab_size + self.audio_pad_token_id = audio_pad_token_id + self.audio_pad_code = audio_pad_token_id + self.im_start_token_id = im_start_token_id + self.im_end_token_id = im_end_token_id + self.audio_start_token_id = audio_start_token_id + self.audio_end_token_id = audio_end_token_id + self.audio_user_slot_token_id = audio_user_slot_token_id + self.audio_assistant_slot_token_id = audio_assistant_slot_token_id + self.sampling_rate = sampling_rate + self.audio_tokenizer_name_or_path = audio_tokenizer_name_or_path + self.local_text_head_mode = local_text_head_mode + + def get_text_config(self, **_: object) -> Qwen3Config: + return self.qwen3_config + + AutoConfig.register("moss_tts_delay", MossTTSDelayConfig) AutoConfig.register("moss_tts_realtime", MossTTSRealtimeConfig) - -__all__ = ["MossTTSDelayConfig", "MossTTSRealtimeConfig", "MossTTSLocalTransformerConfig"] +AutoConfig.register("moss_tts_local", MossTTSLocalConfig) + +__all__ = [ + "MossTTSDelayConfig", + "MossTTSRealtimeConfig", + "MossTTSLocalTransformerConfig", + "MossTTSLocalConfig", +] 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 b1b6063d47c..7c5d47653f2 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,7 +5,6 @@ from __future__ import annotations -import os from collections.abc import Iterable from typing import Any @@ -42,10 +41,8 @@ class MossTTSCodecDecoder(nn.Module): ``n_vq=32`` (MOSS-TTS) or ``n_vq=16`` (all other variants) without swapping weights. - The codec checkpoint path defaults to the value stored in - ``vllm_config.model_config.hf_config.codec_model_name_or_path`` but can - be overridden by setting the environment variable - ``MOSS_TTS_CODEC_PATH``. + The codec checkpoint path comes from + ``vllm_config.model_config.hf_config.codec_model_name_or_path``. """ input_modalities = "audio" @@ -64,11 +61,19 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: cfg = vllm_config.model_config.hf_config self._n_vq: int = int(getattr(cfg, "n_vq", getattr(cfg, "rvq", 16))) - self._codec_path: str = str(getattr(cfg, "codec_model_name_or_path", "OpenMOSS-Team/MOSS-Audio-Tokenizer")) + self._codec_path: str = str( + getattr( + cfg, + "codec_model_name_or_path", + getattr(cfg, "audio_tokenizer_name_or_path", "OpenMOSS-Team/MOSS-Audio-Tokenizer"), + ) + ) - # Resolved at load_weights() time + # Resolved at load_weights() time, once the codec checkpoint's own + # config (sampling rate, channel count) is known. self._codec: MossAudioTokenizerModel | None = None self._cuda_graph_wrapper: MossTTSCUDAGraphCodecWrapper | None = None + self._n_channels: int = 1 self._sr_tensor = torch.tensor(self._OUTPUT_SAMPLE_RATE, dtype=torch.int32) # ------------------------------------------------------------------ @@ -119,7 +124,7 @@ def forward( """ sr_tensor = self._sr_tensor empty = torch.zeros((0,), dtype=torch.float32) - info_list: list[dict[str, Any]] = runtime_additional_information or [] + info_list: list[dict[str, Any]] = runtime_additional_information or [{}] num_req = max(len(info_list), 1) if self._codec is None: @@ -147,7 +152,7 @@ def forward( # kwargs if available, else assume one request. ids_flat = input_ids.reshape(-1).to(dtype=torch.long) num_scheduled_tokens = kwargs.get("num_scheduled_tokens") - if isinstance(num_scheduled_tokens, list) and len(num_scheduled_tokens) == num_req: + 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)) @@ -176,12 +181,8 @@ def forward( codebook_size = self._codec.config.codebook_size codes_nq_t = codes_nq_t.clamp_(0, int(codebook_size) - 1) - # left_context_size: number of left-context frames prepended to - # this chunk by the stage input processor (overlap-add pattern, - # same as Qwen3-TTS code2wav). We decode them together with the - # real chunk for smooth boundaries, then trim them from the output. meta = (info.get("meta", {}) if isinstance(info, dict) else {}) or {} - left_ctx = meta.get("left_context_size", 0) + left_ctx = self._runtime_value(info, meta, "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): @@ -196,28 +197,40 @@ def forward( if out.audio is None: continue - wav = out.audio.reshape(-1).to(dtype=torch.float32).cpu() + # ``out.audio`` is ``(1, C, T)``; keep the channel axis for + # stereo codecs (Local-v1.5) and flatten to ``(T,)`` for mono + # ones (Delay/Realtime) to preserve their existing output shape. + wav = out.audio[0].to(dtype=torch.float32).cpu() if out.audio_lengths is not None: - wav = wav[: int(out.audio_lengths[0].item())] + wav = wav[..., : int(out.audio_lengths[0].item())] - # Trim left-context samples. + # Trim left-context samples (per-channel sample axis, so the + # trim amount is identical for mono and interleaved-stereo). if left_ctx > 0: - trim = min(left_ctx * self._codec.downsample_rate, wav.shape[0]) + trim = min(left_ctx * self._codec.downsample_rate, wav.shape[-1]) if trim < left_ctx * self._codec.downsample_rate: logger.warning( "left_ctx trim (%d samples) exceeds wav length (%d); returning empty audio.", left_ctx * self._codec.downsample_rate, - wav.shape[0], + wav.shape[-1], ) - wav = wav[trim:] + wav = wav[..., trim:] - audios[i] = wav + audios[i] = wav.reshape(-1) if self._n_channels == 1 else wav return OmniOutput( text_hidden_states=None, multimodal_outputs={"model_outputs": audios, "sr": srs}, ) + @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] + return default + # ------------------------------------------------------------------ # Weight loading # ------------------------------------------------------------------ @@ -234,11 +247,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: for _ in weights: pass - codec_path = os.environ.get("MOSS_TTS_CODEC_PATH", self._codec_path) + codec_path = self._codec_path logger.info("Loading MOSS Audio Tokenizer from %s", codec_path) - codec_cfg = MossAudioTokenizerConfig.from_pretrained(codec_path) - codec = MossAudioTokenizerModel(codec_cfg) + codec_cfg, codec = self._build_codec(codec_path) model_loader = DefaultModelLoader(self.vllm_config.load_config) source = DefaultModelLoader.Source( @@ -255,10 +267,19 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # and the rest stay at their random init, which produces noise that # sounds correct in duration but is structurally garbage. _SUFFIX_REMAP: list[tuple[str, str]] = [ + # v1 (MOSS-Audio-Tokenizer) naming. (".self_attn.in_projs.0.", ".attn.in_proj."), (".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."), (".layer_scale_1.", ".ls1."), (".layer_scale_2.", ".ls2."), (".input_proj.", ".in_proj."), @@ -295,11 +316,14 @@ def _remap(name: str) -> str: codec.to(device=device, dtype=torch.float32) codec.eval() self._codec = codec + self._n_channels = int(getattr(codec_cfg, "number_channels", 1) or 1) + self._sr_tensor = torch.tensor(int(codec_cfg.sampling_rate), dtype=torch.int32) logger.info( - "MOSS Audio Tokenizer loaded: sampling_rate=%d, n_vq=%d", + "MOSS Audio Tokenizer loaded: sampling_rate=%d, n_vq=%d, n_channels=%d", codec_cfg.sampling_rate, codec_cfg.num_quantizers, + self._n_channels, ) self._maybe_enable_decoder_cudagraph(device) @@ -310,6 +334,24 @@ def _remap(name: str) -> str: # mirror that here. return {f"_codec.{name}" for name, _ in codec.named_parameters()} + 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) + return codec_cfg, codec + except Exception: + logger.exception( + "Failed to instantiate official MOSS Audio Tokenizer via HF remote code; " + "falling back to vendored codec." + ) + + codec_cfg = MossAudioTokenizerConfig.from_pretrained(codec_path) + codec = MossAudioTokenizerModel(codec_cfg) + return codec_cfg, codec + def _maybe_enable_decoder_cudagraph(self, device: torch.device) -> None: """Capture CUDA Graphs for the codec decoder if enforce_eager is False.""" if getattr(self.vllm_config.model_config, "enforce_eager", True): 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 new file mode 100644 index 00000000000..cb2e95551b5 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py @@ -0,0 +1,223 @@ +# Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). +"""Per-frame depth transformer for MossTTSLocalModel (MOSS-TTS-Local-Transformer-v1.5). + +A 1-layer GPT2-style block that decodes the ``n_vq`` audio codebook codes for +one audio frame, run inside the talker's ``make_omni_output`` independent of +vLLM's main scheduler -- mirrors ``MossTTSRealtimeLocalTransformer`` +(``modeling_moss_tts_local.py``) in role, but the algorithm and numerics are +faithful to the official ``gpt2_decoder.py`` / ``modeling_moss_tts.py`` +(GPT2-style LayerNorm + bias, SiLU MLP, **interleaved/GPT-J-style RoPE** -- +not vLLM's neox-style concat-half rotation) rather than Realtime's +Qwen3-style ``CodePredictorBaseModel``. + +Its KV cache and positions reset to 0 every audio frame: there is no +cross-frame state. Submodule names (``h.0.*`` / ``ln_f``) match the +checkpoint 1:1 so ``load_weights()`` needs no remapping. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local import ( + _sample_token, +) + + +class _MossTTSLocalAttention(nn.Module): + """GPT2-style fused-QKV self-attention with interleaved RoPE. + + Faithful to the official ``MossTTSNanoGPT2Attention``: ``rotate_half`` + operates on even/odd index pairs (GPT-J style), and ``cos``/``sin`` are + built via ``repeat_interleave(2, dim=-1)`` rather than the neox-style + concat-half construction vLLM uses elsewhere. + """ + + def __init__(self, hidden_size: int, n_head: int, rope_base: float) -> None: + super().__init__() + if hidden_size % n_head != 0: + raise ValueError(f"hidden_size={hidden_size} must be divisible by n_head={n_head}") + self.n_head = n_head + self.head_dim = hidden_size // n_head + self.embed_dim = hidden_size + self.c_attn = nn.Linear(hidden_size, 3 * hidden_size, bias=True) + self.c_proj = nn.Linear(hidden_size, hidden_size, bias=True) + inv_freq = 1.0 / (rope_base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32) / self.head_dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def _rope_cos_sin( + self, seq_len: int, device: torch.device, dtype: torch.dtype + ) -> tuple[torch.Tensor, torch.Tensor]: + position_ids = torch.arange(seq_len, device=device, dtype=torch.float32) + freqs = torch.einsum("s,d->sd", position_ids, self.inv_freq.to(device=device)) + cos = freqs.cos().repeat_interleave(2, dim=-1).to(dtype) + sin = freqs.sin().repeat_interleave(2, dim=-1).to(dtype) + return cos.view(1, seq_len, 1, self.head_dim), sin.view(1, seq_len, 1, self.head_dim) + + @staticmethod + def _rotate_half(x: torch.Tensor) -> torch.Tensor: + even = x[..., ::2] + odd = x[..., 1::2] + return torch.stack((-odd, even), dim=-1).reshape_as(x) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """``hidden_states``: ``(B, S, H)``. Re-prefills with a fresh causal mask over ``[0, S)`` every call.""" + batch_size, seq_len, _ = hidden_states.shape + qkv = self.c_attn(hidden_states) + query, key, value = qkv.split(self.embed_dim, dim=-1) + query = query.view(batch_size, seq_len, self.n_head, self.head_dim) + key = key.view(batch_size, seq_len, self.n_head, self.head_dim) + value = value.view(batch_size, seq_len, self.n_head, self.head_dim) + + cos, sin = self._rope_cos_sin(seq_len, hidden_states.device, hidden_states.dtype) + query = query * cos + self._rotate_half(query) * sin + key = key * cos + self._rotate_half(key) * sin + + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + attn_output = F.scaled_dot_product_attention(query, key, value, is_causal=True) + attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.embed_dim) + return self.c_proj(attn_output) + + +class _MossTTSLocalMLP(nn.Module): + def __init__(self, hidden_size: int, inner_size: int) -> None: + super().__init__() + self.fc_in = nn.Linear(hidden_size, inner_size, bias=True) + self.fc_out = nn.Linear(inner_size, hidden_size, bias=True) + self.act = nn.SiLU() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.fc_out(self.act(self.fc_in(hidden_states))) + + +class _MossTTSLocalBlock(nn.Module): + def __init__(self, hidden_size: int, n_head: int, inner_size: int, rope_base: float, eps: float) -> None: + super().__init__() + self.ln_1 = nn.LayerNorm(hidden_size, eps=eps) + self.attn = _MossTTSLocalAttention(hidden_size, n_head, rope_base) + self.ln_2 = nn.LayerNorm(hidden_size, eps=eps) + self.mlp = _MossTTSLocalMLP(hidden_size, inner_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = hidden_states + self.attn(self.ln_1(hidden_states)) + hidden_states = hidden_states + self.mlp(self.ln_2(hidden_states)) + return hidden_states + + +class MossTTSLocalDepthTransformer(nn.Module): + """Per-frame depth transformer for MOSS-TTS-Local-Transformer-v1.5. + + Per frame: + - position 0's input is the backbone's last hidden state; its output + feeds BOTH the binary continue/stop head (``local_text_lm_head``) + and codebook-0's head (``audio_lm_heads[0]``) simultaneously. + - codebooks 1..n_vq-1 are sampled sequentially: each sampled code is + re-embedded (``audio_embeddings[c]``) and appended as the next + position, re-prefilling the block over the growing (<=n_vq) sequence + with a fresh causal mask each call -- mathematically identical to + incremental KV-cache decoding since attention is strictly causal and + only the last position is ever read, but avoids any cache plumbing + given the trivially short sequence length. + """ + + def __init__(self, gpt2_config, hidden_size: int | None = None) -> None: + super().__init__() + self.hidden_size = int(hidden_size if hidden_size is not None else gpt2_config.n_embd) + n_head = int(gpt2_config.n_head) + inner_size = int(gpt2_config.n_inner) + eps = float(getattr(gpt2_config, "layer_norm_epsilon", 1e-5)) + rope_base = float(getattr(gpt2_config, "rope_base", 1_000_000.0)) + self.h = nn.ModuleList([_MossTTSLocalBlock(self.hidden_size, n_head, inner_size, rope_base, eps)]) + self.ln_f = nn.LayerNorm(self.hidden_size, eps=eps) + + def _forward_prefix(self, seq_embeds: torch.Tensor) -> torch.Tensor: + hidden_states = seq_embeds + for block in self.h: + hidden_states = block(hidden_states) + return self.ln_f(hidden_states) + + @torch.no_grad() + def generate_frame( + self, + backbone_last_hidden: torch.Tensor, # (B, H) + audio_lm_heads: nn.ModuleList, # n_vq x Linear(H -> audio_vocab_size) + audio_embeddings: nn.ModuleList, # n_vq x Embedding(audio_vocab_size, H) + local_text_lm_head: nn.Module, # Linear(H -> 2): [continue, stop] + *, + n_vq: int, + do_sample: bool = True, + temperature: float = 1.0, + top_k: int = 50, + top_p: float = 1.0, + text_temperature: float = 1.0, + text_top_k: int = 50, + text_top_p: float = 1.0, + repetition_penalty: float = 1.0, + history_per_codebook: list[list[int]] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Generate one audio frame for batch B. + + Returns ``(should_continue, codes)``: ``should_continue`` is a + ``(B,)`` bool tensor (``True`` iff the binary head picked the + "continue" candidate, i.e. logits index 0); ``codes`` is a + ``(B, n_vq)`` LongTensor of sampled codebook indices. + + ``history_per_codebook[c]`` is a list of recently-emitted token ids + for codebook ``c``; when ``repetition_penalty != 1.0`` those tokens' + logits get scaled down before sampling (mirrors upstream's + ``_apply_repetition_penalty``). + """ + batch_size = backbone_last_hidden.shape[0] + dtype = self.ln_f.weight.dtype + + embeds = backbone_last_hidden.new_zeros((batch_size, n_vq, self.hidden_size), dtype=dtype) + embeds[:, 0, :] = backbone_last_hidden.to(dtype) + + hidden = self._forward_prefix(embeds[:, :1, :]) + 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) + should_continue = binary_choice.eq(0) + import os as _os + + if _os.environ.get("MOSS_TTS_DEBUG_STOP"): + import logging as _logging + + _logging.getLogger("moss_tts_debug").warning( + "binary_logits=%s choice=%s", binary_logits.tolist(), binary_choice.tolist() + ) + + codes = backbone_last_hidden.new_zeros((batch_size, n_vq), dtype=torch.long) + for channel_index in range(n_vq): + channel_logits = audio_lm_heads[channel_index](local_hidden).float() + if ( + repetition_penalty != 1.0 + and history_per_codebook is not None + and channel_index < len(history_per_codebook) + ): + hist = history_per_codebook[channel_index] + if hist: + hist_t = torch.tensor(hist, dtype=torch.long, device=channel_logits.device) + sel = channel_logits.index_select(-1, hist_t) + 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) + codes[:, channel_index] = channel_token + + if channel_index + 1 < n_vq: + embeds[:, channel_index + 1, :] = audio_embeddings[channel_index](channel_token).to(dtype) + hidden = self._forward_prefix(embeds[:, : channel_index + 2, :]) + local_hidden = hidden[:, channel_index + 1, :] + + return should_continue, codes + + +__all__ = ["MossTTSLocalDepthTransformer"] 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 d32af7c9020..2b520336cbc 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 @@ -6,7 +6,7 @@ from __future__ import annotations import copy -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any import torch @@ -22,11 +22,15 @@ from vllm_omni.model_executor.models.moss_tts.configuration_moss_tts import ( MossTTSDelayConfig, + MossTTSLocalConfig, MossTTSRealtimeConfig, ) from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local import ( MossTTSRealtimeLocalTransformer, ) +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import ( + MossTTSLocalDepthTransformer, +) from vllm_omni.model_executor.models.output_templates import OmniOutput logger = init_logger(__name__) @@ -36,6 +40,40 @@ def _maybe_prefix(prefix: str, name: str) -> str: return f"{prefix}.{name}" if prefix else name +def _iter_state_row_spans( + states: Sequence[dict[str, Any]], + spans: Sequence[tuple[int, int]] | None, + num_rows: int, +) -> Iterable[tuple[dict[str, Any], int, int]]: + """Yield state/logit-row ranges without assuming equal request row counts.""" + if num_rows <= 0 or not states: + return + + if spans is not None: + if len(spans) != len(states): + raise RuntimeError(f"request_token_spans has {len(spans)} entries for {len(states)} request states") + if all(row_end <= num_rows for _, row_end in spans): + for state, (row_start, row_end) in zip(states, spans, strict=False): + if row_start < row_end: + yield state, int(row_start), int(row_end) + return + if num_rows != len(states): + raise RuntimeError( + "request_token_spans describe full hidden rows, but current logits " + f"have {num_rows} rows for {len(states)} states and are not one-row-per-request" + ) + + if num_rows == len(states): + for i, state in enumerate(states): + yield state, i, i + 1 + return + + raise RuntimeError( + "MOSS-TTS compute_logits requires request_token_spans or one-row-per-request logits; " + f"got {num_rows} rows for {len(states)} request states" + ) + + # --------------------------------------------------------------------------- # MossTTSDelayTalkerForGeneration # --------------------------------------------------------------------------- @@ -209,14 +247,13 @@ def compute_logits( n_vq = self.n_vq device = logits.device vocab_size = logits.shape[-1] - rows_per_state = max(1, logits.shape[0] // max(1, len(states))) - for i, state in enumerate(states): - if not isinstance(state, dict): - continue - row_start = i * rows_per_state - row_end = min(row_start + rows_per_state, logits.shape[0]) - if row_start >= row_end: + for state, row_start, row_end in _iter_state_row_spans( + states, + getattr(self, "_batch_state_spans", None), + logits.shape[0], + ): + if state is None: continue row = logits[row_start:row_end] @@ -569,6 +606,7 @@ def make_omni_output( """ if isinstance(model_outputs, OmniOutput): self._batch_state = None + self._batch_state_spans = None return model_outputs hidden = model_outputs # (S, H) @@ -590,11 +628,11 @@ def make_omni_output( # Stash the per-row state for compute_logits to apply masks. logits # rows align with hidden rows, which align with info_dicts in order. 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") # Per-request (start, end) hidden-row spans from the runner. In mixed # prefill+decode steps the per-request token counts differ, so an equal - # rows-per-request split would sample codes from the wrong rows; fall - # back to that split only when spans are unavailable. + # rows-per-request split would sample codes from the wrong rows. spans = kwargs.get("request_token_spans") # One accumulated-codes tensor per request, in batch order. The generic @@ -607,18 +645,18 @@ def make_omni_output( if hidden.numel() > 0: num_rows = hidden.shape[0] - rows_per_req = max(1, num_rows // max(1, len(info_dicts) or 1)) + 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): per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long)) continue - if spans is not None and i < len(spans): - row_start, row_end = spans[i] - row_end = min(int(row_end), num_rows) - else: - row_start = i * rows_per_req - row_end = min(row_start + rows_per_req, num_rows) + row_start, row_end = spans[i] + row_end = min(int(row_end), num_rows) if row_start >= row_end: per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long)) continue @@ -864,13 +902,12 @@ def compute_logits( logits = hidden_states.new_full((B, V), float("-inf")) states = self._batch_state or [] - rows_per_state = max(1, B // max(1, len(states) or 1)) - for i, state in enumerate(states): - r0 = i * rows_per_state - r1 = min(r0 + rows_per_state, B) - if r0 >= r1: - continue - if not isinstance(state, dict): + for state, r0, r1 in _iter_state_row_spans( + states, + getattr(self, "_batch_state_spans", None), + B, + ): + if state is None: logits[r0:r1, self.text_pad_id] = 0.0 continue if state.get("is_stopping"): @@ -991,6 +1028,7 @@ def make_omni_output( ) -> OmniOutput: if isinstance(model_outputs, OmniOutput): self._batch_state = None + self._batch_state_spans = None return model_outputs hidden = model_outputs # (S, H) @@ -1004,9 +1042,10 @@ def make_omni_output( info["audio_state"] = {"is_stopping": False, "step": 0} 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") - # See delay talker: real per-request row spans from the runner, with the - # equal-split as fallback when unavailable. + # See delay talker: real per-request row spans from the runner are + # required because mixed prefill+decode steps have unequal row counts. spans = kwargs.get("request_token_spans") # Per-request accumulated codes in batch order, pre-filled with empty # placeholders so every skip path (stopped / eos / bos / pad) keeps @@ -1016,16 +1055,16 @@ def make_omni_output( have_codes = False if hidden.numel() > 0 and info_dicts: num_rows = hidden.shape[0] - rows_per_req = max(1, num_rows // max(1, len(info_dicts) or 1)) + 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 - if spans is not None and i < len(spans): - row_start, row_end = spans[i] - row_end = min(int(row_end), num_rows) - else: - row_start = i * rows_per_req - row_end = min(row_start + rows_per_req, num_rows) + row_start, row_end = spans[i] + row_end = min(int(row_end), num_rows) if row_start >= row_end: continue @@ -1189,7 +1228,430 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loaded +# --------------------------------------------------------------------------- +# MossTTSLocalTalkerForGeneration +# --------------------------------------------------------------------------- + + +class MossTTSLocalTalkerForGeneration(nn.Module): + """Stage-0 talker for MossTTSLocalModel (MOSS-TTS-Local-Transformer-v1.5). + + Architecture differs from both existing variants: + + * Backbone (Qwen3; checkpoint prefix ``transformer.*``, unlike Delay/ + Realtime's ``language_model.*``) consumes additive text+audio fusion + the same way Delay does (real text embeddings + masked per-codebook + audio embeddings from ``info["codes"]["ref"]`` / ``info["audio_codes"] + ["current"]``), but there is no delay-pattern countdown: the prompt + already ends right at the audio_start boundary, and every step from + the first decode onward drives exactly one audio frame. + * The checkpoint's own ``text_lm_head`` is loaded (trained for a + teacher-forced text loss) but never consulted at inference: the + official reference's own ``generate()`` hardcodes the next + text-channel input to ``audio_assistant_slot_token_id`` regardless of + what ``text_lm_head`` would predict, because + ``local_text_head_mode == "binary"`` for this checkpoint. + ``compute_logits`` therefore synthesises a one-hot logit row + (Realtime-style), forced to ``audio_assistant_slot_token_id`` while + continuing or ``im_end_token_id`` once the local transformer's binary + head decides to stop. + * 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 + condition: the local transformer's binary head selects the + "audio_end" candidate (logits index 1). + """ + + have_multimodal_outputs: bool = True + has_preprocess: bool = True + has_postprocess: bool = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.vllm_config = vllm_config + self.config: MossTTSLocalConfig = vllm_config.model_config.hf_config + + qwen3_cfg = self.config.qwen3_config + hidden_size = int(qwen3_cfg.hidden_size) + self.hidden_size = hidden_size + self.n_vq: int = int(self.config.n_vq) + self.audio_vocab_size: int = int(self.config.audio_vocab_size) + self.audio_pad_token_id: int = int(self.config.audio_pad_token_id) + self.text_vocab_size: int = int(qwen3_cfg.vocab_size) + + self.audio_assistant_slot_token_id: int = int(self.config.audio_assistant_slot_token_id) + self.im_end_token_id: int = int(self.config.im_end_token_id) + + # Qwen3 backbone. get_text_config() is sufficient for vLLM to size + # KV cache / heads correctly, as long as AutoConfig resolves to the + # local MossTTSLocalConfig rather than the HF remote-code class. + self.model = Qwen3Model( + vllm_config=vllm_config, + prefix=_maybe_prefix(prefix, "model"), + ) + + # Real text LM head -- loaded from the checkpoint but never consulted + # at inference (see class docstring); kept purely as a destination + # for ``text_lm_head.weight`` so load_weights() loads every tensor. + self.text_lm_head = ParallelLMHead( + self.text_vocab_size, + hidden_size, + bias=False, + prefix=_maybe_prefix(prefix, "text_lm_head"), + ) + self.logits_processor = LogitsProcessor(self.text_vocab_size) + + # Per-codebook audio embeddings/heads. audio_pad_token_id (== + # audio_vocab_size) is outside the embedding table -- masked to zero + # contribution rather than stored as a 1025th row (mirrors upstream + # ``_build_inputs_embeds``'s ``valid_mask`` masking). + self.audio_embeddings = nn.ModuleList( + [nn.Embedding(self.audio_vocab_size, hidden_size) for _ in range(self.n_vq)] + ) + self.audio_lm_heads = nn.ModuleList( + [nn.Linear(hidden_size, self.audio_vocab_size, bias=False) for _ in range(self.n_vq)] + ) + self.local_text_lm_head = nn.Linear(hidden_size, 2, bias=False) + self.local_transformer = MossTTSLocalDepthTransformer(self.config.gpt2_config, hidden_size=hidden_size) + + self._batch_state: list[dict[str, Any]] | None = None + # 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.gpu_resident_buffer_keys: set[tuple[str, str]] = { + ("audio_codes", "current"), + ("audio_codes", "accumulated"), + ("hidden_states", "last"), + } + + # ------------------------------------------------------------------ + # vLLM hooks + # ------------------------------------------------------------------ + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + return self.model.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **_: Any, + ) -> torch.Tensor | IntermediateTensors: + return self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor | OmniOutput, + sampling_metadata: SamplingMetadata | None = None, + ) -> torch.Tensor | None: + """Synthesise a one-hot text logit row per request. + + The text channel is never freely generated here: the upstream + reference hardcodes it to ``audio_assistant_slot_token_id`` while + audio frames keep coming, and only emits ``im_end_token_id`` once the + local transformer's binary head decides to stop (see class + docstring for why the real ``text_lm_head`` is bypassed). + """ + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hidden_states is None or hidden_states.numel() == 0: + return None + + num_rows = hidden_states.shape[0] + logits = hidden_states.new_full((num_rows, self.text_vocab_size), float("-inf")) + + states = self._batch_state or [] + for state, row_start, row_end in _iter_state_row_spans( + states, + getattr(self, "_batch_state_spans", None), + num_rows, + ): + if state is not None and state.get("is_stopping"): + 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 + if not states: + logits[:, self.audio_assistant_slot_token_id] = 0.0 + return logits + + # ------------------------------------------------------------------ + # Embedding (text + audio codebooks, additive) + # ------------------------------------------------------------------ + + def _audio_embed(self, codes: torch.Tensor) -> torch.Tensor: + """``codes``: ``(T, n_vq)`` long, possibly containing ``audio_pad_token_id`` + (outside the valid embedding range). Returns the ``(T, H)`` additive + audio embedding, masking out pad positions per codebook. + """ + device = codes.device + valid_mask = codes.ne(self.audio_pad_token_id) # (T, n_vq) + safe_codes = codes.masked_fill(~valid_mask, 0).clamp(0, self.audio_vocab_size - 1) + if self._stacked_audio_emb_w is not None: + n_vq_idx = torch.arange(self.n_vq, device=device) + codes_t = safe_codes.t() # (n_vq, T) + gathered = self._stacked_audio_emb_w[n_vq_idx.unsqueeze(1), codes_t] # (n_vq, T, H) + gathered = gathered * valid_mask.t().unsqueeze(-1) + return gathered.sum(0) + embeds: torch.Tensor | None = None + for i, emb_layer in enumerate(self.audio_embeddings): + col = emb_layer(safe_codes[:, i]) * valid_mask[:, i].unsqueeze(-1) + embeds = col if embeds is None else embeds + col + return embeds + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Build per-step input embeddings (text + audio additive fusion). + + Prefill: initialise the per-request ``{is_stopping, step}`` state. + Decode: combine the forced text token (from ``compute_logits``) with + the audio codes the local transformer produced last step. + """ + device = input_ids.device + span_len = int(input_ids.shape[0]) + audio_state = info_dict.get("audio_state") + is_first_call = not isinstance(audio_state, dict) + + if span_len > 1 or is_first_call: + embeds = self.model.embed_tokens(input_ids) + + ref_codes = (info_dict.get("codes", {}) or {}).get("ref") + ref_offset = int(info_dict.get("ref_offset", 0)) + last_ref_row = None + if isinstance(ref_codes, torch.Tensor) and ref_codes.numel() > 0: + if ref_codes.dim() == 1 and ref_codes.numel() % self.n_vq == 0: + ref_codes = ref_codes.view(-1, self.n_vq) + if isinstance(ref_codes, torch.Tensor) and ref_codes.dim() == 2: + end_off = ref_offset + span_len + chunk = ref_codes[ref_offset:end_off] + if chunk.numel() > 0 and chunk.shape[0] == span_len: + codes = chunk.to(device=device, dtype=torch.long) + embeds = embeds + self._audio_embed(codes) + last_ref_row = codes[-1].detach() + + current_codes = ( + last_ref_row + if last_ref_row is not None + 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 + + info_update: dict[str, Any] = { + "audio_state": {"is_stopping": False, "step": 0, "max_new_frames": max_new_frames}, + "audio_codes": {"current": current_codes}, + "ref_offset": ref_offset + span_len, + } + return input_ids, embeds, info_update + + # Decode step: input_ids is the forced text token from compute_logits. + 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) + else: + combined = text_embed + return input_ids, combined, {} + + 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()}} + + # ------------------------------------------------------------------ + # Per-step audio generation via local transformer + # ------------------------------------------------------------------ + + def make_omni_output( + self, + model_outputs: torch.Tensor | OmniOutput, + **kwargs: Any, + ) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + self._batch_state = None + self._batch_state_spans = None + return model_outputs + + hidden = model_outputs # (S, H) + info_dicts: list[dict[str, Any]] = ( + kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or [] + ) + + for info in info_dicts: + if isinstance(info, dict) and not isinstance(info.get("audio_state"), dict): + info["audio_state"] = {"is_stopping": False, "step": 0, "max_new_frames": -1} + + 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] + 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 + have_codes = True + + 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 + return OmniOutput( + text_hidden_states=hidden, + multimodal_outputs={ + "codes": {"audio": per_req_codes}, + "meta": {"finished": torch.tensor([True], dtype=torch.bool, device=device)}, + }, + ) + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Map checkpoint weight names to vLLM-Omni module names. + + Checkpoint layout (prefix ``transformer.*`` for the backbone -- + unlike Delay/Realtime's ``language_model.*``): + transformer.* → model.* (Qwen3 backbone) + text_lm_head.weight → text_lm_head.weight + audio_embeddings.{i}.weight → audio_embeddings.{i}.weight + audio_lm_heads.{i}.weight → audio_lm_heads.{i}.weight + local_text_lm_head.weight → local_text_lm_head.weight + local_transformer.h.0.* → local_transformer.h.0.* + local_transformer.ln_f.* → local_transformer.ln_f.* + + Every "tied" pair in the upstream config is stored as a full, + separate tensor in this checkpoint (confirmed via the raw + safetensors header: 438 keys, no dedup), so aside from the backbone + prefix this is pure 1:1 name-based loading. + """ + loaded: set[str] = set() + params_dict = dict(self.named_parameters()) + + backbone_weights: list[tuple[str, torch.Tensor]] = [] + skipped: list[str] = [] + + for name, tensor in weights: + if name.startswith("transformer."): + backbone_weights.append((name[len("transformer.") :], tensor)) + continue + if name in params_dict: + default_weight_loader(params_dict[name], tensor) + loaded.add(name) + else: + skipped.append(name) + + backbone_loaded = self.model.load_weights(iter(backbone_weights)) + for n in backbone_loaded: + loaded.add(f"model.{n}") + + self._stacked_audio_emb_w = torch.stack( + [e.weight.detach() for e in self.audio_embeddings], dim=0 + ) # (n_vq, audio_vocab_size, hidden_size) + + logger.info( + "[MossTTSLocal] loaded %d/%d params; skipped=%d (first 5: %s)", + len(loaded), + len(params_dict), + len(skipped), + skipped[:5], + ) + not_loaded = [n for n in params_dict if n not in loaded] + if not_loaded: + logger.warning("[MossTTSLocal] %d params NOT loaded (first 5: %s)", len(not_loaded), not_loaded[:5]) + return loaded + + __all__ = [ "MossTTSDelayTalkerForGeneration", "MossTTSRealtimeTalkerForGeneration", + "MossTTSLocalTalkerForGeneration", ] diff --git a/vllm_omni/model_executor/models/moss_tts/pipeline.py b/vllm_omni/model_executor/models/moss_tts/pipeline.py index fbbdd8897ff..d0cec08e5b1 100644 --- a/vllm_omni/model_executor/models/moss_tts/pipeline.py +++ b/vllm_omni/model_executor/models/moss_tts/pipeline.py @@ -80,9 +80,39 @@ ), ) -# The pipeline config is otherwise the same for all 5 variants; the per-variant +MOSS_TTS_LOCAL_PIPELINE = PipelineConfig( + model_type="moss_tts_local", + model_arch="MossTTSLocalModel", # different talker class: GPT2-style local depth transformer + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="moss_tts_local", + execution_type=StageExecutionType.LLM_AR, + 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}, + ), + StagePipelineConfig( + stage_id=1, + model_stage="moss_tts_local_codec", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + final_output=True, + final_output_type="audio", + engine_output_type="audio", + model_arch="MossTTSCodecDecoder", + sync_process_input_func=f"{_PROC}.talker2codec", + sampling_constraints={"detokenize": True}, + ), + ), +) + +# The pipeline config is otherwise the same for all variants; the per-variant # differences (n_vq, backbone size, generation strategy) are encoded in the -# HF config.json and the deploy YAML. Realtime is split out because it has a -# different talker architecture (MossTTSRealtime vs MossTTSDelayModel). +# HF config.json and the deploy YAML. Realtime and Local are split out because +# they have different talker architectures from the delay variant +# (MossTTSRealtime / MossTTSLocalModel vs MossTTSDelayModel). -__all__ = ["MOSS_TTS_PIPELINE", "MOSS_TTS_REALTIME_PIPELINE"] +__all__ = ["MOSS_TTS_PIPELINE", "MOSS_TTS_REALTIME_PIPELINE", "MOSS_TTS_LOCAL_PIPELINE"] diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 279de283dc8..6749d5cfc63 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -294,6 +294,12 @@ "modeling_moss_tts_talker", "MossTTSRealtimeTalkerForGeneration", ), + # MossTTSLocalModel: MOSS-TTS-Local-Transformer-v1.5 + "MossTTSLocalModel": ( + "moss_tts", + "modeling_moss_tts_talker", + "MossTTSLocalTalkerForGeneration", + ), # Stage-1 codec decoder (shared by all 5 variants) "MossTTSCodecDecoder": ( "moss_tts",