275 lines
12 KiB
Python
275 lines
12 KiB
Python
"""A minimal diffusion transformer (DiT) used by the example model extension.
|
|
|
|
This file stands in for the situation where diffusers does NOT have your model.
|
|
You vendor the architecture yourself inside your extension's ``src/`` folder and
|
|
load the weights manually in your model class (see ``../example_model.py``).
|
|
|
|
The architecture here is intentionally tiny and boring:
|
|
|
|
latents (B, C, h, w)
|
|
-> patchify with a strided conv (B, N_img, hidden)
|
|
text embeds (B, L, text_dim)
|
|
-> linear projection (B, L, hidden)
|
|
concat [text | image] into one joint sequence (B, L + N_img, hidden)
|
|
-> N transformer blocks (self attention + mlp, adaLN-zero
|
|
modulated by the timestep embedding)
|
|
-> final modulated norm + linear
|
|
take only the image tokens and unpatchify back to (B, C, h, w)
|
|
|
|
Real models add RoPE position embeddings, fancier attention, guidance
|
|
embeddings, etc. For real-world reference implementations in this repo see:
|
|
- ../../chroma/src/model.py (flux-style double/single stream blocks)
|
|
- ../../ernie_image/transformer.py (diffusers ModelMixin based)
|
|
- ../../ideogram4/src/transformer.py (packed single-sequence model)
|
|
|
|
GRADIENT CHECKPOINTING
|
|
======================
|
|
ai-toolkit enables gradient checkpointing on your model from
|
|
``jobs/process/BaseSDTrainProcess.py`` which does, in order of preference:
|
|
|
|
if hasattr(unet, 'enable_gradient_checkpointing'):
|
|
unet.enable_gradient_checkpointing()
|
|
elif hasattr(unet, 'gradient_checkpointing'):
|
|
unet.gradient_checkpointing = True
|
|
|
|
So a custom model only needs:
|
|
1. a ``self.gradient_checkpointing`` flag (default False)
|
|
2. (optionally) an ``enable_gradient_checkpointing()`` method
|
|
3. to wrap each transformer block call in ``torch.utils.checkpoint.checkpoint``
|
|
when the flag is set AND grads are enabled.
|
|
|
|
IMPORTANT: gate on ``torch.is_grad_enabled()``, NOT on ``self.training``.
|
|
Sampling runs under ``torch.no_grad()`` where checkpointing is pure overhead,
|
|
and some training setups (e.g. certain adapters) run the module in eval mode
|
|
while still needing gradients. ``torch.is_grad_enabled()`` handles both.
|
|
"""
|
|
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch import nn
|
|
from torch.utils.checkpoint import checkpoint
|
|
|
|
|
|
def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
|
|
"""Standard sinusoidal embedding.
|
|
|
|
in: t (B,) float tensor, the flow-matching time in [0, 1] (1 = pure noise)
|
|
out: emb (B, dim)
|
|
|
|
We scale t by 1000 before embedding so the sinusoids get a useful range,
|
|
the same trick flux and friends use.
|
|
"""
|
|
t = t.float() * 1000.0
|
|
half = dim // 2
|
|
freqs = torch.exp(
|
|
-math.log(max_period) * torch.arange(half, dtype=torch.float32, device=t.device) / half
|
|
)
|
|
args = t[:, None] * freqs[None]
|
|
return torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
|
|
|
|
|
class ExampleTransformerBlock(nn.Module):
|
|
"""One DiT block: adaLN-zero modulated self-attention + MLP.
|
|
|
|
in: x (B, S, hidden) the joint [text | image] token sequence
|
|
temb (B, hidden) the timestep embedding
|
|
attn_mask (B, 1, 1, S) bool, True = attend, False = padding
|
|
out: x (B, S, hidden)
|
|
"""
|
|
|
|
def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0):
|
|
super().__init__()
|
|
self.num_heads = num_heads
|
|
self.head_dim = hidden_size // num_heads
|
|
|
|
self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
|
self.qkv = nn.Linear(hidden_size, hidden_size * 3)
|
|
self.proj = nn.Linear(hidden_size, hidden_size)
|
|
|
|
self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
|
mlp_hidden = int(hidden_size * mlp_ratio)
|
|
self.mlp = nn.Sequential(
|
|
nn.Linear(hidden_size, mlp_hidden),
|
|
nn.GELU(approximate="tanh"),
|
|
nn.Linear(mlp_hidden, hidden_size),
|
|
)
|
|
|
|
# adaLN-zero: timestep embedding -> shift/scale/gate for attn and mlp.
|
|
# Zero-init so the block starts as identity (standard DiT trick).
|
|
self.adaLN_modulation = nn.Sequential(
|
|
nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size)
|
|
)
|
|
nn.init.zeros_(self.adaLN_modulation[-1].weight)
|
|
nn.init.zeros_(self.adaLN_modulation[-1].bias)
|
|
|
|
def forward(self, x: torch.Tensor, temb: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor:
|
|
b, s, d = x.shape
|
|
shift_a, scale_a, gate_a, shift_m, scale_m, gate_m = (
|
|
self.adaLN_modulation(temb).unsqueeze(1).chunk(6, dim=-1)
|
|
) # each (B, 1, hidden), broadcasts over the sequence
|
|
|
|
# --- attention ---
|
|
# ALWAYS default to torch's built-in scaled_dot_product_attention so the
|
|
# model runs with no extra dependency. If the reference repo you are
|
|
# porting hard-codes flash-attn (or xformers, sage, ...), do NOT carry
|
|
# that requirement over -- make it OPTIONAL. The clean pattern is a
|
|
# per-module ``attention_backend`` flag toggled in bulk from the parent
|
|
# model (e.g. ``set_attention_backend("flash")``), branching to the
|
|
# flash kernel only when explicitly selected AND the package is present.
|
|
# See ../../ideogram4/src/transformer.py and ../../boogu_image/src for
|
|
# working "native" (SDPA) + optional "flash" implementations.
|
|
h = self.norm1(x) * (1 + scale_a) + shift_a
|
|
q, k, v = self.qkv(h).chunk(3, dim=-1)
|
|
q = q.view(b, s, self.num_heads, self.head_dim).transpose(1, 2)
|
|
k = k.view(b, s, self.num_heads, self.head_dim).transpose(1, 2)
|
|
v = v.view(b, s, self.num_heads, self.head_dim).transpose(1, 2)
|
|
h = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
|
|
h = h.transpose(1, 2).reshape(b, s, d)
|
|
x = x + gate_a * self.proj(h)
|
|
|
|
# --- mlp ---
|
|
h = self.norm2(x) * (1 + scale_m) + shift_m
|
|
x = x + gate_m * self.mlp(h)
|
|
return x
|
|
|
|
|
|
class ExampleTransformer2DModel(nn.Module):
|
|
"""The denoiser. Plain ``nn.Module`` on purpose.
|
|
|
|
You could also subclass ``diffusers.ModelMixin``/``ConfigMixin`` (see
|
|
../../ernie_image/transformer.py) to get ``save_pretrained``,
|
|
``_gradient_checkpointing_func`` etc. for free, but a plain module shows
|
|
exactly what ai-toolkit actually requires, which is very little:
|
|
|
|
- a forward pass
|
|
- ``device`` / ``dtype`` properties (BaseModel reads ``self.model.device``
|
|
and ``self.model.dtype`` in a few places, e.g. save_device_state)
|
|
- the gradient checkpointing flag described in the module docstring
|
|
|
|
NOTE: the class NAME matters. ``ExampleModel.target_lora_modules`` lists
|
|
"ExampleTransformer2DModel" -- that string is matched against module class
|
|
names when deciding where to attach LoRA layers.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int = 16, # VAE latent channels
|
|
out_channels: int = 16, # predicted velocity has the same channels
|
|
patch_size: int = 2, # latent pixels per token side
|
|
hidden_size: int = 1024,
|
|
num_heads: int = 16,
|
|
num_layers: int = 12,
|
|
text_dim: int = 2048, # width of the text encoder hidden states
|
|
):
|
|
super().__init__()
|
|
self.in_channels = in_channels
|
|
self.out_channels = out_channels
|
|
self.patch_size = patch_size
|
|
self.hidden_size = hidden_size
|
|
|
|
# latent (B, C, h, w) -> image tokens (B, N_img, hidden)
|
|
self.x_embedder = nn.Conv2d(
|
|
in_channels, hidden_size, kernel_size=patch_size, stride=patch_size
|
|
)
|
|
# text encoder hidden states -> model width
|
|
self.text_proj = nn.Linear(text_dim, hidden_size)
|
|
# sinusoidal timestep embedding -> mlp
|
|
self.t_embedder = nn.Sequential(
|
|
nn.Linear(hidden_size, hidden_size),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_size, hidden_size),
|
|
)
|
|
|
|
# ``blocks`` is the repeated-layer ModuleList. The attribute name is
|
|
# what ExampleModel.get_transformer_block_names() returns, which the
|
|
# LoRA code uses for block targeting / "transformer only" training.
|
|
self.blocks = nn.ModuleList(
|
|
[
|
|
ExampleTransformerBlock(hidden_size, num_heads)
|
|
for _ in range(num_layers)
|
|
]
|
|
)
|
|
|
|
# final adaLN + projection back to patch pixels, zero-init so the
|
|
# untrained model predicts zeros.
|
|
self.norm_out = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
|
self.adaLN_out = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size))
|
|
self.proj_out = nn.Linear(hidden_size, patch_size * patch_size * out_channels)
|
|
nn.init.zeros_(self.adaLN_out[-1].weight)
|
|
nn.init.zeros_(self.adaLN_out[-1].bias)
|
|
nn.init.zeros_(self.proj_out.weight)
|
|
nn.init.zeros_(self.proj_out.bias)
|
|
|
|
# gradient checkpointing flag, flipped on by the trainer (see module
|
|
# docstring). Off by default so inference pays no cost.
|
|
self.gradient_checkpointing = False
|
|
|
|
# the trainer prefers this method if it exists
|
|
def enable_gradient_checkpointing(self, enable: bool = True):
|
|
self.gradient_checkpointing = enable
|
|
|
|
def disable_gradient_checkpointing(self):
|
|
self.gradient_checkpointing = False
|
|
|
|
@property
|
|
def device(self):
|
|
return next(self.parameters()).device
|
|
|
|
@property
|
|
def dtype(self):
|
|
return next(self.parameters()).dtype
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor, # (B, C, h, w) noisy latents
|
|
timestep: torch.Tensor, # (B,) flow time in [0, 1], 1 = pure noise
|
|
encoder_hidden_states: torch.Tensor, # (B, L, text_dim) padded text features
|
|
attention_mask: torch.Tensor, # (B, L) 1 = real text token, 0 = padding
|
|
) -> torch.Tensor:
|
|
"""Predict the flow-matching velocity.
|
|
|
|
out: (B, C, h, w) velocity in the ai-toolkit convention
|
|
(noise - clean), matching ExampleModel.get_loss_target().
|
|
"""
|
|
b, c, h, w = hidden_states.shape
|
|
p = self.patch_size
|
|
gh, gw = h // p, w // p
|
|
n_img = gh * gw
|
|
|
|
# tokens
|
|
img = self.x_embedder(hidden_states) # (B, hidden, gh, gw)
|
|
img = img.flatten(2).transpose(1, 2) # (B, N_img, hidden)
|
|
txt = self.text_proj(encoder_hidden_states) # (B, L, hidden)
|
|
x = torch.cat([txt, img], dim=1) # (B, L + N_img, hidden)
|
|
|
|
# timestep conditioning
|
|
temb = self.t_embedder(timestep_embedding(timestep, self.hidden_size))
|
|
temb = temb.to(x.dtype)
|
|
|
|
# joint attention mask: text padding is masked out, image tokens and
|
|
# real text tokens attend everywhere. (B, 1, 1, S) bool for sdpa.
|
|
img_mask = torch.ones(b, n_img, dtype=torch.bool, device=x.device)
|
|
attn_mask = torch.cat([attention_mask.bool(), img_mask], dim=1)
|
|
attn_mask = attn_mask[:, None, None, :]
|
|
|
|
for block in self.blocks:
|
|
if torch.is_grad_enabled() and self.gradient_checkpointing:
|
|
# Recompute this block's activations during backward instead
|
|
# of storing them -- trades compute for a big VRAM saving.
|
|
# use_reentrant=False is the modern, correct variant.
|
|
x = checkpoint(block, x, temb, attn_mask, use_reentrant=False)
|
|
else:
|
|
x = block(x, temb, attn_mask)
|
|
|
|
# final modulation + project, keep only the image tokens
|
|
shift, scale = self.adaLN_out(temb).unsqueeze(1).chunk(2, dim=-1)
|
|
x = self.norm_out(x) * (1 + scale) + shift
|
|
x = self.proj_out(x)[:, -n_img:] # (B, N_img, p*p*C)
|
|
|
|
# unpatchify back to the latent layout
|
|
x = x.view(b, gh, gw, p, p, self.out_channels)
|
|
x = x.permute(0, 5, 1, 3, 2, 4).reshape(b, self.out_channels, h, w)
|
|
return x
|