[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

@@ -0,0 +1,379 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PretrainedConfig, PreTrainedModel, Qwen2Config, Qwen2Model
+from transformers.utils import is_flash_attn_2_available
+from vllm.logger import init_logger
+
+from .audio_dsp import ISTFTHead
+
+logger = init_logger(__name__)
+
+
+class AudioVAEConfig(PretrainedConfig):
+    model_type = "audio_vae"
+
+    def __init__(
+        self,
+        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)
+
+
+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):
+    def __init__(self, scale_factor=4):
+        super().__init__()
+        self.scale_factor = scale_factor
+        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}
+
+        if x is None and not is_last:
+            return None, state
+
+        if state["is_first"] and is_last:
+            out = self.upsampler(x.transpose(1, 2)).transpose(1, 2)
+            return out, None
+
+        output_chunks = []
+
+        if state["is_first"]:
+            state["prev_chunk"] = x
+            state["is_first"] = False
+            if not is_last:
+                return None, state
+
+        if state["prev_chunk"] is not None:
+            p = state["prev_chunk"].transpose(1, 2)
+
+            if state["history_last"] is None:
+                lookahead = x[:, :1, :].transpose(1, 2)
+                inp = torch.cat([p, lookahead], dim=2)
+                up = self.upsampler(inp)
+                out_prev = up[:, :, : p.size(2) * self.scale_factor]
+            else:
+                lookahead = x[:, :1, :].transpose(1, 2)
+                inp = torch.cat([state["history_last"], p, lookahead], dim=2)
+                up = self.upsampler(inp)
+                start = self.scale_factor
+                end = start + p.size(2) * self.scale_factor
+                out_prev = up[:, :, start:end]
+
+            output_chunks.append(out_prev.transpose(1, 2))
+            state["history_last"] = p[:, :, -1:]
+            state["prev_chunk"] = x
+
+        if is_last:
+            p = state["prev_chunk"].transpose(1, 2)
+            inp = torch.cat([state["history_last"], p], dim=2)
+            up = self.upsampler(inp)
+            out_last = up[:, :, self.scale_factor :]
+            output_chunks.append(out_last.transpose(1, 2))
+            state = None
+
+        final_out = torch.cat(output_chunks, dim=1) if output_chunks else None
+        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 = _qwen2_config(decoder_args)
+        logger.info("AudioVAE Decoder: using attn_implementation=%r", config._attn_implementation)
+        self.decoder = Qwen2Model(config)
+        self.output_dim = output_dim
+        self.latent_dim = latent_dim
+        self.fc1 = nn.Linear(latent_dim, config.hidden_size)
+        self.hop_length = output_dim
+        self.head = ISTFTHead(
+            dim=config.hidden_size, n_fft=self.hop_length * 4, hop_length=self.hop_length, padding="same"
+        )
+        self.patch_size = patch_size
+        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)
+        if self.patch_size != -1:
+            if use_cache:
+                x, upsample_state = self.upsampling(x, state=upsample_state, is_last=last_chunk)
+                if x is None:
+                    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 = []
+
+        if use_cache and getattr(self.decoder.config, "sliding_window", None) is not None:
+            sw_size = self.decoder.config.sliding_window
+            target_len = sw_size - 1
+            if past_key_values is None:
+                past_len = 0
+            elif hasattr(past_key_values, "get_seq_length"):
+                past_len = past_key_values.get_seq_length()
+            elif isinstance(past_key_values, tuple) and len(past_key_values) > 0:
+                past_len = past_key_values[0][0].shape[-2]
+            else:
+                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, :]
+                outputs = self.decoder(inputs_embeds=x_fill, past_key_values=past_key_values, use_cache=use_cache)
+                hidden_states_list.append(outputs.last_hidden_state)
+                past_key_values = outputs.past_key_values
+                x = x[:, fill_len:, :]
+
+        outputs = self.decoder(inputs_embeds=x, past_key_values=past_key_values, use_cache=use_cache)
+        hidden_states_list.append(outputs.last_hidden_state)
+        past_key_values = outputs.past_key_values
+
+        if len(hidden_states_list) > 1:
+            full_hidden_state = torch.cat(hidden_states_list, dim=1)
+        else:
+            full_hidden_state = hidden_states_list[0]
+
+        x_out, _, audio_buffer, window_buffer = self.head(
+            full_hidden_state,
+            streaming=use_cache,
+            audio_buffer=audio_buffer,
+            window_buffer=window_buffer,
+            last_chunk=last_chunk,
+        )
+
+        stream_state = (upsample_state, audio_buffer, window_buffer)
+        return x_out, stream_state, past_key_values
+
+
+class AudioVAE(PreTrainedModel):
+    config_class = 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=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=dec_kwargs["backbone"],
+            output_dim=dec_kwargs["output_dim"],
+            latent_dim=dec_kwargs["latent_dim"],
+            patch_size=config.patch_size,
+        )
+        self.post_init()
+
+    def _init_weights(self, module):
+        std = 0.02
+        if isinstance(module, nn.Linear):
+            if self.config.init_method == "kaiming":
+                nn.init.kaiming_normal_(module.weight, mode="fan_in", nonlinearity="relu")
+            else:
+                module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    @torch.inference_mode()
+    def encode_latent(self, waveform, waveform_length):
+        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, _ = self.encoder(waveform)
+        h = h.transpose(1, 2)
+
+        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,
+            use_cache=use_cache,
+            stream_state=stream_state,
+            last_chunk=last_chunk,
+        )
+        return waveform, stream_state, 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