"""Backbone = diffusers SD3 MMDiT (SD3Transformer2DModel). We deliberately do NOT hand-roll the transformer: diffusers' implementation is more optimized (fused / SDPA = FlashAttention, gradient checkpointing, torch.compile-friendly) and battle-tested. We only wrap it to map our conditioning: - hidden_states = noised line latent [B, C, h, w] - encoder_hidden_states = content + style tokens [B, L+K, context_dim] (joint attention) - pooled_projections = pooled style vector [B, pooled_dim] (AdaLN) - timestep = flow-matching timestep [B] Output: predicted velocity [B, C, h, w]. """ from __future__ import annotations import torch import torch.nn as nn from diffusers import SD3Transformer2DModel from diffusers.models.embeddings import PatchEmbed from ..config import BackboneConfig from .rope import RoPEJointAttnProcessor, build_2d_rope class Backbone(nn.Module): """SD3 Transformer2D wrapper mapping (latent, timestep, content+style tokens) -> velocity. With ``cfg.rope`` (default) the SD3 absolute ``pos_embed`` table is replaced by a position-free ``PatchEmbed`` and a custom 2D-RoPE joint-attention processor; per-call rotary freqs are passed via ``joint_attention_kwargs`` (grad-checkpoint-safe — SD3 threads that dict through every block and through its gradient-checkpointing wrapper). With ``cfg.rope=False`` the stock SD3 pos_embed is kept. """ def __init__( self, cfg: BackboneConfig, in_channels: int, context_dim: int, pooled_dim: int, out_channels: int | None = None, ) -> None: super().__init__() # in_channels may exceed out_channels when a glyph latent is channel-concatenated onto the input # (glyph_concat): the model READS extra channels but still predicts only the latent_channels velocity. out_channels = out_channels if out_channels is not None else in_channels self.transformer = SD3Transformer2DModel( sample_size=cfg.sample_size, patch_size=cfg.patch, in_channels=in_channels, out_channels=out_channels, num_layers=cfg.num_layers, attention_head_dim=cfg.dim // cfg.heads, num_attention_heads=cfg.heads, joint_attention_dim=context_dim, caption_projection_dim=cfg.dim, pooled_projection_dim=pooled_dim, pos_embed_max_size=cfg.pos_embed_max_size, qk_norm=cfg.qk_norm, # SD3.5: RMSNorm on Q/K (training stability) dual_attention_layers=cfg.dual_attention_layers, # SD3.5-medium / MMDiT-X (() = off) ) self.rope = cfg.rope self.patch = cfg.patch self.head_dim = cfg.dim // cfg.heads self.theta = cfg.rope_theta if self.rope: # Swap in a position-free PatchEmbed: drops the absolute sincos table AND the # pos_embed_max_size width cap, so variable-width lines patch-embed without a hard limit. self.transformer.pos_embed = PatchEmbed( height=cfg.sample_size, width=cfg.sample_size, patch_size=cfg.patch, in_channels=in_channels, embed_dim=cfg.dim, pos_embed_type=None, ) self.transformer.set_attn_processor(RoPEJointAttnProcessor()) # Memoized 2D-RoPE freqs keyed by (h_tokens, w_tokens, device, dtype) — see forward(). self._rope_cache: dict[ tuple[int, int, torch.device, torch.dtype], tuple[torch.Tensor, torch.Tensor] ] = {} def _rope( self, h_t: int, w_t: int, device: torch.device, dtype: torch.dtype ) -> tuple[torch.Tensor, torch.Tensor]: """Memoized 2D-RoPE (cos, sin) for an ``h_t x w_t`` grid (constants, so caching is bit-identical).""" key = (h_t, w_t, device, dtype) cached = self._rope_cache.get(key) if cached is None: cos, sin = build_2d_rope(h_t, w_t, self.head_dim, device, self.theta) cached = (cos.to(dtype), sin.to(dtype)) self._rope_cache[key] = cached return cached def enable_optimizations(self) -> None: """Enable diffusers' gradient checkpointing (non-reentrant). Non-reentrant checkpointing keeps the checkpointed blocks' outputs attached to the autograd graph during the forward pass, so REPA's forward-hook capture of a mid-block hidden state can still backprop into the backbone (reentrant checkpointing detaches it -> REPA would train only its projection head). diffusers 0.38 defaults to non-reentrant for torch>=1.11, so the plain call already gives us what REPA needs. (The previous ``gradient_checkpointing_kwargs=`` form was never a diffusers arg — it is a HF Transformers arg — so that try-branch always raised.) """ self.transformer.enable_gradient_checkpointing() def compile_blocks(self) -> None: """Regionally ``torch.compile`` the repeated ``JointTransformerBlock`` (shared by train + inference). Compiles ONE block (not all 24): cold start is cheap and per-step throughput matches a full compile. The change is IN-PLACE (no ``_orig_mod.`` state-dict prefix, no ``OptimizedModule`` proxy) so a REPA forward-hook on an inner block still fires with grad intact. ``dynamic=True`` keeps the variable line width on one symbolic shape — pair with width bucketing (train) / one width per call (inference) so only a handful of shapes occur and recompiles stay bounded. Measured: −70% kernel launches, ~1.9× train throughput (docs/PERF_AUDIT.md §0e). """ import torch._dynamo import torch.fx.experimental._config as fx_config # use_duck_shape=False keeps the line width a free symbol, so a width that coincidentally equals # batch/heads/h_t doesn't trigger a recompile. Process-global; set once per entrypoint. fx_config.use_duck_shape = False # Variable widths + the last block's context_pre_only=True guard trigger several recompiles; the # default recompile_limit (8) is too low and Dynamo ABORTS. Raise it (bucketing keeps the real count # to ~num_buckets; dynamic=True keeps most widths on one symbolic shape). torch._dynamo.config.recompile_limit = 256 torch._dynamo.config.accumulated_recompile_limit = 2048 # SD3 ships `_repeated_blocks` empty, so populate it or compile_repeated_blocks() is a no-op. self.transformer._repeated_blocks = ["JointTransformerBlock"] # fullgraph=True so a graph break in the custom RoPE processor (apply_rotary_emb on a slice + cat, # qk RMSNorm, SDPA) fails LOUDLY at compile time instead of silently un-fusing and erasing the win. self.transformer.compile_repeated_blocks(dynamic=True, fullgraph=True) def forward( self, latent: torch.Tensor, timestep: torch.Tensor, context_tokens: torch.Tensor, pooled: torch.Tensor, n_content: int | None = None, ) -> torch.Tensor: joint_attention_kwargs: dict[str, object] | None = None if self.rope: # The 2D-RoPE (cos, sin) are deterministic constants of the token grid (arange-derived, # non-differentiable), so memoize them per (grid, device, dtype): under width-bucketing only a # handful of grids occur, so this rebuilds arange/get_1d_rotary/cat at most once per bucket # instead of every step. Bit-identical to rebuilding. (apply_rotary_emb upcasts internally; the # dtype cast is just a clean device/dtype match, cached alongside.) h_t = latent.shape[-2] // self.patch w_t = latent.shape[-1] // self.patch joint_attention_kwargs = {"image_rotary_emb": self._rope(h_t, w_t, latent.device, latent.dtype)} if n_content: # Shared-column RoPE for the line-glyph content tokens: 1 row × n_content columns, so # content token j carries the SAME column frequency as image patches in column j (MSRoPE-style). joint_attention_kwargs["content_rotary_emb"] = self._rope(1, n_content, latent.device, latent.dtype) joint_attention_kwargs["n_content"] = n_content return self.transformer( hidden_states=latent, timestep=timestep, encoder_hidden_states=context_tokens, pooled_projections=pooled, joint_attention_kwargs=joint_attention_kwargs, return_dict=False, )[0]