Nick CaoandClaude
9de3806624
Signed-off-by: Nick Cao <ncao@redhat.com> Co-authored-by: Claude <noreply@anthropic.com>
94 lines
3.2 KiB
Python
94 lines
3.2 KiB
Python
# 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)
|