539 lines
30 KiB
Python
539 lines
30 KiB
Python
"""ComfyUI-Krea2Edit — in-context edit forward for the Krea2 model.
|
|
|
|
ComfyUI's native Krea2 `_forward` is text-to-image only: it builds the sequence
|
|
`[text | target]`. The krea2_edit LoRA (trained in ai-toolkit) needs the *appearance
|
|
path*: the VAE-encoded SOURCE latent prepended as a block of clean tokens, distinguished
|
|
from the (noisy) target purely by the 3-axis RoPE frame index (source=1, target=0, h/w
|
|
aligned). This node adds that by wrapping the model's DIFFUSION_MODEL forward and rebuilding
|
|
the sequence as `[text | source(frame=1) | target(frame=0)]`, keeping only the target tokens
|
|
out — mirroring ai-toolkit's `predict_velocity_edit` exactly, using the model's own submodules.
|
|
|
|
Wiring: LoadImage -> VAEEncode(source) --\
|
|
Krea2EditModelPatch(model, source_latent) -> KSampler
|
|
UNETLoader -> LoraLoaderModelOnly -/
|
|
KSampler.latent_image <- EmptySD3LatentImage (noise). Text: NATIVE krea2 CLIP + CLIPTextEncode.
|
|
"""
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from einops import rearrange
|
|
|
|
import comfy.patcher_extension
|
|
import comfy.utils
|
|
import comfy.ldm.common_dit
|
|
from comfy.ldm.flux.layers import timestep_embedding
|
|
from comfy.ldm.flux.math import apply_rope
|
|
from comfy.ldm.modules.attention import optimized_attention_masked
|
|
|
|
|
|
def _imgids(bs, frame, h_, w_, device):
|
|
ids = torch.zeros(h_, w_, 3, device=device, dtype=torch.float32)
|
|
ids[..., 0] = frame
|
|
ids[..., 1] = torch.arange(h_, device=device, dtype=torch.float32)[:, None]
|
|
ids[..., 2] = torch.arange(w_, device=device, dtype=torch.float32)[None, :]
|
|
return ids.reshape(1, h_ * w_, 3).repeat(bs, 1, 1)
|
|
|
|
|
|
def _imgids_offset(bs, frame, gh, gw, th, tw, device):
|
|
"""Stride-1 integer positions at a centered integer offset. For `fit` refs the
|
|
pixels are already resampled to target grid density, so the position grid is
|
|
stride-1 BY CONSTRUCTION — scaling it again only manufactures skip/collision
|
|
artifacts. Requires gh<=th, gw<=tw (guaranteed by the floor+cap in fit)."""
|
|
# fractional center (2026-07-28): integer floor placed odd-gap refs 8px off
|
|
# their true center — RoPE is continuous, half-token positions are exact.
|
|
off_h, off_w = max(0.0, (th - gh) / 2), max(0.0, (tw - gw) / 2)
|
|
ids = torch.zeros(gh, gw, 3, device=device, dtype=torch.float32)
|
|
ids[..., 0] = frame
|
|
ids[..., 1] = (torch.arange(gh, device=device, dtype=torch.float32) + off_h)[:, None]
|
|
ids[..., 2] = (torch.arange(gw, device=device, dtype=torch.float32) + off_w)[None, :]
|
|
return ids.reshape(1, gh * gw, 3).repeat(bs, 1, 1)
|
|
|
|
|
|
REF_POSITIONS = ["aligned (trained)", "beside (right of target)", "bag of features"]
|
|
|
|
|
|
def _place_refs(ref_ids, modes, h_, w_):
|
|
"""Reference placement on the h/w RoPE axes (frame axis untouched: refs keep 1..N).
|
|
aligned (trained): ref token (h, w) sits on target token (h, w) -- how the models were trained.
|
|
beside: each ref is shifted to the right of the target canvas (and of earlier refs): no h/w overlap,
|
|
vertical placement and the ref's internal layout unchanged.
|
|
bag of features (experimental): every token of a ref gets one position (the target's centre), so
|
|
attention to it is decided by content alone; the ref's layout is discarded.
|
|
modes: one per ref (a single string applies to all), e.g. scene aligned + subject beside."""
|
|
if isinstance(modes, str):
|
|
modes = [modes] * len(ref_ids)
|
|
out, x0 = [], float(w_)
|
|
for ids, mode in zip(ref_ids, modes):
|
|
if mode.startswith("aligned"):
|
|
out.append(ids)
|
|
continue
|
|
ids = ids.clone()
|
|
if mode.startswith("beside"):
|
|
ids[..., 2] = ids[..., 2] - ids[..., 2].min() + x0
|
|
x0 = float(ids[..., 2].max()) + 1.0
|
|
elif mode.startswith("bag"):
|
|
ids[..., 1] = (h_ - 1) / 2.0
|
|
ids[..., 2] = (w_ - 1) / 2.0
|
|
else:
|
|
raise ValueError(f"unknown ref_position {mode!r}")
|
|
out.append(ids)
|
|
return out
|
|
|
|
|
|
def _to_4d(v):
|
|
"""(B,C,T,H,W) -> (B*T,C,H,W); pass 4D through. Images use T=1."""
|
|
if v.ndim == 5:
|
|
b, c, t, h, w = v.shape
|
|
return v.reshape(b * t, c, h, w)
|
|
return v
|
|
|
|
|
|
def _fit_src(src, H, W):
|
|
"""Fit a source latent to the target grid the way TRAINING did: center-crop to
|
|
the target aspect ratio, then resize. A plain interpolate (the pre-fix behavior)
|
|
STRETCHES mixed-AR sources — users saw stretched people whenever their input AR
|
|
differed from the output resolution."""
|
|
sh, sw = src.shape[-2:]
|
|
if (sh, sw) == (H, W):
|
|
return src
|
|
s = max(H / sh, W / sw)
|
|
ch, cw = min(sh, int(round(H / s))), min(sw, int(round(W / s)))
|
|
y0, x0 = (sh - ch) // 2, (sw - cw) // 2
|
|
src = src[..., y0:y0 + ch, x0:x0 + cw]
|
|
return F.interpolate(src.float(), size=(H, W), mode="bilinear")
|
|
|
|
|
|
def _fit_encode_image(image, vae, H, W, cache, key, fit_mode="crop"):
|
|
"""Pixel-space source prep (blur-proof path): center-crop the IMAGE to the
|
|
target AR, resize to the exact target pixel grid, VAE-encode. Latent-space
|
|
resizing (the old fallback) softens VAE latents — this path never resizes
|
|
latents at all. Cached per target resolution (encode once, not per step)."""
|
|
key = key + (fit_mode,)
|
|
if key in cache:
|
|
return cache[key]
|
|
print(f"[krea2edit] _fit_encode_image: mode={fit_mode} in={tuple(image.shape)} target_latent={H}x{W}", flush=True)
|
|
px_h, px_w = H * 8, W * 8
|
|
img = image.movedim(-1, 1) # B,H,W,C -> B,C,H,W
|
|
ih, iw = img.shape[-2:]
|
|
if fit_mode == "fit":
|
|
# "bilinear" answer to scale mismatch: resample CONTENT (pixel space, bicubic)
|
|
# to the target's grid density instead of moving positions. AR-preserving
|
|
# fit-inside, no crop, no grey canvas — the forward places it at an integer
|
|
# centered offset (scaled-pos with s=1 -> stride 1, no rounding artifacts).
|
|
sc = min(px_h / ih, px_w / iw)
|
|
# NEAR-MATCHED AR: fill the target grid EXACTLY via a minimal center-crop.
|
|
# Fit-inside margins of 1-2 tokens are not harmless: target edge columns
|
|
# with no ref correspondence get filled by repeating adjacent ref content
|
|
# (2026-07-14 edge-duplication bug: ref (74,54) vs target (74,56)).
|
|
# This also restores the design promise fit == crop at matched AR.
|
|
CROP_TOL = 0.08
|
|
if ih * sc >= px_h * (1 - CROP_TOL) and iw * sc >= px_w * (1 - CROP_TOL):
|
|
s = max(px_h / ih, px_w / iw)
|
|
ch, cw = min(ih, int(round(px_h / s))), min(iw, int(round(px_w / s)))
|
|
y0, x0 = (ih - ch) // 2, (iw - cw) // 2
|
|
img = img[..., y0:y0 + ch, x0:x0 + cw]
|
|
nh, nw = px_h, px_w
|
|
else:
|
|
# genuine AR mismatch: MUST match the trainer's _fit_prep EXACTLY
|
|
# (krea2_edit.py) — /16 floor snap capped at the target's /16 floor.
|
|
# The model is trained on this geometry; a /8-round node grid would
|
|
# produce a different ref latent size -> different centered offset ->
|
|
# a visible margin-boundary seam even from a well-trained model
|
|
# (train/infer geometry must be byte-identical). 2026-07-15 alignment.
|
|
nh = min(max(16, int(ih * sc) // 16 * 16), max(16, px_h // 16 * 16))
|
|
nw = min(max(16, int(iw * sc) // 16 * 16), max(16, px_w // 16 * 16))
|
|
# CROP-TO-GRID (2026-07-28 seam-doubling RCA): resizing ih*sc -> floor16
|
|
# SQUASHES content by up to 15px; the misregistration peaks exactly at
|
|
# the ref band edges — the outpaint seam — and renders as a doubled
|
|
# band (proven causal: 754px vs 753px input A/B, one pixel flips
|
|
# clean<->worst). Center-crop the source so the fitted axis lands on
|
|
# the /16 grid at scale sc EXACTLY: zero squash, stride-1 stays true.
|
|
ch2, cw2 = min(ih, max(1, int(round(nh / sc)))), min(iw, max(1, int(round(nw / sc))))
|
|
y0, x0 = (ih - ch2) // 2, (iw - cw2) // 2
|
|
img = img[..., y0:y0 + ch2, x0:x0 + cw2]
|
|
img = F.interpolate(img.float(), size=(nh, nw), mode="bicubic", antialias=True)
|
|
lat = vae.encode(img.movedim(1, -1)[..., :3].clamp(0, 1))
|
|
cache[key] = lat
|
|
return lat
|
|
# crop (default / "v1 legacy"): center-crop to the target AR, then resize.
|
|
s = max(px_h / ih, px_w / iw)
|
|
ch, cw = min(ih, int(round(px_h / s))), min(iw, int(round(px_w / s)))
|
|
y0, x0 = (ih - ch) // 2, (iw - cw) // 2
|
|
img = img[..., y0:y0 + ch, x0:x0 + cw]
|
|
img = F.interpolate(img.float(), size=(px_h, px_w), mode="bicubic", antialias=True)
|
|
lat = vae.encode(img.movedim(1, -1)[..., :3].clamp(0, 1))
|
|
cache[key] = lat
|
|
return lat
|
|
|
|
|
|
def _ref_attn_bias(boosts, boost_mask, txtlen, slens, tgtlen, mask_hw, device, dtype):
|
|
"""Additive attention-logit bias on the [text | refs... | target] sequence.
|
|
|
|
boosts: per-ref factor on target->ref attention, aligned with the source blocks
|
|
(last entry = last ref = the subject by workflow convention). Equivalent to
|
|
multiplying those keys' post-softmax attention weight before renormalization.
|
|
boost_mask (ComfyUI MASK, ref-image pixel space) restricts the LAST ref's boost
|
|
to a region (e.g. the face).
|
|
"""
|
|
nsrc = len(slens)
|
|
offs = [txtlen]
|
|
for sl in slens:
|
|
offs.append(offs[-1] + sl)
|
|
rows0 = offs[-1]
|
|
L = rows0 + tgtlen
|
|
bias = torch.zeros(1, 1, L, L, device=device, dtype=dtype)
|
|
for i, b in enumerate(boosts):
|
|
if b == 1.0:
|
|
continue
|
|
off, sl = offs[i], slens[i]
|
|
if boost_mask is not None and i == nsrc - 1 and mask_hw is not None:
|
|
mask = boost_mask[:1]
|
|
if mask.ndim == 2:
|
|
mask = mask[None]
|
|
mask = F.interpolate(mask[None].float(), size=mask_hw[i], mode="area")[0, 0]
|
|
cols = off + torch.nonzero(mask.reshape(-1) > 0.5, as_tuple=True)[0].to(device)
|
|
else:
|
|
cols = torch.arange(off, off + sl, device=device)
|
|
bias[:, :, rows0:, cols] = math.log(max(b, 1e-4))
|
|
return bias
|
|
|
|
|
|
def krea2_edit_forward(m, x, timesteps, context, src_latent, transformer_options,
|
|
ref_boost=1.0, ref_boost_a=1.0, ref_boost_mask=None,
|
|
ref_native=False, pos_mode="anchor", ref_position="aligned (trained)",
|
|
ref_position_a="aligned (trained)"):
|
|
"""Krea2 SingleStreamDiT._forward, but with source block(s) prepended.
|
|
|
|
m : the SingleStreamDiT (LoRA-patched at sample time)
|
|
x : (B,C,H,W) or (B,C,T,H,W) noisy TARGET latent
|
|
src_latent : clean SOURCE latent (VAE-encoded), 4D/5D — or a LIST of them
|
|
(multi-ref: [scene, subject], frames 1..N, training-matched)
|
|
context : (B, seq, txtlayers*txtdim) — the 12-layer Qwen3-VL stack
|
|
"""
|
|
patch = m.patch
|
|
|
|
# Mirror ComfyUI _forward: latents may arrive 5D (B,C,T,H,W) for this model.
|
|
temporal = x.ndim == 5
|
|
if temporal:
|
|
b5, c5, t5, h5, w5 = x.shape
|
|
x = _to_4d(x)
|
|
bs, c, H_orig, W_orig = x.shape
|
|
|
|
x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch, patch), padding_mode="replicate")
|
|
H, W = x.shape[-2], x.shape[-1]
|
|
h_, w_ = H // patch, W // patch
|
|
|
|
# source(s) -> (bs, C, H, W): flatten temporal, match batch, fit to the target grid
|
|
# (center-crop to target AR then resize — training-matched; never stretch).
|
|
src_list = src_latent if isinstance(src_latent, (list, tuple)) else [src_latent]
|
|
srcs = []
|
|
for sl in src_list:
|
|
src = _to_4d(sl).to(x.device, x.dtype)
|
|
if src.shape[0] != bs:
|
|
src = src[:1].expand(bs, *src.shape[1:])
|
|
if not ref_native and src.shape[-2:] != (H, W):
|
|
print(f"[krea2edit] LATENT-PATH fit_src (crop): src={tuple(src.shape[-2:])} -> {H}x{W}", flush=True)
|
|
src = _fit_src(src, H, W).to(x.dtype)
|
|
srcs.append(comfy.ldm.common_dit.pad_to_patch_size(src, (patch, patch), padding_mode="replicate"))
|
|
src_grids = [(s_.shape[-2] // patch, s_.shape[-1] // patch) for s_ in srcs]
|
|
|
|
context = m._unpack_context(context) # (B, seq, 12, 2560)
|
|
|
|
tgt_img = m.first(rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch))
|
|
src_imgs = [m.first(rearrange(s_, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch))
|
|
for s_ in srcs]
|
|
|
|
t = m.tmlp(timestep_embedding(timesteps, m.tdim).unsqueeze(1).to(tgt_img.dtype))
|
|
tvec = m.tproj(t)
|
|
|
|
context = m.txtfusion(context, mask=None, transformer_options=transformer_options)
|
|
context = m.txtmlp(context)
|
|
|
|
txtlen, tgtlen = context.shape[1], tgt_img.shape[1]
|
|
srclen = sum(si.shape[1] for si in src_imgs)
|
|
combined = torch.cat([context] + src_imgs + [tgt_img], dim=1) # [text | refs... | target]
|
|
|
|
device = combined.device
|
|
if pos_mode == "stride1" and ref_native:
|
|
print(f"[krea2edit] STRIDE1-POS fit: ref grids {src_grids} centered in ({h_},{w_})", flush=True)
|
|
ref_ids = [_imgids_offset(bs, i + 1, gh, gw, h_, w_, device)
|
|
for i, (gh, gw) in enumerate(src_grids)]
|
|
else:
|
|
ref_ids = [_imgids(bs, i + 1, gh, gw, device) for i, (gh, gw) in enumerate(src_grids)]
|
|
# per-ref placement, same convention as ref_boost/ref_boost_a: last ref = subject (ref_position),
|
|
# earlier refs = scene (ref_position_a)
|
|
ref_ids = _place_refs(ref_ids, [ref_position_a] * (len(ref_ids) - 1) + [ref_position], h_, w_)
|
|
pos = torch.cat([
|
|
torch.zeros(bs, txtlen, 3, device=device, dtype=torch.float32)] # text @ 0
|
|
+ ref_ids
|
|
+ [_imgids(bs, 0, h_, w_, device)], # target frame=0
|
|
dim=1)
|
|
freqs = m.pe_embedder(pos)
|
|
|
|
attn_bias = None
|
|
if ref_boost != 1.0 or ref_boost_a != 1.0:
|
|
# last ref = subject (single-ref: the only ref); earlier refs (scene) get ref_boost_a
|
|
boosts = [ref_boost_a] * (len(src_imgs) - 1) + [ref_boost]
|
|
attn_bias = _ref_attn_bias(boosts, ref_boost_mask, txtlen,
|
|
[si.shape[1] for si in src_imgs], tgtlen,
|
|
src_grids, combined.device, combined.dtype)
|
|
|
|
for block in m.blocks:
|
|
combined = block(combined, tvec, freqs, attn_bias, transformer_options=transformer_options)
|
|
|
|
final = m.last(combined, t)
|
|
out = final[:, txtlen + srclen: txtlen + srclen + tgtlen, :] # target tokens only
|
|
out = rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)",
|
|
h=h_, w=w_, ph=patch, pw=patch, c=m.channels)
|
|
out = out[:, :, :H_orig, :W_orig]
|
|
if temporal:
|
|
out = out.reshape(b5, t5, m.channels, H_orig, W_orig).movedim(1, 2)
|
|
return out
|
|
|
|
|
|
class Krea2EditModelPatch:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {
|
|
"model": ("MODEL",),
|
|
"source_latent": ("LATENT",),
|
|
}, "optional": {
|
|
"source_latent_b": ("LATENT", {"tooltip": "2nd reference (subject photo) for multi-ref LoRAs -> RoPE frame=2, training-matched order: scene first, subject second"}),
|
|
"ref_boost": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1000.0, "step": 0.01, "round": 0.001,
|
|
"tooltip": "reference-fidelity dial: multiplies target->reference attention. Applies to the LAST ref (= the subject in two-ref workflows, the only ref in single-ref). 1.0 = off, >1 pulls harder toward the reference's appearance, <1 loosens. Optimal value is model-specific"}),
|
|
"ref_boost_a": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1000.0, "step": 0.01, "round": 0.001,
|
|
"tooltip": "same dial for the FIRST ref (= the scene in two-ref workflows). No effect in single-ref workflows. 1.0 = off"}),
|
|
"fit_mode": (["fit", "crop (legacy)"], {"default": "fit",
|
|
"tooltip": "how an image source fits a mismatched output aspect ratio (needs vae + source_image connected): fit = resample the source to the target grid at a centered offset — matches how this model was trained (default, use this); crop (legacy) = center-crop to the target AR then resize (v1/v1.1 geometry, only for older weights)"}),
|
|
"ref_boost_mask": ("MASK", {"tooltip": "optional region on the (last) reference to boost, e.g. the face; empty = whole reference"}),
|
|
"vae": ("VAE", {"tooltip": "RECOMMENDED with source_image: enables the blur-proof pixel-space path (crop+resize in pixels, encode internally) — immune to input/output resolution mismatches"}),
|
|
"source_image": ("IMAGE", {"tooltip": "source as IMAGE (with vae connected): overrides source_latent with exact pixel-space fitting — fixes blurry results from mismatched resolutions"}),
|
|
"source_image_b": ("IMAGE", {"tooltip": "2nd reference as IMAGE (with vae)"}),
|
|
# declared last on purpose: appending keeps every existing socket index put,
|
|
# so workflows saved before this input still reload with their links intact.
|
|
"target_latent": ("LATENT", {"tooltip": "RECOMMENDED with vae + source_image: wire the SAME latent you feed KSampler.latent_image. Lets the node VAE-encode the source here, before sampling starts, instead of on the first step — otherwise the VAE is pulled onto the GPU mid-sampling and can evict part of the diffusion model, slowing every remaining step on VRAM-tight setups"}),
|
|
# appended last, same reason as target_latent
|
|
"ref_position": (REF_POSITIONS, {"default": REF_POSITIONS[0],
|
|
"tooltip": "where the LAST reference (the person in two-image workflows, or the only image) sits in position space. aligned (trained) = pinned onto the output layout, how the models were trained (default). beside = placed next to the canvas: identity still flows but the layout pin is gone, so the person can be restaged more freely. bag of features = experimental, content without layout"}),
|
|
"ref_position_a": (REF_POSITIONS, {"default": REF_POSITIONS[0],
|
|
"tooltip": "the same choice for the FIRST reference (the scene in two-image workflows); no effect with a single image. For putting a person into a scene: scene aligned + person beside"}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "patch"
|
|
CATEGORY = "krea2edit"
|
|
DESCRIPTION = "Adds the krea2_edit in-context source-preservation path (source latent as frame=1 tokens) to a Krea2 model."
|
|
|
|
def patch(self, model, source_latent, source_latent_b=None, ref_boost=1.0, ref_boost_a=1.0,
|
|
ref_boost_mask=None, vae=None, source_image=None,
|
|
source_image_b=None, fit_mode="fit", target_latent=None,
|
|
ref_position="aligned (trained)", ref_position_a="aligned (trained)", **_future):
|
|
if _future:
|
|
print(f"[krea2edit] WARNING: workflow provides inputs this node version does not "
|
|
f"know ({', '.join(sorted(_future))}). The workflow is newer than the "
|
|
f"installed node pack. Update comfyui-krea2edit (Manager -> Update, or git "
|
|
f"pull) and restart ComfyUI. Continuing without them.", flush=True)
|
|
if not (ref_position.startswith("aligned") and ref_position_a.startswith("aligned")):
|
|
print(f"[krea2edit] reference placement: last ref={ref_position!r}, "
|
|
f"first ref={ref_position_a!r}", flush=True)
|
|
m = model.clone()
|
|
# The target latent reaches the diffusion model already scaled (process_latent_in);
|
|
# scale the source(s) the same way so all share one latent space.
|
|
src_samples = model.model.process_latent_in(source_latent["samples"])
|
|
if source_latent_b is not None:
|
|
src_samples = [src_samples, model.model.process_latent_in(source_latent_b["samples"])]
|
|
|
|
px_cache = {} # pixel-path encoded sources, keyed per target resolution
|
|
mm = model.model # for process_latent_in on the pixel path
|
|
state = {"announced": False}
|
|
|
|
if fit_mode == "fit" and (vae is None or source_image is None):
|
|
print(f"[krea2edit] WARNING: fit_mode='fit' has NO EFFECT — it needs both "
|
|
f"'vae' and 'source_image' connected (the pixel path). Falling back to the "
|
|
f"latent crop path.", flush=True)
|
|
|
|
# Pre-encode OUTSIDE the sampling window. vae.encode -> load_models_gpu ->
|
|
# free_memory(keep_loaded=[]), which partially unloads whatever is resident —
|
|
# including the diffusion model, if the first call lands inside the sampler.
|
|
# Nothing re-expands it (sampler_helpers loads once, before the loop), so the
|
|
# rest of the run streams weights from CPU every step. Running the encode here,
|
|
# at node-execution time, restores the ordinary VAEEncode -> KSampler order
|
|
# where the sampler evicts the VAE instead of the reverse.
|
|
primed = None
|
|
if vae is not None and source_image is not None:
|
|
if target_latent is not None:
|
|
Hh, Ww = target_latent["samples"].shape[-2], target_latent["samples"].shape[-1]
|
|
print(f"[krea2edit] pre-encoding sources at target {Hh * 8}x{Ww * 8}px "
|
|
f"(before sampling, fit_mode={fit_mode})", flush=True)
|
|
_fit_encode_image(source_image, vae, Hh, Ww, px_cache, ("a", Hh, Ww), fit_mode)
|
|
if source_image_b is not None:
|
|
_fit_encode_image(source_image_b, vae, Hh, Ww, px_cache, ("b", Hh, Ww), fit_mode)
|
|
primed = (Hh, Ww)
|
|
else:
|
|
print("[krea2edit] NOTE: connect 'target_latent' (the same latent that feeds "
|
|
"KSampler.latent_image) to pre-encode the source here instead of on the "
|
|
"first sampling step. Without it the VAE is loaded mid-sampling and can "
|
|
"evict part of the diffusion model, slowing every remaining step.",
|
|
flush=True)
|
|
|
|
def wrapper(executor, x, timesteps, context, *wargs, **kwargs):
|
|
# ComfyUI signature drift (2026-07-19, commit c9602625 adds ref_latents):
|
|
# old: execute(x, t, ctx, attention_mask, transformer_options)
|
|
# new: execute(x, t, ctx, attention_mask, ref_latents, transformer_options)
|
|
# Accept both: transformer_options is the trailing dict; any native
|
|
# ref_latents are ignored (this patch supplies its own source path).
|
|
transformer_options = kwargs.pop("transformer_options", None)
|
|
if transformer_options is None:
|
|
transformer_options = {}
|
|
for a in reversed(wargs):
|
|
if isinstance(a, dict):
|
|
transformer_options = a
|
|
break
|
|
dm = executor.class_obj # the SingleStreamDiT instance
|
|
src = src_samples
|
|
if vae is not None and source_image is not None:
|
|
xx = _to_4d(x)
|
|
Hh, Ww = xx.shape[-2], xx.shape[-1]
|
|
# not `if not px_cache` — the cache is already primed when target_latent
|
|
# is wired, so an explicit one-shot flag is needed to announce once.
|
|
if not state["announced"]:
|
|
state["announced"] = True
|
|
print(f"[krea2edit] pixel path ACTIVE (fit_mode={fit_mode})", flush=True)
|
|
if primed is not None and primed != (Hh, Ww):
|
|
print(f"[krea2edit] WARNING: 'target_latent' is {primed[0] * 8}x{primed[1] * 8}px but "
|
|
f"sampling is at {Hh * 8}x{Ww * 8}px — the pre-encode is unused and the VAE "
|
|
f"will run mid-sampling. Wire the SAME latent that feeds KSampler.",
|
|
flush=True)
|
|
lat = mm.process_latent_in(_fit_encode_image(source_image, vae, Hh, Ww, px_cache, ("a", Hh, Ww), fit_mode))
|
|
if source_image_b is not None:
|
|
lat = [lat, mm.process_latent_in(_fit_encode_image(source_image_b, vae, Hh, Ww, px_cache, ("b", Hh, Ww), fit_mode))]
|
|
src = lat
|
|
v = krea2_edit_forward(dm, x, timesteps, context, src, transformer_options,
|
|
ref_boost=ref_boost, ref_boost_a=ref_boost_a,
|
|
ref_boost_mask=ref_boost_mask,
|
|
ref_native=(fit_mode == "fit" and vae is not None
|
|
and source_image is not None),
|
|
pos_mode=("stride1" if fit_mode == "fit" else "anchor"),
|
|
ref_position=ref_position, ref_position_a=ref_position_a)
|
|
return v
|
|
|
|
to = m.model_options.setdefault("transformer_options", {})
|
|
comfy.patcher_extension.add_wrapper_with_key(
|
|
comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, "krea2_edit", wrapper, to
|
|
)
|
|
return (m,)
|
|
|
|
|
|
class Krea2EditGroundedEncode:
|
|
"""Image-grounded instruction encode — the SEMANTIC path of krea2_edit.
|
|
|
|
Training always encodes the instruction TOGETHER with the source image through
|
|
Qwen3-VL (user turn = <vision tokens: source> + instruction) and taps 12 layers.
|
|
Stock CLIPTextEncode is text-only, so inference was running with the grounding
|
|
half of the recipe missing (the VAE source tokens carry appearance; THIS carries
|
|
scene semantics: "the man on the left", "the sign in the back").
|
|
|
|
Requires a qwen3vl TE checkpoint WITH the vision tower (all local ones have it).
|
|
grounding_px caps the longest side fed to the VLM — the 2026-07-02 LoRA trained
|
|
with 384-768px jitter, so 640-768 is in-distribution; 0 = native res (the jitter
|
|
training makes that tolerable too). For CFG, ground the NEGATIVE too: second node,
|
|
empty prompt, same image (matches training's unconditional).
|
|
"""
|
|
DEFAULT_SYSTEM = (
|
|
"Describe the image by detailing the color, shape, size, "
|
|
"texture, quantity, text, spatial relationships of the objects and background:"
|
|
)
|
|
|
|
KREA2_EDIT_TEMPLATE = (
|
|
"<|im_start|>system\nDescribe the image by detailing the color, shape, size, "
|
|
"texture, quantity, text, spatial relationships of the objects and background:"
|
|
"<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
|
|
"{}<|im_end|>\n<|im_start|>assistant\n"
|
|
)
|
|
|
|
@classmethod
|
|
def _template(cls, nimg, system_prompt=""):
|
|
sp = system_prompt.strip() or cls.DEFAULT_SYSTEM
|
|
vis = "<|vision_start|><|image_pad|><|vision_end|>" * nimg
|
|
return ("<|im_start|>system\n" + sp + "<|im_end|>\n<|im_start|>user\n"
|
|
+ vis + "{}<|im_end|>\n<|im_start|>assistant\n")
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP",),
|
|
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE",),
|
|
"image_b": ("IMAGE", {"tooltip": "2nd reference (subject) for multi-ref LoRAs; vision blocks in training order: scene, subject"}),
|
|
"grounding_px": ("INT", {"default": 768, "min": 0, "max": 4096, "step": 64,
|
|
"tooltip": "cap longest side fed to Qwen3-VL; 0 = native"}),
|
|
"system_prompt": ("STRING", {"multiline": True, "default": "",
|
|
"tooltip": "advanced (optional): override the grounding system prompt (empty = training default). Steers what the vision encoder attends to, e.g. facial identity detail."}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "encode"
|
|
CATEGORY = "krea2edit"
|
|
DESCRIPTION = "Encodes the edit instruction grounded on the source image (training-matched semantic path)."
|
|
|
|
KREA2_EDIT_TEMPLATE_2REF = (
|
|
"<|im_start|>system\nDescribe the image by detailing the color, shape, size, "
|
|
"texture, quantity, text, spatial relationships of the objects and background:"
|
|
"<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
|
|
"<|vision_start|><|image_pad|><|vision_end|>"
|
|
"{}<|im_end|>\n<|im_start|>assistant\n"
|
|
)
|
|
|
|
def _prep(self, image, grounding_px):
|
|
samples = image.movedim(-1, 1) # B,H,W,C -> B,C,H,W
|
|
h, w = samples.shape[2], samples.shape[3]
|
|
if grounding_px and max(h, w) > grounding_px:
|
|
s = grounding_px / max(h, w)
|
|
samples = comfy.utils.common_upscale(samples, round(w * s), round(h * s), "area", "disabled")
|
|
return samples.movedim(1, -1)[:, :, :, :3]
|
|
|
|
def encode(self, clip, prompt, image=None, image_b=None, grounding_px=768, system_prompt="", **_future):
|
|
if _future:
|
|
print(f"[krea2edit] WARNING: workflow provides inputs this node version does not "
|
|
f"know ({', '.join(sorted(_future))}). Update comfyui-krea2edit and restart "
|
|
f"ComfyUI. Continuing without them.", flush=True)
|
|
if image is None: # text-only fallback = old behavior
|
|
tokens = clip.tokenize(prompt)
|
|
return (clip.encode_from_tokens_scheduled(tokens),)
|
|
imgs = [self._prep(image, grounding_px)]
|
|
if image_b is not None:
|
|
imgs.append(self._prep(image_b, grounding_px))
|
|
template = self._template(len(imgs), system_prompt)
|
|
tokens = clip.tokenize(prompt, images=imgs, llama_template=template)
|
|
return (clip.encode_from_tokens_scheduled(tokens),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Krea2EditModelPatch": Krea2EditModelPatch,
|
|
"Krea2EditGroundedEncode": Krea2EditGroundedEncode,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Krea2EditModelPatch": "Krea2 Edit (source patch)",
|
|
"Krea2EditGroundedEncode": "Krea2 Edit (grounded encode)",
|
|
}
|
|
|
|
|
|
def _pack_version():
|
|
# single source of truth = pyproject.toml, so this never drifts from the release
|
|
try:
|
|
import os, re
|
|
p = os.path.join(os.path.dirname(__file__), "pyproject.toml")
|
|
with open(p) as f:
|
|
m = re.search(r'^version\s*=\s*"([^"]+)"', f.read(), re.M)
|
|
return m.group(1) if m else "unknown"
|
|
except Exception:
|
|
return "unknown"
|
|
|
|
|
|
print(f"[krea2edit] nodes v{_pack_version()} loaded", flush=True)
|