"""Anima style stream — decoupled cross-attention adapter (stage 2). Design (spec §6): the character stream is in-context latent concatenation (detail copying is the point); the *style* stream must NOT copy details, so it uses statistical transfer via embedding injection instead: encoder SigLIP 2 patch tokens from the last K hidden layers, aggregated with learnable softmax layer weights and projected to style tokens (AnimeAdapter-style). injection each DiT block's text cross-attention output gains a decoupled attention term: out = out_text + style_weight * gamma_i * Attn(Q, K_s, V_s) where Q is the block's frozen query, K_s/V_s are new trainable projections of the style tokens, and gamma_i is a per-block learnable gate initialized to 0 — the adapter is an exact no-op at init, so training starts from the base model's behaviour. This module is pure torch (no ComfyUI imports) so the training code can import the same definition. ComfyUI integration lives in style_nodes.py. """ import torch import torch.nn.functional as F from torch import nn # Anima 2B DiT geometry ANIMA_X_DIM = 2048 ANIMA_N_HEADS = 16 ANIMA_N_BLOCKS = 28 class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x): norm = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps) return (norm * self.weight.float()).type_as(x) class StyleTokenAggregator(nn.Module): """Aggregate SigLIP patch tokens from the last K hidden layers into style tokens: softmax-weighted layer mix -> LayerNorm -> projection.""" def __init__(self, siglip_dim=1152, style_dim=1024, n_layers=6): super().__init__() self.n_layers = n_layers self.layer_weights = nn.Parameter(torch.zeros(n_layers)) self.norm = nn.LayerNorm(siglip_dim) self.proj = nn.Linear(siglip_dim, style_dim) def forward(self, hidden_states_B_K_N_D): """hidden_states: (B, K, N, D) — the last K hidden layers of the vision tower (K == n_layers), N patch tokens of dim D.""" assert hidden_states_B_K_N_D.shape[1] == self.n_layers, ( f"expected {self.n_layers} layers, got {hidden_states_B_K_N_D.shape[1]}" ) w = torch.softmax(self.layer_weights, dim=0) mixed = (hidden_states_B_K_N_D * w[None, :, None, None].to(hidden_states_B_K_N_D)).sum(dim=1) return self.proj(self.norm(mixed)) # (B, N, style_dim) class StyleBlockKV(nn.Module): """Per-DiT-block decoupled K/V projections + gate.""" def __init__(self, x_dim=ANIMA_X_DIM, style_dim=1024, n_heads=ANIMA_N_HEADS): super().__init__() self.n_heads = n_heads self.head_dim = x_dim // n_heads self.k_proj = nn.Linear(style_dim, x_dim, bias=False) self.v_proj = nn.Linear(style_dim, x_dim, bias=False) # match the base attention's K normalization (RMSNorm per head) self.k_norm = RMSNorm(self.head_dim) # gate init 0 -> adapter is a no-op until trained self.gate = nn.Parameter(torch.zeros(1)) def kv(self, style_tokens_B_N_D): B, N, _ = style_tokens_B_N_D.shape k = self.k_proj(style_tokens_B_N_D).view(B, N, self.n_heads, self.head_dim) v = self.v_proj(style_tokens_B_N_D).view(B, N, self.n_heads, self.head_dim) k = self.k_norm(k) return k, v class AnimaStyleAdapter(nn.Module): """Aggregator + one StyleBlockKV per DiT block.""" def __init__(self, siglip_dim=1152, style_dim=1024, n_layers=6, x_dim=ANIMA_X_DIM, n_heads=ANIMA_N_HEADS, n_blocks=ANIMA_N_BLOCKS): super().__init__() self.config = { "siglip_dim": siglip_dim, "style_dim": style_dim, "n_layers": n_layers, "x_dim": x_dim, "n_heads": n_heads, "n_blocks": n_blocks, } self.aggregator = StyleTokenAggregator(siglip_dim, style_dim, n_layers) self.blocks = nn.ModuleList( [StyleBlockKV(x_dim, style_dim, n_heads) for _ in range(n_blocks)] ) @classmethod def from_state_dict(cls, sd): """Instantiate with dimensions inferred from a checkpoint.""" siglip_dim = sd["aggregator.proj.weight"].shape[1] style_dim = sd["aggregator.proj.weight"].shape[0] n_layers = sd["aggregator.layer_weights"].shape[0] x_dim = sd["blocks.0.k_proj.weight"].shape[0] n_blocks = 0 while f"blocks.{n_blocks}.k_proj.weight" in sd: n_blocks += 1 head_dim = sd["blocks.0.k_norm.weight"].shape[0] adapter = cls(siglip_dim, style_dim, n_layers, x_dim, x_dim // head_dim, n_blocks) adapter.load_state_dict(sd) return adapter class StyleState: """Runtime state shared by all patched cross-attention wrappers. Armed by the style diffusion wrapper before each forward.""" def __init__(self): self.active = False # style K/V per block, computed once per forward: list of (k, v) self.kv_per_block = None # per-sample multiplier (0 masks a sample, e.g. the uncond chunk) self.sample_scale_B = None self.weight = 1.0 def style_attention(q_B_S_H_D, k_B_N_H_D, v_B_N_H_D): """Decoupled attention with the block's frozen query. Returns (B, S, H*D) to match the text branch pre-output_proj layout.""" q = q_B_S_H_D.transpose(1, 2) k = k_B_N_H_D.transpose(1, 2) v = v_B_N_H_D.transpose(1, 2) out = F.scaled_dot_product_attention(q, k, v) return out.transpose(1, 2).reshape(q_B_S_H_D.shape[0], q_B_S_H_D.shape[1], -1) class StyleCrossAttention(nn.Module): """Replacement for a DiT block's cross_attn module. Reimplements the original attention using the original (frozen) projections, then adds the decoupled style term *before* output_proj — the IP-Adapter formulation: out = output_proj( Attn(Q, K_text, V_text) + weight * gamma * Attn(Q, K_style, V_style) ) When the state is inactive it computes exactly the original attention (the reimplementation is numerically identical: same modules, same attn_op). """ def __init__(self, orig_attn, block_kv, state, block_index): super().__init__() self.orig = orig_attn self.block_kv = block_kv self.state = state self.block_index = block_index def forward(self, x, context=None, rope_emb=None, transformer_options={}): state = self.state if not state.active or state.kv_per_block is None: return self.orig(x, context, rope_emb=rope_emb, transformer_options=transformer_options) q, k, v = self.orig.compute_qkv(x, context, rope_emb=rope_emb) text_out = self.orig.attn_op(q, k, v, transformer_options=transformer_options) ks, vs = state.kv_per_block[self.block_index] B = q.shape[0] if ks.shape[0] != B: # single style batch broadcast over CFG batch ks = ks.expand(B, -1, -1, -1) vs = vs.expand(B, -1, -1, -1) style_out = style_attention(q, ks.to(q.dtype), vs.to(q.dtype)) gate = self.block_kv.gate.to(q.dtype) scale = state.weight * gate if state.sample_scale_B is not None: scale = scale * state.sample_scale_B.to(q.dtype).view(B, 1, 1) text_out = text_out + scale * style_out return self.orig.output_dropout(self.orig.output_proj(text_out))