Spaces:
Running on Zero
Running on Zero
File size: 7,002 Bytes
4fe31ed | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | # SPDX-License-Identifier: Apache-2.0
#
# Independent reimplementation of standard sinusoidal position- and
# timestep-embedding utilities, written from their public mathematical
# definitions (see the per-symbol references below). This file is NOT derived
# from the DiT source tree and carries no third-party (CC BY-NC) copyright; it
# is released under Apache-2.0 like the rest of this package.
#
# References for the underlying formulas (math only, no code reuse):
# * Sinusoidal position encoding — Vaswani et al., "Attention Is All You
# Need" (2017), Section 3.5.
# * Sinusoidal timestep embedding — Ho et al., "Denoising Diffusion
# Probabilistic Models" (2020); the same closed form is reused across
# diffusion and flow-matching models.
#
# The module layout, parameter names, and numerical outputs are intentionally
# identical to the previous version so that existing checkpoints load unchanged
# and inference results are bit-for-bit reproducible.
import math
import numpy as np
import torch
from torch import nn
# --------------------------------------------------------------------------- #
# Sinusoidal (sin-cos) position embeddings
#
# For a position p and channel index i the embedding interleaves
# sin(p / 10000^(2i/d)) and cos(p / 10000^(2i/d)),
# i.e. the classic Transformer positional encoding. The 2D variant encodes the
# height and width axes with half the channels each and concatenates them.
# --------------------------------------------------------------------------- #
def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
"""1D sin-cos embedding for a flat array of positions.
Args:
embed_dim: even output dimension per position.
pos: array of positions, any shape; flattened to (M,).
Returns:
(M, embed_dim) array, [sin | cos] halves concatenated.
"""
assert embed_dim % 2 == 0
# Inverse frequencies 1 / 10000^(k/(embed_dim/2)) for k = 0 .. embed_dim/2-1.
inv_freq = np.arange(embed_dim // 2, dtype=np.float64)
inv_freq /= embed_dim / 2.0
inv_freq = 1.0 / 10000**inv_freq # (embed_dim/2,)
angles = np.einsum("m,d->md", pos.reshape(-1), inv_freq) # outer product (M, D/2)
return np.concatenate([np.sin(angles), np.cos(angles)], axis=1) # (M, D)
def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
"""2D sin-cos embedding: half the channels encode each spatial axis."""
assert embed_dim % 2 == 0
emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2)
emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2)
return np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0):
"""Sin-cos table for a ``grid_size x grid_size`` patch grid.
Returns an ``(grid_size**2 [+ extra_tokens], embed_dim)`` array. When
``cls_token`` is set and ``extra_tokens > 0``, that many zero rows are
prepended.
"""
axis_h = np.arange(grid_size, dtype=np.float32)
axis_w = np.arange(grid_size, dtype=np.float32)
grid = np.meshgrid(axis_w, axis_h) # width varies fastest
grid = np.stack(grid, axis=0).reshape([2, 1, grid_size, grid_size])
pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
if cls_token and extra_tokens > 0:
pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
return pos_embed
# --------------------------------------------------------------------------- #
# Timestep embedding
# --------------------------------------------------------------------------- #
class TimestepEmbedder(nn.Module):
"""Embed scalar (possibly fractional) timesteps into vectors.
Sinusoidal frequency features are fed through a two-layer MLP with a SiLU
non-linearity. This is the standard timestep conditioning used by diffusion
and flow-matching models.
"""
def __init__(self, hidden_size, frequency_embedding_size=256):
super().__init__()
self.mlp = nn.Sequential(
nn.Linear(frequency_embedding_size, hidden_size, bias=True),
nn.SiLU(),
nn.Linear(hidden_size, hidden_size, bias=True),
)
self.frequency_embedding_size = frequency_embedding_size
@staticmethod
def timestep_embedding(t, dim, max_period=10000):
"""Sinusoidal features for a 1D tensor of (possibly fractional) timesteps.
Args:
t: (N,) tensor of timestep values.
dim: output feature dimension.
max_period: lowest angular frequency (longest period).
Returns:
(N, dim) tensor of [cos | sin] features (zero-padded if ``dim`` is odd).
"""
half = dim // 2
# Geometrically spaced frequencies over [1, 1/max_period].
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):
freq_features = self.timestep_embedding(t, self.frequency_embedding_size)
freq_features = freq_features.to(next(self.mlp.parameters()).dtype)
return self.mlp(freq_features)
class PositionEmbedding(nn.Module):
"""Fixed 2D sin-cos position table exposed as a lookup by position id."""
def __init__(self, max_num_patch_per_side, hidden_size):
super().__init__()
self.max_num_patch_per_side = max_num_patch_per_side
self.hidden_size = hidden_size
# Build the (non-trainable) table eagerly. If it were created lazily it
# could stay on the meta device through transformers'
# from_pretrained(dtype=...) path and later materialize as uninitialized
# memory when moved to the GPU.
table = get_2d_sincos_pos_embed(hidden_size, max_num_patch_per_side)
self.pos_embed = nn.Parameter(
torch.from_numpy(table).float(),
requires_grad=False,
)
def _reset_parameters(self):
"""Recompute the table after a meta-init path (call post-from_pretrained)."""
table = get_2d_sincos_pos_embed(self.hidden_size, self.max_num_patch_per_side)
on_meta = self.pos_embed.is_meta or self.pos_embed.device.type == "meta"
materialized = torch.from_numpy(table).to(
device="cpu" if on_meta else self.pos_embed.device,
dtype=torch.float32 if on_meta else self.pos_embed.dtype,
)
if on_meta:
self.pos_embed = nn.Parameter(materialized.float(), requires_grad=False)
else:
self.pos_embed.data.copy_(materialized.to(self.pos_embed.dtype))
def forward(self, position_ids):
return self.pos_embed[position_ids]
|