[Model] Add Ming-omni-tts dense 0.5B pipeline (#2906)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/model_executor/models/common/ming/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/model_executor/models/common/ming/audio_dsp.py

@@ -0,0 +1,209 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/audio_tokenizer/istft.py
+
+import torch
+import torch.nn as nn
+
+
+class ISTFT(nn.Module):
+    """
+    Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with
+    windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges.
+    See issue: https://github.com/pytorch/pytorch/issues/62323
+    Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs.
+    The NOLA constraint is met as we trim padded samples anyway.
+
+    Args:
+        n_fft (int): Size of Fourier transform.
+        hop_length (int): The distance between neighboring sliding window frames.
+        win_length (int): The size of window frame and STFT filter.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"):
+        super().__init__()
+        if padding not in ["center", "same"]:
+            raise ValueError("Padding must be 'center' or 'same'.")
+        self.padding = padding
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        window = torch.hann_window(win_length)
+        self.register_buffer("window", window)
+
+        self.buffer_len = self.win_length - self.hop_length
+
+    def _buffer_process(
+        self,
+        x: torch.Tensor,
+        buffer: torch.Tensor | None,
+        pad: int,
+        last_chunk: bool = False,
+        streaming: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if streaming:
+            if buffer is None:
+                # first chunk
+                x = x[:, pad:]
+            if buffer is not None:
+                # next chunk
+                x[:, : self.buffer_len] += buffer
+            buffer = x[:, -self.buffer_len :]
+            if not last_chunk:
+                x = x[:, : -self.buffer_len]
+            else:
+                x = x[:, :-pad]
+        else:
+            x = x[:, pad:-pad]
+
+        return x, buffer
+
+    def forward(
+        self,
+        spec: torch.Tensor,
+        audio_buffer: torch.Tensor | None = None,
+        window_buffer: torch.Tensor | None = None,
+        streaming: bool = False,
+        last_chunk: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+        """
+        Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram.
+
+        Args:
+            spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size,
+                            N is the number of frequency bins, and T is the number of time frames.
+            audio_buffer (Tensor): [Streaming Input/State] The audio overlap buffer from the previous chunk.
+                            Shape: (B, win_length - hop_length)
+            window_buffer (Tensor): [Streaming Input/State] The window overlap buffer from the previous chunk.
+            streaming: If `True`, the function operates in streaming mode, processing `spec` as a single chunk.
+            last_chunk: When `streaming=True` and `last_chunk=True`, the function can perform final "flush" operations
+
+        Returns:
+            Reconstructed signal, plus streaming buffers when `padding="same"`.
+        """
+        if self.padding == "center":
+            # Fallback to pytorch native implementation
+            return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True)
+        elif self.padding == "same":
+            pad = (self.win_length - self.hop_length) // 2
+        else:
+            raise ValueError("Padding must be 'center' or 'same'.")
+
+        if spec.dim() != 3:
+            raise ValueError(f"Expected spec rank-3 [Batch, Freq, Time], got {tuple(spec.shape)}")
+        _, _, T = spec.shape
+
+        # Inverse FFT
+        ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
+        ifft = ifft * self.window[None, :, None]
+
+        # Overlap and Add
+        output_size = (T - 1) * self.hop_length + self.win_length
+        y = torch.nn.functional.fold(
+            ifft,
+            output_size=(1, output_size),
+            kernel_size=(1, self.win_length),
+            stride=(1, self.hop_length),
+        )[:, 0, 0, :]
+
+        y, audio_buffer = self._buffer_process(y, audio_buffer, pad, last_chunk=last_chunk, streaming=streaming)
+
+        # Window envelope
+        window_sq = self.window.square().expand(1, T, -1).transpose(1, 2)
+        window_envelope = (
+            torch.nn.functional.fold(
+                window_sq,
+                output_size=(1, output_size),
+                kernel_size=(1, self.win_length),
+                stride=(1, self.hop_length),
+            )
+            .squeeze(0)
+            .squeeze(0)
+        )
+
+        window_envelope, window_buffer = self._buffer_process(
+            window_envelope, window_buffer, pad, last_chunk=last_chunk, streaming=streaming
+        )
+        window_envelope = window_envelope.squeeze()
+
+        # Normalize
+        if not (window_envelope > 1e-11).all():
+            raise RuntimeError("ISTFT window envelope underflowed; invalid overlap-add state.")
+        y = y / window_envelope
+
+        return y, audio_buffer, window_buffer
+
+
+class FourierHead(nn.Module):
+    """Base class for inverse fourier modules."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+        """
+        raise NotImplementedError("Subclasses must implement the forward method.")
+
+
+class ISTFTHead(FourierHead):
+    """
+    ISTFT Head module for predicting STFT complex coefficients.
+
+    Args:
+        dim (int): Hidden dimension of the model.
+        n_fft (int): Size of Fourier transform.
+        hop_length (int): The distance between neighboring sliding window frames, which should align with
+                          the resolution of the input features.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"):
+        super().__init__()
+        out_dim = n_fft + 2
+        self.out = torch.nn.Linear(dim, out_dim)
+        self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        audio_buffer: torch.Tensor | None = None,
+        window_buffer: torch.Tensor | None = None,
+        streaming: bool = False,
+        last_chunk: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+        """
+        Forward pass of the ISTFTHead module.
+
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Audio, predicted spectrogram coefficients, and streaming buffers.
+        """
+        x_pred = self.out(x)
+        x_pred = x_pred.transpose(1, 2)
+        mag, p = x_pred.chunk(2, dim=1)
+        mag = torch.exp(mag)
+        mag = torch.clip(mag, max=1e2)  # safeguard to prevent excessively large magnitudes
+        # wrapping happens here. These two lines produce real and imaginary value
+        x = torch.cos(p)
+        y = torch.sin(p)
+        # recalculating phase here does not produce anything new
+        # only costs time
+        # phase = torch.atan2(y, x)
+        # S = mag * torch.exp(phase * 1j)
+        # better directly produce the complex value
+        S = mag * (x + 1j * y)
+        audio, audio_buffer, window_buffer = self.istft(
+            S, audio_buffer=audio_buffer, window_buffer=window_buffer, streaming=streaming, last_chunk=last_chunk
+        )
+        return audio.unsqueeze(1), x_pred, audio_buffer, window_buffer
+
+
+__all__ = ["FourierHead", "ISTFT", "ISTFTHead"]

vllm_omni/model_executor/models/common/ming/audio_vae.py

@@ -1,10 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# Copyright 2025 The vLLM-Omni team.
-# Copyright (c) Ant Group. All rights reserved.
-# Adapted from:
-# https://github.com/inclusionAI/Ming/tree/e58533db227031990c5a6864dcf5f08fb53ed0d2/AudioVAE
-
-from __future__ import annotations
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import torch
 import torch.nn as nn
@@ -13,6 +8,8 @@ from transformers import PretrainedConfig, PreTrainedModel, Qwen2Config, Qwen2Mo
 from transformers.utils import is_flash_attn_2_available
 from vllm.logger import init_logger
 
+from .audio_dsp import ISTFTHead
+
 logger = init_logger(__name__)
 
 
@@ -21,114 +18,50 @@ class AudioVAEConfig(PretrainedConfig):
 
     def __init__(
         self,
-        sample_rate: int = 44100,
-        enc_kwargs: dict | None = None,
-        dec_kwargs: dict | None = None,
-        init_method: str = "kaiming",
-        patch_size: int = 4,
+        sample_rate=44100,
+        enc_kwargs=None,
+        semantic_module_kwargs=None,
+        dec_kwargs=None,
+        hifi_gan_disc_kwargs=None,
+        spec_disc_kwargs=None,
+        lambda_disc=1.0,
+        lambda_mel_loss=15,
+        lambda_adv=1.0,
+        lambda_feat_match_loss=1.0,
+        lambda_semantic=5.0,
+        init_method="kaiming",
+        patch_size=4,
         **kwargs,
     ):
         self.sample_rate = sample_rate
         self.enc_kwargs = enc_kwargs or {}
+        self.semantic_module_kwargs = semantic_module_kwargs
         self.dec_kwargs = dec_kwargs or {}
+        self.hifi_gan_disc_kwargs = hifi_gan_disc_kwargs
+        self.spec_disc_kwargs = spec_disc_kwargs
+        self.lambda_disc = lambda_disc
+        self.lambda_mel_loss = lambda_mel_loss
+        self.lambda_adv = lambda_adv
+        self.lambda_feat_match_loss = lambda_feat_match_loss
+        self.lambda_semantic = lambda_semantic
         self.init_method = init_method
         self.patch_size = patch_size
         super().__init__(**kwargs)
 
 
-class ISTFT(nn.Module):
-    def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"):
-        super().__init__()
-        if padding not in ["center", "same"]:
-            raise ValueError("Padding must be 'center' or 'same'.")
-        self.padding = padding
-        self.n_fft = n_fft
-        self.hop_length = hop_length
-        self.win_length = win_length
-        window = torch.hann_window(win_length)
-        self.register_buffer("window", window)
-        self.buffer_len = self.win_length - self.hop_length
-
-    def _buffer_process(self, x, buffer, pad, last_chunk=False, streaming=False):
-        if streaming:
-            if buffer is None:
-                x = x[:, pad:]
-            if buffer is not None:
-                x[:, : self.buffer_len] += buffer
-            buffer = x[:, -self.buffer_len :]
-            if not last_chunk:
-                x = x[:, : -self.buffer_len]
-            else:
-                x = x[:, :-pad]
-        else:
-            x = x[:, pad:-pad]
-        return x, buffer
-
-    def forward(self, spec, audio_buffer=None, window_buffer=None, streaming=False, last_chunk=False):
-        if self.padding == "center":
-            return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True)
-        elif self.padding == "same":
-            pad = (self.win_length - self.hop_length) // 2
-        else:
-            raise ValueError("Padding must be 'center' or 'same'.")
-
-        B, N, T = spec.shape
-        ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
-        ifft = ifft * self.window[None, :, None]
-
-        output_size = (T - 1) * self.hop_length + self.win_length
-        y = torch.nn.functional.fold(
-            ifft,
-            output_size=(1, output_size),
-            kernel_size=(1, self.win_length),
-            stride=(1, self.hop_length),
-        )[:, 0, 0, :]
-
-        y, audio_buffer = self._buffer_process(y, audio_buffer, pad, last_chunk=last_chunk, streaming=streaming)
-
-        window_sq = self.window.square().expand(1, T, -1).transpose(1, 2)
-        window_envelope = (
-            torch.nn.functional.fold(
-                window_sq,
-                output_size=(1, output_size),
-                kernel_size=(1, self.win_length),
-                stride=(1, self.hop_length),
-            )
-            .squeeze(0)
-            .squeeze(0)
-        )
-
-        window_envelope, window_buffer = self._buffer_process(
-            window_envelope, window_buffer, pad, last_chunk=last_chunk, streaming=streaming
-        )
-        window_envelope = window_envelope.squeeze()
-
-        assert (window_envelope > 1e-11).all()
-        y = y / window_envelope
-
-        return y, audio_buffer, window_buffer
-
-
-class ISTFTHead(nn.Module):
-    def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"):
-        super().__init__()
-        out_dim = n_fft + 2
-        self.out = nn.Linear(dim, out_dim)
-        self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding)
-
-    def forward(self, x, audio_buffer=None, window_buffer=None, streaming=False, last_chunk=False):
-        x_pred = self.out(x)
-        x_pred = x_pred.transpose(1, 2)
-        mag, p = x_pred.chunk(2, dim=1)
-        mag = torch.exp(mag)
-        mag = torch.clip(mag, max=1e2)
-        x = torch.cos(p)
-        y = torch.sin(p)
-        S = mag * (x + 1j * y)
-        audio, audio_buffer, window_buffer = self.istft(
-            S, audio_buffer=audio_buffer, window_buffer=window_buffer, streaming=streaming, last_chunk=last_chunk
-        )
-        return audio.unsqueeze(1), x_pred, audio_buffer, window_buffer
+def _qwen2_config(backbone):
+    config = Qwen2Config.from_dict(config_dict=backbone)
+    if (getattr(config, "_attn_implementation", None) or getattr(config, "attn_implementation", None)) not in (
+        None,
+        "flash_attention_2",
+    ):
+        return config
+    if is_flash_attn_2_available():
+        config._attn_implementation_autoset = True
+        config._attn_implementation = "flash_attention_2"
+    else:
+        config._attn_implementation = "sdpa"
+    return config
 
 
 class StreamingLinearUpsample(nn.Module):
@@ -138,6 +71,8 @@ class StreamingLinearUpsample(nn.Module):
         self.upsampler = nn.Upsample(scale_factor=scale_factor, mode="linear", align_corners=False)
 
     def forward(self, x, state=None, is_last=False):
+        if x is None and is_last and (state is None or state.get("prev_chunk") is None):
+            raise ValueError("Received end-of-stream without any latent chunk to upsample.")
         if state is None:
             state = {"prev_chunk": None, "history_last": None, "is_first": True}
 
@@ -188,16 +123,72 @@ class StreamingLinearUpsample(nn.Module):
         return final_out, state
 
 
+class Encoder(nn.Module):
+    def __init__(self, encoder_args, input_dim=320, hop_size=320, latent_dim=64, patch_size=-1):
+        super().__init__()
+        config = _qwen2_config(encoder_args)
+        logger.info("AudioVAE Encoder: using attn_implementation=%r", config._attn_implementation)
+        self.encoder = Qwen2Model(config)
+        self.input_dim = input_dim
+        self.hop_size = hop_size
+        self.latent_dim = latent_dim
+        self.fc1 = nn.Linear(input_dim, config.hidden_size, bias=False)
+        self.fc2 = nn.Linear(config.hidden_size, config.hidden_size)
+        self.fc3 = nn.Linear(config.hidden_size, latent_dim * 2)
+        self.norm = nn.LayerNorm(config.hidden_size)
+        self.patch_size = patch_size
+        if patch_size != -1:
+            aggregator_config = _qwen2_config({**encoder_args, "num_hidden_layers": 4})
+            self.aggregator = Qwen2Model(aggregator_config)
+            self.cls_embed = nn.Parameter(torch.rand(1, 1, config.hidden_size))
+            self.cls_embed.data.normal_(0, 0.02)
+
+    def get_frames(self, x):
+        num_frames_total = (x.size(-1) + self.hop_size - 1) // self.hop_size
+        expected_len = (num_frames_total - 1) * self.hop_size + self.input_dim
+        padding_needed = expected_len - x.size(-1)
+        waveform = F.pad(x, (0, padding_needed), value=0.0)
+        frames = waveform.unfold(dimension=-1, size=self.input_dim, step=self.hop_size)
+        return frames
+
+    def pad_patch_insert_cls(self, x):
+        bsz, _, dim = x.size()
+        num_frame = x.size(1)
+        r = num_frame % self.patch_size
+        pad_num = self.patch_size - r if r else 0
+        x = F.pad(x, (0, 0, 0, pad_num), value=0.0)
+        # [Batch, Time, Dimension] -> [Batch*PatchGroups, Patch, Dimension].
+        x = x.reshape(-1, self.patch_size, dim)
+        x = torch.cat((x, self.cls_embed.expand(x.size(0), -1, -1)), dim=1)
+        # [Batch*PatchGroups, Patch+1, Dimension] -> [Batch, Time, Dimension].
+        x = x.reshape(bsz, -1, dim)
+        return x
+
+    def forward(self, waveform):
+        x = self.get_frames(waveform)
+        x = self.fc1(x)
+        x = self.fc2(x)
+        x = self.encoder(inputs_embeds=x)
+        x = x.last_hidden_state
+
+        if self.patch_size != -1:
+            x = self.pad_patch_insert_cls(x)
+            x = self.aggregator(inputs_embeds=x)
+            x = x.last_hidden_state
+            bsz, _, dim = x.size()
+            # [Batch, Time, Dimension] -> [Batch*PatchGroups, Patch+1, Dimension].
+            x = x.reshape(-1, self.patch_size + 1, dim)
+            # [Batch*PatchGroups, 1, Dimension] -> [Batch, PatchGroups, Dimension].
+            x = x[:, -1:, :].reshape(bsz, -1, dim)
+
+        x = self.fc3(x)
+        return x, waveform.unsqueeze(1)
+
+
 class Decoder(nn.Module):
     def __init__(self, decoder_args, output_dim=320, latent_dim=64, patch_size=-1):
         super().__init__()
-        config = Qwen2Config.from_dict(config_dict=decoder_args)
-        if is_flash_attn_2_available():
-            config._attn_implementation_autoset = True
-            config._attn_implementation = "flash_attention_2"
-        else:
-            config._attn_implementation = "sdpa"
-
+        config = _qwen2_config(decoder_args)
         logger.info("AudioVAE Decoder: using attn_implementation=%r", config._attn_implementation)
         self.decoder = Qwen2Model(config)
         self.output_dim = output_dim
@@ -211,7 +202,21 @@ class Decoder(nn.Module):
         if self.patch_size != -1:
             self.upsampling = StreamingLinearUpsample(scale_factor=patch_size)
 
+    def forward(self, x):
+        x = self.fc1(x)
+
+        if self.patch_size != -1:
+            # [Batch, Time, Dimension] -> [Batch, Dimension, Time] -> [Batch, Time, Dimension].
+            x = self.upsampling.upsampler(x.transpose(1, 2)).transpose(1, 2)
+
+        x = self.decoder(inputs_embeds=x)
+        x = x.last_hidden_state
+        x, _ = self.head(x)
+        return x, None
+
     def low_level_reconstruct(self, x, past_key_values=None, use_cache=False, stream_state=None, last_chunk=False):
+        if stream_state is None:
+            stream_state = (None, None, None)
         upsample_state, audio_buffer, window_buffer = stream_state
         bsz, device, dtype = x.size(0), x.device, x.dtype
         x = self.fc1(x)
@@ -222,6 +227,7 @@ class Decoder(nn.Module):
                     stream_state = (upsample_state, audio_buffer, window_buffer)
                     return torch.empty(bsz, 1, 0, device=device, dtype=dtype), stream_state, past_key_values
             else:
+                # [Batch, Time, Dimension] -> [Batch, Dimension, Time] -> [Batch, Time, Dimension].
                 x = self.upsampling.upsampler(x.transpose(1, 2)).transpose(1, 2)
 
         hidden_states_list = []
@@ -239,7 +245,6 @@ class Decoder(nn.Module):
                 past_len = 0
 
             curr_len = x.shape[1]
-
             if past_len < target_len and (past_len + curr_len) >= sw_size:
                 fill_len = target_len - past_len
                 x_fill = x[:, :fill_len, :]
@@ -269,86 +274,39 @@ class Decoder(nn.Module):
         return x_out, stream_state, past_key_values
 
 
-class Encoder(nn.Module):
-    def __init__(self, encoder_args, input_dim=320, hop_size=320, latent_dim=64, patch_size=-1):
-        super().__init__()
-        config = Qwen2Config.from_dict(config_dict=encoder_args)
-        if is_flash_attn_2_available():
-            config._attn_implementation_autoset = True
-            config._attn_implementation = "flash_attention_2"
-        else:
-            config._attn_implementation = "sdpa"
-
-        logger.info("AudioVAE Encoder: using attn_implementation=%r", config._attn_implementation)
-        self.encoder = Qwen2Model(config)
-        self.input_dim = input_dim
-        self.hop_size = hop_size
-        self.latent_dim = latent_dim
-        self.fc1 = nn.Linear(input_dim, config.hidden_size, bias=False)
-        self.fc2 = nn.Linear(config.hidden_size, config.hidden_size)
-        self.fc3 = nn.Linear(config.hidden_size, latent_dim * 2)
-        self.norm = nn.LayerNorm(config.hidden_size)
-        self.patch_size = patch_size
-        if patch_size != -1:
-            config.num_hidden_layers = 4
-            self.aggregator = Qwen2Model(config)
-            self.cls_embed = nn.Parameter(torch.rand(1, 1, config.hidden_size))
-            self.cls_embed.data.normal_(0, 0.02)
-
-    def get_frames(self, x):
-        num_frames_total = (x.size(-1) + self.hop_size - 1) // self.hop_size
-        expected_len = (num_frames_total - 1) * self.hop_size + self.input_dim
-        padding_needed = expected_len - x.size(-1)
-        waveform = F.pad(x, (0, padding_needed), value=0.0)
-        frames = waveform.unfold(dimension=-1, size=self.input_dim, step=self.hop_size)
-        return frames
-
-    def pad_patch_insert_cls(self, x):
-        bsz, _, dim = x.size()
-        num_frame = x.size(1)
-        r = num_frame % self.patch_size
-        pad_num = self.patch_size - r if r else 0
-        x = F.pad(x, (0, 0, 0, pad_num), value=0.0)
-        x = x.reshape(-1, self.patch_size, dim)
-        x = torch.cat((x, self.cls_embed.expand(x.size(0), -1, -1)), dim=1)
-        x = x.reshape(bsz, -1, dim)
-        return x
-
-    def forward(self, waveform):
-        x = self.get_frames(waveform)
-        x = self.fc1(x)
-        x = self.fc2(x)
-        x = self.encoder(inputs_embeds=x)
-        x = x.last_hidden_state
-
-        if self.patch_size != -1:
-            x = self.pad_patch_insert_cls(x)
-            x = self.aggregator(inputs_embeds=x)
-            x = x.last_hidden_state
-            bsz, _, dim = x.size()
-            x = x.reshape(-1, self.patch_size + 1, dim)
-            x = x[:, -1:, :].reshape(bsz, -1, dim)
-
-        x = self.fc3(x)
-        return x, waveform.unsqueeze(1)
-
-
 class AudioVAE(PreTrainedModel):
     config_class = AudioVAEConfig
 
-    def __init__(self, config: AudioVAEConfig):
+    def __init__(self, config):
         super().__init__(config)
+        self.config = config
+        enc_kwargs = config.enc_kwargs
+        dec_kwargs = config.dec_kwargs
+        for key in ("backbone", "input_dim", "latent_dim"):
+            if key not in enc_kwargs:
+                raise ValueError(f"AudioVAE.enc_kwargs missing required key: {key}")
+        for key in ("backbone", "output_dim", "latent_dim"):
+            if key not in dec_kwargs:
+                raise ValueError(f"AudioVAE.dec_kwargs missing required key: {key}")
+
+        hop_size = enc_kwargs.get("hop_size", enc_kwargs["input_dim"])
+        if enc_kwargs["input_dim"] != hop_size:
+            raise ValueError(f"AudioVAE encoder input_dim ({enc_kwargs['input_dim']}) != hop_size ({hop_size}).")
+        if hop_size != dec_kwargs["output_dim"]:
+            raise ValueError(
+                f"AudioVAE encoder hop_size ({hop_size}) != decoder output_dim ({dec_kwargs['output_dim']})."
+            )
         self.encoder = Encoder(
-            encoder_args=config.enc_kwargs["backbone"],
-            input_dim=config.enc_kwargs["input_dim"],
-            hop_size=config.enc_kwargs.get("hop_size", 320),
-            latent_dim=config.enc_kwargs["latent_dim"],
+            encoder_args=enc_kwargs["backbone"],
+            input_dim=enc_kwargs["input_dim"],
+            hop_size=hop_size,
+            latent_dim=enc_kwargs["latent_dim"],
             patch_size=config.patch_size,
         )
         self.decoder = Decoder(
-            decoder_args=config.dec_kwargs["backbone"],
-            output_dim=config.dec_kwargs["output_dim"],
-            latent_dim=config.dec_kwargs["latent_dim"],
+            decoder_args=dec_kwargs["backbone"],
+            output_dim=dec_kwargs["output_dim"],
+            latent_dim=dec_kwargs["latent_dim"],
             patch_size=config.patch_size,
         )
         self.post_init()
@@ -367,21 +325,50 @@ class AudioVAE(PreTrainedModel):
             if module.padding_idx is not None:
                 module.weight.data[module.padding_idx].zero_()
 
+    @torch.inference_mode()
     def encode_latent(self, waveform, waveform_length):
-        from diffusers.models.autoencoders.autoencoder_oobleck import OobleckDiagonalGaussianDistribution
+        if waveform.ndim != 2:
+            raise ValueError(f"Expected waveform rank-2 [Batch, Time], got {tuple(waveform.shape)}")
+        if waveform_length.ndim != 1:
+            raise ValueError(f"Expected waveform_length rank-1 [Batch], got {tuple(waveform_length.shape)}")
+        if waveform.shape[0] != waveform_length.shape[0]:
+            raise ValueError(
+                "Batch mismatch: "
+                f"waveform batch={waveform.shape[0]} vs "
+                f"waveform_length batch={waveform_length.shape[0]}"
+            )
+        if torch.any(waveform_length <= 0):
+            raise ValueError("waveform_length must be strictly positive.")
 
         frame_num = torch.ceil(waveform_length / self.config.enc_kwargs["input_dim"]).to(torch.int32)
         if self.config.patch_size != -1:
             frame_num = torch.ceil(frame_num / self.config.patch_size)
-        h, y = self.encoder(waveform)
+        h, _ = self.encoder(waveform)
         h = h.transpose(1, 2)
 
-        posterior = OobleckDiagonalGaussianDistribution(h)
-        latent = posterior.sample()
+        mean, logvar = torch.chunk(h, 2, dim=1)
+        logvar = torch.clamp(logvar, -30.0, 20.0)
+        std = torch.exp(0.5 * logvar)
+        latent = mean + std * torch.randn_like(mean)
         latent = latent.transpose(1, 2)
         return latent, frame_num
 
+    @torch.inference_mode()
     def decode(self, latent, past_key_values=None, use_cache=False, stream_state=(None, None, None), last_chunk=False):
+        if latent.dim() != 3:
+            raise ValueError(f"Expected latent rank-3 [B,T,D], got shape={tuple(latent.shape)}")
+        if latent.shape[0] <= 0:
+            raise ValueError("latent batch size must be positive.")
+
+        target_dtype = next(self.decoder.parameters()).dtype
+        target_device = next(self.decoder.parameters()).device
+        if latent.dtype != target_dtype or latent.device != target_device:
+            latent = latent.to(device=target_device, dtype=target_dtype)
+
+        expected_latent_dim = self.config.dec_kwargs["latent_dim"]
+        if latent.shape[-1] != expected_latent_dim:
+            raise ValueError(f"Latent dim mismatch in decode(): got {latent.shape[-1]}, expected {expected_latent_dim}")
+
         waveform, stream_state, past_key_values = self.decoder.low_level_reconstruct(
             latent,
             past_key_values=past_key_values,

vllm_omni/model_executor/models/common/ming/dit.py

@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+import torch.nn.functional as F
+from torch import nn
+from x_transformers.x_transformers import apply_rotary_pos_emb
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        if self.weight.dtype in [torch.float16, torch.bfloat16]:
+            x = x.to(self.weight.dtype)
+        x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps)
+        return x
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate="none"):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+
+        activation = nn.GELU(approximate=approximate)
+        project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation)
+        self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out))
+
+    def forward(self, x):
+        return self.ff(x)
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        heads=8,
+        dim_head=64,
+        dropout=0.0,
+        qk_norm=None,
+        pe_attn_head=None,
+        attn_mask_enabled=True,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.heads = heads
+        self.inner_dim = dim_head * heads
+        self.dropout = dropout
+        self.to_q = nn.Linear(dim, self.inner_dim)
+        self.to_k = nn.Linear(dim, self.inner_dim)
+        self.to_v = nn.Linear(dim, self.inner_dim)
+        if qk_norm is None:
+            self.q_norm = None
+            self.k_norm = None
+        elif qk_norm == "rms_norm":
+            self.q_norm = RMSNorm(dim_head)
+            self.k_norm = RMSNorm(dim_head)
+        else:
+            raise ValueError(f"Unimplemented qk_norm: {qk_norm}")
+
+        self.to_out = nn.ModuleList([])
+        self.to_out.append(nn.Linear(self.inner_dim, dim))
+        self.to_out.append(nn.Dropout(dropout))
+        self.pe_attn_head = pe_attn_head
+        self.attn_mask_enabled = attn_mask_enabled
+
+    def forward(self, x, mask=None, rope=None):
+        batch_size = x.shape[0]
+        query = self.to_q(x)
+        key = self.to_k(x)
+        value = self.to_v(x)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // self.heads
+        query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+
+        if self.q_norm is not None:
+            query = self.q_norm(query)
+        if self.k_norm is not None:
+            key = self.k_norm(key)
+
+        if rope is not None:
+            freqs, xpos_scale = rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+            if self.pe_attn_head is not None:
+                on = self.pe_attn_head
+                query[:, :on, :, :] = apply_rotary_pos_emb(query[:, :on, :, :], freqs, q_xpos_scale)
+                key[:, :on, :, :] = apply_rotary_pos_emb(key[:, :on, :, :], freqs, k_xpos_scale)
+            else:
+                query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+                key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+
+        if self.attn_mask_enabled and mask is not None:
+            valid_sample_indices = mask.any(dim=1)
+            final_output = torch.zeros_like(query).to(query.device)
+            attn_mask = mask[valid_sample_indices]
+            query = query[valid_sample_indices]
+            key = key[valid_sample_indices]
+            value = value[valid_sample_indices]
+            attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)
+            attn_mask = attn_mask.expand(valid_sample_indices.sum().item(), self.heads, query.shape[-2], key.shape[-2])
+        else:
+            attn_mask = None
+
+        x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+        if self.attn_mask_enabled and mask is not None:
+            final_output[valid_sample_indices] = x
+            x = final_output
+
+        x = x.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim)
+        x = x.to(query.dtype)
+        x = self.to_out[0](x)
+        x = self.to_out[1](x)
+
+        if mask is not None:
+            mask = mask.unsqueeze(-1)
+            x = x.masked_fill(~mask, 0.0)
+
+        return x
+
+
+class DiTBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        num_heads,
+        mlp_ratio=4.0,
+        dropout=0.1,
+        qk_norm=None,
+        pe_attn_head=None,
+        attn_mask_enabled=True,
+        **kwargs,
+    ):
+        super().__init__()
+        self.norm1 = RMSNorm(hidden_size)
+        self.attn = Attention(
+            dim=hidden_size,
+            heads=num_heads,
+            dim_head=hidden_size // num_heads,
+            dropout=dropout,
+            qk_norm=qk_norm,
+            pe_attn_head=pe_attn_head,
+            attn_mask_enabled=attn_mask_enabled,
+        )
+        self.norm2 = RMSNorm(hidden_size)
+        self.mlp = FeedForward(dim=hidden_size, mult=mlp_ratio, dropout=dropout, approximate="tanh")
+
+    def forward(self, x, mask, rope):
+        x = x + self.attn(self.norm1(x), mask=mask, rope=rope)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, hidden_size, out_channels):
+        super().__init__()
+        self.norm_final = RMSNorm(hidden_size)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+
+    def forward(self, x):
+        x = self.norm_final(x)
+        x = self.linear(x)
+        return x
+
+
+class CondEmbedder(nn.Module):
+    def __init__(self, input_feature_size, hidden_size):
+        super().__init__()
+        self.cond_embedder = nn.Linear(input_feature_size, hidden_size)
+
+    def forward(self, llm_cond):
+        return self.cond_embedder(llm_cond)
+
+
+def get_epss_timesteps(n, device, dtype):
+    dt = 1 / 32
+    predefined_timesteps = {
+        5: [0, 2, 4, 8, 16, 32],
+        6: [0, 2, 4, 6, 8, 16, 32],
+        7: [0, 2, 4, 6, 8, 16, 24, 32],
+        10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32],
+        12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32],
+        16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32],
+    }
+    t = predefined_timesteps.get(n, [])
+    if not t:
+        return torch.linspace(0, 1, n + 1, device=device, dtype=dtype)
+    return dt * torch.tensor(t, device=device, dtype=dtype)

vllm_omni/model_executor/models/common/ming/fm.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from .dit import get_epss_timesteps
+
+
+def apply_sway_sampling(t, sway_sampling_coef):
+    if sway_sampling_coef is None:
+        return t
+    return t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)
+
+
+def build_timesteps(steps, device, dtype, use_epss=True, sway_sampling_coef=-1.0):
+    if steps <= 0:
+        raise ValueError(f"steps must be positive, got {steps}")
+    if use_epss:
+        t = get_epss_timesteps(steps, device=device, dtype=dtype)
+    else:
+        t = torch.linspace(0, 1, steps + 1, device=device, dtype=dtype)
+    return apply_sway_sampling(t, sway_sampling_coef)
+
+
+class Solver:
+    def __init__(self, func, y0, sigma=0.25, temperature=1.5) -> None:
+        self.func = func
+        self.y0 = y0
+        self.sigma = sigma
+        self.temperature = temperature
+
+    def integrate(self, t):
+        solution = torch.empty(len(t), *self.y0.shape, dtype=self.y0.dtype, device=self.y0.device)
+        solution[0] = self.y0
+
+        j = 1
+        y0 = self.y0
+        for t0, t1 in zip(t[:-1], t[1:]):
+            dt = t1 - t0
+            f0 = self.func(t0, y0)
+            dy = dt * f0
+            y1 = y0 + dy
+
+            while j < len(t) and t1 >= t[j]:
+                solution[j] = self._linear_interp(t0, t1, y0, y1, t[j])
+                j += 1
+
+            noise = torch.randn_like(y0)
+            shift = self.sigma * (self.temperature**0.5) * (abs(dt) ** 0.5) * noise
+            y0 = y1 + shift
+
+        return solution
+
+    def _linear_interp(self, t0, t1, y0, y1, t):
+        if t == t0:
+            return y0
+        if t == t1:
+            return y1
+        slope = (t - t0) / (t1 - t0)
+        return y0 + slope * (y1 - y0)
+
+
+def integrate_cfm_steps(fn, y0, t, sde_args, sde_rnd, steps):
+    for step in range(steps):
+        dt = t[step + 1] - t[step]
+        y0 = y0 + fn(t[step], y0) * dt
+        y0 = y0 + sde_args[1] * (sde_args[2] ** 0.5) * (dt.abs() ** 0.5) * sde_rnd[step]
+    return y0

vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py

@@ -26,10 +26,10 @@ from vllm.sequence import IntermediateTensors
 
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE, AudioVAEConfig
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.transformers_utils.configs.ming_flash_omni import MingFlashOmniTalkerConfig
 
-from .audio_vae import AudioVAE, AudioVAEConfig
 from .prompt_utils import DEFAULT_PROMPT as MING_DEFAULT_PROMPT
 from .talker_module import CFM, Aggregator, DiT, MingAudioGenerator, build_tts_input
 from .text_processing import segment_and_normalize

vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py

@@ -697,6 +697,9 @@ class MingFlashOmniThinkerForConditionalGeneration(
             f"image_gen_scales={list(self.query_tokens_dict.keys())}"
         )
 
+    def get_language_model(self):
+        return self.language_model
+
     def _load_image_gen_query_tokens(self, model_path: str) -> None:
         """Load learnable query tokens used for image generation.
 

vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py

@@ -110,12 +110,18 @@ BASE_CAPTION_TEMPLATE: dict[str, Any] = {
 }
 
 
-def create_instruction(user_input: dict[str, Any]) -> str:
+def create_instruction(user_input: Any) -> str | None:
     """Return a JSON caption string for ``audio_sequence[0]``.
 
     Only keys already present on the base template are merged in; unknown
     keys are silently ignored to keep the output schema stable.
     """
+    if user_input is None:
+        return None
+    if isinstance(user_input, str):
+        return user_input
+    if not isinstance(user_input, dict):
+        raise ValueError(f"Ming instruction must be str, dict, or None; got {type(user_input).__name__}")
     caption = copy.deepcopy(BASE_CAPTION_TEMPLATE)
     item = caption["audio_sequence"][0]
     for key, value in user_input.items():

vllm_omni/model_executor/models/ming_flash_omni/talker_module.py

@@ -27,223 +27,19 @@ from threading import Lock
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from transformers import PreTrainedTokenizerBase, Qwen2Config, Qwen2Model, StaticCache
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
-from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb
+from x_transformers.x_transformers import RotaryEmbedding
 
 from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding
-
-from .audio_vae import AudioVAE
+from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE
+from vllm_omni.model_executor.models.common.ming.dit import CondEmbedder, DiTBlock, FinalLayer, get_epss_timesteps
+from vllm_omni.model_executor.models.common.ming.fm import apply_sway_sampling, integrate_cfm_steps
 
 logger = init_logger(__name__)
 
 
-########################################################################
-# DiT Modules
-# Ported from:
-# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/modules.py
-# Ported from:
-# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/dit.py
-########################################################################
-
-
-class RMSNorm(nn.Module):
-    def __init__(self, dim: int, eps: float = 1e-6):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim))
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if self.weight.dtype in [torch.float16, torch.bfloat16]:
-            x = x.to(self.weight.dtype)
-        x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps)
-        return x
-
-
-class FeedForward(nn.Module):
-    def __init__(
-        self, dim: int, dim_out: int | None = None, mult: float = 4, dropout: float = 0.0, approximate: str = "none"
-    ):
-        super().__init__()
-        inner_dim = int(dim * mult)
-        dim_out = dim_out if dim_out is not None else dim
-
-        activation = nn.GELU(approximate=approximate)
-        project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation)
-        self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out))
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.ff(x)
-
-
-class Attention(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        heads: int = 8,
-        dim_head: int = 64,
-        dropout: float = 0.0,
-        qk_norm: str | None = None,
-        pe_attn_head: int | None = None,
-        attn_mask_enabled: bool = True,
-    ):
-        super().__init__()
-        self.dim = dim
-        self.heads = heads
-        self.inner_dim = dim_head * heads
-        self.dropout = dropout
-
-        self.to_q = nn.Linear(dim, self.inner_dim)
-        self.to_k = nn.Linear(dim, self.inner_dim)
-        self.to_v = nn.Linear(dim, self.inner_dim)
-        if qk_norm is None:
-            self.q_norm = None
-            self.k_norm = None
-        elif qk_norm == "rms_norm":
-            self.q_norm = RMSNorm(dim_head)
-            self.k_norm = RMSNorm(dim_head)
-        else:
-            raise ValueError(f"Unimplemented qk_norm: {qk_norm}")
-
-        self.to_out = nn.ModuleList([])
-        self.to_out.append(nn.Linear(self.inner_dim, dim))
-        self.to_out.append(nn.Dropout(dropout))
-
-        self.pe_attn_head = pe_attn_head
-        self.attn_mask_enabled = attn_mask_enabled
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        mask: torch.Tensor | None = None,
-        rope: tuple[torch.Tensor, torch.Tensor | None] | None = None,
-    ) -> torch.Tensor:
-        batch_size = x.shape[0]
-
-        query = self.to_q(x)
-        key = self.to_k(x)
-        value = self.to_v(x)
-
-        inner_dim = key.shape[-1]
-        head_dim = inner_dim // self.heads
-        query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
-        key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
-        value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
-
-        if self.q_norm is not None:
-            query = self.q_norm(query)
-        if self.k_norm is not None:
-            key = self.k_norm(key)
-
-        if rope is not None:
-            freqs, xpos_scale = rope
-            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
-
-            if self.pe_attn_head is not None:
-                on = self.pe_attn_head
-                query[:, :on, :, :] = apply_rotary_pos_emb(query[:, :on, :, :], freqs, q_xpos_scale)
-                key[:, :on, :, :] = apply_rotary_pos_emb(key[:, :on, :, :], freqs, k_xpos_scale)
-            else:
-                query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
-                key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
-
-        if self.attn_mask_enabled and mask is not None:
-            valid_sample_indices = mask.any(dim=1)
-            final_output = torch.zeros_like(query).to(query.device)
-
-            attn_mask = mask[valid_sample_indices]
-            query = query[valid_sample_indices]
-            key = key[valid_sample_indices]
-            value = value[valid_sample_indices]
-            attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)
-            attn_mask = attn_mask.expand(valid_sample_indices.sum().item(), self.heads, query.shape[-2], key.shape[-2])
-        else:
-            attn_mask = None
-
-        x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
-        if self.attn_mask_enabled and mask is not None:
-            final_output[valid_sample_indices] = x
-            x = final_output
-
-        x = x.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim)
-        x = x.to(query.dtype)
-
-        x = self.to_out[0](x)
-        x = self.to_out[1](x)
-
-        if mask is not None:
-            mask = mask.unsqueeze(-1)
-            x = x.masked_fill(~mask, 0.0)
-
-        return x
-
-
-class DiTBlock(nn.Module):
-    """A DiT block with pre-norm and residual connections."""
-
-    def __init__(
-        self,
-        hidden_size: int,
-        num_heads: int,
-        mlp_ratio: float = 4.0,
-        dropout: float = 0.1,
-        qk_norm: str | None = None,
-        pe_attn_head: int | None = None,
-        attn_mask_enabled: bool = True,
-        **kwargs,
-    ):
-        super().__init__()
-        self.norm1 = RMSNorm(hidden_size)
-        self.attn = Attention(
-            dim=hidden_size,
-            heads=num_heads,
-            dim_head=hidden_size // num_heads,
-            dropout=dropout,
-            qk_norm=qk_norm,
-            pe_attn_head=pe_attn_head,
-            attn_mask_enabled=attn_mask_enabled,
-        )
-        self.norm2 = RMSNorm(hidden_size)
-        self.mlp = FeedForward(dim=hidden_size, mult=mlp_ratio, dropout=dropout, approximate="tanh")
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        mask: torch.Tensor | None,
-        rope: tuple[torch.Tensor, torch.Tensor | None] | None,
-    ) -> torch.Tensor:
-        x = x + self.attn(self.norm1(x), mask=mask, rope=rope)
-        x = x + self.mlp(self.norm2(x))
-        return x
-
-
-class FinalLayer(nn.Module):
-    """The final layer of DiT."""
-
-    def __init__(self, hidden_size: int, out_channels: int):
-        super().__init__()
-        self.norm_final = RMSNorm(hidden_size)
-        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.norm_final(x)
-        x = self.linear(x)
-        return x
-
-
-class CondEmbedder(nn.Module):
-    """Embeds LLM hidden states with optional CFG dropout."""
-
-    def __init__(self, input_feature_size: int, hidden_size: int):
-        super().__init__()
-        self.cond_embedder = nn.Linear(input_feature_size, hidden_size)
-
-    def forward(self, llm_cond: torch.Tensor) -> torch.Tensor:
-        return self.cond_embedder(llm_cond)
-
-
 class DiT(nn.Module):
     """Diffusion model with a Transformer backbone for audio latent generation."""
 
@@ -324,29 +120,6 @@ class DiT(nn.Module):
         return model_out[:, -x.shape[1] :, :]
 
 
-#########################################################################################
-# CFM
-# Ported from:
-# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/cfm.py
-#########################################################################################
-
-
-def get_epss_timesteps(n, device, dtype):
-    dt = 1 / 32
-    predefined_timesteps = {
-        5: [0, 2, 4, 8, 16, 32],
-        6: [0, 2, 4, 6, 8, 16, 32],
-        7: [0, 2, 4, 6, 8, 16, 24, 32],
-        10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32],
-        12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32],
-        16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32],
-    }
-    t = predefined_timesteps.get(n, [])
-    if not t:
-        return torch.linspace(0, 1, n + 1, device=device, dtype=dtype)
-    return dt * torch.tensor(t, device=device, dtype=dtype)
-
-
 class CFM(nn.Module):
     """Conditional Flow Matching module for audio latent generation."""
 
@@ -391,15 +164,8 @@ class CFM(nn.Module):
             pred, null_pred = torch.chunk(pred_cfg, 2, dim=0)
             return pred + (pred - null_pred) * sde_args[0]
 
-        if self.sway_sampling_coef is not None:
-            t = t + self.sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)
-
-        for step in range(self.steps):
-            dt = t[step + 1] - t[step]
-            y0 = y0 + fn(t[step], y0) * dt
-            y0 = y0 + sde_args[1] * (sde_args[2] ** 0.5) * (dt.abs() ** 0.5) * sde_rnd[step]
-
-        return y0
+        t = apply_sway_sampling(t, self.sway_sampling_coef)
+        return integrate_cfm_steps(fn, y0, t, sde_args, sde_rnd, self.steps)
 
 
 class CFMGraphExecutor:

vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py

@@ -22,7 +22,8 @@ from .spk_embedding import SpkembExtractor
 from .talker_module import resample
 
 if TYPE_CHECKING:
-    from .audio_vae import AudioVAE
+    from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE
+
     from .talker_module import Aggregator
 
 logger = init_logger(__name__)

vllm_omni/model_executor/models/ming_tts/__init__.py

@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from .config_ming_tts import MingDenseConfig
+from .ming_tts import MingTTSForConditionalGeneration
+from .ming_tts_audio_vae import MingAudioVAEModel
+from .ming_tts_llm import MingLLMModel
+
+__all__ = [
+    "MingDenseConfig",
+    "MingTTSForConditionalGeneration",
+    "MingLLMModel",
+    "MingAudioVAEModel",
+]

vllm_omni/model_executor/models/ming_tts/aggregator.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/fm/dit.py
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from x_transformers.x_transformers import RotaryEmbedding
+
+from vllm_omni.model_executor.models.common.ming.dit import DiTBlock, FinalLayer
+
+
+class Aggregator(nn.Module):
+    def __init__(
+        self,
+        in_channels=4,
+        hidden_size=1152,
+        depth=28,
+        num_heads=16,
+        mlp_ratio=4.0,
+        llm_input_dim=896,
+        **kwargs,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.num_heads = num_heads
+
+        self.word_embedder = nn.Embedding(1, hidden_size)
+        self.x_embedder = nn.Linear(in_channels, hidden_size)
+        self.hidden_size = hidden_size
+        self.rotary_embed = RotaryEmbedding(hidden_size // num_heads)
+        self.blocks = nn.ModuleList(
+            [DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **kwargs) for _ in range(depth)]
+        )
+        self.final_layer = FinalLayer(hidden_size, llm_input_dim)
+
+    def forward(self, x, mask=None):
+        if x.ndim != 3:
+            raise ValueError(f"Expected x rank-3 [Batch, Time, Dimension], got {tuple(x.shape)}")
+        if x.shape[-1] != self.in_channels:
+            raise ValueError(f"x feature dim mismatch: got {x.shape[-1]}, expected {self.in_channels}")
+
+        # [Batch, Time, Dimension] -> [Batch, Time, Hidden].
+        x = self.x_embedder(x)
+        cls_embed = self.word_embedder(torch.zeros((x.shape[0], 1), dtype=torch.long, device=x.device))
+        # Prepend a learned CLS token: [Batch, Time, Hidden] -> [Batch, Time + 1, Hidden].
+        x = torch.cat([cls_embed, x], dim=1)
+
+        rope = self.rotary_embed.forward_from_seq_len(x.shape[1])
+        if mask is not None:
+            if mask.ndim != 2:
+                raise ValueError(f"Expected mask rank-2 [Batch, Time], got {tuple(mask.shape)}")
+            if mask.shape[0] != x.shape[0] or mask.shape[1] != x.shape[1] - 1:
+                raise ValueError(
+                    f"Mask shape mismatch: got {tuple(mask.shape)}, expected {(x.shape[0], x.shape[1] - 1)}"
+                )
+            mask_pad = mask.clone().detach()[:, :1]
+            mask = torch.cat([mask_pad, mask], dim=-1)
+        for block in self.blocks:
+            x = block(x, mask, rope)
+        x = self.final_layer(x)
+        # Keep the CLS projection only: [Batch, Time + 1, Hidden] -> [Batch, 1, Hidden].
+        return x[:, :1, :]
+
+
+__all__ = ["Aggregator"]

vllm_omni/model_executor/models/ming_tts/audio_prep.py

@@ -0,0 +1,261 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from __future__ import annotations
+
+import math
+from io import BytesIO
+from typing import Any
+
+import torch
+
+from .config_ming_tts import (
+    AUDIO_FRAME_HOP,
+    LATENT_DIM,
+    PATCH_SIZE,
+    SAMPLE_RATE,
+    VAE_PATCH_SIZE,
+    VISION_START_TOKEN_ID,
+    MingTTSConfig,
+)
+
+
+def pad_prompt_waveform(
+    waveform: Any,
+    *,
+    patch_size: int = PATCH_SIZE,
+    sample_rate: int = SAMPLE_RATE,
+) -> torch.Tensor:
+    tensor = coerce_prompt_waveform(waveform)
+    pad_align = int((float(sample_rate) / 12.5) * int(patch_size))
+    new_len = ((int(tensor.shape[-1]) + pad_align - 1) // pad_align) * pad_align
+    if new_len == int(tensor.shape[-1]):
+        return tensor
+    padded = torch.zeros((1, new_len), dtype=tensor.dtype, device=tensor.device)
+    padded[:, : tensor.shape[-1]] = tensor
+    return padded
+
+
+def coerce_prompt_waveform(value: Any) -> torch.Tensor:
+    if value is None:
+        raise ValueError("prompt waveform cannot be None")
+    if isinstance(value, torch.Tensor):
+        tensor = value.detach()
+        if tensor.ndim == 1:
+            return tensor.unsqueeze(0).to(torch.float32)
+        if tensor.ndim == 2:
+            if tensor.shape[0] != 1:
+                return tensor.reshape(1, -1).to(torch.float32)
+            return tensor.to(torch.float32)
+        raise ValueError(f"Unsupported Ming prompt waveform rank: {tuple(tensor.shape)}")
+    if isinstance(value, (list, tuple)):
+        parts = [coerce_prompt_waveform(item) for item in value if item is not None]
+        if not parts:
+            raise ValueError("prompt waveform list was empty")
+        return torch.cat(parts, dim=-1)
+    return coerce_prompt_waveform(torch.as_tensor(value))
+
+
+def coerce_speaker_embeddings(value: Any, *, use_zero_spk_emb: bool = False) -> list[torch.Tensor] | None:
+    if value is None:
+        return [torch.zeros((192,), dtype=torch.float32)] if use_zero_spk_emb else None
+    if isinstance(value, torch.Tensor):
+        tensor = value.detach()
+        if tensor.ndim == 1:
+            tensor = tensor.unsqueeze(0)
+        if tensor.ndim != 2:
+            raise ValueError(f"Unsupported Ming speaker embedding shape: {tuple(tensor.shape)}")
+        items = [row.reshape(-1).to(torch.float32).cpu() for row in tensor]
+    elif isinstance(value, (list, tuple)):
+        if value and all(not isinstance(item, (list, tuple, torch.Tensor)) for item in value):
+            items = [torch.as_tensor(value).detach().reshape(-1).to(torch.float32).cpu()]
+        else:
+            items = []
+            for item in value:
+                if item is None:
+                    continue
+                if not isinstance(item, torch.Tensor):
+                    item = torch.as_tensor(item)
+                items.append(item.detach().reshape(-1).to(torch.float32).cpu())
+    else:
+        return coerce_speaker_embeddings(torch.as_tensor(value), use_zero_spk_emb=use_zero_spk_emb)
+    if not items:
+        return [torch.zeros((192,), dtype=torch.float32)] if use_zero_spk_emb else None
+    for item in items:
+        if int(item.numel()) != 192:
+            raise ValueError(f"Ming speaker embedding must have 192 dims, got {int(item.numel())}")
+    return items
+
+
+def count_prompt_latent_patches(
+    value: Any,
+    *,
+    patch_size: int = PATCH_SIZE,
+    latent_dim: int = LATENT_DIM,
+) -> int:
+    if value is None:
+        return 0
+    if not isinstance(value, torch.Tensor):
+        value = torch.as_tensor(value)
+    latents = value.detach()
+    if latents.ndim == 3 and latents.shape[0] == 1:
+        latents = latents.squeeze(0)
+    if latents.ndim == 3 and latents.shape[-2:] == (patch_size, latent_dim):
+        return int(latents.shape[0])
+    if latents.ndim != 2 or latents.shape[-1] != latent_dim:
+        raise ValueError(f"Unsupported Ming prompt_latents shape: {tuple(latents.shape)}")
+    if latents.shape[0] % patch_size != 0:
+        raise ValueError(
+            f"Ming prompt_latents frame count must be divisible by patch_size={patch_size}, "
+            f"got frames={int(latents.shape[0])}"
+        )
+    return int(latents.shape[0] // patch_size)
+
+
+def count_prompt_waveform_patches(
+    value: Any,
+    *,
+    patch_size: int = PATCH_SIZE,
+    frame_hop: int = AUDIO_FRAME_HOP,
+    vae_patch_size: int = VAE_PATCH_SIZE,
+) -> int:
+    if value is None:
+        return 0
+    waveform = pad_prompt_waveform(value, patch_size=patch_size)
+    frame_count = int(math.ceil(float(waveform.shape[-1]) / float(frame_hop)))
+    latent_frames = int(math.ceil(float(frame_count) / float(vae_patch_size)))
+    if latent_frames % int(patch_size) != 0:
+        raise ValueError(
+            f"Ming prompt waveform produced latent frame count not divisible by patch_size={patch_size}: "
+            f"frames={latent_frames}"
+        )
+    return int(latent_frames // int(patch_size))
+
+
+def _normalize_prompt_waveform(value: Any, *, target_sr: int) -> torch.Tensor:
+    if isinstance(value, bytes):
+        import torchaudio
+
+        waveform, sr = torchaudio.load(BytesIO(value))
+        waveform = waveform[:1].to(torch.float32)
+        if int(sr) != int(target_sr):
+            from torchaudio.functional import resample as resample_audio
+
+            waveform = resample_audio(waveform, int(sr), int(target_sr))
+        return waveform
+
+    if isinstance(value, tuple) and len(value) == 2 and isinstance(value[1], int):
+        waveform = coerce_prompt_waveform(value[0])
+        if int(value[1]) != int(target_sr):
+            from torchaudio.functional import resample as resample_audio
+
+            waveform = resample_audio(waveform, int(value[1]), int(target_sr))
+        return waveform
+
+    if isinstance(value, dict):
+        samples = value.get("samples", value.get("array", value.get("waveform")))
+        sr = value.get("sample_rate", value.get("sr", target_sr))
+        return _normalize_prompt_waveform((samples, int(sr)), target_sr=target_sr)
+
+    return coerce_prompt_waveform(value)
+
+
+def _coerce_prompt_latents(
+    value: Any,
+    *,
+    patch_size: int,
+    latent_dim: int,
+) -> dict[str, torch.Tensor] | None:
+    if value is None:
+        return None
+    if not isinstance(value, torch.Tensor):
+        value = torch.as_tensor(value)
+
+    latents = value.detach()
+    if latents.ndim == 3 and latents.shape[0] == 1:
+        latents = latents.squeeze(0)
+
+    if latents.ndim == 3 and latents.shape[-2:] == (patch_size, latent_dim):
+        patches = latents
+        # [Patch, Time, Dimension] -> [Frame, Dimension] for history seeding.
+        frames = patches.reshape(-1, latent_dim)
+        return {"patches": patches, "frames": frames}
+
+    if latents.ndim != 2 or latents.shape[-1] != latent_dim:
+        raise ValueError(f"Unsupported prompt latent shape: {tuple(latents.shape)}")
+    if latents.shape[0] % patch_size != 0:
+        raise ValueError(
+            f"Prompt latent frame count must be divisible by patch_size={patch_size}, "
+            f"got frames={int(latents.shape[0])}"
+        )
+    # [Frame, Dimension] -> [Patch, Time, Dimension] for Aggregator prompt slots.
+    patches = latents.reshape(-1, patch_size, latent_dim) if latents.shape[0] > 0 else None
+    return {"patches": patches, "frames": latents}
+
+
+def _initial_history(
+    frames: torch.Tensor | None,
+    *,
+    history_size: int,
+    latent_dim: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    history = torch.zeros((history_size, latent_dim), device=device, dtype=dtype)
+    if frames is None or frames.numel() == 0:
+        return history
+    frames = frames.to(device=device, dtype=dtype)
+    take = min(history_size, int(frames.shape[0]))
+    history[-take:] = frames[-take:]
+    return history
+
+
+def _take_scalar(value: Any, idx: int) -> float | None:
+    if not isinstance(value, torch.Tensor) or value.numel() == 0:
+        return None
+    return float(value.reshape(-1)[idx].item())
+
+
+def _find_audio_placeholder_positions(input_ids: torch.Tensor, cfg: MingTTSConfig) -> torch.Tensor:
+    dummy_pos = (input_ids == cfg.audio_dummy_token_id).nonzero(as_tuple=True)[0]
+    if dummy_pos.numel() == 0:
+        return dummy_pos
+
+    audio_start_pos = (input_ids == cfg.audio_start_token_id).nonzero(as_tuple=True)[0]
+    audio_end_pos = (input_ids == cfg.audio_end_token_id).nonzero(as_tuple=True)[0]
+    if audio_start_pos.numel() == 0:
+        return dummy_pos
+
+    start = int(audio_start_pos[0].item())
+    end = int(audio_end_pos[0].item()) if audio_end_pos.numel() > 0 else int(input_ids.shape[0])
+    keep = (dummy_pos > start) & (dummy_pos < end)
+    filtered = dummy_pos[keep]
+    return filtered if filtered.numel() > 0 else dummy_pos
+
+
+def _find_speaker_placeholder_positions(input_ids: torch.Tensor, hf_config: Any) -> list[int]:
+    vision_start_token_id = getattr(hf_config, "vision_start_token_id", VISION_START_TOKEN_ID)
+    vision_start_pos = (input_ids == int(vision_start_token_id)).nonzero(as_tuple=True)[0]
+    if vision_start_pos.numel() == 0:
+        return []
+
+    slots = []
+    for pos in vision_start_pos:
+        slot = int(pos.item()) + 1
+        if slot < int(input_ids.shape[0]):
+            slots.append(slot)
+    return slots
+
+
+__all__ = [
+    "coerce_prompt_waveform",
+    "coerce_speaker_embeddings",
+    "count_prompt_latent_patches",
+    "count_prompt_waveform_patches",
+    "pad_prompt_waveform",
+    "_coerce_prompt_latents",
+    "_find_audio_placeholder_positions",
+    "_find_speaker_placeholder_positions",
+    "_initial_history",
+    "_normalize_prompt_waveform",
+    "_take_scalar",
+]

vllm_omni/model_executor/models/ming_tts/config_ming_tts.py

@@ -0,0 +1,209 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Ming dense checkpoint config adapters."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from transformers import PretrainedConfig, Qwen2Config
+
+from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig
+
+from .constants import (
+    AGGREGATOR_HIDDEN_SIZE,
+    AUDIO_DUMMY_TOKEN_ID,
+    AUDIO_END_TOKEN_ID,
+    AUDIO_EOS_TOKEN_ID,
+    AUDIO_FRAME_HOP,
+    AUDIO_START_TOKEN_ID,
+    DEFAULT_CFG,
+    DEFAULT_SIGMA,
+    DEFAULT_TEMPERATURE,
+    HISTORY_PATCH_SIZE,
+    KEY_CFG,
+    KEY_CHUNK_ID,
+    KEY_DECODE_STEP,
+    KEY_LAST_STOP_PROB,
+    KEY_LATENT_HISTORY,
+    KEY_MAX_DECODE_STEPS,
+    KEY_MIN_DECODE_STEPS,
+    KEY_NEXT_EMBEDS,
+    KEY_PROMPT_LATENT_TAIL,
+    KEY_PROMPT_LATENTS,
+    KEY_REQUEST_ID,
+    KEY_SIGMA,
+    KEY_SPEAKER_EMBEDDING,
+    KEY_TEMPERATURE,
+    KEY_TEXT_MODE,
+    LATENT_CHUNK_SIZE,
+    LATENT_DIM,
+    LATENT_LEFT_CONTEXT,
+    LLM_HIDDEN_SIZE,
+    LLM_VOCAB_SIZE,
+    MAX_DECODE_STEPS,
+    PATCH_SIZE,
+    SAMPLE_RATE,
+    STOP_HEAD_MIN_STEPS,
+    STOP_HEAD_THRESHOLD,
+    TEXT_EOS_TOKEN_ID,
+    VAE_PATCH_SIZE,
+    VISION_START_TOKEN_ID,
+)
+from .validation import _coerce_audio_vae_config, _nested_get, _to_plain_dict, validate_ming_tts_config
+
+
+def _coerce_qwen2_config(value: Any) -> Qwen2Config:
+    if isinstance(value, Qwen2Config):
+        return value
+    if isinstance(value, PretrainedConfig):
+        return Qwen2Config.from_dict(value.to_dict())
+    if isinstance(value, dict):
+        return Qwen2Config.from_dict(dict(value))
+    raise TypeError(f"Unsupported llm_config type for Ming dense config: {type(value)!r}")
+
+
+class MingDenseConfig(PretrainedConfig):
+    # The upstream checkpoint declares model_type="dense". Keep it for HF
+    # config compatibility; deploy/ming_tts.yaml selects the vLLM-Omni
+    # pipeline via pipeline: ming_tts.
+    model_type = "dense"
+
+    def __init__(
+        self,
+        llm_config: Qwen2Config | dict[str, Any] | None = None,
+        ditar_config: dict[str, Any] | None = None,
+        aggregator_config: dict[str, Any] | None = None,
+        audio_tokenizer_config: AudioVAEConfig | dict[str, Any] | None = None,
+        architectures: list[str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        self.llm_config = _coerce_qwen2_config(llm_config or {})
+        self.ditar_config = dict(ditar_config or {})
+        self.aggregator_config = dict(aggregator_config or {})
+        self.audio_tokenizer_config = _coerce_audio_vae_config(audio_tokenizer_config)
+        super().__init__(architectures=architectures, **kwargs)
+
+    def get_text_config(self, decoder: bool = False, **kwargs: Any) -> Qwen2Config:
+        del decoder, kwargs
+        return self.llm_config
+
+
+@dataclass
+class MingTTSConfig:
+    """Flat config object shared by Stage-1 and Stage-2. Build via from_hf_config()."""
+
+    llm_hidden_size: int = LLM_HIDDEN_SIZE
+    llm_vocab_size: int = LLM_VOCAB_SIZE
+    llm_config: dict[str, Any] = field(default_factory=dict)
+
+    latent_dim: int = LATENT_DIM
+    patch_size: int = PATCH_SIZE
+    history_patch_size: int = HISTORY_PATCH_SIZE
+
+    ditar_config: dict[str, Any] = field(default_factory=dict)
+    aggregator_config: dict[str, Any] = field(default_factory=dict)
+
+    audio_tokenizer_config: AudioVAEConfig | None = None
+    vae_patch_size: int = VAE_PATCH_SIZE
+    sample_rate: int = SAMPLE_RATE
+    audio_frame_hop: int = AUDIO_FRAME_HOP
+
+    cfg: float = DEFAULT_CFG
+    sigma: float = DEFAULT_SIGMA
+    temperature: float = DEFAULT_TEMPERATURE
+    stop_head_min_steps: int = STOP_HEAD_MIN_STEPS
+    stop_head_threshold: float = STOP_HEAD_THRESHOLD
+    max_decode_steps: int = MAX_DECODE_STEPS
+
+    latent_chunk_size: int = LATENT_CHUNK_SIZE
+    latent_left_context: int = LATENT_LEFT_CONTEXT
+
+    text_eos_token_id: int = TEXT_EOS_TOKEN_ID
+    audio_dummy_token_id: int = AUDIO_DUMMY_TOKEN_ID
+    audio_start_token_id: int = AUDIO_START_TOKEN_ID
+    audio_end_token_id: int = AUDIO_END_TOKEN_ID
+    audio_eos_token_id: int = AUDIO_EOS_TOKEN_ID
+
+    @classmethod
+    def from_hf_config(cls, hf_config: PretrainedConfig) -> MingTTSConfig:
+        llm_raw = getattr(hf_config, "llm_config", {}) or {}
+        ditar_raw = getattr(hf_config, "ditar_config", {}) or {}
+        agg_raw = getattr(hf_config, "aggregator_config", {}) or {}
+        atc_raw = getattr(hf_config, "audio_tokenizer_config", None)
+
+        llm_dict = _to_plain_dict(llm_raw)
+        ditar = _to_plain_dict(ditar_raw)
+        agg = _to_plain_dict(agg_raw)
+        ditar.setdefault("attn_backend", "torch")
+
+        atc = _coerce_audio_vae_config(atc_raw)
+        atc_enc_latent_dim = _nested_get(atc, "enc_kwargs", "latent_dim", default=LATENT_DIM)
+        atc_patch_size = _nested_get(atc, "patch_size", default=VAE_PATCH_SIZE)
+        atc_sample_rate = _nested_get(atc, "sample_rate", default=SAMPLE_RATE)
+
+        enc_hop_size = _nested_get(atc, "enc_kwargs", "hop_size", default=AUDIO_FRAME_HOP)
+
+        cfg = cls(
+            llm_hidden_size=llm_dict.get("hidden_size", LLM_HIDDEN_SIZE),
+            llm_vocab_size=llm_dict.get("vocab_size", LLM_VOCAB_SIZE),
+            llm_config=llm_dict,
+            latent_dim=atc_enc_latent_dim,
+            patch_size=ditar.get("patch_size", PATCH_SIZE),
+            history_patch_size=ditar.get("history_patch_size", HISTORY_PATCH_SIZE),
+            ditar_config=ditar,
+            aggregator_config=agg,
+            audio_tokenizer_config=atc,
+            vae_patch_size=atc_patch_size,
+            sample_rate=atc_sample_rate,
+            audio_frame_hop=enc_hop_size if enc_hop_size is not None else AUDIO_FRAME_HOP,
+        )
+        return cfg
+
+    def validate(self) -> None:
+        validate_ming_tts_config(self)
+
+
+__all__ = [
+    "AGGREGATOR_HIDDEN_SIZE",
+    "AUDIO_DUMMY_TOKEN_ID",
+    "AUDIO_END_TOKEN_ID",
+    "AUDIO_EOS_TOKEN_ID",
+    "AUDIO_FRAME_HOP",
+    "AUDIO_START_TOKEN_ID",
+    "DEFAULT_CFG",
+    "DEFAULT_SIGMA",
+    "DEFAULT_TEMPERATURE",
+    "HISTORY_PATCH_SIZE",
+    "KEY_CFG",
+    "KEY_CHUNK_ID",
+    "KEY_DECODE_STEP",
+    "KEY_LAST_STOP_PROB",
+    "KEY_LATENT_HISTORY",
+    "KEY_MAX_DECODE_STEPS",
+    "KEY_MIN_DECODE_STEPS",
+    "KEY_NEXT_EMBEDS",
+    "KEY_PROMPT_LATENT_TAIL",
+    "KEY_PROMPT_LATENTS",
+    "KEY_REQUEST_ID",
+    "KEY_SIGMA",
+    "KEY_SPEAKER_EMBEDDING",
+    "KEY_TEMPERATURE",
+    "KEY_TEXT_MODE",
+    "LATENT_CHUNK_SIZE",
+    "LATENT_DIM",
+    "LATENT_LEFT_CONTEXT",
+    "LLM_HIDDEN_SIZE",
+    "LLM_VOCAB_SIZE",
+    "MAX_DECODE_STEPS",
+    "MingDenseConfig",
+    "MingTTSConfig",
+    "PATCH_SIZE",
+    "SAMPLE_RATE",
+    "STOP_HEAD_MIN_STEPS",
+    "STOP_HEAD_THRESHOLD",
+    "TEXT_EOS_TOKEN_ID",
+    "VAE_PATCH_SIZE",
+    "VISION_START_TOKEN_ID",
+]

vllm_omni/model_executor/models/ming_tts/constants.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from __future__ import annotations
+
+# ---------------------------------------------------------------------------
+# Token IDs (confirmed from tokenizer_config.json)
+# ---------------------------------------------------------------------------
+
+AUDIO_DUMMY_TOKEN_ID = 151705  # <audioPatch>
+AUDIO_START_TOKEN_ID = 151706  # <audio>
+AUDIO_END_TOKEN_ID = 151707  # </audio>
+AUDIO_EOS_TOKEN_ID = 151704  # <end_of_audio>
+VISION_START_TOKEN_ID = 151652  # <|vision_start|>
+
+TEXT_EOS_TOKEN_ID = 151669  # <text_eos>
+
+
+# ---------------------------------------------------------------------------
+# Architectural constants (confirmed from original config.json)
+# ---------------------------------------------------------------------------
+
+LATENT_DIM = 64
+PATCH_SIZE = 4
+HISTORY_PATCH_SIZE = 32
+LLM_HIDDEN_SIZE = 896
+LLM_VOCAB_SIZE = 151936
+AGGREGATOR_HIDDEN_SIZE = 1024
+VAE_PATCH_SIZE = 4
+SAMPLE_RATE = 44100
+
+# AudioVAE frame/hop geometry (confirmed)
+AUDIO_FRAME_HOP = 882  # enc input_dim / hop_size / dec output_dim
+
+# stop_head defaults
+STOP_HEAD_MIN_STEPS = 3
+STOP_HEAD_THRESHOLD = 0.5
+
+# FlowLoss sampling defaults
+DEFAULT_CFG = 2.0
+DEFAULT_SIGMA = 0.25
+DEFAULT_TEMPERATURE = 0.0
+
+# Connector / Stage-2 streaming defaults (runtime tuning)
+LATENT_CHUNK_SIZE = 25
+LATENT_LEFT_CONTEXT = 0
+MAX_DECODE_STEPS = 200
+
+
+# ---------------------------------------------------------------------------
+# seq_data.extra_data keys
+# ---------------------------------------------------------------------------
+
+KEY_LATENT_HISTORY = "ming_latent_history"
+KEY_DECODE_STEP = "ming_decode_step"
+KEY_LAST_STOP_PROB = "ming_last_stop_prob"
+KEY_NEXT_EMBEDS = "ming_next_embeds"
+KEY_PROMPT_LATENTS = "ming_prompt_latents"
+KEY_PROMPT_LATENT_TAIL = "ming_prompt_latent_tail"
+KEY_SPEAKER_EMBEDDING = "ming_speaker_embedding"
+KEY_REQUEST_ID = "ming_request_id"
+KEY_CHUNK_ID = "ming_chunk_id"
+KEY_CFG = "ming_cfg"
+KEY_SIGMA = "ming_sigma"
+KEY_TEMPERATURE = "ming_temperature"
+KEY_MAX_DECODE_STEPS = "ming_max_decode_steps"
+KEY_MIN_DECODE_STEPS = "ming_min_decode_steps"
+KEY_TEXT_MODE = "ming_text_mode"

vllm_omni/model_executor/models/ming_tts/flowloss_head.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/fm/flowloss.py
+
+import torch
+import torch.nn as nn
+
+from vllm_omni.model_executor.models.common.ming.fm import Solver, build_timesteps
+
+from .fm.dit import DiT
+
+
+class FlowLoss(nn.Module):
+    """Diffusion Loss"""
+
+    def __init__(self, z_channels, llm_cond_dim, **kwargs):
+        super().__init__()
+        self.z_channels = z_channels
+        # Preserve checkpoint parameter names under flowloss.cfm.model.* while
+        # keeping the inference-only CFM sampling logic local to FlowLoss.
+        self.cfm = nn.Module()
+        self.cfm.add_module("model", DiT(in_channels=z_channels, llm_cond_dim=llm_cond_dim, **kwargs))
+
+    @torch.no_grad()
+    def sample(self, z, latent_history, cfg=2.0, patch_size=1, sigma=0.25, temperature=0):
+        # z: [Batch, Time=1, Dimension]
+        # latent_history: [Batch, History_Time, Dimension]
+        if z.ndim != 3:
+            raise ValueError(f"Expected conditioning rank-3 [Batch, Time, Dimension], got {tuple(z.shape)}")
+        if latent_history.ndim != 3:
+            raise ValueError(
+                f"Expected latent_history rank-3 [Batch, Time, Dimension], got {tuple(latent_history.shape)}"
+            )
+        if z.shape[0] != latent_history.shape[0]:
+            raise ValueError(
+                f"Batch mismatch across conditioning and latent_history: {z.shape[0]}, {latent_history.shape[0]}"
+            )
+        dtype = z.dtype
+        noise = torch.randn(z.shape[0], self.z_channels, patch_size, device=z.device, dtype=torch.float32)
+
+        def fn(t, x):
+            x = x.to(dtype)
+            t = t.to(dtype)
+            if cfg < 1e-5:
+                if t.ndim == 0:
+                    t = t.repeat(x.shape[0])
+                pred = self.cfm.model(x=x, t=t, c=torch.zeros_like(z), latent_history=latent_history)
+                return pred[:, -patch_size:, :].float()
+
+            pred_cfg = self.cfm.model.forward_with_cfg(
+                x=x, t=t, c=z, latent_history=latent_history, cfg_scale=cfg, patch_size=patch_size
+            )
+            pred, null_pred = torch.chunk(pred_cfg, 2, dim=0)
+            return (pred + (pred - null_pred) * cfg).float()
+
+        t = build_timesteps(
+            10,
+            device=noise.device,
+            dtype=noise.dtype,
+            use_epss=True,
+            sway_sampling_coef=-1.0,
+        )
+        return Solver(fn, noise.transpose(1, 2), sigma=sigma, temperature=temperature).integrate(t)[-1].to(dtype)

vllm_omni/model_executor/models/ming_tts/fm/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/model_executor/models/ming_tts/fm/dit.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/fm/dit.py
+
+import torch
+import torch.nn as nn
+from x_transformers.x_transformers import RotaryEmbedding
+
+from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding
+from vllm_omni.model_executor.models.common.ming.dit import CondEmbedder, DiTBlock, FinalLayer
+
+
+class DiT(nn.Module):
+    def __init__(
+        self,
+        in_channels=4,
+        hidden_size=1024,
+        depth=28,
+        num_heads=16,
+        mlp_ratio=4.0,
+        llm_cond_dim=896,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.num_heads = num_heads
+        self.t_embedder = DiTTimestepEmbedding(hidden_size)
+        self.x_embedder = nn.Linear(in_channels, hidden_size)
+        self.c_embedder = CondEmbedder(llm_cond_dim, hidden_size)
+        self.hidden_size = hidden_size
+        self.rotary_embed = RotaryEmbedding(hidden_size // num_heads)
+        self.blocks = nn.ModuleList(
+            [DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **kwargs) for _ in range(depth)]
+        )
+        self.final_layer = FinalLayer(hidden_size, self.out_channels)
+
+    def forward(self, x, t, c, latent_history, mask=None):
+        if x.ndim != 3:
+            raise ValueError(f"Expected x rank-3 [Batch, Time, Dimension], got {tuple(x.shape)}")
+        if latent_history.ndim != 3:
+            raise ValueError(
+                f"Expected latent_history rank-3 [Batch, Time, Dimension], got {tuple(latent_history.shape)}"
+            )
+        if c.ndim != 3:
+            raise ValueError(f"Expected conditioning rank-3 [Batch, Time, Dimension], got {tuple(c.shape)}")
+        if x.shape[0] != latent_history.shape[0] or x.shape[0] != c.shape[0]:
+            raise ValueError(
+                "Batch mismatch across x, conditioning, and latent_history: "
+                f"{x.shape[0]}, {c.shape[0]}, {latent_history.shape[0]}"
+            )
+        if x.shape[-1] != self.in_channels:
+            raise ValueError(f"x feature dim mismatch: got {x.shape[-1]}, expected {self.in_channels}")
+        if latent_history.shape[-1] != self.in_channels:
+            raise ValueError(
+                f"latent_history feature dim mismatch: got {latent_history.shape[-1]}, expected {self.in_channels}"
+            )
+        if t.ndim == 0:
+            t = t.reshape(1)
+        if t.ndim != 1:
+            raise ValueError(f"Expected timestep rank-1 [Batch], got {tuple(t.shape)}")
+        if t.shape[0] != x.shape[0]:
+            raise ValueError(f"Timestep batch mismatch: got {t.shape[0]}, expected {x.shape[0]}")
+
+        t = self.t_embedder(t).unsqueeze(1)
+        x_now = self.x_embedder(x)
+        x_history = self.x_embedder(latent_history)
+        x = torch.cat([x_history, x_now], dim=1)
+        c = self.c_embedder(c)
+        y = t + c
+        x = torch.cat([y, x], dim=1)
+        rope = self.rotary_embed.forward_from_seq_len(x.shape[1])
+
+        if mask is not None:
+            if mask.ndim != 2:
+                raise ValueError(f"Expected mask rank-2 [Batch, Time], got {tuple(mask.shape)}")
+            if mask.shape[0] != x_now.shape[0] or mask.shape[1] != x_now.shape[1]:
+                raise ValueError(
+                    f"Mask shape mismatch: got {tuple(mask.shape)}, expected {(x_now.shape[0], x_now.shape[1])}"
+                )
+            mask_pad = mask.clone().detach()[:, :1].expand(-1, x_history.shape[1] + c.shape[1])
+            mask = torch.cat([mask_pad, mask], dim=-1)
+        for block in self.blocks:
+            x = block(x, mask, rope)
+        x = self.final_layer(x)
+        return x
+
+    def forward_with_cfg(self, x, t, c, cfg_scale, latent_history, patch_size):
+        if not cfg_scale == 1:
+            x = torch.cat([x, x], dim=0)
+            latent_history = torch.cat([latent_history, latent_history], dim=0)
+            fake_latent = torch.zeros_like(c)
+            c = torch.cat([c, fake_latent], dim=0)
+        if t.ndim == 0:
+            t = t.repeat(x.shape[0])
+        model_out = self.forward(x, t, c, latent_history)
+        return model_out[:, -patch_size:, :]

vllm_omni/model_executor/models/ming_tts/ming_tts.py

@@ -0,0 +1,269 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from __future__ import annotations
+
+from collections.abc import Iterable
+from functools import cached_property
+from typing import Any
+
+import torch
+import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.model_executor.models import SupportsPP
+from vllm.model_executor.models.utils import init_vllm_registered_model
+from vllm.v1.sample.sampler import Sampler
+
+from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
+
+from .audio_prep import (
+    _coerce_prompt_latents,
+    _find_audio_placeholder_positions,
+    _find_speaker_placeholder_positions,
+    _initial_history,
+    _take_scalar,
+    coerce_speaker_embeddings,
+)
+from .config_ming_tts import (
+    AUDIO_START_TOKEN_ID,
+    KEY_CFG,
+    KEY_DECODE_STEP,
+    KEY_LAST_STOP_PROB,
+    KEY_LATENT_HISTORY,
+    KEY_MAX_DECODE_STEPS,
+    KEY_MIN_DECODE_STEPS,
+    KEY_NEXT_EMBEDS,
+    KEY_PROMPT_LATENT_TAIL,
+    KEY_REQUEST_ID,
+    KEY_SIGMA,
+    KEY_SPEAKER_EMBEDDING,
+    KEY_TEMPERATURE,
+    KEY_TEXT_MODE,
+    MingTTSConfig,
+)
+from .patch_emission import MING_STOP_REASON_KEY
+from .prompt_encoder import _resolve_prompt_latents
+
+
+class _ModelSampleAdapter(nn.Module):
+    def __init__(self, model: nn.Module):
+        super().__init__()
+        self.model = model
+
+    def forward(self, logits, sampling_metadata):
+        return self.model.sample(logits, sampling_metadata)
+
+
+class MingTTSForConditionalGeneration(nn.Module, SupportsPP, CustomProcessMixin):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        del prefix
+        self.vllm_config = vllm_config
+        self.ming_config = MingTTSConfig.from_hf_config(vllm_config.model_config.hf_config)
+        self.ming_config.validate()
+        self.have_multimodal_outputs = True
+        self.has_preprocess = False
+        self.has_postprocess = False
+        self.requires_raw_input_tokens = False
+        self.model_stage = vllm_config.model_config.model_stage
+        self._prompt_encoder = None
+
+        if self.model_stage == "llm":
+            self.model = init_vllm_registered_model(vllm_config=vllm_config, architectures=["MingLLMModel"])
+            self.has_preprocess = True
+            self.has_postprocess = True
+            self.set_custom_preprocess(self.preprocess)
+            self.set_custom_postprocess(self.postprocess)
+        elif self.model_stage == "audio_vae":
+            self.model = init_vllm_registered_model(vllm_config=vllm_config, architectures=["MingAudioVAEModel"])
+            self.requires_raw_input_tokens = True
+        else:
+            raise ValueError(f"Invalid Ming model_stage={self.model_stage}")
+        self.make_empty_intermediate_tensors = getattr(self.model, "make_empty_intermediate_tensors", lambda: None)
+
+    @cached_property
+    def sampler(self):
+        if hasattr(self.model, "sample"):
+            return _ModelSampleAdapter(self.model)
+        if hasattr(self.model, "sampler"):
+            return self.model.sampler
+        return Sampler()
+
+    def embed_input_ids(self, input_ids: torch.Tensor, **kwargs: Any) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids=input_ids, **kwargs)
+
+    def forward(self, *args: Any, **kwargs: Any):
+        return self.model(*args, **kwargs)
+
+    def compute_logits(self, hidden_states, sampling_metadata=None):
+        return self.model.compute_logits(hidden_states, sampling_metadata=sampling_metadata)
+
+    def sample(self, logits, sampling_metadata):
+        return self.model.sample(logits, sampling_metadata) if hasattr(self.model, "sample") else None
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        weights = list(weights)
+        if self.model_stage == "llm":
+            allowed = ("model.", "linear_proj_audio.", "flowloss.", "stop_head.", "spk_head.")
+            llm_weights = [(k, v) for k, v in weights if k.startswith(allowed)]
+            if not llm_weights:
+                raise RuntimeError(
+                    "Ming Stage-0 received no loadable checkpoint weights. "
+                    "Expected prefixes: model.*, linear_proj_audio.*, flowloss.*, stop_head.*, spk_head.*"
+                )
+            loaded = self.model.load_weights(llm_weights)
+            return {f"model.{name}" for name in loaded}
+
+        audio_weights = [(k, v) for k, v in weights if k.startswith("audio.")]
+        if not audio_weights:
+            raise RuntimeError("Ming Stage-1 received no loadable checkpoint weights. Expected prefix: audio.*")
+        loaded = self.model.load_weights(audio_weights)
+        return {f"model.{name}" for name in loaded}
+
+    def _resolve_prompt_latents(self, info_dict: dict[str, Any]):
+        return _resolve_prompt_latents(self, info_dict)
+
+    def preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor | None, **info_dict: Any):
+        if self.model_stage != "llm":
+            return input_ids, input_embeds, {}
+        input_embeds = self.model.embed_input_ids(input_ids).clone()
+        return (
+            self._prefill_preprocess(input_ids, input_embeds, **info_dict)
+            if int(input_ids.shape[0]) > 1
+            else self._decode_preprocess(input_ids, input_embeds, **info_dict)
+        )
+
+    def postprocess(self, hidden_states: torch.Tensor, **info_dict: Any) -> dict[str, Any]:
+        if self.model_stage != "llm" or hidden_states.numel() == 0:
+            return {}
+        req_id = info_dict.get(KEY_REQUEST_ID, info_dict.get("request_id"))
+        pending = self.model.pop_postprocess_update(req_id)
+        if not pending or not isinstance(pending.get("ming_latent_patch"), torch.Tensor):
+            return {}
+
+        update = {
+            KEY_LATENT_HISTORY: pending[KEY_LATENT_HISTORY].detach().to("cpu").contiguous(),
+            KEY_NEXT_EMBEDS: pending[KEY_NEXT_EMBEDS].detach().to("cpu").contiguous(),
+            KEY_DECODE_STEP: int(info_dict.get(KEY_DECODE_STEP, 0)) + 1,
+        }
+        stop_prob = _take_scalar(pending.get("ming_stop_prob"), 0)
+        if stop_prob is not None:
+            update[KEY_LAST_STOP_PROB] = stop_prob
+        stop_reason = pending.get(MING_STOP_REASON_KEY)
+        if isinstance(stop_reason, str):
+            update[MING_STOP_REASON_KEY] = stop_reason
+        if isinstance(req_id, str):
+            update[KEY_REQUEST_ID] = req_id
+        return update
+
+    def _prefill_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: Any):
+        if bool(info_dict.get(KEY_TEXT_MODE, False)):
+            update: dict[str, Any] = {KEY_TEXT_MODE: True}
+            request_id = info_dict.get(KEY_REQUEST_ID, info_dict.get("request_id"))
+            if request_id is not None:
+                update[KEY_REQUEST_ID] = request_id
+            if int(input_ids.shape[0]) > 1 and int(input_ids[-1].item()) == AUDIO_START_TOKEN_ID:
+                return input_ids[:-1], input_embeds[:-1], update
+            return input_ids, input_embeds, update
+
+        update: dict[str, Any] = {KEY_DECODE_STEP: int(info_dict.get(KEY_DECODE_STEP, 0))}
+        prompt_latents = _resolve_prompt_latents(self, info_dict)
+        history = _initial_history(
+            prompt_latents["frames"] if prompt_latents is not None else None,
+            history_size=self.ming_config.history_patch_size,
+            latent_dim=self.ming_config.latent_dim,
+            device=input_embeds.device,
+            dtype=torch.float32,
+        )
+        update[KEY_LATENT_HISTORY] = history.detach().to("cpu").contiguous()
+        update[KEY_PROMPT_LATENT_TAIL] = update[KEY_LATENT_HISTORY]
+
+        speaker_embedding = info_dict.get(KEY_SPEAKER_EMBEDDING, info_dict.get("speaker_embedding"))
+        speaker_embeddings = None
+        if speaker_embedding is not None:
+            speaker_embeddings = coerce_speaker_embeddings(
+                speaker_embedding,
+                use_zero_spk_emb=bool(info_dict.get("use_zero_spk_emb", False)),
+            )
+        if speaker_embeddings is not None and len(speaker_embeddings) > 0:
+            speaker_slots = _find_speaker_placeholder_positions(input_ids, self.vllm_config.model_config.hf_config)
+            if len(speaker_slots) < len(speaker_embeddings):
+                raise RuntimeError(
+                    "Could not locate enough speaker placeholder slots: "
+                    f"found {len(speaker_slots)}, need {len(speaker_embeddings)}"
+                )
+            for speaker_slot, spk in zip(speaker_slots, speaker_embeddings):
+                spk_proj = self.model.project_speaker_embedding(
+                    spk.to(device=input_embeds.device, dtype=input_embeds.dtype).unsqueeze(0)
+                ).squeeze(0)
+                input_embeds[speaker_slot] = spk_proj
+
+        if prompt_latents is not None and prompt_latents["patches"] is not None:
+            prompt_patches = prompt_latents["patches"].to(dtype=getattr(self.model, "fm_dtype", torch.float32))
+            prompt_embeds = self.model.linear_proj_audio(prompt_patches).squeeze(1)
+            placeholder_pos = _find_audio_placeholder_positions(input_ids, self.ming_config)
+            take = min(int(placeholder_pos.numel()), int(prompt_embeds.shape[0]))
+            if take > 0:
+                input_embeds[placeholder_pos[:take]] = prompt_embeds[:take].to(dtype=input_embeds.dtype)
+
+        request_id = info_dict.get(KEY_REQUEST_ID, info_dict.get("request_id"))
+        if request_id is not None:
+            update[KEY_REQUEST_ID] = request_id
+        _copy_runtime_controls(update, info_dict)
+        return input_ids, input_embeds, update
+
+    def _decode_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: Any):
+        if bool(info_dict.get(KEY_TEXT_MODE, False)):
+            update: dict[str, Any] = {KEY_TEXT_MODE: True}
+            request_id = info_dict.get(KEY_REQUEST_ID, info_dict.get("request_id"))
+            if request_id is not None:
+                update[KEY_REQUEST_ID] = request_id
+            return input_ids, input_embeds, update
+
+        update: dict[str, Any] = {KEY_DECODE_STEP: int(info_dict.get(KEY_DECODE_STEP, 0))}
+        history = info_dict.get(KEY_LATENT_HISTORY)
+        if isinstance(history, torch.Tensor):
+            update[KEY_LATENT_HISTORY] = history.detach().to("cpu").contiguous()
+        else:
+            zero_history = torch.zeros(
+                (self.ming_config.history_patch_size, self.ming_config.latent_dim),
+                device=input_embeds.device,
+                dtype=torch.float32,
+            )
+            update[KEY_LATENT_HISTORY] = zero_history.detach().to("cpu").contiguous()
+
+        next_embeds = info_dict.get(KEY_NEXT_EMBEDS)
+        if isinstance(next_embeds, torch.Tensor) and input_ids.numel() == 1:
+            if not torch.isfinite(next_embeds).all():
+                raise RuntimeError("Non-finite next_embeds before decode preprocess write.")
+            input_embeds[0] = (
+                next_embeds.detach()
+                .reshape(-1, self.ming_config.llm_hidden_size)[0]
+                .to(
+                    device=input_embeds.device,
+                    dtype=input_embeds.dtype,
+                )
+            )
+            if not torch.isfinite(input_embeds[0]).all():
+                raise RuntimeError("Non-finite backbone input_embeds after decode preprocess write.")
+
+        request_id = info_dict.get(KEY_REQUEST_ID, info_dict.get("request_id"))
+        if request_id is not None:
+            update[KEY_REQUEST_ID] = request_id
+        _copy_runtime_controls(update, info_dict)
+        return input_ids, input_embeds, update
+
+
+def _copy_runtime_controls(update: dict[str, Any], info_dict: dict[str, Any]) -> None:
+    for key in (KEY_CFG, KEY_SIGMA, KEY_TEMPERATURE, KEY_MAX_DECODE_STEPS, KEY_MIN_DECODE_STEPS):
+        if key in info_dict:
+            update[key] = info_dict[key]
+
+
+__all__ = [
+    "MingTTSForConditionalGeneration",
+    "_ModelSampleAdapter",
+    "_coerce_prompt_latents",
+    "_find_audio_placeholder_positions",
+    "_find_speaker_placeholder_positions",
+    "_initial_history",
+]

vllm_omni/model_executor/models/ming_tts/ming_tts_audio_vae.py

@@ -0,0 +1,389 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/audio_tokenizer/modeling_audio_vae.py
+from __future__ import annotations
+
+import os
+import time
+import warnings
+from collections import OrderedDict
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+from .config_ming_tts import KEY_CHUNK_ID, KEY_REQUEST_ID, MingTTSConfig
+from .patch_emission import MING_STOP_REASON_KEY
+
+logger = init_logger(__name__)
+
+MING_FINAL_DECODE_STEP_KEY = "ming_final_decode_step"
+_STREAM_STATE_TTL_SECONDS = 15 * 60
+_MAX_STREAM_STATES = 1024
+
+
+class MingAudioVAEModel(nn.Module):
+    input_modalities = "audio"
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        del prefix
+        self.vllm_config = vllm_config
+        self.ming_config = MingTTSConfig.from_hf_config(vllm_config.model_config.hf_config)
+        if self.ming_config.audio_tokenizer_config is None:
+            raise ValueError("MingAudioVAEModel requires audio_tokenizer_config")
+        self.ming_config.validate()
+
+        self.audio = AudioVAE(self.ming_config.audio_tokenizer_config)
+        self.have_multimodal_outputs = True
+        self.has_preprocess = False
+        self.has_postprocess = False
+        self.enable_update_additional_information = True
+        self.requires_raw_input_tokens = True
+
+        self._past_key_values: dict[str, Any] = {}
+        self._stream_state: dict[str, tuple[Any, Any, Any]] = {}
+        self._patch_totals: dict[str, int] = {}
+        self._sample_totals: dict[str, int] = {}
+        self._state_access_times: OrderedDict[str, float] = OrderedDict()
+
+    def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
+        hidden_size = int(self.ming_config.llm_hidden_size)
+        if input_ids is None or input_ids.numel() == 0:
+            return torch.empty((0, hidden_size), device=self.vllm_config.device_config.device, dtype=torch.float32)
+        return torch.zeros((input_ids.shape[0], hidden_size), device=input_ids.device, dtype=torch.float32)
+
+    def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> None:
+        del hidden_states, sampling_metadata
+        return None
+
+    def chunked_decode_streaming(
+        self,
+        latent_chunk: torch.Tensor,
+        *,
+        request_id: str,
+        finished: bool,
+    ) -> tuple[torch.Tensor, Any, Any, bool, bool]:
+        had_past_key_values = request_id in self._past_key_values
+        had_stream_state = _has_stream_state(self._stream_state.get(request_id))
+        stream_state = self._stream_state.get(request_id, (None, None, None))
+        past_key_values = self._past_key_values.get(request_id)
+        waveform_parts: list[torch.Tensor] = []
+
+        patch_count = int(latent_chunk.shape[0])
+        for patch_idx in range(patch_count):
+            # [Batch, Time, Dimension] = [1, patch_size, latent_dim]
+            latent_patch = latent_chunk[patch_idx : patch_idx + 1]
+            is_last_patch = finished and patch_idx == patch_count - 1
+            waveform, stream_state, past_key_values = self.audio.decode(
+                latent_patch,
+                past_key_values=past_key_values,
+                use_cache=True,
+                stream_state=stream_state,
+                last_chunk=is_last_patch,
+            )
+            waveform_parts.append(waveform.reshape(-1).to(torch.float32))
+
+        waveform_flat = (
+            torch.cat(waveform_parts, dim=0)
+            if waveform_parts
+            else torch.zeros((0,), dtype=torch.float32, device=latent_chunk.device)
+        )
+        return waveform_flat, stream_state, past_key_values, had_past_key_values, had_stream_state
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
+        intermediate_tensors: Any = None,
+        inputs_embeds: torch.Tensor | None = None,
+        model_intermediate_buffer: list[dict[str, Any]] | None = None,
+        **kwargs: Any,
+    ) -> OmniOutput:
+        del input_ids, positions, intermediate_tensors, inputs_embeds
+
+        info_list = model_intermediate_buffer if isinstance(model_intermediate_buffer, list) else None
+        if info_list is None:
+            runtime_infos = kwargs.get("runtime_additional_information")
+            info_list = runtime_infos if isinstance(runtime_infos, list) else None
+        info_list = info_list or [{}]
+        num_reqs = max(len(info_list), 1)
+        sr_tensor = torch.tensor(int(self.ming_config.sample_rate), dtype=torch.int32)
+        empty = torch.zeros((0,), dtype=torch.float32, device=self.vllm_config.device_config.device)
+
+        outputs: list[torch.Tensor] = []
+        srs: list[torch.Tensor] = []
+
+        for idx in range(num_reqs):
+            info = info_list[idx] if idx < len(info_list) and isinstance(info_list[idx], dict) else {}
+            meta = info.get("meta") or {}
+            kv_meta = info.get("kv_metadata") or {}
+            has_ming_context = _has_ming_context(info)
+            if has_ming_context and KEY_REQUEST_ID not in info and KEY_REQUEST_ID not in kv_meta:
+                raise RuntimeError(
+                    f"Ming Stage-2 received a payload without {KEY_REQUEST_ID}. keys={sorted(info.keys())}"
+                )
+            request_id = _resolve_request_id(info, idx)
+            self._touch_request_state(request_id)
+            chunk_id = _coerce_optional_int(_first_present(info.get(KEY_CHUNK_ID), kv_meta.get(KEY_CHUNK_ID)))
+            finished = _coerce_finished(
+                _first_present(
+                    meta.get("stream_finished"),
+                    meta.get("finished"),
+                    info.get("stream_finished"),
+                    kv_meta.get("stream_finished"),
+                    torch.tensor(True),
+                )
+            )
+            latent = _first_present(
+                info.get("latent"),
+                info.get("ming_latent_patches"),
+                kv_meta.get("ming_latent_patches"),
+            )
+            stripped = bool(
+                _first_present(
+                    meta.get("_ming_payload_stripped"),
+                    info.get("_ming_payload_stripped"),
+                    kv_meta.get("_ming_payload_stripped"),
+                    False,
+                )
+            )
+            if stripped:
+                raise RuntimeError(
+                    "Ming Stage-2 payload was stripped before model entry. "
+                    f"request_id={request_id} chunk_id={chunk_id} keys={sorted(info.keys())}"
+                )
+            if latent is None:
+                if has_ming_context and not finished:
+                    raise RuntimeError(
+                        "Ming Stage-2 received no latent chunk for an unfinished request. "
+                        f"request_id={request_id} chunk_id={chunk_id} keys={sorted(info.keys())}"
+                    )
+                if finished:
+                    self._clear_request_state(request_id)
+                outputs.append(empty)
+                srs.append(sr_tensor)
+                continue
+
+            latent_tensor = _coerce_latent_chunk(
+                latent,
+                device=self.vllm_config.device_config.device,
+                dtype=next(self.audio.parameters()).dtype,
+                latent_dim=self.ming_config.latent_dim,
+                patch_size=self.ming_config.patch_size,
+            )
+            if latent_tensor is None or latent_tensor.numel() == 0:
+                if not finished:
+                    raise RuntimeError(
+                        "Ming Stage-2 received an empty latent chunk before final flush. "
+                        f"request_id={request_id} chunk_id={chunk_id} latent_shape={_shape_of(latent_tensor)}"
+                    )
+                if finished:
+                    self._clear_request_state(request_id)
+                outputs.append(empty)
+                srs.append(sr_tensor)
+                continue
+
+            patch_count = int(latent_tensor.shape[0])
+            waveform_flat, stream_state, past_key_values, had_past_key_values, had_stream_state = (
+                self.chunked_decode_streaming(
+                    latent_tensor,
+                    request_id=request_id,
+                    finished=finished,
+                )
+            )
+            total_patch_count = self._patch_totals.get(request_id, 0) + patch_count
+            total_waveform_numel = self._sample_totals.get(request_id, 0) + int(waveform_flat.numel())
+            self._patch_totals[request_id] = total_patch_count
+            self._sample_totals[request_id] = total_waveform_numel
+            if (had_past_key_values or had_stream_state) and waveform_flat.numel() == 0 and not finished:
+                raise RuntimeError(
+                    "Ming Stage-2 produced an empty waveform after cached streaming state already existed. "
+                    f"request_id={request_id} chunk_id={chunk_id} latent_shape={tuple(latent_tensor.shape)} "
+                    f"had_past_key_values={had_past_key_values} had_stream_state={had_stream_state}"
+                )
+
+            if finished:
+                if os.environ.get("MING_TTS_STAGE1_FINAL_LOG") == "1":
+                    logger.info(
+                        "MING_STAGE1_FINAL %s",
+                        {
+                            "request_id": request_id,
+                            "chunk_id": chunk_id,
+                            "stop_reason": _first_present(
+                                info.get(MING_STOP_REASON_KEY),
+                                kv_meta.get(MING_STOP_REASON_KEY),
+                            ),
+                            "final_decode_step": _coerce_optional_int(
+                                _first_present(
+                                    info.get(MING_FINAL_DECODE_STEP_KEY),
+                                    kv_meta.get(MING_FINAL_DECODE_STEP_KEY),
+                                )
+                            ),
+                            "final_chunk_patch_count": patch_count,
+                            "total_patch_count": total_patch_count,
+                            "final_chunk_waveform_numel": int(waveform_flat.numel()),
+                            "total_waveform_numel": total_waveform_numel,
+                        },
+                    )
+                self._clear_request_state(request_id)
+            else:
+                self._past_key_values[request_id] = past_key_values
+                self._stream_state[request_id] = stream_state
+
+            outputs.append(waveform_flat)
+            srs.append(sr_tensor)
+
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={
+                "model_outputs": outputs,
+                "sr": srs,
+            },
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        weights = list(weights)
+        if not weights:
+            raise RuntimeError("MingAudioVAEModel received no checkpoint weights.")
+
+        params_dict = dict(self.named_parameters(remove_duplicate=False))
+        loaded: set[str] = set()
+        skipped: list[str] = []
+
+        for name, loaded_weight in weights:
+            if name not in params_dict:
+                skipped.append(name)
+                continue
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, loaded_weight)
+            loaded.add(name)
+
+        missing = {name for name in params_dict if name.startswith("audio.")} - loaded
+        if missing:
+            raise RuntimeError(f"MingAudioVAEModel: {len(missing)} params not loaded. First few: {sorted(missing)[:5]}")
+        if skipped:
+            warnings.warn(
+                f"MingAudioVAEModel: skipped {len(skipped)} checkpoint keys during load. First few: {skipped[:8]}",
+                stacklevel=2,
+            )
+        return loaded
+
+    def _clear_request_state(self, request_id: str) -> None:
+        self._past_key_values.pop(request_id, None)
+        self._stream_state.pop(request_id, None)
+        self._patch_totals.pop(request_id, None)
+        self._sample_totals.pop(request_id, None)
+        self._state_access_times.pop(request_id, None)
+
+    def _touch_request_state(self, request_id: str) -> None:
+        now = time.monotonic()
+        self._evict_expired_request_states(now)
+        self._state_access_times[request_id] = now
+        self._state_access_times.move_to_end(request_id)
+        while len(self._state_access_times) > _MAX_STREAM_STATES:
+            evicted_request_id, _ = self._state_access_times.popitem(last=False)
+            self._clear_request_state(evicted_request_id)
+
+    def _evict_expired_request_states(self, now: float) -> None:
+        cutoff = now - _STREAM_STATE_TTL_SECONDS
+        while self._state_access_times:
+            request_id, last_access = next(iter(self._state_access_times.items()))
+            if last_access >= cutoff:
+                break
+            self._clear_request_state(request_id)
+
+
+def _coerce_finished(value: Any) -> bool:
+    if isinstance(value, torch.Tensor):
+        if value.numel() == 0:
+            return False
+        return bool(value.reshape(-1)[0].item())
+    return bool(value)
+
+
+def _first_present(*values: Any) -> Any:
+    for value in values:
+        if value is not None:
+            return value
+    return None
+
+
+def _coerce_optional_int(value: Any, default: int | None = None) -> int | None:
+    if value is None:
+        return default
+    if isinstance(value, list):
+        value = value[0] if value else default
+    if isinstance(value, torch.Tensor):
+        if value.numel() == 0:
+            return default
+        return int(value.reshape(-1)[0].item())
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return default
+
+
+def _resolve_request_id(info: dict[str, Any], idx: int) -> str:
+    kv_meta = info.get("kv_metadata") or {}
+    request_id = info.get(KEY_REQUEST_ID) or kv_meta.get(KEY_REQUEST_ID)
+    if request_id is None:
+        return str(idx)
+    if not isinstance(request_id, str) or not request_id:
+        raise RuntimeError(f"Ming Stage-2 received invalid request id: {request_id!r}")
+    return request_id
+
+
+def _has_ming_context(info: dict[str, Any]) -> bool:
+    kv_meta = info.get("kv_metadata") or {}
+    keys = (
+        KEY_REQUEST_ID,
+        KEY_CHUNK_ID,
+        "ming_latent_patches",
+        "_ming_payload_stripped",
+    )
+    return any(key in info or key in kv_meta for key in keys) or "latent" in info
+
+
+def _shape_of(value: Any) -> tuple[int, ...] | None:
+    if isinstance(value, torch.Tensor):
+        return tuple(value.shape)
+    return None
+
+
+def _has_stream_state(value: Any) -> bool:
+    if not isinstance(value, tuple):
+        return False
+    return any(item is not None for item in value)
+
+
+def _coerce_latent_chunk(
+    value: Any,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+    latent_dim: int,
+    patch_size: int,
+) -> torch.Tensor | None:
+    if not isinstance(value, torch.Tensor):
+        value = torch.as_tensor(value)
+
+    latent = value.detach()
+    if latent.ndim == 2:
+        latent = latent.unsqueeze(0)
+    if latent.ndim != 3:
+        raise ValueError(f"Expected latent chunk rank-3 [Batch, Time, Dimension], got {tuple(latent.shape)}")
+    if latent.shape[-2] != patch_size:
+        raise ValueError(f"Latent patch size mismatch: got {latent.shape[-2]}, expected {patch_size}")
+    if latent.shape[-1] != latent_dim:
+        raise ValueError(f"Latent dim mismatch: got {latent.shape[-1]}, expected {latent_dim}")
+    return latent.to(device=device, dtype=dtype)

vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py

@@ -0,0 +1,414 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/modeling_bailingmm.py
+from __future__ import annotations
+
+import warnings
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.models.qwen2 import Qwen2Model
+from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
+from vllm.sequence import IntermediateTensors
+from vllm.v1.outputs import SamplerOutput
+from vllm.v1.sample.metadata import SamplingMetadata
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+from .aggregator import Aggregator
+from .config_ming_tts import (
+    KEY_CFG,
+    KEY_DECODE_STEP,
+    KEY_LATENT_HISTORY,
+    KEY_MAX_DECODE_STEPS,
+    KEY_MIN_DECODE_STEPS,
+    KEY_NEXT_EMBEDS,
+    KEY_REQUEST_ID,
+    KEY_SIGMA,
+    KEY_TEMPERATURE,
+    KEY_TEXT_MODE,
+    MingTTSConfig,
+)
+from .flowloss_head import FlowLoss
+from .patch_emission import (
+    MING_STOP_REASON_CODES,
+    MING_STOP_REASON_CONTINUE,
+    MING_STOP_REASON_KEY,
+    MING_STOP_REASON_MAX_DECODE_STEPS,
+    MING_STOP_REASON_STOP_HEAD,
+    _coerce_latent_history,
+    _get_request_token_counts,
+    _normalize_request_infos,
+    _resolve_ming_stop_decision,
+    _resolve_optional_runtime_int,
+    _resolve_runtime_float,
+    _resolve_runtime_int,
+    _validate_ming_decode_window,
+)
+
+logger = init_logger(__name__)
+
+
+class MingLLMModel(nn.Module):
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            "model.model.": "model.",
+        }
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.ming_config = MingTTSConfig.from_hf_config(vllm_config.model_config.hf_config)
+        self.ming_config.validate()
+        self.vllm_config = vllm_config
+        self.prefix = prefix
+        self.fm_dtype = _resolve_ming_runtime_dtype(vllm_config)
+        self.model = Qwen2Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
+        self.linear_proj_audio = Aggregator(
+            in_channels=self.ming_config.latent_dim,
+            llm_input_dim=self.ming_config.llm_hidden_size,
+            **self.ming_config.aggregator_config,
+        )
+        self.flowloss = FlowLoss(
+            z_channels=self.ming_config.latent_dim,
+            llm_cond_dim=self.ming_config.llm_hidden_size,
+            **self.ming_config.ditar_config,
+        )
+        self.stop_head = nn.Linear(self.ming_config.llm_hidden_size, 2, bias=True)
+        self.spk_head = nn.Linear(192, self.ming_config.llm_hidden_size, bias=True)
+        self.flowloss.to(dtype=self.fm_dtype)
+        self.linear_proj_audio.to(dtype=self.fm_dtype)
+        self.stop_head.to(dtype=self.fm_dtype)
+        self.spk_head.to(dtype=self.fm_dtype)
+        self._pending_postprocess_updates: dict[str, dict[str, Any]] = {}
+        self._last_ming_next_token_ids = None
+        self._last_text_mode = False
+
+    def embed_input_ids(
+        self, input_ids: torch.Tensor, inputs_embeds: torch.Tensor | None = None, **_: Any
+    ) -> torch.Tensor:
+        if inputs_embeds is not None:
+            return inputs_embeds
+        return self.model.embed_input_ids(input_ids)
+
+    def project_speaker_embedding(self, spk_emb: torch.Tensor) -> torch.Tensor:
+        return self.spk_head(spk_emb)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        latent_history: torch.Tensor | None = None,
+        model_intermediate_buffer: list[dict[str, Any]] | None = None,
+        seq_token_counts: list[int] | None = None,
+        **kwargs: object,
+    ) -> OmniOutput | IntermediateTensors | torch.Tensor:
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_input_ids(input_ids)
+        if model_intermediate_buffer is None:
+            model_intermediate_buffer = kwargs.get("runtime_additional_information")
+        request_infos = _normalize_request_infos(model_intermediate_buffer)
+        _validate_ming_decode_window(
+            request_infos,
+            min_stop_step=int(self.ming_config.stop_head_min_steps),
+            default_max_decode_steps=self.ming_config.max_decode_steps,
+        )
+        backbone_out = self.model(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )
+        if isinstance(backbone_out, IntermediateTensors):
+            self._last_ming_next_token_ids = None
+            return backbone_out
+        hidden_states = _extract_hidden_states(backbone_out)
+        token_counts = _get_request_token_counts(hidden_states, request_infos, seq_token_counts)
+        text_mode = bool(request_infos) and all(bool(info.get(KEY_TEXT_MODE, False)) for info in request_infos)
+        if request_infos and any(bool(info.get(KEY_TEXT_MODE, False)) for info in request_infos) and not text_mode:
+            raise RuntimeError("Mixed Ming text/audio modes in one Stage-0 batch are unsupported.")
+        if text_mode:
+            self._last_ming_next_token_ids = None
+            self._last_text_mode = True
+            return OmniOutput(
+                text_hidden_states=hidden_states,
+                multimodal_outputs={KEY_TEXT_MODE: True},
+                intermediate_tensors=intermediate_tensors,
+            )
+        self._last_text_mode = False
+        if latent_history is None and not token_counts:
+            self._last_ming_next_token_ids = None
+            return OmniOutput(
+                text_hidden_states=hidden_states, multimodal_outputs=None, intermediate_tensors=intermediate_tensors
+            )
+        if latent_history is not None and not token_counts:
+            token_counts, request_infos = [hidden_states.shape[0]], [{KEY_LATENT_HISTORY: latent_history}]
+
+        total_tokens = hidden_states.shape[0]
+        latent_patch_tokens = next_embed_tokens = new_history_tokens = None
+        decode_step_tokens = has_patch = stop_reason_code_tokens = None
+        next_token_ids = []
+        pending_updates: dict[str, dict[str, Any]] = {}
+        cursor = 0
+        any_decode = False
+        for req_idx, token_count in enumerate(token_counts):
+            end = min(cursor + token_count, total_tokens)
+            if end <= cursor:
+                continue
+            req_info = request_infos[req_idx] if req_idx < len(request_infos) else {}
+            req_id = req_info.get(KEY_REQUEST_ID)
+            req_history = _coerce_latent_history(
+                req_info.get(KEY_LATENT_HISTORY), device=hidden_states.device, dtype=self.fm_dtype, cfg=self.ming_config
+            )
+            if req_history is None:
+                cursor = end
+                continue
+            decode_step = int(req_info.get(KEY_DECODE_STEP, req_info.get("generated_len", 0)))
+            decode_hidden, output_index = (
+                (hidden_states[cursor:end], cursor) if token_count == 1 else (hidden_states[end - 1 : end], end - 1)
+            )
+            sampled_token_latent, next_embeds, new_history, stop_probs = self._decode_one_step(
+                hidden_states=decode_hidden,
+                latent_history=req_history,
+                cfg_scale=_resolve_runtime_float(req_info, KEY_CFG, self.ming_config.cfg),
+                sigma=_resolve_runtime_float(req_info, KEY_SIGMA, self.ming_config.sigma),
+                temperature=_resolve_runtime_float(req_info, KEY_TEMPERATURE, self.ming_config.temperature),
+            )
+            req_max_decode_steps = _resolve_runtime_int(
+                req_info, KEY_MAX_DECODE_STEPS, self.ming_config.max_decode_steps
+            )
+            req_min_decode_steps = _resolve_optional_runtime_int(req_info, KEY_MIN_DECODE_STEPS, 0)
+            if latent_patch_tokens is None:
+                latent_patch_tokens = sampled_token_latent.new_zeros(
+                    (total_tokens, self.ming_config.patch_size, self.ming_config.latent_dim)
+                )
+                next_embed_tokens = next_embeds.new_zeros((total_tokens, 1, self.ming_config.llm_hidden_size))
+                new_history_tokens = new_history.new_zeros(
+                    (total_tokens, self.ming_config.history_patch_size, self.ming_config.latent_dim)
+                )
+                decode_step_tokens = torch.zeros((total_tokens,), dtype=torch.int32, device=hidden_states.device)
+                has_patch = torch.zeros((total_tokens,), dtype=torch.bool, device=hidden_states.device)
+                stop_reason_code_tokens = torch.zeros((total_tokens,), dtype=torch.int32, device=hidden_states.device)
+            latent_patch_tokens[output_index : output_index + 1] = sampled_token_latent
+            next_embed_tokens[output_index : output_index + 1] = next_embeds
+            new_history_tokens[output_index : output_index + 1] = new_history
+            decode_step_tokens[output_index : output_index + 1] = decode_step
+            has_patch[output_index : output_index + 1] = True
+            stop_reason, _, _, _, next_token_id = _resolve_ming_stop_decision(
+                step=decode_step,
+                stop_prob=float(stop_probs.reshape(-1)[0].item()),
+                stop_threshold=float(self.ming_config.stop_head_threshold),
+                min_stop_step=int(self.ming_config.stop_head_min_steps),
+                min_decode_steps=req_min_decode_steps,
+                max_decode_steps=req_max_decode_steps,
+                audio_dummy_token_id=int(self.ming_config.audio_dummy_token_id),
+                text_eos_token_id=int(self.ming_config.text_eos_token_id),
+            )
+            stop_reason_code_tokens[output_index : output_index + 1] = MING_STOP_REASON_CODES[stop_reason]
+            next_token_ids.append(int(next_token_id))
+            if isinstance(req_id, str):
+                pending_updates[req_id] = {
+                    KEY_LATENT_HISTORY: new_history,
+                    KEY_NEXT_EMBEDS: next_embeds,
+                    "ming_latent_patch": sampled_token_latent,
+                    "ming_stop_prob": stop_probs,
+                    MING_STOP_REASON_KEY: stop_reason,
+                }
+            any_decode = True
+            cursor = end
+
+        self._pending_postprocess_updates = pending_updates
+        if not any_decode:
+            self._last_ming_next_token_ids = None
+            return OmniOutput(
+                text_hidden_states=hidden_states, multimodal_outputs=None, intermediate_tensors=intermediate_tensors
+            )
+        self._last_ming_next_token_ids = next_token_ids
+        return OmniOutput(
+            text_hidden_states=hidden_states,
+            multimodal_outputs={
+                "ming_latent_patch": latent_patch_tokens,
+                "ming_next_embeds": next_embed_tokens,
+                "ming_new_history": new_history_tokens,
+                "ming_decode_step": decode_step_tokens,
+                "ming_has_patch": has_patch,
+                MING_STOP_REASON_KEY: stop_reason_code_tokens,
+            },
+            intermediate_tensors=intermediate_tensors,
+        )
+
+    def pop_postprocess_update(self, req_id: str) -> dict[str, Any]:
+        return self._pending_postprocess_updates.pop(req_id, {}) if isinstance(req_id, str) else {}
+
+    def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata: SamplingMetadata) -> torch.Tensor | None:
+        del sampling_metadata
+        if self._last_text_mode:
+            return (
+                None
+                if hidden_states is None or hidden_states.numel() == 0
+                else self.model.compute_logits(hidden_states)
+            )
+        if hidden_states is None or hidden_states.numel() == 0:
+            return None
+        if hidden_states.dim() != 2:
+            raise RuntimeError(
+                f"Expected hidden_states rank-2 [B,H] in compute_logits, got {tuple(hidden_states.shape)}"
+            )
+        batch_size = hidden_states.shape[0]
+        next_token_ids = self._last_ming_next_token_ids
+        self._last_ming_next_token_ids = None
+        if next_token_ids is None:
+            logger.debug(
+                "Missing Ming forced next-token ids before compute_logits. "
+                "Using dummy next-token IDs (likely during a profiling or dummy run)."
+            )
+            next_token_ids = [int(self.ming_config.text_eos_token_id)] * batch_size
+        if len(next_token_ids) != batch_size:
+            raise RuntimeError(
+                "Ming forced next-token batch mismatch: "
+                f"got {len(next_token_ids)} ids for hidden_states batch {batch_size}."
+            )
+        logits = torch.full(
+            (batch_size, self.ming_config.llm_vocab_size),
+            float("-inf"),
+            device=hidden_states.device,
+            dtype=torch.float32,
+        )
+        for i, next_token_id in enumerate(next_token_ids):
+            logits[i, int(next_token_id)] = 0.0
+        return logits
+
+    def sample(self, logits, sampling_metadata):
+        if logits is None:
+            return None
+        if self._last_text_mode:
+            return self.model.sample(logits, sampling_metadata)
+        del sampling_metadata
+        return SamplerOutput(
+            sampled_token_ids=logits.argmax(dim=-1, keepdim=True).to(dtype=torch.int32),
+            logprobs_tensors=None,
+        )
+
+    def _decode_one_step(
+        self,
+        *,
+        hidden_states: torch.Tensor,
+        latent_history: torch.Tensor,
+        cfg_scale: float,
+        sigma: float,
+        temperature: float,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        if hidden_states.dim() != 2:
+            raise RuntimeError(f"Expected decode hidden_states rank-2 [B,H], got {tuple(hidden_states.shape)}")
+        if latent_history.dim() != 3:
+            raise RuntimeError(f"Expected latent_history rank-3 [B,T,D], got {tuple(latent_history.shape)}")
+        if hidden_states.shape[0] != latent_history.shape[0]:
+            raise RuntimeError(
+                "Batch mismatch: "
+                f"hidden_states B={hidden_states.shape[0]} "
+                f"vs latent_history B={latent_history.shape[0]}"
+            )
+        # [Batch, Hidden] -> [Batch, Time, Hidden] = [B, 1, H] for FlowLoss conditioning.
+        z_diff_cond = hidden_states.to(dtype=self.fm_dtype).unsqueeze(1)
+        if not torch.isfinite(z_diff_cond).all():
+            raise RuntimeError("Non-finite z_diff_cond before FlowLoss.sample().")
+        sampled_token_latent = self.flowloss.sample(
+            z=z_diff_cond,
+            latent_history=latent_history,
+            cfg=cfg_scale,
+            patch_size=self.ming_config.patch_size,
+            sigma=sigma,
+            temperature=temperature,
+        )
+        expected_shape = (hidden_states.shape[0], self.ming_config.patch_size, self.ming_config.latent_dim)
+        if tuple(sampled_token_latent.shape) != expected_shape:
+            raise RuntimeError(
+                f"FlowLoss output shape mismatch: got {tuple(sampled_token_latent.shape)}, expected {expected_shape}"
+            )
+        new_history = torch.cat([latent_history[:, self.ming_config.patch_size :, :], sampled_token_latent], dim=1)
+        # Aggregator expects [Batch, Time, Dimension] = [B, 4, 64] and returns [B, 1, H].
+        next_embeds = self.linear_proj_audio(sampled_token_latent)
+        stop_probs = self.stop_head(hidden_states.to(dtype=self.fm_dtype)).softmax(dim=-1)[:, 1]
+        if not torch.isfinite(sampled_token_latent).all():
+            raise RuntimeError("Non-finite sampled_token_latent in Ming decode step.")
+        if not torch.isfinite(next_embeds).all():
+            raise RuntimeError("Non-finite next_embeds in Ming decode step.")
+        if not torch.isfinite(stop_probs).all():
+            raise RuntimeError("Non-finite stop_probs in Ming decode step.")
+        return sampled_token_latent, next_embeds, new_history, stop_probs
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        params_dict = dict(self.named_parameters(remove_duplicate=False))
+        loaded_params = AutoWeightsLoader(self).load_weights(weights, mapper=self.hf_to_vllm_mapper)
+        _warn_missing_prefix("flowloss", params_dict, loaded_params, prefix="flowloss.", fatal=True)
+        _warn_missing_prefix("linear_proj_audio", params_dict, loaded_params, prefix="linear_proj_audio.", fatal=True)
+        _warn_missing_prefix("stop_head", params_dict, loaded_params, prefix="stop_head.", fatal=True)
+        _warn_missing_prefix("spk_head", params_dict, loaded_params, prefix="spk_head.", fatal=True)
+        return loaded_params
+
+
+def _extract_hidden_states(backbone_out: object) -> torch.Tensor:
+    if isinstance(backbone_out, torch.Tensor):
+        return backbone_out
+    if hasattr(backbone_out, "last_hidden_state"):
+        return backbone_out.last_hidden_state
+    if isinstance(backbone_out, (tuple, list)) and backbone_out and isinstance(backbone_out[0], torch.Tensor):
+        return backbone_out[0]
+    raise TypeError(f"Unsupported backbone forward output type: {type(backbone_out)}")
+
+
+def _resolve_ming_runtime_dtype(vllm_config: VllmConfig) -> torch.dtype:
+    dtype = getattr(vllm_config.model_config, "dtype", None)
+    if isinstance(dtype, torch.dtype):
+        return dtype
+    if isinstance(dtype, str):
+        normalized = dtype.strip().lower()
+        if normalized in ("float16", "half", "torch.float16"):
+            return torch.float16
+        if normalized in ("bfloat16", "bf16", "torch.bfloat16"):
+            return torch.bfloat16
+        if normalized in ("float32", "fp32", "torch.float32"):
+            return torch.float32
+    return torch.float32
+
+
+def _warn_missing_prefix(
+    module_name: str,
+    params_dict: dict[str, nn.Parameter],
+    loaded_params: set[str],
+    prefix: str,
+    fatal: bool = False,
+) -> None:
+    missing = {key for key in params_dict if key.startswith(prefix)} - loaded_params
+    if not missing:
+        return
+    msg = (
+        f"MingLLMModel: {len(missing)} {module_name} params not loaded "
+        f"(prefix={prefix}). First few: {sorted(missing)[:5]}"
+    )
+    if fatal:
+        raise RuntimeError(msg)
+    warnings.warn(msg, stacklevel=3)
+
+
+__all__ = [
+    "Aggregator",
+    "FlowLoss",
+    "MING_STOP_REASON_CODES",
+    "MING_STOP_REASON_CONTINUE",
+    "MING_STOP_REASON_KEY",
+    "MING_STOP_REASON_MAX_DECODE_STEPS",
+    "MING_STOP_REASON_STOP_HEAD",
+    "MingLLMModel",
+    "_coerce_latent_history",
+    "_extract_hidden_states",
+    "_resolve_ming_runtime_dtype",
+    "_resolve_ming_stop_decision",
+    "_warn_missing_prefix",
+]

vllm_omni/model_executor/models/ming_tts/patch_emission.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/modeling_bailingmm.py
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from vllm.forward_context import get_forward_context, is_forward_context_available
+
+from .config_ming_tts import KEY_MAX_DECODE_STEPS, KEY_MIN_DECODE_STEPS, KEY_REQUEST_ID, KEY_TEXT_MODE, MingTTSConfig
+
+MING_STOP_REASON_CONTINUE = "continue"
+MING_STOP_REASON_STOP_HEAD = "stop_head"
+MING_STOP_REASON_MAX_DECODE_STEPS = "max_decode_steps"
+MING_STOP_REASON_KEY = "ming_stop_reason"
+MING_STOP_REASON_CODES = {
+    MING_STOP_REASON_CONTINUE: 0,
+    MING_STOP_REASON_STOP_HEAD: 1,
+    MING_STOP_REASON_MAX_DECODE_STEPS: 2,
+}
+
+
+def _normalize_request_infos(model_intermediate_buffer: object) -> list[dict[str, Any]]:
+    if not isinstance(model_intermediate_buffer, list):
+        return []
+    infos: list[dict[str, Any]] = []
+    for item in model_intermediate_buffer:
+        infos.append(item if isinstance(item, dict) else {})
+    return infos
+
+
+def _get_request_token_counts(
+    hidden_states: torch.Tensor,
+    request_infos: list[dict[str, Any]],
+    seq_token_counts: list[int] | None,
+) -> list[int]:
+    if seq_token_counts:
+        return [int(x) for x in seq_token_counts]
+
+    if is_forward_context_available():
+        slices = getattr(get_forward_context(), "ubatch_slices", None)
+        if slices is not None and len(slices) > 0:
+            counts: list[int] = []
+            for item in slices:
+                if isinstance(item, int):
+                    counts.append(int(item))
+                elif hasattr(item, "stop") and hasattr(item, "start"):
+                    counts.append(int(item.stop) - int(item.start))
+            if counts:
+                return counts
+
+    if request_infos:
+        if len(request_infos) == hidden_states.shape[0]:
+            return [1] * hidden_states.shape[0]
+        return [hidden_states.shape[0]]
+
+    return []
+
+
+def _coerce_latent_history(
+    value: object,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+    cfg: MingTTSConfig,
+) -> torch.Tensor | None:
+    if value is None:
+        return None
+    if not isinstance(value, torch.Tensor):
+        value = torch.as_tensor(value)
+
+    history = value.detach()
+    if history.ndim == 2:
+        history = history.unsqueeze(0)
+    if history.ndim != 3:
+        raise RuntimeError(f"Expected latent_history rank-3 [B,T,D], got {tuple(history.shape)}")
+    if history.shape[1] != cfg.history_patch_size or history.shape[2] != cfg.latent_dim:
+        raise RuntimeError(
+            f"latent_history shape mismatch: got {tuple(history.shape)}, "
+            f"expected [B,{cfg.history_patch_size},{cfg.latent_dim}]"
+        )
+    return history.to(device=device, dtype=dtype)
+
+
+def _resolve_runtime_float(req_info: dict[str, Any], key: str, default_value: float) -> float:
+    raw = req_info.get(key, default_value)
+    try:
+        value = float(raw)
+    except (TypeError, ValueError) as exc:
+        raise RuntimeError(f"Invalid {key}: expected float-like value, got {raw!r}") from exc
+    if not value >= 0.0:
+        raise RuntimeError(f"Invalid {key}: expected non-negative value, got {value}")
+    return value
+
+
+def _resolve_runtime_int(req_info: dict[str, Any], key: str, default_value: int) -> int:
+    raw = req_info.get(key, default_value)
+    try:
+        value = int(raw)
+    except (TypeError, ValueError) as exc:
+        raise RuntimeError(f"Invalid {key}: expected int-like value, got {raw!r}") from exc
+    if value <= 0:
+        raise RuntimeError(f"Invalid {key}: expected positive value, got {value}")
+    return value
+
+
+def _resolve_optional_runtime_int(req_info: dict[str, Any], key: str, default_value: int) -> int:
+    raw = req_info.get(key, default_value)
+    try:
+        value = int(raw)
+    except (TypeError, ValueError) as exc:
+        raise RuntimeError(f"Invalid {key}: expected int-like value, got {raw!r}") from exc
+    if value < 0:
+        raise RuntimeError(f"Invalid {key}: expected non-negative value, got {value}")
+    return value
+
+
+def _validate_ming_decode_window(
+    request_infos: list[dict[str, Any]],
+    *,
+    min_stop_step: int,
+    default_max_decode_steps: int,
+) -> None:
+    for i, info in enumerate(request_infos):
+        if info.get(KEY_TEXT_MODE):
+            continue
+        max_steps = _resolve_runtime_int(info, KEY_MAX_DECODE_STEPS, default_max_decode_steps)
+        min_steps = _resolve_optional_runtime_int(info, KEY_MIN_DECODE_STEPS, 0)
+        min_required = max(min_stop_step + 1, min_steps)
+        if max_steps < min_required:
+            req_id = info.get(KEY_REQUEST_ID, f"idx={i}")
+            raise ValueError(
+                f"Ming request {req_id!r}: max_decode_steps={max_steps} < "
+                f"min_required_decode_steps={min_required} "
+                f"(min_stop_step={min_stop_step}, min_decode_steps={min_steps})"
+            )
+
+
+def _resolve_ming_stop_decision(
+    *,
+    step: int,
+    stop_prob: float,
+    stop_threshold: float,
+    min_stop_step: int,
+    min_decode_steps: int,
+    max_decode_steps: int,
+    audio_dummy_token_id: int,
+    text_eos_token_id: int,
+) -> tuple[str, bool, bool, int, int]:
+    min_required_decode_steps = max(min_stop_step + 1, min_decode_steps)
+    if max_decode_steps < min_required_decode_steps:
+        raise RuntimeError(
+            "Invalid Ming decode window: "
+            f"max_decode_steps={max_decode_steps} is smaller than "
+            f"min_required_decode_steps={min_required_decode_steps}"
+        )
+    should_force_stop = (step + 1) >= max_decode_steps
+    should_stop_head = ((step + 1) >= min_required_decode_steps) and stop_prob > stop_threshold
+
+    if should_force_stop:
+        return (
+            MING_STOP_REASON_MAX_DECODE_STEPS,
+            True,
+            True,
+            min_required_decode_steps,
+            text_eos_token_id,
+        )
+    if should_stop_head:
+        return (
+            MING_STOP_REASON_STOP_HEAD,
+            True,
+            False,
+            min_required_decode_steps,
+            text_eos_token_id,
+        )
+    return (
+        MING_STOP_REASON_CONTINUE,
+        False,
+        False,
+        min_required_decode_steps,
+        audio_dummy_token_id,
+    )

vllm_omni/model_executor/models/ming_tts/pipeline.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Ming TTS pipeline: Stage-0 LLM+flow -> Stage-1 audio VAE."""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.ming_tts"
+
+MING_TTS_PIPELINE = PipelineConfig(
+    model_type="ming_tts",
+    model_arch="MingTTSForConditionalGeneration",
+    hf_architectures=("MingTTSForConditionalGeneration",),
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="llm",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            owns_tokenizer=True,
+            hf_config_name="llm_config",
+            engine_output_type="latent",
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.llm2audio_vae_async_chunk"),
+            sampling_constraints={
+                "temperature": 0.0,
+                "top_p": 1.0,
+                "top_k": -1,
+                "max_tokens": 512,
+                "detokenize": True,
+            },
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="audio_vae",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="audio",
+            hf_config_name="llm_config",
+            engine_output_type="audio",
+            sync_process_input_func=f"{_PROC}.llm2audio_vae",
+            sampling_constraints={
+                "temperature": 0.0,
+                "top_p": 1.0,
+                "top_k": -1,
+                "max_tokens": 1,
+                "detokenize": False,
+            },
+        ),
+    ),
+)

vllm_omni/model_executor/models/ming_tts/prompt_assembly.py

@@ -0,0 +1,274 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from __future__ import annotations
+
+import copy
+import json
+import math
+import re
+from typing import Any
+
+import torch
+
+from .audio_prep import (
+    coerce_speaker_embeddings,
+    count_prompt_latent_patches,
+    count_prompt_waveform_patches,
+    pad_prompt_waveform,
+)
+from .config_ming_tts import (
+    AUDIO_FRAME_HOP,
+    KEY_MAX_DECODE_STEPS,
+    KEY_MIN_DECODE_STEPS,
+    KEY_PROMPT_LATENTS,
+    KEY_REQUEST_ID,
+    KEY_SPEAKER_EMBEDDING,
+    LATENT_DIM,
+    PATCH_SIZE,
+    SAMPLE_RATE,
+    VAE_PATCH_SIZE,
+)
+
+DEFAULT_PROMPT = "Please generate speech based on the following description.\n"
+BASE_CAPTION_TEMPLATE: dict[str, Any] = {
+    "audio_sequence": [
+        {
+            "序号": 1,
+            "说话人": "speaker_1",
+            "方言": None,
+            "风格": None,
+            "语速": None,
+            "基频": None,
+            "音量": None,
+            "情感": None,
+            "BGM": {
+                "Genre": None,
+                "Mood": None,
+                "Instrument": None,
+                "Theme": None,
+                "ENV": None,
+                "SNR": None,
+            },
+            "IP": None,
+        }
+    ]
+}
+_DURATION_SECONDS_RE = re.compile(r"Duration:\s*([0-9]+(?:\.[0-9]+)?)\s*s\b", re.IGNORECASE)
+
+
+def create_instruction(user_input: Any) -> str | None:
+    if user_input is None:
+        return None
+    if isinstance(user_input, str):
+        return user_input
+    if not isinstance(user_input, dict):
+        raise ValueError(f"Ming instruction must be str, dict, or None; got {type(user_input).__name__}")
+    caption = copy.deepcopy(BASE_CAPTION_TEMPLATE)
+    item = caption["audio_sequence"][0]
+    for key, value in user_input.items():
+        if key in item:
+            item[key] = value
+    return json.dumps(caption, ensure_ascii=False)
+
+
+def parse_duration_seconds(text: str | None) -> float | None:
+    if not isinstance(text, str):
+        return None
+    match = _DURATION_SECONDS_RE.search(text)
+    if match is None:
+        return None
+    try:
+        value = float(match.group(1))
+    except ValueError:
+        return None
+    if value <= 0.0:
+        return None
+    return value
+
+
+def estimate_decode_steps_for_duration(
+    duration_seconds: float,
+    *,
+    sample_rate: int = SAMPLE_RATE,
+    frame_hop: int = AUDIO_FRAME_HOP,
+    patch_size: int = PATCH_SIZE,
+    vae_patch_size: int = VAE_PATCH_SIZE,
+) -> int:
+    if duration_seconds <= 0.0:
+        return 0
+    samples_per_decode_step = int(frame_hop) * int(patch_size) * int(vae_patch_size)
+    required_samples = float(duration_seconds) * float(sample_rate)
+    return max(1, int(math.ceil(required_samples / float(samples_per_decode_step))))
+
+
+def estimate_decode_step_window_for_duration(duration_seconds: float) -> tuple[int, int]:
+    target_steps = estimate_decode_steps_for_duration(duration_seconds)
+    min_steps = max(1, target_steps - 3)
+    max_steps = max(min_steps, target_steps + 3)
+    return min_steps, max_steps
+
+
+def resolve_effective_runtime_controls(
+    *,
+    text: str,
+    runtime_controls: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    controls = {} if runtime_controls is None else dict(runtime_controls)
+    has_explicit_min = KEY_MIN_DECODE_STEPS in controls and controls[KEY_MIN_DECODE_STEPS] is not None
+    has_explicit_max = KEY_MAX_DECODE_STEPS in controls and controls[KEY_MAX_DECODE_STEPS] is not None
+    if has_explicit_min or has_explicit_max:
+        return controls
+    duration_seconds = parse_duration_seconds(text)
+    if duration_seconds is None:
+        return controls
+    min_decode_steps, max_decode_steps = estimate_decode_step_window_for_duration(duration_seconds)
+    controls[KEY_MIN_DECODE_STEPS] = min_decode_steps
+    controls[KEY_MAX_DECODE_STEPS] = max_decode_steps
+    return controls
+
+
+def build_dense_prompt_token_ids(
+    tokenizer: Any,
+    *,
+    prompt: str,
+    text: str,
+    instruction: str | None = None,
+    prompt_text: str | None = None,
+    speaker_count: int = 0,
+    prompt_patch_count: int = 0,
+) -> list[int]:
+    speaker_prompt = []
+    for idx in range(int(speaker_count)):
+        speaker_prompt.extend(
+            tokenizer.encode(f"  speaker_{idx + 1}:")
+            + tokenizer.encode("<|vision_start|>")
+            + tokenizer.encode("<|vision_pad|>")
+            + tokenizer.encode("<|vision_end|>\n")
+        )
+    instruction_prompt = (
+        tokenizer.encode(instruction) + tokenizer.encode("<|endoftext|>") if instruction is not None else []
+    )
+    prompt_text_tokens = (
+        tokenizer.encode(prompt_text) if int(prompt_patch_count) > 0 and prompt_text is not None else []
+    )
+    audio_patch_token_id = tokenizer.convert_tokens_to_ids("<audioPatch>")
+    if audio_patch_token_id == tokenizer.unk_token_id:
+        raise ValueError("Ming tokenizer is missing required <audioPatch> token.")
+    prompt_latent_tokens = [audio_patch_token_id] * int(prompt_patch_count)
+    text_input_prefix = (
+        []
+        if all(token in text for token in ("Genre: ", "Mood: ", "Instrument: ", "Theme: ", "Duration: "))
+        else tokenizer.encode(" Text input:\n")
+    )
+    return (
+        tokenizer.encode("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n")
+        + tokenizer.encode("<|im_start|>user\n")
+        + tokenizer.encode(prompt)
+        + speaker_prompt
+        + text_input_prefix
+        + prompt_text_tokens
+        + tokenizer.encode(text)
+        + tokenizer.encode("<|im_end|>\n")
+        + tokenizer.encode("<|im_start|>assistant\n")
+        + instruction_prompt
+        + tokenizer.encode("<audio>")
+        + prompt_latent_tokens
+    )
+
+
+def build_ming_dense_prompt(
+    tokenizer: Any,
+    *,
+    prompt: str,
+    text: str,
+    runtime_controls: dict[str, Any] | None = None,
+    instruction: Any = None,
+    prompt_text: str | None = None,
+    prompt_waveform: Any = None,
+    prompt_latents: Any = None,
+    speaker_embedding: Any = None,
+    use_zero_spk_emb: bool = False,
+    request_id: str | None = None,
+) -> dict[str, Any]:
+    instruction_text = create_instruction(instruction)
+    speaker_embeddings = coerce_speaker_embeddings(speaker_embedding, use_zero_spk_emb=use_zero_spk_emb)
+    effective_runtime_controls = resolve_effective_runtime_controls(text=text, runtime_controls=runtime_controls)
+
+    prompt_waveform_tensor = None
+    prompt_patch_count = 0
+    if prompt_waveform is not None:
+        prompt_waveform_tensor = pad_prompt_waveform(prompt_waveform)
+        prompt_patch_count = count_prompt_waveform_patches(prompt_waveform_tensor)
+    if prompt_waveform_tensor is not None and prompt_latents is not None:
+        raise ValueError(
+            "Ming waveform cloning request provided both raw prompt_waveform and explicit prompt_latents. "
+            "Choose exactly one source of truth."
+        )
+
+    prompt_latent_value = None
+    if prompt_waveform_tensor is not None and prompt_text is None:
+        raise ValueError(
+            "Ming prompt_waveform requires prompt_text for prompt-latent conditioning. "
+            "Use speaker_embedding for reference-audio-only speaker conditioning."
+        )
+    if prompt_latents is not None:
+        prompt_latent_value = torch.as_tensor(prompt_latents)
+        prompt_patch_count = count_prompt_latent_patches(
+            prompt_latent_value, patch_size=PATCH_SIZE, latent_dim=LATENT_DIM
+        )
+
+    prompt_token_ids = build_dense_prompt_token_ids(
+        tokenizer,
+        prompt=prompt,
+        text=text,
+        instruction=instruction_text,
+        prompt_text=prompt_text if prompt_patch_count > 0 else None,
+        speaker_count=0 if speaker_embeddings is None else len(speaker_embeddings),
+        prompt_patch_count=prompt_patch_count,
+    )
+
+    additional_information = {}
+    for key, value in effective_runtime_controls.items():
+        if isinstance(value, torch.Tensor):
+            additional_information[key] = value
+        elif key in (KEY_MIN_DECODE_STEPS, KEY_MAX_DECODE_STEPS):
+            additional_information[key] = torch.tensor(int(value), dtype=torch.int32)
+        else:
+            additional_information[key] = torch.tensor(float(value), dtype=torch.float32)
+    if request_id is not None:
+        additional_information[KEY_REQUEST_ID] = request_id
+    if instruction_text is not None:
+        additional_information["instruction"] = instruction_text
+    if prompt_text is not None:
+        additional_information["prompt_text"] = prompt_text
+    if prompt_waveform_tensor is not None:
+        additional_information["prompt_waveform"] = prompt_waveform_tensor
+        additional_information["prompt_waveform_length"] = torch.tensor(
+            [int(prompt_waveform_tensor.shape[-1])], dtype=torch.int32
+        )
+    if prompt_latent_value is not None:
+        additional_information[KEY_PROMPT_LATENTS] = prompt_latent_value
+    if speaker_embeddings is not None:
+        additional_information[KEY_SPEAKER_EMBEDDING] = (
+            speaker_embeddings[0] if len(speaker_embeddings) == 1 else torch.stack(speaker_embeddings, dim=0)
+        )
+    if use_zero_spk_emb:
+        additional_information["use_zero_spk_emb"] = True
+    return {
+        "prompt": prompt,
+        "text": text,
+        "prompt_token_ids": prompt_token_ids,
+        "additional_information": additional_information,
+    }
+
+
+__all__ = [
+    "DEFAULT_PROMPT",
+    "build_dense_prompt_token_ids",
+    "build_ming_dense_prompt",
+    "create_instruction",
+    "estimate_decode_step_window_for_duration",
+    "estimate_decode_steps_for_duration",
+    "parse_duration_seconds",
+    "resolve_effective_runtime_controls",
+]

vllm_omni/model_executor/models/ming_tts/prompt_encoder.py

@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from __future__ import annotations
+
+import json
+import threading
+import time
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors import safe_open
+from vllm.logger import init_logger
+
+from vllm_omni.engine.stage_init_utils import _resolve_model_to_local_path
+from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE
+
+from .audio_prep import (
+    _coerce_prompt_latents,
+    _normalize_prompt_waveform,
+    count_prompt_latent_patches,
+    pad_prompt_waveform,
+)
+from .config_ming_tts import KEY_PROMPT_LATENTS
+
+logger = init_logger(__name__)
+_PROMPT_ENCODER_LOAD_LOCK = threading.Lock()
+
+
+def _resolve_prompt_latents(wrapper: Any, info_dict: dict[str, Any]) -> dict[str, torch.Tensor] | None:
+    raw_latents = info_dict.get(KEY_PROMPT_LATENTS, info_dict.get("prompt_latents"))
+    raw_waveform = info_dict.get("prompt_waveform", info_dict.get("prompt_waveforms"))
+    if raw_latents is not None and raw_waveform is not None:
+        raise ValueError(
+            "Ming waveform cloning request provided both raw prompt_waveform and explicit prompt_latents. "
+            "Choose exactly one source of truth."
+        )
+
+    direct_latents = _coerce_prompt_latents(
+        raw_latents,
+        patch_size=wrapper.ming_config.patch_size,
+        latent_dim=wrapper.ming_config.latent_dim,
+    )
+    if direct_latents is not None:
+        return direct_latents
+    if raw_waveform is None:
+        return None
+
+    encode_fn = getattr(wrapper, "_encode_prompt_waveform_to_latents", None)
+    if callable(encode_fn):
+        latents = encode_fn(raw_waveform, info_dict.get("prompt_waveform_length"))
+    else:
+        latents = _encode_prompt_waveform_to_latents(
+            wrapper,
+            raw_waveform,
+            info_dict.get("prompt_waveform_length"),
+        )
+    return _coerce_prompt_latents(
+        latents,
+        patch_size=wrapper.ming_config.patch_size,
+        latent_dim=wrapper.ming_config.latent_dim,
+    )
+
+
+def _load_prompt_encoder(wrapper: Any) -> AudioVAE:
+    if wrapper._prompt_encoder is not None:
+        return wrapper._prompt_encoder
+    with _PROMPT_ENCODER_LOAD_LOCK:
+        if wrapper._prompt_encoder is not None:
+            return wrapper._prompt_encoder
+        if wrapper.ming_config.audio_tokenizer_config is None:
+            raise RuntimeError("Ming Stage-0 requires audio_tokenizer_config to encode prompt audio.")
+
+        load_start = time.perf_counter()
+        encoder = AudioVAE(wrapper.ming_config.audio_tokenizer_config).eval()
+        state_dict = encoder.state_dict()
+        loaded = 0
+        loaded_encoder_params = set()
+        with torch.no_grad():
+            for shard_path in _iter_model_safetensors(
+                _resolve_model_to_local_path(str(wrapper.vllm_config.model_config.model))
+            ):
+                with safe_open(str(shard_path), framework="pt", device="cpu") as handle:
+                    for key in handle.keys():
+                        if not key.startswith("audio.encoder."):
+                            continue
+                        name = key[len("audio.") :]
+                        if name not in state_dict:
+                            continue
+                        target = state_dict[name]
+                        target.copy_(handle.get_tensor(key).to(device=target.device, dtype=target.dtype))
+                        loaded += 1
+                        loaded_encoder_params.add(name)
+        if loaded == 0:
+            raise RuntimeError("Ming prompt encoder received no audio.encoder.* weights from checkpoint.")
+
+        expected_encoder_params = {f"encoder.{name}" for name, _ in encoder.encoder.named_parameters()}
+        missing = expected_encoder_params - loaded_encoder_params
+        if missing:
+            raise RuntimeError(
+                f"Ming prompt encoder: {len(missing)} params not loaded. First few: {sorted(missing)[:5]}"
+            )
+
+        dev = next(wrapper.parameters()).device
+        try:
+            del encoder.decoder
+            encoder.decoder = None
+            if dev.type != "cpu":
+                encoder.encoder.to(dev, dtype=getattr(wrapper.model, "fm_dtype", torch.bfloat16))
+            else:
+                encoder.encoder.to(dev)
+        except Exception as exc:
+            raise RuntimeError(f"Failed to move Ming prompt encoder to {dev}: {exc}") from exc
+        wrapper._prompt_encoder = encoder
+        logger.info("Ming prompt encoder cold-loaded in %.3f ms", (time.perf_counter() - load_start) * 1000.0)
+        return encoder
+
+
+@torch.inference_mode()
+def _encode_prompt_waveform_to_latents(wrapper: Any, waveform: Any, waveform_length: Any = None) -> torch.Tensor:
+    encoder = _load_prompt_encoder(wrapper)
+    waveform = _normalize_prompt_waveform(waveform, target_sr=wrapper.ming_config.sample_rate)
+    waveform = pad_prompt_waveform(
+        waveform,
+        patch_size=wrapper.ming_config.patch_size,
+        sample_rate=wrapper.ming_config.sample_rate,
+    )
+    dev = next(encoder.encoder.parameters()).device
+    waveform = waveform.to(device=dev, dtype=next(encoder.encoder.parameters()).dtype)
+    if waveform_length is None:
+        waveform_length = torch.full((waveform.shape[0],), waveform.shape[-1], dtype=torch.int32, device=dev)
+    elif not isinstance(waveform_length, torch.Tensor):
+        waveform_length = torch.as_tensor(waveform_length, dtype=torch.int32, device=dev)
+    else:
+        waveform_length = waveform_length.to(device=dev, dtype=torch.int32)
+
+    latents, _ = encoder.encode_latent(waveform, waveform_length)
+    if latents.ndim == 3 and latents.shape[0] == 1:
+        latents = latents.squeeze(0)
+    count_prompt_latent_patches(
+        latents,
+        patch_size=wrapper.ming_config.patch_size,
+        latent_dim=wrapper.ming_config.latent_dim,
+    )
+    return latents.detach().to(dtype=torch.float32).contiguous()
+
+
+def _iter_model_safetensors(local_model_path: str) -> list[Path]:
+    model_root = Path(local_model_path)
+    index_path = model_root / "model.safetensors.index.json"
+    if index_path.exists():
+        with index_path.open("r", encoding="utf-8") as handle:
+            index_data = json.load(handle)
+        filenames = sorted(set(index_data.get("weight_map", {}).values()))
+        if not filenames:
+            raise RuntimeError(f"No checkpoint shards listed in {index_path}")
+        return [model_root / filename for filename in filenames]
+
+    single_file = model_root / "model.safetensors"
+    if single_file.exists():
+        return [single_file]
+
+    files = sorted(model_root.glob("*.safetensors"))
+    if not files:
+        raise RuntimeError(f"No .safetensors checkpoint found under {local_model_path}")
+    return files
+
+
+__all__ = [
+    "_encode_prompt_waveform_to_latents",
+    "_iter_model_safetensors",
+    "_load_prompt_encoder",
+    "_resolve_prompt_latents",
+]

vllm_omni/model_executor/models/ming_tts/speaker_extractor.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/spkemb_extractor.py
+import os
+
+import torch
+import torchaudio
+from vllm.multimodal.media.audio import load_audio
+
+from vllm_omni.model_executor.models.ming_flash_omni.spk_embedding import SpkembExtractor
+
+
+def resolve_model_to_local_path(model):
+    if os.path.isdir(model):
+        return model
+
+    from huggingface_hub import snapshot_download
+
+    return snapshot_download(model)
+
+
+class MingSpeakerEmbeddingExtractor:
+    def __init__(self, model, target_sr=16000):
+        local_model_path = resolve_model_to_local_path(model)
+        campplus_path = os.path.join(local_model_path, "campplus.onnx")
+        if not os.path.exists(campplus_path):
+            raise RuntimeError(f"Missing Ming speaker extractor model: {campplus_path}")
+
+        self.target_sr = int(target_sr)
+        self._core = SpkembExtractor(campplus_path, target_sr=self.target_sr)
+
+    def extract_from_waveform(self, waveform, sample_rate):
+        if not isinstance(waveform, torch.Tensor):
+            waveform = torch.as_tensor(waveform)
+
+        tensor = waveform.detach().to(torch.float32)
+        if tensor.ndim == 1:
+            tensor = tensor.unsqueeze(0)
+        if int(sample_rate) != self.target_sr:
+            tensor = torchaudio.transforms.Resample(orig_freq=int(sample_rate), new_freq=self.target_sr)(tensor)
+
+        embedding = self._core._extract_spk_embedding(tensor)
+        return embedding.squeeze(0).to(dtype=torch.float32)
+
+    def extract_from_file(self, audio_path):
+        audio, sample_rate = load_audio(audio_path, sr=None, mono=False)
+        waveform = torch.from_numpy(audio).contiguous()
+        return self.extract_from_waveform(waveform, sample_rate)
+
+    def extract_many(self, audio_paths):
+        return [self.extract_from_file(path) for path in audio_paths]

vllm_omni/model_executor/models/ming_tts/validation.py

@@ -0,0 +1,172 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from __future__ import annotations
+
+from typing import Any
+
+from transformers import PretrainedConfig
+
+from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig
+
+from .constants import (
+    AGGREGATOR_HIDDEN_SIZE,
+    HISTORY_PATCH_SIZE,
+    LATENT_DIM,
+    LLM_HIDDEN_SIZE,
+    LLM_VOCAB_SIZE,
+    PATCH_SIZE,
+    SAMPLE_RATE,
+)
+
+
+def _to_plain_dict(obj: Any) -> dict[str, Any]:
+    """Normalize nested config objects into plain dicts when possible."""
+    if obj is None:
+        return {}
+    if isinstance(obj, dict):
+        return dict(obj)
+    if isinstance(obj, PretrainedConfig):
+        return obj.to_dict()
+    if hasattr(obj, "to_dict") and callable(obj.to_dict):
+        try:
+            return dict(obj.to_dict())
+        except Exception:
+            pass
+    try:
+        return dict(vars(obj))
+    except Exception:
+        return {}
+
+
+def _coerce_audio_vae_config(atc_raw: Any) -> AudioVAEConfig | None:
+    if atc_raw is None:
+        return None
+    if isinstance(atc_raw, AudioVAEConfig):
+        return atc_raw
+    if isinstance(atc_raw, PretrainedConfig):
+        atc_dict = atc_raw.to_dict()
+    elif isinstance(atc_raw, dict):
+        atc_dict = dict(atc_raw)
+    elif hasattr(atc_raw, "to_dict") and callable(atc_raw.to_dict):
+        atc_dict = dict(atc_raw.to_dict())
+    else:
+        raise TypeError(f"Unsupported audio_tokenizer_config type for Ming dense config: {type(atc_raw)!r}")
+
+    return AudioVAEConfig(**atc_dict)
+
+
+def _nested_get(obj: Any, *keys: str, default: Any = None) -> Any:
+    """Safe nested attribute/key access for dicts and config-like objects."""
+    cur = obj
+    for key in keys:
+        if cur is None:
+            return default
+        if isinstance(cur, dict):
+            cur = cur.get(key)
+        else:
+            cur = getattr(cur, key, None)
+    return cur if cur is not None else default
+
+
+def validate_ming_tts_config(cfg: Any) -> None:
+    """Run before GPU allocation/weight loading. Raises ValueError on mismatches."""
+    if cfg.audio_dummy_token_id != 151705:
+        raise ValueError(
+            f"audio_dummy_token_id={cfg.audio_dummy_token_id}, expected 151705 (<audioPatch>). "
+            "Wrong tokenizer/checkpoint?"
+        )
+    if cfg.audio_eos_token_id != 151704:
+        raise ValueError(
+            f"audio_eos_token_id={cfg.audio_eos_token_id}, expected 151704 (<end_of_audio>). "
+            "Wrong tokenizer/checkpoint?"
+        )
+    if cfg.text_eos_token_id != 151669:
+        raise ValueError(
+            f"text_eos_token_id={cfg.text_eos_token_id}, expected 151669 (<text_eos>). Wrong tokenizer/checkpoint?"
+        )
+
+    if cfg.audio_tokenizer_config is None:
+        raise ValueError("audio_tokenizer_config is None. Nested AudioVAE config was not deserialized correctly.")
+
+    if cfg.latent_dim != LATENT_DIM:
+        raise ValueError(
+            f"latent_dim mismatch: got {cfg.latent_dim}, expected {LATENT_DIM}. "
+            "Check audio_tokenizer_config.enc_kwargs.latent_dim."
+        )
+    if cfg.patch_size != PATCH_SIZE:
+        raise ValueError(
+            f"patch_size mismatch: got {cfg.patch_size}, expected {PATCH_SIZE}. Check ditar_config.patch_size."
+        )
+    if cfg.history_patch_size != HISTORY_PATCH_SIZE:
+        raise ValueError(
+            f"history_patch_size mismatch: got {cfg.history_patch_size}, expected {HISTORY_PATCH_SIZE}. "
+            "Check ditar_config.history_patch_size."
+        )
+    if cfg.llm_hidden_size != LLM_HIDDEN_SIZE:
+        raise ValueError(
+            f"llm_hidden_size mismatch: got {cfg.llm_hidden_size}, expected {LLM_HIDDEN_SIZE}. "
+            "Check llm_config.hidden_size."
+        )
+    if cfg.llm_vocab_size != LLM_VOCAB_SIZE:
+        raise ValueError(f"llm_vocab_size mismatch: got {cfg.llm_vocab_size}, expected {LLM_VOCAB_SIZE}.")
+    if cfg.sample_rate != SAMPLE_RATE:
+        raise ValueError(f"sample_rate mismatch: got {cfg.sample_rate}, expected {SAMPLE_RATE}.")
+
+    if cfg.vae_patch_size != cfg.patch_size:
+        raise ValueError(f"VAE patch size ({cfg.vae_patch_size}) != flow/DiT patch size ({cfg.patch_size}).")
+
+    llm_hidden_from_cfg = cfg.llm_config.get("hidden_size")
+    if llm_hidden_from_cfg is not None and llm_hidden_from_cfg != cfg.llm_hidden_size:
+        raise ValueError(f"llm_hidden_size ({cfg.llm_hidden_size}) != llm_config.hidden_size ({llm_hidden_from_cfg}).")
+
+    agg_h = cfg.aggregator_config.get("hidden_size")
+    dit_h = cfg.ditar_config.get("hidden_size")
+    if agg_h is not None and dit_h is not None and agg_h != dit_h:
+        raise ValueError(f"aggregator_config.hidden_size ({agg_h}) != ditar_config.hidden_size ({dit_h}).")
+    if agg_h is not None and agg_h != AGGREGATOR_HIDDEN_SIZE:
+        raise ValueError(f"aggregator hidden_size mismatch: got {agg_h}, expected {AGGREGATOR_HIDDEN_SIZE}.")
+    if dit_h is not None and dit_h != AGGREGATOR_HIDDEN_SIZE:
+        raise ValueError(f"ditar hidden_size mismatch: got {dit_h}, expected {AGGREGATOR_HIDDEN_SIZE}.")
+
+    atc = cfg.audio_tokenizer_config
+    semantic_module_kwargs = _nested_get(atc, "semantic_module_kwargs", default=None)
+    if semantic_module_kwargs is not None:
+        raise ValueError("Ming dense 0.5B expects audio_tokenizer_config.semantic_module_kwargs to be null.")
+
+    enc_latent = _nested_get(atc, "enc_kwargs", "latent_dim", default=None)
+    dec_latent = _nested_get(atc, "dec_kwargs", "latent_dim", default=None)
+    if enc_latent is not None and enc_latent != cfg.latent_dim:
+        raise ValueError(f"audio enc latent_dim ({enc_latent}) != Ming latent_dim ({cfg.latent_dim}).")
+    if dec_latent is not None and dec_latent != cfg.latent_dim:
+        raise ValueError(f"audio dec latent_dim ({dec_latent}) != Ming latent_dim ({cfg.latent_dim}).")
+
+    atc_patch = _nested_get(atc, "patch_size", default=None)
+    if atc_patch is not None and atc_patch != cfg.vae_patch_size:
+        raise ValueError(f"audio_tokenizer_config.patch_size ({atc_patch}) != vae_patch_size ({cfg.vae_patch_size}).")
+
+    atc_sr = _nested_get(atc, "sample_rate", default=None)
+    if atc_sr is not None and atc_sr != cfg.sample_rate:
+        raise ValueError(f"audio_tokenizer_config.sample_rate ({atc_sr}) != sample_rate ({cfg.sample_rate}).")
+
+    enc_input_dim = _nested_get(atc, "enc_kwargs", "input_dim", default=None)
+    enc_hop_size = _nested_get(atc, "enc_kwargs", "hop_size", default=None)
+    dec_output_dim = _nested_get(atc, "dec_kwargs", "output_dim", default=None)
+
+    if enc_input_dim is not None and enc_hop_size is not None and enc_input_dim != enc_hop_size:
+        raise ValueError(f"AudioVAE encoder input_dim ({enc_input_dim}) != hop_size ({enc_hop_size}).")
+    if enc_hop_size is not None and dec_output_dim is not None and enc_hop_size != dec_output_dim:
+        raise ValueError(
+            f"AudioVAE encoder hop_size ({enc_hop_size}) != decoder output_dim ({dec_output_dim}). "
+            "Expected 882 in this checkpoint family."
+        )
+
+    if cfg.latent_chunk_size <= 0:
+        raise ValueError(f"latent_chunk_size must be > 0, got {cfg.latent_chunk_size}.")
+    if cfg.latent_left_context < 0:
+        raise ValueError(f"latent_left_context must be >= 0, got {cfg.latent_left_context}.")
+    if cfg.max_decode_steps <= 0:
+        raise ValueError(f"max_decode_steps must be > 0, got {cfg.max_decode_steps}.")
+    if not (0.0 <= cfg.stop_head_threshold <= 1.0):
+        raise ValueError(f"stop_head_threshold must be in [0,1], got {cfg.stop_head_threshold}.")
+    if cfg.stop_head_min_steps < 0:
+        raise ValueError(f"stop_head_min_steps must be >= 0, got {cfg.stop_head_min_steps}.")

vllm_omni/model_executor/models/registry.py

@@ -156,6 +156,22 @@ _OMNI_MODELS = {
         "mimo_audio_code2wav",
         "MiMoAudioToken2WavForConditionalGenerationVLLM",
     ),
+    ## ming-tts
+    "MingTTSForConditionalGeneration": (
+        "ming_tts",
+        "ming_tts",
+        "MingTTSForConditionalGeneration",
+    ),
+    "MingLLMModel": (
+        "ming_tts",
+        "ming_tts_llm",
+        "MingLLMModel",
+    ),
+    "MingAudioVAEModel": (
+        "ming_tts",
+        "ming_tts_audio_vae",
+        "MingAudioVAEModel",
+    ),
     ## glm_image
     "GlmImageForConditionalGeneration": (
         "glm_image",