diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index 38de0e0c5..c01782790 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -7,7 +7,6 @@ # Original file was released under Apache-2.0, with the full license text # available at https://github.com/huggingface/transformers/blob/main/LICENSE. -import math from collections.abc import Iterable from dataclasses import dataclass @@ -49,6 +48,7 @@ from vllm_omni.diffusion.distributed.parallel_state import ( ) from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available from vllm_omni.diffusion.layers.rope import RotaryEmbedding +from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding def patchify(imgs, p): @@ -1041,28 +1041,8 @@ class TimestepEmbedder(nn.Module): ) self.frequency_embedding_size = frequency_embedding_size - @staticmethod - def timestep_embedding(t, dim, max_period=10000): - """ - Create sinusoidal timestep embeddings. - :param t: a 1-D Tensor of N indices, one per batch element. - These may be fractional. - :param dim: the dimension of the output. - :param max_period: controls the minimum frequency of the embeddings. - :return: an (N, D) Tensor of positional embeddings. - """ - half = dim // 2 - freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( - device=t.device - ) - args = t[:, None].float() * freqs[None] - embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) - if dim % 2: - embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) - return embedding - def forward(self, t): - t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + t_freq = timestep_embedding(t, self.frequency_embedding_size) t_emb = self.mlp(t_freq) return t_emb diff --git a/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py b/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py index 9c5af8f33..6333819b0 100644 --- a/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py +++ b/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py @@ -16,6 +16,7 @@ from vllm.logger import init_logger from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb from vllm_omni.diffusion.attention.layer import Attention as DiffusionAttention +from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding logger = init_logger(__name__) @@ -277,38 +278,6 @@ class ConvNeXtV2Block(nn.Module): return residual + x -class SinusPositionEmbedding(nn.Module): - """Sinusoidal position embedding.""" - - def __init__(self, dim): - super().__init__() - self.dim = dim - - def forward(self, x, scale=1000): - device = x.device - half_dim = self.dim // 2 - emb = math.log(10000) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb) - emb = scale * x.unsqueeze(1) * emb.unsqueeze(0) - emb = torch.cat((emb.sin(), emb.cos()), dim=-1) - return emb - - -class TimestepEmbedding(nn.Module): - """Timestep embedding with MLP.""" - - def __init__(self, dim, freq_embed_dim=256): - super().__init__() - self.time_embed = SinusPositionEmbedding(freq_embed_dim) - self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) - - def forward(self, timestep: torch.Tensor): - time_hidden = self.time_embed(timestep) - time_hidden = time_hidden.to(timestep.dtype) - time = self.time_mlp(time_hidden) - return time - - class TextEmbedding(nn.Module): """Text embedding with optional ConvNeXt modeling.""" @@ -394,7 +363,7 @@ class DiT(nn.Module): ): super().__init__() - self.time_embed = TimestepEmbedding(dim) + self.time_embed = DiTTimestepEmbedding(dim) if mu_dim is None: mu_dim = mel_dim self.input_embed = InputEmbedding(mel_dim, mu_dim, dim, spk_dim) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index 0f3c33389..e1acc4cdc 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py @@ -3,7 +3,6 @@ import inspect import logging -import math from collections.abc import Callable, Iterable from dataclasses import dataclass from typing import Any, cast @@ -75,6 +74,7 @@ from vllm_omni.diffusion.distributed.sp_plan import ( from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.layers.rope import RotaryEmbedding from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import HunyuanFusedMoE +from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding logger = logging.getLogger(__name__) @@ -2878,31 +2878,6 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline): return HunyuanImage3Text2ImagePipelineOutput(samples=image) -def timestep_embedding(t, dim, max_period=10000): - """ - Create sinusoidal timestep embeddings. - - Args: - t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional. - dim (int): the dimension of the output. - max_period (int): controls the minimum frequency of the embeddings. - - Returns: - embedding (torch.Tensor): An (N, D) Tensor of positional embeddings. - - .. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py - """ - half = dim // 2 - freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( - device=t.device - ) - args = t[:, None].float() * freqs[None] - embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) - if dim % 2: - embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) - return embedding - - class TimestepEmbedder(nn.Module): """ Embeds scalar timesteps into vector representations. diff --git a/vllm_omni/diffusion/models/nextstep_1_1/modeling_nextstep_heads.py b/vllm_omni/diffusion/models/nextstep_1_1/modeling_nextstep_heads.py index 7768ac765..3a7a76633 100644 --- a/vllm_omni/diffusion/models/nextstep_1_1/modeling_nextstep_heads.py +++ b/vllm_omni/diffusion/models/nextstep_1_1/modeling_nextstep_heads.py @@ -6,11 +6,11 @@ from __future__ import annotations -import math - import torch import torch.nn as nn +from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding + # --------------------------------------------------------------------------- # Utilities (inlined from remote utils/model_utils.py) # --------------------------------------------------------------------------- @@ -109,20 +109,8 @@ class TimestepEmbedder(nn.Module): ) self.frequency_embedding_size = frequency_embedding_size - @staticmethod - def timestep_embedding(t: torch.Tensor, dim: int, max_period: float = 10000.0) -> torch.Tensor: - half = dim // 2 - freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( - device=t.device - ) - args = t[:, None].float() * freqs[None] - embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) - if dim % 2: - embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) - return embedding - def forward(self, t: torch.Tensor) -> torch.Tensor: - t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + t_freq = timestep_embedding(t, self.frequency_embedding_size) return self.mlp(t_freq.to(self.mlp[0].weight.dtype)) diff --git a/vllm_omni/diffusion/models/z_image/z_image_transformer.py b/vllm_omni/diffusion/models/z_image/z_image_transformer.py index c36ea7466..80daec6b8 100644 --- a/vllm_omni/diffusion/models/z_image/z_image_transformer.py +++ b/vllm_omni/diffusion/models/z_image/z_image_transformer.py @@ -50,6 +50,7 @@ from vllm_omni.diffusion.forward_context import ( is_forward_context_available, ) from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk +from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding ADALN_EMBED_DIM = 256 SEQ_MULTI_OF = 32 @@ -236,20 +237,8 @@ class TimestepEmbedder(nn.Module): self.frequency_embedding_size = frequency_embedding_size - @staticmethod - def timestep_embedding(t, dim, max_period=10000): - half = dim // 2 - freqs = torch.exp( - -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half - ) - args = t[:, None].float() * freqs[None] - embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) - if dim % 2: - embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) - return embedding - def forward(self, t): - t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + t_freq = timestep_embedding(t, self.frequency_embedding_size) weight_dtype = self.mlp[0].bias.dtype if weight_dtype.is_floating_point: t_freq = t_freq.to(weight_dtype) diff --git a/vllm_omni/model_executor/layers/timestep_embedding.py b/vllm_omni/model_executor/layers/timestep_embedding.py new file mode 100644 index 000000000..8c6324c22 --- /dev/null +++ b/vllm_omni/model_executor/layers/timestep_embedding.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared timestep embedding primitives for diffusion models. + +- ``SinusPositionEmbedding``: sin/cos positional encoding (TTS DiT/CFM). + Used by Qwen3-TTS, Qwen2.5-Omni, CosyVoice3, Ming-Flash-Omni. +- ``DiTTimestepEmbedding``: SinusPosEmb + Linear + SiLU + Linear MLP. + Used by the same models as above. +- ``timestep_embedding()``: standalone function (GLIDE/DiT convention). + Used by Bagel, NextStep, Z-Image, HunyuanImage3. +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn + + +def timestep_embedding(t: torch.Tensor, dim: int, max_period: float = 10000.0) -> torch.Tensor: + """Create sinusoidal timestep embeddings (GLIDE/DiT convention). + + Produces cos-then-sin embeddings with log-spaced frequencies. + + Args: + t: (N,) 1-D tensor of timestep indices (may be fractional). + dim: Output embedding dimension. + max_period: Controls the minimum frequency. + + Returns: + (N, dim) tensor of positional embeddings. + """ + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + +class SinusPositionEmbedding(nn.Module): + """Sinusoidal position embedding for scalar timesteps. + + Maps scalar timestep values to ``dim``-dimensional embeddings using + the standard log-spaced frequency formula from DDPM/DiT. + + Args: + dim: Output embedding dimension (must be even). + """ + + def __init__(self, dim: int) -> None: + super().__init__() + self.dim = dim + + def forward(self, x: torch.Tensor, scale: float = 1000.0) -> torch.Tensor: + """ + Args: + x: (N,) scalar timesteps. + scale: Frequency scaling factor. + + Returns: + (N, dim) sinusoidal embeddings, cast to the input dtype. + """ + device = x.device + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb) + emb = scale * x.unsqueeze(1) * emb.unsqueeze(0) + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb.to(x.dtype) + + +class DiTTimestepEmbedding(nn.Module): + """Timestep conditioning: SinusPositionEmbedding + Linear + SiLU + Linear. + + Args: + dim: Hidden dimension (output size). + freq_embed_dim: Sinusoidal embedding dimension (input to MLP). + """ + + def __init__(self, dim: int, freq_embed_dim: int = 256) -> None: + super().__init__() + self.time_embed = SinusPositionEmbedding(freq_embed_dim) + self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + + def forward(self, timestep: torch.Tensor) -> torch.Tensor: + time_hidden = self.time_embed(timestep) + time_hidden = time_hidden.to(timestep.dtype) + return self.time_mlp(time_hidden) diff --git a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py index 80acbaad0..ae34eb865 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py @@ -21,7 +21,6 @@ # MAE: https://github.com/facebookresearch/mae/blob/main/models_mae.py # -------------------------------------------------------- import logging -import math from functools import cached_property from queue import Queue from threading import Lock @@ -33,6 +32,8 @@ from transformers import PreTrainedTokenizerBase, Qwen2Config, Qwen2Model, Stati from vllm.logger import init_logger from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb +from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding + from .audio_vae import AudioVAE logger = init_logger(__name__) @@ -231,34 +232,6 @@ class FinalLayer(nn.Module): return x -class SinusPositionEmbedding(nn.Module): - def __init__(self, dim: int): - super().__init__() - self.dim = dim - - def forward(self, x: torch.Tensor, scale: float = 1000) -> torch.Tensor: - device = x.device - half_dim = self.dim // 2 - emb = math.log(10000) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb) - emb = scale * x.unsqueeze(1) * emb.unsqueeze(0) - emb = torch.cat((emb.sin(), emb.cos()), dim=-1) - return emb - - -class TimestepEmbedder(nn.Module): - def __init__(self, dim: int, freq_embed_dim: int = 256): - super().__init__() - self.time_embed = SinusPositionEmbedding(freq_embed_dim) - self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) - - def forward(self, timestep: torch.Tensor) -> torch.Tensor: - time_hidden = self.time_embed(timestep) - time_hidden = time_hidden.to(timestep.dtype) - time = self.time_mlp(time_hidden) - return time - - class CondEmbedder(nn.Module): """Embeds LLM hidden states with optional CFG dropout.""" @@ -288,7 +261,7 @@ class DiT(nn.Module): self.out_channels = in_channels self.num_heads = num_heads - self.t_embedder = TimestepEmbedder(hidden_size) + self.t_embedder = DiTTimestepEmbedding(hidden_size) self.x_embedder = nn.Linear(in_channels, hidden_size) self.c_embedder = CondEmbedder(llm_cond_dim, hidden_size) if "spk_dim" in kwargs: diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py index 84ff8b55f..34c88bad9 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py @@ -34,6 +34,7 @@ from vllm.v1.outputs import SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler +from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import cap_and_align_mel_length, resolve_max_mel_frames from vllm_omni.platforms import current_omni_platform @@ -598,36 +599,6 @@ class DiTAttention(nn.Module): return attention_output -# time step conditioning embedding -class SinusPositionEmbedding(nn.Module): - def __init__(self, dim): - super().__init__() - self.dim = dim - - def forward(self, hidden_states, scale=1000): - device = hidden_states.device - half_dim = self.dim // 2 - emb = math.log(10000) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb) - emb = scale * hidden_states.unsqueeze(1) * emb.unsqueeze(0) - emb = torch.cat((emb.sin(), emb.cos()), dim=-1) - return emb.type_as(hidden_states) - - -class DiTTimestepEmbedding(nn.Module): - def __init__(self, dim, freq_embed_dim=256): - super().__init__() - self.time_embed = SinusPositionEmbedding(freq_embed_dim) - self.time_mlp = nn.ModuleList([nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)]) - - def forward(self, timestep): # noqa: F821 - time_hidden = self.time_embed(timestep) - time_hidden = time_hidden.to(timestep.dtype) - for layer in self.time_mlp: - time_hidden = layer(time_hidden) # b d - return time_hidden - - class DiTDecoderLayer(nn.Module): def __init__(self, config: Qwen2_5OmniDiTConfig, look_ahead_block=0, look_backward_block=0): super().__init__() diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py index bed108ff6..7bc0030bc 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py @@ -26,6 +26,8 @@ from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from transformers.utils import ModelOutput, auto_docstring, logging from transformers.utils.hub import cached_file +from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding + from .configuration_qwen3_tts_tokenizer_v1 import ( Qwen3TTSTokenizerV1Config, Qwen3TTSTokenizerV1DecoderBigVGANConfig, @@ -626,36 +628,6 @@ class DiTAttention(nn.Module): return attention_output -# time step conditioning embedding -class SinusPositionEmbedding(nn.Module): - def __init__(self, dim): - super().__init__() - self.dim = dim - - def forward(self, hidden_states, scale=1000): - device = hidden_states.device - half_dim = self.dim // 2 - emb = math.log(10000) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb) - emb = scale * hidden_states.unsqueeze(1) * emb.unsqueeze(0) - emb = torch.cat((emb.sin(), emb.cos()), dim=-1) - return emb.type_as(hidden_states) - - -class DiTTimestepEmbedding(nn.Module): - def __init__(self, dim, freq_embed_dim=256): - super().__init__() - self.time_embed = SinusPositionEmbedding(freq_embed_dim) - self.time_mlp = nn.ModuleList([nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)]) - - def forward(self, timestep): - time_hidden = self.time_embed(timestep) - time_hidden = time_hidden.to(timestep.dtype) - for layer in self.time_mlp: - time_hidden = layer(time_hidden) # b d - return time_hidden - - class DiTDecoderLayer(nn.Module): def __init__(self, config: Qwen3TTSTokenizerV1DecoderBigVGANConfig, look_ahead_block=0, look_backward_block=0): super().__init__()