commit cb1b535ffad75ce5c41f4a26797c7afa1bfa8df6 Author: Seunghyun Ji Date: Tue Jun 2 19:54:05 2026 +0900 asdf diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c6844d5 --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +node.zip +*.pyc +__pycache__/ +.ruff_cache/ diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..854a266 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Seunghyun Ji + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..e165314 --- /dev/null +++ b/README.md @@ -0,0 +1,105 @@ +# ComfyUI-EasyControl-KSamplerCompat + +KSampler-compatible **Anima EasyControl** image conditioning for ComfyUI. + +One node — **`Anima EasyControl (KSampler)`** — takes `MODEL + VAE + IMAGE` and +returns a `MODEL`. The reference-image conditioning rides on the returned model +as a set of ModelPatcher object-patches, so it flows straight into the **stock +KSampler** and composes with the rest of your graph (schedulers, ControlNet, +other adapters). + +This is deliberately unlike the upstream Flux / Qwen EasyControl ComfyUI nodes, +which make you load a model through a dedicated loader and sample with a +**dedicated sampler node**. Here there is no custom sampler — just a model patch. + +## What it's for + +EasyControl is a frozen-DiT reference-image adapter for the Anima model. The +flagship use case is **colorization**: feed a grayscale / manga line page into +the `IMAGE` socket and a color prompt into your normal text encode, and the DiT +colorizes the reference while respecting its structure. + +## Install + +1. Clone into `ComfyUI/custom_nodes/`: + ``` + git clone https://github.com/sorryhyun/ComfyUI-EasyControl-KSamplerCompat + ``` +2. Put your EasyControl checkpoint (`*.safetensors` trained with + `networks.methods.easycontrol`, e.g. the colorize adapter) in + `ComfyUI/models/loras/`. +3. Restart ComfyUI. + +## Wiring + +``` +UNETLoader ──► MODEL ─┐ + ├─► Anima EasyControl (KSampler) ──► MODEL ──► KSampler ──► ... +VAELoader ───► VAE ───┤ +LoadImage ──► IMAGE ──┘ ▲ + easycontrol_lora = + strength = 1.0 +``` + +- **`image`** — the conditioning image (the grayscale / lineart page for + colorization). It is VAE-encoded once; resolution should match what you + sample at (the cond stream runs at the image's own token count). +- **`strength`** — scales the trained cond effect. `1.0` = as trained. Raise to + push harder toward the reference, lower to loosen. +- **`cond_scale_override`** (optional) — replace the checkpoint's trained + `cond_scale` outright (before `strength`); `0` = keep the trained value. + +Text prompt: encode it the normal way and feed the KSampler's positive/negative +as usual. For the colorize adapter, the prompt carries **color** facts the +lineart can't (`"pink hair, blue eyes, white dress"`); structure comes from the +reference image. + +## How it works + +EasyControl extends each DiT block's self-attention to attend over a +**reference-image key/value stream** in addition to the target tokens, gated by +a trained per-block scalar bias `b_cond`. On apply: + +1. The reference image is VAE-encoded and mapped into the DiT input space + (`process_latent_in`). +2. On the first sampling step, the cond stream is walked once through all blocks + to build a per-block `(cond_k, cond_v)` cache (deterministic across steps — + the cond t-embedding is fixed at `t=0`). Reused for every step and CFG + branch. +3. Each block's `forward` is replaced (via reversible object-patch) with one + that runs extended self-attention `[target_k ; cond_k]` with the `b_cond` + bias on the cond columns. Cross-attention and MLP run baseline. + +It targets ComfyUI's native Anima/Cosmos backbone +(`comfy/ldm/cosmos/predict2.py`) directly — a from-scratch reimplementation of +anima_lora's inference path against ComfyUI's **split** `q_proj/k_proj/v_proj` +layout. The trained cond-LoRA delta (a fused `D→3D` tensor) is sliced into q/k/v +thirds and added onto the split projections; the base weights are numerically +identical between the two DiTs (same pretrained model, fused-vs-split layout), +so this reproduces what training saw. No anima_lora vendoring required — it uses +only ComfyUI's own modules plus the checkpoint tensors. + +## Known limitations / notes + +- **Latent space.** The node assumes `VAE.encode` → `process_latent_in` yields + the same latent the DiT receives for the noisy target. This is the principled + match for the native Anima model; if a future build changes the latent + pipeline, the cond stream would need the same change. +- **Block-compile ordering.** If you also use a node that *rebuilds* the DiT + (e.g. Anima block-compile), apply **this node after** it in the chain. The + per-block patches resolve their block from the live diffusion_model each + forward (rebuild-tolerant), but as with the other Anima nodes, mixing + DiT-rebuilding patches is best avoided. +- **Attention.** Extended attention uses `scaled_dot_product_attention` with an + additive bias mask (memory-efficient backends handle the bias). This differs + from anima_lora's training-time flash-LSE path numerically only at the ulp + level. +- **One reference image.** If a batch is fed to `IMAGE`, the first image is used. + +## Relationship to the Anima Adapter Loader + +This is a separate repo from +[`ComfyUI-Anima_lora-Adapter`](https://github.com/sorryhyun/ComfyUI-Anima_lora-Adapter) +(LoRA / HydraLoRA / ReFT / FeRA / Soft Tokens). Those add residuals to whole +Linear/block outputs and can ride generic forward hooks; EasyControl reaches +inside self-attention, so it needs this dedicated reimplementation. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..ba37be9 --- /dev/null +++ b/__init__.py @@ -0,0 +1,18 @@ +"""ComfyUI-EasyControl-KSamplerCompat. + +A single model-patch node, ``AnimaEasyControlPatch`` (display: "Anima +EasyControl (KSampler)"), that applies Anima EasyControl reference-image +conditioning as a ModelPatcher patch and returns a MODEL for the stock KSampler. + +Unlike the upstream Flux/Qwen EasyControl ComfyUI nodes — which require a +dedicated model loader and a custom sampler node — this node takes MODEL + VAE + +IMAGE and returns a MODEL, so the conditioning rides through the normal KSampler +path and composes with the rest of the graph. + +Targets ComfyUI's native Anima/Cosmos DiT backbone (comfy/ldm/cosmos/predict2.py). +EasyControl checkpoints (networks.methods.easycontrol) go in models/loras/. +""" + +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/easycontrol_patch.py b/easycontrol_patch.py new file mode 100644 index 0000000..6fb630a --- /dev/null +++ b/easycontrol_patch.py @@ -0,0 +1,410 @@ +"""EasyControl image conditioning for ComfyUI — KSampler-compatible. + +This is a from-scratch reimplementation of anima_lora's EasyControl *inference* +path, rewritten against ComfyUI's native Anima/Cosmos DiT backbone +(``comfy/ldm/cosmos/predict2.py``) so it runs under the stock KSampler with no +dedicated sampler node. + +Why a reimplementation rather than vendoring anima_lora's network module: +EasyControl reaches *inside* each DiT block's self-attention (it concatenates a +reference-image key/value stream into the target's attention). anima_lora's +training DiT uses a **fused** ``self_attn.qkv_proj`` (D->3D) and an +``adaln_fused_down`` / ``adaln_up_*`` AdaLN-LoRA topology; ComfyUI's backbone +uses **split** ``q_proj`` / ``k_proj`` / ``v_proj`` and ``adaln_modulation_*`` +Sequentials. So the in-block splice cannot be reused as-is. + +Correctness hinge: the trained cond-LoRA delta is a fused ``D->3D`` tensor whose +three chunks are (q, k, v) in that order — exactly the split layout ComfyUI +uses. The *base* projection weights are numerically identical between the two +DiTs (same pretrained Anima model, just fused-vs-split layout — see anima_lora +``networks/attn_fuse.py``), so adding the sliced fused delta onto ComfyUI's +split q/k/v projections reproduces what training saw. + +Mechanism, per denoising step: + - The reference image is VAE-encoded once (node side) and the cond stream is + walked once through all blocks (lazily, on the first DiT forward) to produce + a per-block (cond_k, cond_v) cache. The cond stream is deterministic across + steps (cond t-embedding at t=0, frozen weights), so this is computed once + and reused for every step and every CFG branch. + - Each block's self-attention is replaced by an *extended* attention over + ``[target_k ; cond_k]`` / ``[target_v ; cond_v]`` with a per-block additive + logit bias ``b_cond`` on the cond columns (the trained gate). Cross-attn and + MLP run baseline. + +Robustness: the per-block forward closures resolve their block from +``state.live_blocks`` (refreshed by the pre-hook from the live diffusion_model +each forward) rather than capturing a block reference, so a downstream node that +rebuilds the DiT (e.g. block-compile) doesn't strand the patch on a dead +instance — same hazard the in-tree Anima nodes document. +""" + +from __future__ import annotations + +import logging +from collections import OrderedDict +from typing import Optional + +import torch +import torch.nn.functional as F +from einops import rearrange + +logger = logging.getLogger(__name__) + +# ComfyUI backbone helpers. Imported at module load — this node only makes sense +# inside a ComfyUI install, where these are always importable. +from comfy.ldm.cosmos.predict2 import apply_rotary_pos_emb + + +# --------------------------------------------------------------------------- +# Weight loading +# --------------------------------------------------------------------------- + + +def _load_weights_and_meta(path: str): + from safetensors import safe_open + from safetensors.torch import load_file + + sd = load_file(path) + meta = {} + if path.endswith(".safetensors"): + with safe_open(path, framework="pt") as f: + meta = f.metadata() or {} + return sd, meta + + +class _LoRA: + """A single ``D->r->out`` cond-LoRA delta with an fp32 bottleneck. + + Mirrors anima_lora ``networks/methods/easycontrol.py::_LoRAProj``: the + bottleneck runs in fp32 for bf16 stability, the ``alpha/r`` scale is folded + in here, and the caller-side ``cond_scale * strength`` is applied via + ``eff_scale`` at the call site. + """ + + __slots__ = ("down", "up", "scale") + + def __init__(self, down: torch.Tensor, up: torch.Tensor, alpha_over_r: float): + self.down = down + self.up = up + self.scale = alpha_over_r + + def to(self, device, dtype): + # Keep master copies in fp32 (the forward casts to fp32 anyway). + self.down = self.down.to(device=device, dtype=torch.float32) + self.up = self.up.to(device=device, dtype=torch.float32) + return self + + def __call__(self, x: torch.Tensor) -> torch.Tensor: + h = F.linear(x.float(), self.down) + h = F.linear(h, self.up) + return (h * self.scale).to(x.dtype) + + +class EasyControlState: + """Holds the trained cond-LoRA tensors, the reference latent, and the lazily + built per-block KV cache. One per applied node instance; carried on the + ModelPatcher via a strong ref so it isn't GC'd.""" + + def __init__(self, sd: dict, meta: dict, strength: float, + cond_scale_override: Optional[float] = None): + self.num_blocks = int(meta.get("ss_num_blocks", 28)) + self.r = int(meta.get("ss_cond_lora_dim", 16)) + self.alpha = float(meta.get("ss_cond_lora_alpha", float(self.r))) + self.apply_ffn = bool(int(meta.get("ss_apply_ffn_lora", 1))) + cond_scale = ( + float(cond_scale_override) + if cond_scale_override is not None + else float(meta.get("ss_cond_scale", 1.0)) + ) + # alpha/r folds into each _LoRA; cond_scale * strength is the call-site + # effective scale (matches anima_lora get_effective_scale + multiplier). + a_over_r = self.alpha / self.r if self.r > 0 else 1.0 + self.eff_scale = cond_scale * float(strength) + + def grab(prefix, idx): + d = sd.get(f"{prefix}.{idx}.lora_down.weight") + u = sd.get(f"{prefix}.{idx}.lora_up.weight") + if d is None or u is None: + return None + return _LoRA(d, u, a_over_r) + + self.lora_qkv = [grab("cond_lora_qkv", i) for i in range(self.num_blocks)] + self.lora_o = [grab("cond_lora_o", i) for i in range(self.num_blocks)] + if self.apply_ffn: + self.lora_ffn1 = [grab("cond_lora_ffn1", i) for i in range(self.num_blocks)] + self.lora_ffn2 = [grab("cond_lora_ffn2", i) for i in range(self.num_blocks)] + else: + self.lora_ffn1 = [None] * self.num_blocks + self.lora_ffn2 = [None] * self.num_blocks + + # b_cond stored as a ParameterList → keys "b_cond.{i}" (0-d tensors). + self.b_cond = [ + sd.get(f"b_cond.{i}", torch.tensor(-10.0)).float().reshape(()) + for i in range(self.num_blocks) + ] + + if any(q is None for q in self.lora_qkv): + missing = [i for i, q in enumerate(self.lora_qkv) if q is None] + raise ValueError( + f"EasyControl checkpoint missing cond_lora_qkv for blocks {missing}. " + "Is this an EasyControl (networks.methods.easycontrol) checkpoint?" + ) + + # Reference latent in *model space* (process_latent_in already applied by + # the node). 4D [1, C, H, W]; given a device/dtype at prefill time. + self.cond_latent: Optional[torch.Tensor] = None + + # Built lazily on first DiT forward. + self.cond_kv_cache: Optional[list] = None + self._cache_device = None + # Refreshed each forward by the pre-hook from the live diffusion_model. + self.live_blocks = None + + def _to(self, device, dtype): + for lst in (self.lora_qkv, self.lora_o, self.lora_ffn1, self.lora_ffn2): + for m in lst: + if m is not None: + m.to(device, dtype) + self.b_cond = [b.to(device=device) for b in self.b_cond] + + +# --------------------------------------------------------------------------- +# Cond stream prefill (build per-block KV cache) +# --------------------------------------------------------------------------- + + +def _cond_qkv(attn, cond_normed, lora_qkv, eff_scale, rope): + """Split-projection cond Q/K/V with the fused LoRA delta sliced into thirds, + then per-head norm + RoPE. Returns q, k, v as [B, S, H, D].""" + q = attn.q_proj(cond_normed) + k = attn.k_proj(cond_normed) + v = attn.v_proj(cond_normed) + if lora_qkv is not None: + dq, dk, dv = (eff_scale * lora_qkv(cond_normed)).chunk(3, dim=-1) + q = q + dq + k = k + dk + v = v + dv + h, d = attn.n_heads, attn.head_dim + q = rearrange(q, "b s (h d) -> b s h d", h=h, d=d) + k = rearrange(k, "b s (h d) -> b s h d", h=h, d=d) + v = rearrange(v, "b s (h d) -> b s h d", h=h, d=d) + q = attn.q_norm(q) + k = attn.k_norm(k) + v = attn.v_norm(v) + if rope is not None: + q = apply_rotary_pos_emb(q, rope) + k = apply_rotary_pos_emb(k, rope) + return q, k, v + + +def _self_attn_bshd(q, k, v): + """Plain self-attention on [B, S, H, D] tensors → [B, S, H, D].""" + qt, kt, vt = (t.transpose(1, 2) for t in (q, k, v)) # [B,H,S,D] + out = F.scaled_dot_product_attention(qt, kt, vt) + return out.transpose(1, 2) + + +@torch.no_grad() +def _build_cond_kv(dit, state: EasyControlState, device, dtype): + """Walk the cond stream once; cache per-block (cond_k, cond_v). + + The cond stream is self-attention only (cross-attn is dropped on the cond + side, matching anima_lora's two-stream variant), evolving cond_x block by + block exactly as anima_lora ``precompute_cond_kv`` does, but against + ComfyUI's split-projection block API. + """ + state._to(device, dtype) + cond_latent = state.cond_latent.to(device=device, dtype=dtype) + if cond_latent.ndim == 4: + cond_latent = cond_latent.unsqueeze(2) # [B,C,1,H,W] + + cond_x5d, cond_rope_raw, cond_extra = dit.prepare_embedded_sequence( + cond_latent, fps=None, padding_mask=None + ) + cond_x = rearrange(cond_x5d, "b t h w d -> b (t h w) d") + if cond_extra is not None: + # Anima ships extra_per_block_abs_pos_emb=False, so this is None. If a + # future config enables it, the per-block add would need mirroring here. + logger.warning( + "EasyControl: extra_per_block_abs_pos_emb is set; cond stream does " + "not replicate it. Results may drift." + ) + cond_rope = cond_rope_raw.unsqueeze(1).unsqueeze(0) if cond_rope_raw is not None else None + + # cond t-embedding at t=0 — identical to MiniTrainDIT._forward's t path. + B = cond_x.shape[0] + zeros = torch.zeros(B, 1, device=device, dtype=dtype) + cemb, cadaln = dit.t_embedder[1](dit.t_embedder[0](zeros).to(dtype)) + cemb = dit.t_embedding_norm(cemb) + + eff = state.eff_scale + cache = [] + for idx in range(state.num_blocks): + block = dit.blocks[idx] + attn = block.self_attn + + if block.use_adaln_lora: + s_sa, sc_sa, g_sa = ( + block.adaln_modulation_self_attn(cemb) + cadaln + ).chunk(3, dim=-1) + s_mlp, sc_mlp, g_mlp = ( + block.adaln_modulation_mlp(cemb) + cadaln + ).chunk(3, dim=-1) + else: + s_sa, sc_sa, g_sa = block.adaln_modulation_self_attn(cemb).chunk(3, dim=-1) + s_mlp, sc_mlp, g_mlp = block.adaln_modulation_mlp(cemb).chunk(3, dim=-1) + + # ---- self-attn (this is what we cache) ---- + cond_normed = block.layer_norm_self_attn(cond_x) * (1 + sc_sa) + s_sa + cq, ck, cv = _cond_qkv(attn, cond_normed, state.lora_qkv[idx], eff, cond_rope) + cache.append((ck.detach(), cv.detach())) + + # ---- evolve cond_x to feed the next block ---- + attn_out = _self_attn_bshd(cq, ck, cv) + attn_out = rearrange(attn_out, "b s h d -> b s (h d)") + proj = attn.output_proj(attn_out) + if state.lora_o[idx] is not None: + proj = proj + eff * state.lora_o[idx](attn_out) + cond_x = cond_x + g_sa * proj + + # ---- MLP ---- + mlp_normed = block.layer_norm_mlp(cond_x) * (1 + sc_mlp) + s_mlp + hmid = block.mlp.layer1(mlp_normed) + if state.lora_ffn1[idx] is not None: + hmid = hmid + eff * state.lora_ffn1[idx](mlp_normed) + hmid = block.mlp.activation(hmid) + mlp_out = block.mlp.layer2(hmid) + if state.lora_ffn2[idx] is not None: + mlp_out = mlp_out + eff * state.lora_ffn2[idx](hmid) + cond_x = cond_x + g_mlp * mlp_out + + state.cond_kv_cache = cache + state._cache_device = device + nbytes = sum(k.numel() + v.numel() for k, v in cache) * cache[0][0].element_size() + logger.info( + f"EasyControl: built cond KV cache ({len(cache)} blocks, " + f"S_c={cache[0][0].shape[1]}, {nbytes / 1e6:.0f} MB)" + ) + + +# --------------------------------------------------------------------------- +# Extended self-attention (target keys + cached cond keys, b_cond gate) +# --------------------------------------------------------------------------- + + +def _extended_self_attn(tq, tk, tv, ck, cv, b_cond): + """Attention of target queries over [target_k; cond_k] with an additive + ``b_cond`` logit bias on the cond columns. All inputs [B, S, H, D] (cond may + be batch-1 and is broadcast). Returns [B, S, H, D].""" + B, S = tq.shape[0], tq.shape[1] + if ck.shape[0] != B: + ck = ck.expand(B, -1, -1, -1) + cv = cv.expand(B, -1, -1, -1) + Sc = ck.shape[1] + k = torch.cat([tk, ck], dim=1).transpose(1, 2) # [B,H,S+Sc,D] + v = torch.cat([tv, cv], dim=1).transpose(1, 2) + q = tq.transpose(1, 2) + mask = torch.zeros(S + Sc, device=q.device, dtype=q.dtype) + mask[S:] = b_cond.to(q.dtype) + out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask.view(1, 1, 1, S + Sc)) + return out.transpose(1, 2) # [B,S,H,D] + + +def _make_self_attn_forward(state: EasyControlState, idx: int): + """Replacement for one block's ``self_attn.forward`` (an ``Attention``). + + Overriding only the self-attention — not the whole Block — keeps every other + part of ComfyUI's block native (AdaLN, cross-attn, MLP, residuals, dtype + handling), so there is nothing to drift and no block-level weights to strand + (``Attention`` holds no direct parameters; its q/k/v/o projections and norms + are submodules that keep their own cast-weights forward). + + Resolves the live ``Attention`` from ``state.live_blocks[idx].self_attn`` + each call, so a downstream DiT rebuild can't strand the patch on a dead + instance. + """ + + def forward(x, context=None, rope_emb=None, transformer_options={}): + attn = state.live_blocks[idx].self_attn + cache = state.cond_kv_cache + # Cond not ready, or this somehow ran as cross-attn → native attention. + if cache is None or context is not None: + return type(attn).forward(attn, x, context, rope_emb, transformer_options) + + tq, tk, tv = attn.compute_qkv(x, None, rope_emb=rope_emb) # [B,S,H,D] + ck, cv = cache[idx] + out = _extended_self_attn( + tq, tk, tv, ck.to(tk.dtype), cv.to(tv.dtype), state.b_cond[idx] + ) + out = rearrange(out, "b s h d -> b s (h d)") + return attn.output_dropout(attn.output_proj(out)) + + return forward + + +# --------------------------------------------------------------------------- +# Pre-hook: lazy cond-KV prefill + live-block refresh +# --------------------------------------------------------------------------- + + +def _make_prefill_pre_hook(state: EasyControlState): + def pre_hook(module, args): + # module = live diffusion_model; args[0] = x (model-space latent). + state.live_blocks = module.blocks + x = args[0] + device, dtype = x.device, x.dtype + if state.cond_kv_cache is None or state._cache_device != device: + _build_cond_kv(module, state, device, dtype) + return None + + return pre_hook + + +# --------------------------------------------------------------------------- +# Public entry point — called by the node +# --------------------------------------------------------------------------- + + +def install_easycontrol(model, weight_path, cond_latent, strength, cond_scale_override=None): + """Install EasyControl onto a cloned ModelPatcher. + + ``cond_latent`` is the VAE-encoded reference in **model space** (the node + applies ``process_latent_in``), 4D ``[1, C, H, W]``. + """ + sd, meta = _load_weights_and_meta(weight_path) + state = EasyControlState(sd, meta, strength, cond_scale_override) + state.cond_latent = cond_latent + + diffusion_model = model.get_model_object("diffusion_model") + nblocks = len(diffusion_model.blocks) + if nblocks != state.num_blocks: + raise ValueError( + f"EasyControl checkpoint is for {state.num_blocks} blocks but the " + f"loaded DiT has {nblocks}. Wrong base model?" + ) + + # Lazy-prefill pre-hook on diffusion_model (object-patched OrderedDict so it + # reverts on unpatch and composes with other pre-hooks). + pre_hook = _make_prefill_pre_hook(state) + new_pre = OrderedDict(diffusion_model._forward_pre_hooks) + new_pre[id(pre_hook)] = pre_hook + model.add_object_patch("diffusion_model._forward_pre_hooks", new_pre) + + # Per-block self-attention replacement (object-patched → path-resolved to + # the live blocks; reverted on unpatch). Only self_attn is swapped; the rest + # of each block stays native. + for idx in range(state.num_blocks): + model.add_object_patch( + f"diffusion_model.blocks.{idx}.self_attn.forward", + _make_self_attn_forward(state, idx), + ) + + # Keep the state alive for the patcher's lifetime. + if not hasattr(model, "_easycontrol_states"): + model._easycontrol_states = [] + model._easycontrol_states.append(state) + logger.info( + f"EasyControl: installed ({state.num_blocks} blocks, r={state.r}, " + f"ffn_lora={state.apply_ffn}, eff_scale={state.eff_scale:.3f})" + ) + return state diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..11ebf72 --- /dev/null +++ b/nodes.py @@ -0,0 +1,136 @@ +"""EasyControl image-conditioning node for the Anima DiT — KSampler-compatible. + +A single model-patch node: ``MODEL + VAE + IMAGE -> MODEL``. The returned MODEL +carries the reference-image conditioning as a set of object patches, so it flows +straight into the stock KSampler (and composes with ControlNet, schedulers, +other adapters) — no dedicated sampler node, unlike the upstream Flux/Qwen +EasyControl ComfyUI nodes. + +The EasyControl ``.safetensors`` goes in ``ComfyUI/models/loras/`` (same folder +the LoRA loaders read). Feed the conditioning image (e.g. the grayscale / manga +line page for colorization) into the IMAGE socket and a VAE that matches the +Anima latent space into the VAE socket. +""" + +import folder_paths + +from .easycontrol_patch import install_easycontrol + + +class AnimaEasyControlPatch: + """Apply Anima EasyControl reference-image conditioning to a MODEL. + + EasyControl is a frozen-DiT adapter: a per-block cond LoRA on self-attn + (q/k/v/o) + FFN plus a scalar ``b_cond`` logit-bias gate. At inference the + reference image is VAE-encoded, the cond stream is walked once to build a + per-block key/value cache, and each block's self-attention is extended to + attend over those cached cond keys. Output is a MODEL for the stock KSampler. + """ + + @classmethod + def INPUT_TYPES(cls): + loras = folder_paths.get_filename_list("loras") + return { + "required": { + "model": ("MODEL",), + "vae": ("VAE",), + "image": ("IMAGE",), + "base_model": ( + ["anima"], + { + "default": "anima", + "tooltip": ( + "Base model family this node supports. Currently " + "Anima only — the EasyControl reimplementation targets " + "the Anima/Cosmos DiT backbone. More families may be " + "added later." + ), + }, + ), + "easycontrol_lora": ( + loras, + { + "tooltip": ( + "EasyControl checkpoint (networks.methods.easycontrol " + "— per-block cond LoRA + b_cond gate). Lives in " + "ComfyUI/models/loras/." + ) + }, + ), + "strength": ( + "FLOAT", + { + "default": 1.0, + "min": 0.0, + "max": 3.0, + "step": 0.05, + "tooltip": ( + "Conditioning strength. Scales the cond-LoRA deltas " + "(multiplies the trained cond_scale). 1.0 = as trained; " + "raise to push harder toward the reference, lower to " + "loosen." + ), + }, + ), + }, + "optional": { + "cond_scale_override": ( + "FLOAT", + { + "default": 0.0, + "min": 0.0, + "max": 4.0, + "step": 0.05, + "tooltip": ( + "If > 0, replaces the checkpoint's trained cond_scale " + "(before strength). 0 = use the trained value." + ), + }, + ), + }, + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "apply" + CATEGORY = "loaders" + DESCRIPTION = ( + "Anima EasyControl image conditioning, KSampler-compatible. Feeds a " + "reference image (VAE-encoded) into a frozen-DiT extended self-attention " + "stream and returns a patched MODEL for the stock KSampler." + ) + + def apply(self, model, vae, image, base_model, easycontrol_lora, strength, + cond_scale_override=0.0): + # base_model is an informational guard — only "anima" is offered today. + del base_model + new_model = model.clone() + weight_path = folder_paths.get_full_path("loras", easycontrol_lora) + if weight_path is None: + raise FileNotFoundError( + f"EasyControl checkpoint not found in loras folder: {easycontrol_lora}" + ) + + # IMAGE is [B, H, W, C] in [0, 1]; use the first image as the reference. + pixels = image[:1, :, :, :3] + vae_latent = vae.encode(pixels) # ComfyUI latent space [1, C, h, w] + + # Map into the DiT's input space — the same transform apply_model applies + # to the noisy latent before calling diffusion_model. Without this the + # cond stream sees a differently-scaled latent than the target stream. + cond_latent = new_model.model.process_latent_in(vae_latent) + + override = float(cond_scale_override) if cond_scale_override and cond_scale_override > 0 else None + install_easycontrol( + new_model, weight_path, cond_latent, float(strength), + cond_scale_override=override, + ) + return (new_model,) + + +NODE_CLASS_MAPPINGS = { + "AnimaEasyControlPatch": AnimaEasyControlPatch, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AnimaEasyControlPatch": "Anima EasyControl (KSampler)", +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..0085ec1 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "comfyui-easycontrol-ksamplercompat" +description = "KSampler-compatible Anima EasyControl image conditioning for ComfyUI. A single model-patch node (MODEL + VAE + IMAGE -> MODEL) that feeds a reference image into the Anima DiT's frozen extended self-attention stream and returns a patched MODEL for the stock KSampler — no dedicated sampler node. Reimplements anima_lora's EasyControl inference path against ComfyUI's native Anima/Cosmos DiT backbone (split q/k/v projections). EasyControl checkpoints go in models/loras/. Typical use: grayscale/lineart -> color (colorization)." +version = "0.1.0" +license = { file = "LICENSE" } +dependencies = [] + +[project.urls] +Repository = "https://github.com/sorryhyun/ComfyUI-EasyControl-KSamplerCompat" + +[tool.comfy] +PublisherId = "sorryhyun" +DisplayName = "Anima EasyControl (KSampler)" +Icon = "" diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_smoke.py b/tests/test_smoke.py new file mode 100644 index 0000000..e4be8f4 --- /dev/null +++ b/tests/test_smoke.py @@ -0,0 +1,160 @@ +"""Numeric checks for the EasyControl KSampler-compat reimplementation. + +Two independent checks, no ComfyUI runtime / no Anima checkpoint needed: + + test_extended_attn_math + The core correctness claim, tested directly on the attention math: + - with a very negative b_cond the cond keys are gated out, so extended + attention == plain target-only self-attention (the step-0 equivalence + that makes the splice safe to drop on a trained model); + - with b_cond = 0 the cond keys contribute, so the output differs; + - the cond batch (1) broadcasts to the target batch. + + test_plumbing + End-to-end wiring on a tiny real predict2 DiT (the backbone Anima + inherits): build the cond-KV cache, override each block's self_attn, run a + full DiT forward, and assert it's finite and shape-preserving. The DiT is + stabilized (adaln gates near-zero) only so an *untrained* net doesn't blow + up — this check is about wiring, not values. + +Run: python tests/test_smoke.py (with the comfy repo importable) +""" + +import os +import sys + +# Make `comfy` importable: prefer an already-installed comfy, else COMFY_PATH, +# else a couple of common dev locations relative to this checkout. +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +try: + import comfy.ops # noqa: F401 +except ImportError: + for cand in ( + os.environ.get("COMFY_PATH"), + os.path.expanduser("~/anima/comfy"), + os.path.expanduser("~/comfy"), + ): + if cand and os.path.isdir(os.path.join(cand, "comfy")): + sys.path.insert(0, cand) + break + +import torch + +import comfy.ops as ops +from comfy.ldm.cosmos.predict2 import MiniTrainDIT + +import easycontrol_patch as ec + + +def test_extended_attn_math(): + torch.manual_seed(0) + B, S, Sc, H, D = 2, 12, 5, 3, 16 + tq = torch.randn(B, S, H, D) + tk = torch.randn(B, S, H, D) + tv = torch.randn(B, S, H, D) + ck = torch.randn(1, Sc, H, D) # cond is batch-1 → must broadcast + cv = torch.randn(1, Sc, H, D) + + # Reference: plain target-only self-attention. + base = torch.nn.functional.scaled_dot_product_attention( + tq.transpose(1, 2), tk.transpose(1, 2), tv.transpose(1, 2) + ).transpose(1, 2) + + gated = ec._extended_self_attn(tq, tk, tv, ck, cv, torch.tensor(-1e4)) + rel = (gated - base).abs().max().item() / base.abs().max().item() + print(f" [math] gated rel-diff vs target-only = {rel:.2e}") + assert gated.shape == base.shape + assert rel < 1e-5, f"very-negative b_cond must reduce to target-only, rel={rel}" + + active = ec._extended_self_attn(tq, tk, tv, ck, cv, torch.tensor(0.0)) + delta = (active - base).abs().max().item() + print(f" [math] active abs-diff vs target-only = {delta:.2e}") + assert delta > 1e-2, "b_cond=0 must let cond keys change the output" + assert torch.isfinite(active).all() + print("test_extended_attn_math PASSED") + + +def _stable_dit(): + torch.manual_seed(0) + dit = MiniTrainDIT( + max_img_h=16, max_img_w=16, max_frames=1, + in_channels=4, out_channels=4, + patch_spatial=2, patch_temporal=1, + concat_padding_mask=True, + model_channels=32, num_blocks=2, num_heads=2, mlp_ratio=4.0, + crossattn_emb_channels=16, + pos_emb_cls="rope3d", use_adaln_lora=False, + operations=ops.disable_weight_init, + ).eval().float() + # disable_weight_init leaves Linear weights UNINITIALIZED (ComfyUI expects a + # checkpoint), so init everything deterministically: norm weights → 1, all + # other weights → small normal. Then shrink the adaln_modulation output so an + # untrained net's residual stream stays bounded. + torch.manual_seed(1) + with torch.no_grad(): + for name, p in dit.named_parameters(): + if p.ndim == 1: # RMSNorm weights (and any bias) + p.fill_(1.0 if "norm" in name else 0.0) + else: + p.normal_(0.0, 0.05) + for blk in dit.blocks: + for name in ("adaln_modulation_self_attn", + "adaln_modulation_cross_attn", + "adaln_modulation_mlp"): + getattr(blk, name)[-1].weight.mul_(0.1) + return dit + + +def _build_state(num_blocks, D, ffn, r=4): + sd = {} + for i in range(num_blocks): + sd[f"cond_lora_qkv.{i}.lora_down.weight"] = torch.randn(r, D) * 0.05 + sd[f"cond_lora_qkv.{i}.lora_up.weight"] = torch.randn(3 * D, r) * 0.05 + sd[f"cond_lora_o.{i}.lora_down.weight"] = torch.randn(r, D) * 0.05 + sd[f"cond_lora_o.{i}.lora_up.weight"] = torch.randn(D, r) * 0.05 + sd[f"cond_lora_ffn1.{i}.lora_down.weight"] = torch.randn(r, D) * 0.05 + sd[f"cond_lora_ffn1.{i}.lora_up.weight"] = torch.randn(ffn, r) * 0.05 + sd[f"cond_lora_ffn2.{i}.lora_down.weight"] = torch.randn(r, ffn) * 0.05 + sd[f"cond_lora_ffn2.{i}.lora_up.weight"] = torch.randn(D, r) * 0.05 + sd[f"b_cond.{i}"] = torch.tensor(0.0) + meta = { + "ss_num_blocks": str(num_blocks), "ss_cond_lora_dim": str(r), + "ss_cond_lora_alpha": str(r), "ss_apply_ffn_lora": "1", "ss_cond_scale": "1.0", + } + return ec.EasyControlState(sd, meta, strength=1.0) + + +def test_plumbing(): + dit = _stable_dit() + D = dit.model_channels + state = _build_state(len(dit.blocks), D, int(D * 4.0)) + + # Build the cond-KV cache from a random reference latent. + state.cond_latent = torch.randn(1, 4, 8, 8) * 0.5 + state.live_blocks = dit.blocks + ec._build_cond_kv(dit, state, torch.device("cpu"), torch.float32) + assert state.cond_kv_cache is not None and len(state.cond_kv_cache) == len(dit.blocks) + for k, v in state.cond_kv_cache: + assert torch.isfinite(k).all() and torch.isfinite(v).all() + print(f" [plumb] cond KV cache: {len(state.cond_kv_cache)} blocks, " + f"S_c={state.cond_kv_cache[0][0].shape[1]}") + + # Override each block's self_attn and run a full forward. + for idx in range(len(dit.blocks)): + dit.blocks[idx].self_attn.forward = ec._make_self_attn_forward(state, idx) + + x = torch.randn(1, 4, 1, 8, 8) * 0.5 + t = torch.tensor([0.5]) + ctx = torch.randn(1, 7, 16) * 0.5 + with torch.no_grad(): + out = dit(x, t, ctx) + assert out.shape == (1, 4, 1, 8, 8), out.shape + assert torch.isfinite(out).all(), "patched DiT forward produced non-finite output" + print(f" [plumb] full patched forward: shape={tuple(out.shape)} finite=OK") + print("test_plumbing PASSED") + + +if __name__ == "__main__": + test_extended_attn_math() + test_plumbing() + print("\nALL CHECKS PASSED")