model/model_utils.py

# 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]