comfyui-anima-incontext/style_adapter.py

"""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))