""" Anima In-Context Reference — core logic. Strategy -------- Anima's DiT (Cosmos-Predict2 MiniTrainDIT) is a *video* architecture: latents flow through the blocks as (B, T, H, W, D) and self-attention is computed over the flattened (t h w) sequence with 3D RoPE (max_frames=128, patch_temporal=1). We exploit this: the reference image latent is concatenated as an extra *frame* along the T axis. This gives us, for free: * a distinct temporal RoPE coordinate for reference tokens (no spatial position collision with the generated frame), * per-frame timestep conditioning — MiniTrainDIT accepts timesteps of shape (B, T), so the reference frame can be conditioned at t=0 (clean image) while the generated frame follows the sampler's sigma. This matches the OminiControl-style "clean condition token" recipe. The generated frame's self-attention can then attend to reference tokens (shared attention / in-context conditioning). Reference frames are sliced off the output before returning to the sampler. Strength control is implemented by patching each block's `self_attn.attn_op` (a plain attribute, cleanly replaceable via ModelPatcher.add_object_patch) with a version that adds a per-sample additive bias on reference-token key columns: * log(strength) amplifies/attenuates reference attention, * cond_only masks reference keys for the uncond half of the CFG batch (equivalent to not concatenating the reference at all for the uncond forward — this matches the training contract, where the reference is dropped ~10% of the time to form the ref-free distribution). NOTE: the base model was finetuned as a T2I model (single frame), so zero-shot behaviour without a trained in-context LoRA is expected to be weak. This module defines the exact inference-time contract that the LoRA training code must replicate: - reference frames appended after generated frames on the T axis - reference frames receive timestep 0 - reference latents are latent_format-normalized (process_latent_in) - text conditioning unchanged """ import math import torch import torch.nn.functional as F import comfy.patcher_extension from comfy.patcher_extension import WrappersMP WRAPPER_KEY = "anima_incontext_ref" NEG_BIAS = -1e9 # finite mask value; softmax subtracts the row max so this is NaN-safe class RefState: """Mutable state shared between the diffusion-model wrapper and the patched attention ops. The wrapper fills in per-forward token counts and the per-sample reference bias (they depend on resolution and on the CFG batch layout), the attention ops read them.""" def __init__(self): self.active = False self.total_tokens = -1 self.gen_tokens = -1 # per-sample additive bias on reference key columns, shape (B,). # None means "all zero" (neutral -> attn ops fall back). self.bias_B = None # lazily-built full bias tensor (B, 1, 1, S), cached across the # 28 blocks of one forward pass self._bias_cache = None def bias_for(self, device, dtype): if self._bias_cache is None or self._bias_cache.device != device or self._bias_cache.dtype != dtype: b = torch.zeros((self.bias_B.shape[0], 1, 1, self.total_tokens), device=device, dtype=dtype) b[:, 0, 0, self.gen_tokens:] = self.bias_B.to(device=device, dtype=dtype).unsqueeze(1) self._bias_cache = b return self._bias_cache def _tokens_per_frame(h, w, patch_spatial): # pad_to_patch_size pads H and W up to a multiple of patch_spatial hp = math.ceil(h / patch_spatial) wp = math.ceil(w / patch_spatial) return hp * wp def _fit_latent(r, H, W, mode): """Fit reference latent frames (N, C, h, w) to the generation latent size (H, W). stretch: plain bilinear resize (aspect distortion) pad: aspect-preserving resize + edge-replicate center pad. Replicate keeps a white-background reference white at the borders instead of introducing a mean-gray frame. crop: aspect-filling resize + center crop """ h, w = r.shape[-2:] if (h, w) == (H, W): return r if mode == "stretch": return F.interpolate(r, size=(H, W), mode="bilinear", align_corners=False) if mode == "pad": scale = min(H / h, W / w) nh = max(1, min(H, round(h * scale))) nw = max(1, min(W, round(w * scale))) r = F.interpolate(r, size=(nh, nw), mode="bilinear", align_corners=False) pt = (H - nh) // 2 pl = (W - nw) // 2 return F.pad(r, (pl, W - nw - pl, pt, H - nh - pt), mode="replicate") if mode == "crop": scale = max(H / h, W / w) nh = max(H, round(h * scale)) nw = max(W, round(w * scale)) r = F.interpolate(r, size=(nh, nw), mode="bilinear", align_corners=False) ot = (nh - H) // 2 ol = (nw - W) // 2 return r[:, :, ot:ot + H, ol:ol + W] raise ValueError(f"unknown fit mode: {mode}") def make_ref_attn_op(state, fallback_op): """Replacement for Attention.attn_op on self-attention modules. Adds the per-sample reference bias to attention logits for reference-token key columns. Falls back to the original op whenever the reference is not active, the bias is neutral, or the sequence length does not match the full (gen + ref) self-attention sequence (which excludes cross-attention and the LLMAdapter, whose K length differs). """ def ref_attn_op(q_B_S_H_D, k_B_S_H_D, v_B_S_H_D, transformer_options={}): if ( not state.active or state.bias_B is None or k_B_S_H_D.shape[1] != state.total_tokens or q_B_S_H_D.shape[1] != state.total_tokens ): return fallback_op(q_B_S_H_D, k_B_S_H_D, v_B_S_H_D, transformer_options=transformer_options) # (B, S, H, D) -> (B, H, S, D) q = q_B_S_H_D.transpose(1, 2) k = k_B_S_H_D.transpose(1, 2) v = v_B_S_H_D.transpose(1, 2) bias = state.bias_for(q.device, q.dtype) out = F.scaled_dot_product_attention(q, k, v, attn_mask=bias) # (B, H, S, D) -> (B, S, H*D) out = out.transpose(1, 2).reshape(q_B_S_H_D.shape[0], q_B_S_H_D.shape[1], -1) return out return ref_attn_op def _per_sample_bias(B, strength, cond_only, cond_or_uncond, device): """Build the per-sample reference-key bias, shape (B,). cond samples get log(strength) (0 at strength=1); uncond samples get the same unless cond_only, in which case their reference keys are masked out entirely. strength <= 0 masks the reference everywhere. Returns None when every entry is zero (neutral -> no attn patch). """ if strength <= 0.0: base = NEG_BIAS else: base = math.log(strength) bias = torch.full((B,), base, device=device, dtype=torch.float32) if cond_only and cond_or_uncond is not None and len(cond_or_uncond) > 0 and B % len(cond_or_uncond) == 0: # calc_cond_batch concatenates equal-sized chunks along B, one # per entry of cond_or_uncond (0 = cond, 1 = uncond). chunk = B // len(cond_or_uncond) for i, kind in enumerate(cond_or_uncond): if kind == 1: bias[i * chunk:(i + 1) * chunk] = NEG_BIAS if torch.count_nonzero(bias) == 0: return None return bias def make_diffusion_wrapper(opts): """DIFFUSION_MODEL wrapper around MiniTrainDIT._forward. opts is a dict with: ref_latent: (N, C, 1, H, W) latent tensor, already latent_format-normalized (process_latent_in) state: RefState instance shared with the attn ops strength: attention bias multiplier for reference tokens cond_only: mask reference keys for the uncond CFG half fit_mode: stretch | pad | crop (reference latent resize) sigma_start: apply when current sigma <= sigma_start sigma_end: ... and sigma >= sigma_end patch_spatial: DiT spatial patch size (2 for Anima) ref_timestep: timestep value for reference frames (default 0.0) """ def wrapper(executor, x, timesteps, context, fps=None, padding_mask=None, **kwargs): state = opts["state"] ref = opts["ref_latent"] to = kwargs.get("transformer_options", {}) # strength <= 0 fully masks the reference for every sample, which # is mathematically identical to not attaching it — skip the # concat so the output is bit-exact with the reference-free # forward (and faster). Verified on-device: with the concat, the # masked-attention kernel差 compounds over sampling steps. if opts["strength"] <= 0.0: return executor(x, timesteps, context, fps, padding_mask, **kwargs) # ---- sigma window gating ---- sigmas = to.get("sigmas", None) if sigmas is not None: s = float(sigmas.max()) if s > opts["sigma_start"] or s < opts["sigma_end"]: return executor(x, timesteps, context, fps, padding_mask, **kwargs) squeeze_t = False if x.ndim == 4: # (B, C, H, W) -> (B, C, 1, H, W) x = x.unsqueeze(2) squeeze_t = True B, C, T, H, W = x.shape n_ref = ref.shape[0] # ---- prepare reference frames ---- r = ref.to(device=x.device, dtype=x.dtype) # (N, C, 1, H', W') r = r.squeeze(2) # (N, C, H', W') r = _fit_latent(r, H, W, opts.get("fit_mode", "pad")) # (N, C, H, W) -> (1, C, N, H, W) -> (B, C, N, H, W) r = r.permute(1, 0, 2, 3).unsqueeze(0).expand(B, -1, -1, -1, -1) x_cat = torch.cat([x, r], dim=2) # (B, C, T + N, H, W) # ---- per-frame timesteps: generated frames keep the sampler's t, # reference frames get ref_timestep (0 = clean) ---- t = timesteps if t.ndim == 1: t = t.unsqueeze(1) # (B, 1) t = t.expand(B, T) t_ref = torch.full((B, n_ref), opts.get("ref_timestep", 0.0), device=t.device, dtype=t.dtype) t_cat = torch.cat([t, t_ref], dim=1) # (B, T + N) # ---- arm the attention-op state ---- tpf = _tokens_per_frame(H, W, opts["patch_spatial"]) state.gen_tokens = T * tpf state.total_tokens = (T + n_ref) * tpf state.bias_B = _per_sample_bias( B, opts["strength"], opts.get("cond_only", False), to.get("cond_or_uncond", None), x.device ) state._bias_cache = None state.active = True try: out = executor(x_cat, t_cat, context, fps, padding_mask, **kwargs) finally: state.active = False state.bias_B = None state._bias_cache = None out = out[:, :, :T] # drop reference frames if squeeze_t: out = out.squeeze(2) return out return wrapper def apply_incontext_ref( model_patcher, ref_latent, strength, start_percent, end_percent, cond_only=True, fit_mode="pad", ref_timestep=0.0, ): """Clone the ModelPatcher and install the in-context reference patches. ref_latent: raw LATENT samples tensor from a VAE encode, (N, C, H, W) or (N, C, 1, H, W). """ m = model_patcher.clone() lat = ref_latent if lat.ndim == 4: lat = lat.unsqueeze(2) # (N, C, 1, H, W) # Normalize into the model's latent space (Wan21 per-channel # mean/std). The sampler does this for the generated latent via # process_latent_in; we must match it for reference frames. lat = m.model.process_latent_in(lat.clone()) ms = m.get_model_object("model_sampling") sigma_start = ms.percent_to_sigma(start_percent) sigma_end = ms.percent_to_sigma(end_percent) dm = m.get_model_object("diffusion_model") patch_spatial = getattr(dm, "patch_spatial", 2) state = RefState() opts = { "ref_latent": lat, "state": state, "strength": strength, "cond_only": cond_only, "fit_mode": fit_mode, "sigma_start": sigma_start, "sigma_end": sigma_end, "patch_spatial": patch_spatial, "ref_timestep": ref_timestep, } wrapper = make_diffusion_wrapper(opts) if hasattr(m, "add_wrapper_with_key"): m.add_wrapper_with_key(WrappersMP.DIFFUSION_MODEL, WRAPPER_KEY, wrapper) else: comfy.patcher_extension.add_wrapper_with_key( WrappersMP.DIFFUSION_MODEL, WRAPPER_KEY, wrapper, m.model_options, is_model_options=True ) # Patch every block's self-attention op for strength control. # TODO: reference-side K/V is constant across steps and could be # cached; skipped for now (2B model, minor win). for i, block in enumerate(dm.blocks): orig_op = block.self_attn.attn_op m.add_object_patch( "diffusion_model.blocks.{}.self_attn.attn_op".format(i), make_ref_attn_op(state, orig_op), ) return m