File size: 7,627 Bytes
e545366 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | """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))
|