Spaces:
Running on Zero
Running on Zero
| """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 | |
| from toolkit.models.v2._mixin import OstrisModelMixin | |
| 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, OstrisModelMixin): | |
| """The denoiser. Plain ``nn.Module`` plus ``OstrisModelMixin``. | |
| The mixin supplies the universal component API every v2 module shares: | |
| ``load()`` / ``load_model()`` / ``aitk_post_load()`` (quantize, layer | |
| offloading and device placement driven by the holder's model_config via | |
| ``BaseModel.component_load_kwargs(role)``), plus comfy-format save/load. | |
| Override its class hooks (``get_transformer_block_names``, | |
| ``get_quantization_exclude_modules``, ``get_offload_ignore_modules``, | |
| ``convert_state_dict_on_load/save``) as needed. | |
| 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 get_transformer_block_names(cls): | |
| # attribute name(s) of the repeated-block ModuleList(s); the quantizer | |
| # streams these blocks through the GPU one at a time | |
| return ["blocks"] | |
| 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 get_transformer_block_names() returns, which the LoRA code uses | |
| # for block targeting and the quantizer uses for block streaming. | |
| 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 | |
| def device(self): | |
| return next(self.parameters()).device | |
| 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 | |