[Cleanup] Extract shared SinusPositionEmbedding/DiTTimestepEmbedding/timestep_embedding layer (#3285)

Signed-off-by: Nick Cao <ncao@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>

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
 

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)

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.

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))
 
 

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)

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)

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:

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__()

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__()