nodes/sorryhyun-ComfyUI-EasyControl-KSamplerCompat
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

asdf

Browse Source
This commit is contained in:
Seunghyun Ji
2026-06-02 19:54:05 +09:00
commit cb1b535ffa
9 changed files with 868 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+4
View File
@@ -0,0 +1,4 @@
node.zip
*.pyc
__pycache__/
.ruff_cache/
LICENSE
+21
View File
@@ -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.
README.md
+105
View File
@@ -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 = <your ckpt>
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.
__init__.py
+18
View File
@@ -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"]
easycontrol_patch.py
+410
View File
@@ -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
nodes.py
+136
View File
@@ -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)",
}
pyproject.toml
+14
View File
@@ -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 = ""
tests/__init__.py
View File
tests/test_smoke.py
+160
View File
@@ -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")