"""SpatialDiffuseSlot — Semanticist arch + tok_L init + OUR spatial conditioning. Design (2026-07-15, user-confirmed): - encoder : Semanticist ViT (vit_base_patch16, tok_L init, num_slots=256 kept so the ckpt loads 100%). We IGNORE its slots and take the 256 PATCH tokens (patches never attend slots -> patch path is clean). - pool : OUR AttnPool port (per-level learnable queries cross-attend the 16x16 patch grid) -> 85 multi-res tokens (8x8+4x4+2x2+1x1). - dit : Semanticist DiT-L (tok_L init; null_cond re-init at K=85) with OUR spatial-align mask applied inside its concat self-attention: latent->latent full | latent->cond xa-rule | cond->latent BLOCK | cond->cond identity (see spatial_mask.py, unit-tested) - warmup : freeze_dit=True freezes the pretrained DiT trunk (blocks/embedders/ final_layer) so encoder+pool+cond-embedder adapt first; resume with freeze_dit=False (+low lr) to open the trunk. Preserves tok_L init. """ import math import torch import torch.nn as nn import torch.nn.functional as F from semanticist.stage1.diffuse_slot import DiffuseSlot, DiT_with_autoenc_cond from semanticist.stage1 import vision_transformer import sys, os sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) from spatial_mask import build_concat_self_attn_mask def _sincos_2d(dim, s): """(s*s, dim) 2D sin-cos grid embedding (query localization head-start).""" yy, xx = torch.meshgrid(torch.arange(s), torch.arange(s), indexing="ij") grid = torch.stack([yy, xx], 0).float().reshape(2, -1) # (2, s*s) d4 = dim // 4 omega = 1.0 / (10000 ** (torch.arange(d4).float() / d4)) out = [] for g in grid: # y then x ang = g[:, None] * omega[None] out += [torch.sin(ang), torch.cos(ang)] emb = torch.cat(out, dim=1) # (s*s, 4*d4) if emb.shape[1] < dim: emb = torch.cat([emb, torch.zeros(emb.shape[0], dim - emb.shape[1])], 1) return emb class LevelNestedSampler(nn.Module): """Nested level-drop for our multi-res cond tokens (OUR method, port). Cond token order is FINE-first: [8x8(0..63), 4x4(64..79), 2x2(80..83), 1x1(84)]. We drop COARSE-to-FINE keep order: keep 1x1 first, drop 8x8 first (finest, and 8x8 is 1:1 with DiT image tokens so it's the pure per-patch detail level). Drop is WHOLE-LEVEL, coarse-first: keep_levels ~ uniform(1, num_levels) so each of {1x1}, {1x1,2x2}, {1x1,2x2,4x4}, {all} is 25% likely (num_levels=4). This trains every granularity equally (coarse-only levels are properly learned, unlike a token-proportional cut which almost never fully drops 8x8). 1x1(global) always kept -- full drop is CFG's job (separate uncond_drop). Returns (B, M) bool, True = KEEP. inference_with_n_slots: if given, interpreted as a TOKEN count and mapped to the largest whole keep_levels whose token total <= it (so trainer's test_num_slots=85 -> all levels; smaller -> coarse prefix). """ def __init__(self, level_sizes=(8, 4, 2, 1)): super().__init__() self.M = sum(s * s for s in level_sizes) self.num_levels = len(level_sizes) # coarse-first levels and cumulative token counts coarse = sorted(level_sizes) # [1,2,4,8] off, base = {}, 0 for s in level_sizes: # fine-first layout offsets off[s] = base; base += s * s # ranks: token -> which coarse-first LEVEL it belongs to (0=coarsest 1x1) ranks = torch.empty(self.M, dtype=torch.long) cum = [] acc = 0 for lvl, s in enumerate(coarse): for t in range(off[s], off[s] + s * s): ranks[t] = lvl acc += s * s; cum.append(acc) # cumulative tokens keeping lvl+1 levels self.register_buffer("level_of_token", ranks) # (M,) level index per token self.register_buffer("cum_tokens", torch.tensor(cum)) # [1,5,21,85] self.num_slots = self.M def forward(self, batch_size, device, inference_with_n_slots=-1): if self.training: keep_levels = torch.randint(1, self.num_levels + 1, (batch_size,), device=device) else: if inference_with_n_slots != -1: # map a token budget -> largest whole level-prefix that fits cum = self.cum_tokens.to(device) kl = int((cum <= inference_with_n_slots).sum().clamp(min=1)) keep_levels = torch.full((batch_size,), kl, device=device) else: keep_levels = torch.full((batch_size,), self.num_levels, device=device) # keep token if its coarse-first level index < keep_levels return self.level_of_token.to(device)[None, :] < keep_levels[:, None] # (B,M) True=keep class SpatialAttnPool(nn.Module): """Port of OUR AttnPool: per-level queries cross-attend shared patch K/V. Per level s in level_sizes: s*s learnable queries (2D sin-cos init). depth stacked (cross-attn + FFN) layers, K/V shared across levels/layers. Output: (B, sum(s*s), enc_d) in level order given (finest first to match our cond-token order: 8x8, 4x4, 2x2, 1x1 -> 85 tokens). """ def __init__(self, enc_d=768, level_sizes=(8, 4, 2, 1), num_heads=12, depth=2, use_ffn=True): super().__init__() self.level_sizes = tuple(level_sizes) self.enc_d, self.h = enc_d, num_heads self.hd = enc_d // num_heads self.depth, self.use_ffn = depth, use_ffn self.kv = nn.Linear(enc_d, 2 * enc_d) self.kv_norm = nn.LayerNorm(enc_d) self.queries = nn.ParameterList() self.q_projs = nn.ModuleList() self.out_projs = nn.ModuleList() self.q_norms = nn.ModuleList() self.ffns = nn.ModuleList() for s in self.level_sizes: q = nn.Parameter(_sincos_2d(enc_d, s) * 0.02) self.queries.append(q) self.q_projs.append(nn.ModuleList( [nn.Linear(enc_d, enc_d) for _ in range(depth)])) self.out_projs.append(nn.ModuleList( [nn.Linear(enc_d, enc_d) for _ in range(depth)])) self.q_norms.append(nn.ModuleList( [nn.LayerNorm(enc_d) for _ in range(depth)])) self.ffns.append(nn.ModuleList([ nn.Sequential(nn.LayerNorm(enc_d), nn.Linear(enc_d, 4 * enc_d), nn.GELU(), nn.Linear(4 * enc_d, enc_d)) if use_ffn else nn.Identity() for _ in range(depth)])) def forward(self, feats): # feats: (B, P, enc_d) patch tokens B, P, D = feats.shape kv = self.kv(self.kv_norm(feats)).reshape(B, P, 2, self.h, self.hd) k = kv[:, :, 0].transpose(1, 2) # (B,h,P,hd) v = kv[:, :, 1].transpose(1, 2) outs = [] for li, s in enumerate(self.level_sizes): q = self.queries[li].unsqueeze(0).expand(B, -1, -1) # (B,s*s,D) for d in range(self.depth): qh = self.q_projs[li][d](self.q_norms[li][d](q)) qh = qh.reshape(B, s * s, self.h, self.hd).transpose(1, 2) att = F.scaled_dot_product_attention(qh, k, v) # (B,h,s*s,hd) att = att.transpose(1, 2).reshape(B, s * s, D) q = q + self.out_projs[li][d](att) if self.use_ffn: q = q + self.ffns[li][d](q) outs.append(q) return torch.cat(outs, dim=1) # (B, 85, D) class DiTSpatial(DiT_with_autoenc_cond): """Semanticist DiT + our spatial mask in the concat self-attention. forward/forward_with_cfg mirror the parent but pass `mask` to every block.""" def __init__(self, *args, level_sizes=(8, 4, 2, 1), **kwargs): super().__init__(*args, **kwargs) G = self.x_embedder.grid_size[0] if hasattr(self.x_embedder, "grid_size") \ else int(self.x_embedder.num_patches ** 0.5) mask = build_concat_self_attn_mask(G, level_sizes) # (N+M, N+M) bool self.register_buffer("spatial_mask", mask, persistent=False) def forward(self, x, t, autoenc_cond, drop_mask=None): x = self.x_embedder(x) + self.pos_embed num_tokens = x.shape[1] c = self.t_embedder(t) autoenc = self.embed_cond(autoenc_cond, drop_mask) x = torch.cat((x, autoenc), dim=1) assert x.shape[1] == self.spatial_mask.shape[0], \ f"seq {x.shape[1]} != mask {self.spatial_mask.shape[0]}" for i, block in enumerate(self.blocks): if (i + 1) == self.encoder_depth and self.use_repa: projected = self.projector(x) self._repa_hook = projected[:, :num_tokens] x = block(x, c, self.spatial_mask) x = x[:, :num_tokens] x = self.final_layer(x, c) return self.unpatchify(x) def forward_with_cfg(self, x, t, autoenc_cond, drop_mask=None, y=None, cfg_scale=1.0): # parent's CFG wrapper calls self.forward -> mask applied automatically. half = x[: len(x) // 2] combined = torch.cat([half, half], dim=0) model_out = self.forward(combined, t, autoenc_cond, drop_mask) eps, rest = model_out[:, : self.in_channels], model_out[:, self.in_channels:] cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0) half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps) eps = torch.cat([half_eps, half_eps], dim=0) return torch.cat([eps, rest], dim=1) class SpatialDiffuseSlot(DiffuseSlot): """Semanticist DiffuseSlot with OUR spatial conditioning (85 multi-res tokens). Extra params (config): level_sizes : (8,4,2,1) -> 85 cond tokens pool_depth : AttnPool depth (default 2, our L_repa choice) freeze_dit : freeze pretrained DiT trunk (warmup phase 1) NOTE: pass num_slots=256 in config so the tok_L ENCODER loads 100%; the DiT is rebuilt with num_autoenc=85 (null_cond re-init is expected/normal). """ def __init__(self, *, level_sizes=(8, 4, 2, 1), pool_depth=2, freeze_dit=False, init_from=None, dit_lr_scale=None, dit_model="DiT-L-2", **kwargs): super().__init__(dit_model=dit_model, **kwargs) # Phase-2 knob: when the trunk is UNFROZEN, create_optimizer puts dit.* # (except cond-embedder/null_cond) into `lr_scale`-scaled param groups — # timm scheduler applies the scale to warmup AND cosine, so the # pretrained trunk moves ~dit_lr_scale× slower (protects tok_L knowledge). self.dit_lr_scale = dit_lr_scale self.level_sizes = tuple(level_sizes) self.num_cond_tokens = sum(s * s for s in self.level_sizes) # 85 # rebuild DiT with num_autoenc=85 + spatial mask (same size/class family) import semanticist.stage1.diffuse_slot as ds cfgmap = { # mirror DiT_with_autoenc_cond_*_2 constructors "DiT-L-2": dict(depth=24, hidden_size=1024, patch_size=2, num_heads=16), "DiT-XL-2": dict(depth=28, hidden_size=1152, patch_size=2, num_heads=16), "DiT-B-2": dict(depth=12, hidden_size=768, patch_size=2, num_heads=12), }[dit_model] old = self.dit self.dit = DiTSpatial( input_size=self.dit_input_size, in_channels=self.dit_in_channels, num_autoenc=self.num_cond_tokens, autoenc_dim=kwargs.get("slot_dim", 16), use_repa=self.use_repa, encoder_depth=old.encoder_depth if hasattr(old, "encoder_depth") else 8, level_sizes=self.level_sizes, **cfgmap, ) del old # our multi-res pool on top of the (frozen-init) Semanticist ViT patches enc_d = self.encoder.embed_dim if hasattr(self.encoder, "embed_dim") else 768 self.spatial_pool = SpatialAttnPool( enc_d=enc_d, level_sizes=self.level_sizes, depth=pool_depth) # tok_L init (encoder 100% + DiT trunk; null_cond K-mismatch handled). # NOTE: strict=False still ERRORS on shape mismatch -> pop mismatched keys. if init_from: # init_from can be: # (a) tok_L .pkl -> Phase-1 pretrained init # (b) an accelerate ckpt DIR (has model.safetensors) -> Phase-2: # load the whole SpatialDiffuseSlot weights (encoder+pool+DiT+cond) # from a prior frozen-run checkpoint, then continue with DiT unfrozen. if os.path.isdir(init_from): from safetensors.torch import load_file sf = os.path.join(init_from, "model.safetensors") ck = load_file(sf) phase = "PHASE-2 (resume weights from ckpt dir)" else: ck = torch.load(init_from, map_location="cpu") phase = "PHASE-1 (tok_L pretrained)" ck = {k.replace("._orig_mod", ""): v for k, v in ck.items()} own = self.state_dict() dropped = [k for k, v in list(ck.items()) if k in own and own[k].shape != v.shape] for k in dropped: ck.pop(k) ret = self.load_state_dict(ck, strict=False) enc_miss = [k for k in ret.missing_keys if k.startswith("encoder.")] dit_miss = [k for k in ret.missing_keys if k.startswith("dit.") and "null_cond" not in k and "pos_embed" not in k and "spatial_mask" not in k] pool_miss = [k for k in ret.missing_keys if k.startswith("spatial_pool.")] print(f"[SpatialDiffuseSlot] init_from={init_from} [{phase}]\n" f" encoder missing={len(enc_miss)} dit-trunk missing={len(dit_miss)} " f"pool missing={len(pool_miss)} shape-dropped={dropped}") assert not enc_miss and not dit_miss, "init incomplete!" # Phase-2 must also restore the trained attn-pool (Phase-1 doesn't have it) if os.path.isdir(init_from): assert not pool_miss, "Phase-2 resume: spatial_pool weights missing!" # BUG#8 FIX: cond length is 85, but DiffuseSlot kept num_slots=256 (for the # encoder ckpt). Trainer eval + nested_sampler size their drop_mask from # num_slots -> (256) vs cond (85) crash at eval. Re-point BOTH to 85. # (encoder keeps its own internal num_slots=256 attr -> ckpt still loads.) self.num_slots = self.num_cond_tokens # 85 (trainer reads this) # OUR level-drop (coarse-first keep, 8x8 dropped first, token-proportional). self.nested_sampler = LevelNestedSampler(self.level_sizes) self.freeze_dit = freeze_dit if freeze_dit: self._set_dit_trunk_grad(False) def _set_dit_trunk_grad(self, flag: bool): # trunk = pretrained parts (keep cond-embedder/null_cond trainable) for name, p in self.dit.named_parameters(): if name.startswith(("autoenc_cond_embedder", "null_cond")): p.requires_grad = True else: p.requires_grad = flag # ---- encoder path: Semanticist ViT patches -> our pool -> 85 tokens ---- def encode_patches(self, x): enc = self.encoder h = enc.prepare_tokens(x) # patches must not see slots (their causal mask already enforces this); # replicate the mask so patch features match tok_L's training regime. T = h.shape[1] attn_mask = torch.ones(T, T, device=h.device, dtype=torch.bool) ns = enc.num_slots causal = torch.ones(ns, ns, device=h.device, dtype=torch.bool).tril(0) attn_mask[-ns:, -ns:] = causal attn_mask[:-ns, -ns:] = False for blk in enc.blocks: h = blk(h, attn_mask=attn_mask) h = enc.norm(h) num_patches = T - 1 - ns return h[:, 1:1 + num_patches] # (B, 256, enc_d) drop cls/slots def encode_slots(self, x): feats = self.encode_patches(x) tokens = self.spatial_pool(feats) # (B, 85, enc_d) slots = self.encoder2slot(tokens) # reuse Linear(enc_d, slot_dim) if self.norm_slots: std = torch.std(slots, dim=-1, keepdim=True) mean = torch.mean(slots, dim=-1, keepdim=True) slots = (slots - mean) / (std + 1e-6) return slots