541 lines
19 KiB
Python
541 lines
19 KiB
Python
#!/usr/bin/env python3
|
||
|
||
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
|
||
#
|
||
# NVIDIA CORPORATION and its licensors retain all intellectual property
|
||
# and proprietary rights in and to this software, related documentation
|
||
# and any modifications thereto. Any use, reproduction, disclosure or
|
||
# distribution of this software and related documentation without an express
|
||
# license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||
|
||
"""
|
||
DiffiT: Diffusion Vision Transformers for Image Generation.
|
||
Code by Ali Hatamizadeh.
|
||
"""
|
||
|
||
import math
|
||
from typing import Optional
|
||
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn as nn
|
||
from timm.models.layers import trunc_normal_
|
||
from timm.models.vision_transformer import Mlp, PatchEmbed
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Positional embedding utilities
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray) -> np.ndarray:
|
||
"""
|
||
Compute 1-D sinusoidal positional embeddings.
|
||
|
||
Args:
|
||
embed_dim: Output dimension for each position (must be even).
|
||
pos: Positions to encode, arbitrary shape – will be flattened.
|
||
|
||
Returns:
|
||
Embedding array of shape ``(M, embed_dim)``.
|
||
"""
|
||
assert embed_dim % 2 == 0
|
||
omega = np.arange(embed_dim // 2, dtype=np.float64)
|
||
omega /= embed_dim / 2.0
|
||
omega = 1.0 / 10000**omega # (D/2,)
|
||
|
||
pos = pos.reshape(-1) # (M,)
|
||
out = np.einsum("m,d->md", pos, omega) # (M, D/2)
|
||
|
||
return np.concatenate([np.sin(out), np.cos(out)], axis=1) # (M, D)
|
||
|
||
|
||
def get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid: np.ndarray) -> np.ndarray:
|
||
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: int,
|
||
grid_size: int,
|
||
cls_token: bool = False,
|
||
extra_tokens: int = 0,
|
||
) -> np.ndarray:
|
||
"""
|
||
Generate 2-D sinusoidal positional embeddings on a square grid.
|
||
|
||
Args:
|
||
embed_dim: Embedding dimension.
|
||
grid_size: Height (= width) of the grid.
|
||
cls_token: If *True*, prepend ``extra_tokens`` zero rows.
|
||
extra_tokens: Number of extra token slots to prepend.
|
||
|
||
Returns:
|
||
Array of shape ``[grid_size*grid_size, embed_dim]`` (or with
|
||
``extra_tokens`` additional leading rows when requested).
|
||
"""
|
||
grid_h = np.arange(grid_size, dtype=np.float32)
|
||
grid_w = np.arange(grid_size, dtype=np.float32)
|
||
grid = np.meshgrid(grid_w, grid_h) # w first
|
||
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 & label embedders
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TimestepEmbedder(nn.Module):
|
||
"""Embeds scalar diffusion timesteps into vector representations."""
|
||
|
||
def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
|
||
super().__init__()
|
||
self.frequency_embedding_size = frequency_embedding_size
|
||
self.mlp = nn.Sequential(
|
||
nn.Linear(frequency_embedding_size, hidden_size, bias=True),
|
||
nn.SiLU(),
|
||
nn.Linear(hidden_size, hidden_size, bias=True),
|
||
)
|
||
|
||
@staticmethod
|
||
def timestep_embedding(
|
||
t: torch.Tensor, dim: int, max_period: int = 10000
|
||
) -> torch.Tensor:
|
||
"""
|
||
Sinusoidal timestep embeddings following Vaswani et al.
|
||
|
||
Args:
|
||
t: 1-D tensor of ``N`` indices (may be fractional).
|
||
dim: Embedding dimension.
|
||
max_period: Controls the minimum frequency.
|
||
|
||
Returns:
|
||
Tensor of shape ``(N, dim)``.
|
||
"""
|
||
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: torch.Tensor) -> torch.Tensor:
|
||
t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
|
||
return self.mlp(t_freq)
|
||
|
||
|
||
class LabelEmbedder(nn.Module):
|
||
"""
|
||
Embeds class labels into vector representations.
|
||
Supports label dropout for classifier-free guidance.
|
||
"""
|
||
|
||
def __init__(self, num_classes: int, hidden_size: int, dropout_prob: float):
|
||
super().__init__()
|
||
use_cfg_embedding = int(dropout_prob > 0)
|
||
self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)
|
||
self.num_classes = num_classes
|
||
self.dropout_prob = dropout_prob
|
||
|
||
def token_drop(
|
||
self,
|
||
labels: torch.Tensor,
|
||
force_drop_ids: Optional[torch.Tensor] = None,
|
||
) -> torch.Tensor:
|
||
"""Replace a random subset of labels with the null class for CFG."""
|
||
if force_drop_ids is None:
|
||
drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
|
||
else:
|
||
drop_ids = force_drop_ids == 1
|
||
return torch.where(drop_ids, self.num_classes, labels)
|
||
|
||
def forward(
|
||
self,
|
||
labels: torch.Tensor,
|
||
train: bool,
|
||
force_drop_ids: Optional[torch.Tensor] = None,
|
||
) -> torch.Tensor:
|
||
use_dropout = self.dropout_prob > 0
|
||
if (train and use_dropout) or (force_drop_ids is not None):
|
||
labels = self.token_drop(labels, force_drop_ids)
|
||
return self.embedding_table(labels)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# DiffiT attention with relative position bias
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class DiffiTAttention(nn.Module):
|
||
"""
|
||
Window-based multi-head self-attention with time-embedding modulation
|
||
and learnable relative position bias (Swin-style).
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
dim: int,
|
||
temb_dim: Optional[int],
|
||
num_heads: int,
|
||
window_size: int = 16,
|
||
qkv_bias: bool = True,
|
||
qk_scale: Optional[float] = None,
|
||
attn_drop: float = 0.0,
|
||
proj_drop: float = 0.0,
|
||
):
|
||
super().__init__()
|
||
self.num_heads = num_heads
|
||
head_dim = dim // num_heads
|
||
self.scale = qk_scale or head_dim ** -0.5
|
||
|
||
# Time-embedding projection ------------------------------------------
|
||
if temb_dim is not None:
|
||
self.qkv_temb = nn.Linear(temb_dim, dim * 3)
|
||
|
||
# Window / relative position bias ------------------------------------
|
||
self.window_size = (window_size, window_size)
|
||
ws = self.window_size
|
||
|
||
self.relative_position_bias_table = nn.Parameter(
|
||
torch.zeros((2 * ws[0] - 1) * (2 * ws[1] - 1), num_heads)
|
||
)
|
||
trunc_normal_(self.relative_position_bias_table, std=0.02)
|
||
|
||
coords_h = torch.arange(ws[0])
|
||
coords_w = torch.arange(ws[1])
|
||
coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"))
|
||
coords_flat = torch.flatten(coords, 1)
|
||
|
||
relative_coords = coords_flat[:, :, None] - coords_flat[:, None, :]
|
||
relative_coords = relative_coords.permute(1, 2, 0).contiguous()
|
||
relative_coords[:, :, 0] += ws[0] - 1
|
||
relative_coords[:, :, 1] += ws[1] - 1
|
||
relative_coords[:, :, 0] *= 2 * ws[1] - 1
|
||
self.register_buffer("relative_position_index", relative_coords.sum(-1))
|
||
|
||
# Linear projections -------------------------------------------------
|
||
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
||
self.proj = nn.Linear(dim, dim)
|
||
|
||
self.attn_drop = nn.Dropout(attn_drop)
|
||
self.proj_drop = nn.Dropout(proj_drop)
|
||
self.softmax = nn.Softmax(dim=-1)
|
||
|
||
def forward(self, x: torch.Tensor, temb: Optional[torch.Tensor]) -> torch.Tensor:
|
||
B, N, C = x.shape
|
||
|
||
if temb is not None:
|
||
qkv_temb = self.qkv_temb(temb).unsqueeze(1).to(x.dtype)
|
||
qkv = qkv_temb + self.qkv(x)
|
||
else:
|
||
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
|
||
|
||
qkv = qkv.reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
|
||
q, k, v = qkv.unbind(0) # each: (B, heads, N, head_dim)
|
||
|
||
attn = (q * self.scale) @ k.transpose(-2, -1) # (B, heads, N, N)
|
||
|
||
# Add relative position bias
|
||
ws = self.window_size
|
||
rpb = self.relative_position_bias_table[
|
||
self.relative_position_index.view(-1)
|
||
].view(ws[0] * ws[1], ws[0] * ws[1], -1)
|
||
attn = attn + rpb.permute(2, 0, 1).contiguous().unsqueeze(0)
|
||
|
||
attn = self.attn_drop(self.softmax(attn))
|
||
|
||
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
|
||
x = self.proj_drop(self.proj(x))
|
||
return x
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Transformer block
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class DiffiTBlock(nn.Module):
|
||
"""Single DiffiT transformer block: Norm → Attention → Norm → MLP."""
|
||
|
||
def __init__(
|
||
self,
|
||
hidden_size: int,
|
||
num_heads: int,
|
||
mlp_ratio: float = 4.0,
|
||
**block_kwargs,
|
||
):
|
||
super().__init__()
|
||
self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||
self.attn = DiffiTAttention(
|
||
dim=hidden_size,
|
||
temb_dim=hidden_size,
|
||
num_heads=num_heads,
|
||
qkv_bias=True,
|
||
**block_kwargs,
|
||
)
|
||
self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||
self.mlp = Mlp(
|
||
in_features=hidden_size,
|
||
hidden_features=int(hidden_size * mlp_ratio),
|
||
act_layer=lambda: nn.GELU(approximate="tanh"),
|
||
drop=0,
|
||
)
|
||
|
||
def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
|
||
x = x + self.attn(self.norm1(x), c)
|
||
x = x + self.mlp(self.norm2(x))
|
||
return x
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Final projection layer
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class FinalLayer(nn.Module):
|
||
"""Final projection: LayerNorm → SiLU → Linear → patch-level predictions."""
|
||
|
||
def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
|
||
super().__init__()
|
||
self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||
self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
|
||
self.silu = nn.SiLU()
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
return self.linear(self.silu(self.norm_final(x)))
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# DiffiT model
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# Latent-size threshold that separates the two CFG strategies.
|
||
# image_size 256 → latent_size 32 (power-cosine CFG)
|
||
# image_size 512 → latent_size 64 (linear CFG)
|
||
_LATENT_SIZE_THRESHOLD = 32
|
||
|
||
|
||
class DiffiT(nn.Module):
|
||
"""
|
||
DiffiT: Diffusion Vision Transformers for Image Generation.
|
||
|
||
A class-conditional latent diffusion model built from DiffiTBlocks
|
||
with window-based relative-position self-attention and sinusoidal
|
||
time/position embeddings.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
input_size: int = 32,
|
||
patch_size: int = 2,
|
||
in_channels: int = 4,
|
||
hidden_size: int = 1152,
|
||
depth: int = 30,
|
||
num_heads: int = 16,
|
||
mlp_ratio: float = 4.0,
|
||
mask_ratio=None,
|
||
decode_layer=None,
|
||
class_dropout_prob: float = 0.1,
|
||
num_classes: int = 1000,
|
||
learn_sigma: bool = True,
|
||
):
|
||
super().__init__()
|
||
self.input_size = input_size
|
||
self.learn_sigma = learn_sigma
|
||
self.in_channels = in_channels
|
||
self.out_channels = in_channels * 2 if learn_sigma else in_channels
|
||
self.patch_size = patch_size
|
||
self.num_heads = num_heads
|
||
|
||
# window_size must match the spatial grid so the relative
|
||
# position bias table covers every token pair.
|
||
window_size = input_size // patch_size # 256→16, 512→32
|
||
|
||
# Embedders -----------------------------------------------------------
|
||
self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True)
|
||
self.t_embedder = TimestepEmbedder(hidden_size)
|
||
self.y_embedder = LabelEmbedder(num_classes, hidden_size, class_dropout_prob)
|
||
|
||
num_patches = self.x_embedder.num_patches
|
||
self.pos_embed = nn.Parameter(
|
||
torch.zeros(1, num_patches, hidden_size), requires_grad=False
|
||
)
|
||
|
||
# Transformer backbone ------------------------------------------------
|
||
self.blocks = nn.ModuleList([
|
||
DiffiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, window_size=window_size)
|
||
for _ in range(depth)
|
||
])
|
||
self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels)
|
||
|
||
self._initialize_weights()
|
||
|
||
# ----- weight init -------------------------------------------------------
|
||
|
||
def _initialize_weights(self) -> None:
|
||
def _basic_init(module: nn.Module) -> None:
|
||
if isinstance(module, nn.Linear):
|
||
nn.init.xavier_uniform_(module.weight)
|
||
if module.bias is not None:
|
||
nn.init.constant_(module.bias, 0)
|
||
|
||
self.apply(_basic_init)
|
||
|
||
# Fixed sinusoidal position embedding
|
||
pos_embed = get_2d_sincos_pos_embed(
|
||
self.pos_embed.shape[-1], int(self.x_embedder.num_patches ** 0.5)
|
||
)
|
||
self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
|
||
|
||
# Patch embedding projection
|
||
w = self.x_embedder.proj.weight.data
|
||
nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
|
||
nn.init.constant_(self.x_embedder.proj.bias, 0)
|
||
|
||
# Label embedding
|
||
nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
|
||
|
||
# Timestep MLP
|
||
nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
|
||
nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
|
||
|
||
# Zero-init final projection so the model starts near identity
|
||
nn.init.constant_(self.final_layer.linear.weight, 0)
|
||
nn.init.constant_(self.final_layer.linear.bias, 0)
|
||
|
||
# ----- helpers -----------------------------------------------------------
|
||
|
||
def unpatchify(self, x: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
Rearrange patch tokens back into a spatial feature map.
|
||
|
||
Args:
|
||
x: ``(N, T, patch_size**2 * C)``
|
||
|
||
Returns:
|
||
``(N, C, H, W)``
|
||
"""
|
||
c = self.out_channels
|
||
p = self.x_embedder.patch_size[0]
|
||
h = w = int(x.shape[1] ** 0.5)
|
||
assert h * w == x.shape[1]
|
||
|
||
x = x.reshape(x.shape[0], h, w, p, p, c)
|
||
x = torch.einsum("nhwpqc->nchpwq", x)
|
||
return x.reshape(x.shape[0], c, h * p, h * p)
|
||
|
||
def convert_to_fp16(self) -> None:
|
||
"""Switch the full model to fp16 for inference."""
|
||
self.to(dtype=torch.float16)
|
||
|
||
def convert_to_fp32(self) -> None:
|
||
"""Switch the full model back to fp32 for inference."""
|
||
self.to(dtype=torch.float32)
|
||
|
||
# ----- forward -----------------------------------------------------------
|
||
|
||
def forward(
|
||
self,
|
||
x: torch.Tensor,
|
||
t: torch.Tensor,
|
||
y: torch.Tensor,
|
||
enable_mask: bool = False,
|
||
) -> torch.Tensor:
|
||
"""
|
||
Forward pass of DiffiT.
|
||
|
||
Args:
|
||
x: ``(N, C, H, W)`` spatial inputs (images or latent codes).
|
||
t: ``(N,)`` diffusion timesteps.
|
||
y: ``(N,)`` class labels.
|
||
enable_mask: Reserved for future use.
|
||
|
||
Returns:
|
||
``(N, out_C, H, W)`` noise (and optionally variance) prediction.
|
||
"""
|
||
x = self.x_embedder(x) + self.pos_embed
|
||
c = self.t_embedder(t) + self.y_embedder(y, self.training)
|
||
|
||
for block in self.blocks:
|
||
x = block(x, c)
|
||
|
||
x = self.final_layer(x)
|
||
return self.unpatchify(x)
|
||
|
||
def forward_with_cfg(
|
||
self,
|
||
x: torch.Tensor,
|
||
t: torch.Tensor,
|
||
y: torch.Tensor,
|
||
cfg_scale: Optional[float] = None,
|
||
diffusion_steps: int = 1000,
|
||
scale_pow: float = 4.0,
|
||
cfg_eps_channels: Optional[int] = None,
|
||
) -> torch.Tensor:
|
||
"""
|
||
Forward pass with (optional) classifier-free guidance.
|
||
|
||
When ``cfg_scale`` is given the batch is assumed to be
|
||
``[cond_half, uncond_half]`` and guidance is applied to the
|
||
predicted noise.
|
||
|
||
CFG strategy varies by resolution:
|
||
- **image_size 256** (``input_size <= 32``): power-cosine schedule
|
||
that ramps the effective scale over the diffusion process.
|
||
- **image_size 512** (``input_size > 32``): constant linear CFG
|
||
scaling.
|
||
"""
|
||
eps_channels = self.in_channels if cfg_eps_channels is None else int(cfg_eps_channels)
|
||
if eps_channels < 1 or eps_channels > self.out_channels:
|
||
raise ValueError(
|
||
f"cfg_eps_channels must be between 1 and {self.out_channels}. "
|
||
f"Received {eps_channels}."
|
||
)
|
||
|
||
if cfg_scale is None:
|
||
model_out = self.forward(x, t, y)
|
||
eps, rest = model_out[:, :eps_channels], model_out[:, eps_channels:]
|
||
return torch.cat([eps, rest], dim=1)
|
||
|
||
# Duplicate the conditional half for paired evaluation
|
||
half = x[: len(x) // 2]
|
||
combined = torch.cat([half, half], dim=0)
|
||
model_out = self.forward(combined, t, y)
|
||
|
||
eps, rest = model_out[:, :eps_channels], model_out[:, eps_channels:]
|
||
cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
|
||
|
||
if self.input_size <= _LATENT_SIZE_THRESHOLD:
|
||
# ----- image_size 256: power-cosine CFG schedule -----------------
|
||
phase = ((1 - t / diffusion_steps) ** scale_pow) * math.pi
|
||
scale_step = 0.5 * (1 - torch.cos(phase))
|
||
real_cfg_scale = (cfg_scale - 1) * scale_step + 1
|
||
real_cfg_scale = real_cfg_scale[: len(x) // 2].view(-1, 1, 1, 1)
|
||
else:
|
||
# ----- image_size 512: constant linear CFG -----------------------
|
||
real_cfg_scale = cfg_scale
|
||
|
||
half_eps = uncond_eps + real_cfg_scale * (cond_eps - uncond_eps)
|
||
eps = torch.cat([half_eps, half_eps], dim=0)
|
||
return torch.cat([eps, rest], dim=1)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Model constructor
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def Diffit(**kwargs):
|
||
"""DiffiT-XL/2 configuration (depth 28, dim 1152, patch 2, 16 heads)."""
|
||
return DiffiT(depth=28, hidden_size=1152, patch_size=2, num_heads=16, **kwargs)
|