API tools faq
paste
Guest User

Patch for Krea2 Edit Node

a guest
Aug 4th, 2026
1,151
0
Never
1
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 20.87 KB | None | 0 0
raw download clone embed print report
  1. """ [COPY AND PASTE THE ENTIRE CODE INTO __init__.py in the Krea 2 Edit Custom Node folder ]
  2.  
  3. ComfyUI-Krea2Edit — in-context edit forward for the Krea2 model (upgraded).
  4.  
  5. ComfyUI's native Krea2 `_forward` is text-to-image only: it builds the sequence
  6. `[text | target]`. The krea2_edit LoRA (trained in ai-toolkit) needs the *appearance
  7. path*: the VAE-encoded SOURCE latent prepended as a block of clean tokens, distinguished
  8. from the (noisy) target purely by the 3-axis RoPE frame index (source=1..N, target=0, h/w
  9. aligned). This node adds that by wrapping the model's DIFFUSION_MODEL forward and rebuilding
  10. the sequence as `[text | source(s) | target]`, keeping only the target tokens out.
  11.  
  12. Upgraded with Project-2 latent handling:
  13. - Pixel-space fit geometry (blur-proof, AR-safe) matching v1.2 training
  14. - Stride-1 centered RoPE offsets for native-fit refs
  15. - Clean multi-ref ordering (scene first, subject last)
  16. - Regional ref_boost_mask (last ref only)
  17.  
  18. Wiring: LoadImage -> VAEEncode(source) --\
  19. Krea2EditModelPatch(model, source_latent) -> KSampler
  20. UNETLoader -> LoraLoaderModelOnly -/
  21. KSampler.latent_image <- EmptySD3LatentImage (noise). Text: NATIVE krea2 CLIP + CLIPTextEncode
  22. (or Krea2EditGroundedEncode for the semantic path).
  23. """
  24. import math
  25.  
  26. import torch
  27. import torch.nn.functional as F
  28. from einops import rearrange
  29.  
  30. import comfy.patcher_extension
  31. import comfy.utils
  32. import comfy.ldm.common_dit
  33. from comfy.ldm.flux.layers import timestep_embedding
  34. from comfy.ldm.flux.math import apply_rope
  35. from comfy.ldm.modules.attention import optimized_attention_masked
  36.  
  37.  
  38. CROP_TOL = 0.08 # near-matched-AR tolerance for the "fit" geometry (v1.2 value)
  39.  
  40.  
  41. def _imgids(bs, frame, h_, w_, device):
  42. ids = torch.zeros(h_, w_, 3, device=device, dtype=torch.float32)
  43. ids[..., 0] = frame
  44. ids[..., 1] = torch.arange(h_, device=device, dtype=torch.float32)[:, None]
  45. ids[..., 2] = torch.arange(w_, device=device, dtype=torch.float32)[None, :]
  46. return ids.reshape(1, h_ * w_, 3).repeat(bs, 1, 1)
  47.  
  48.  
  49. def _imgids_offset(bs, frame, gh, gw, th, tw, device):
  50. """Stride-1 integer positions at a centered integer offset.
  51. For fit-mode refs the pixels are already resampled to target grid density,
  52. so the position grid stays stride-1 by construction."""
  53. off_h, off_w = max(0, (th - gh) // 2), max(0, (tw - gw) // 2)
  54. ids = torch.zeros(gh, gw, 3, device=device, dtype=torch.float32)
  55. ids[..., 0] = frame
  56. ids[..., 1] = (torch.arange(gh, device=device, dtype=torch.float32) + off_h)[:, None]
  57. ids[..., 2] = (torch.arange(gw, device=device, dtype=torch.float32) + off_w)[None, :]
  58. return ids.reshape(1, gh * gw, 3).repeat(bs, 1, 1)
  59.  
  60.  
  61. def _to_4d(v):
  62. """(B,C,T,H,W) -> (B*T,C,H,W); pass 4D through. Images use T=1."""
  63. if v.ndim == 5:
  64. b, c, t, h, w = v.shape
  65. return v.reshape(b * t, c, h, w)
  66. return v
  67.  
  68.  
  69. def _fit_latent(src, H, W):
  70. """Latent-space fallback: center-crop to target AR then resize.
  71. Equals plain bilinear at matched AR; crops instead of stretching at mismatch."""
  72. sh, sw = src.shape[-2:]
  73. if (sh, sw) == (H, W):
  74. return src
  75. s = max(H / sh, W / sw)
  76. ch, cw = min(sh, int(round(H / s))), min(sw, int(round(W / s)))
  77. y0, x0 = (sh - ch) // 2, (sw - cw) // 2
  78. src = src[..., y0:y0 + ch, x0:x0 + cw]
  79. return F.interpolate(src.float(), size=(H, W), mode="bilinear")
  80.  
  81.  
  82. def fit_image_to_latent(image, lat_h, lat_w, fit_mode="fit"):
  83. """Pixel-space source prep (v1.2 'fit' geometry — the blur/stretch fix).
  84. Resample the IMAGE to the target grid BEFORE VAE-encoding so latents are
  85. never resized.
  86.  
  87. fit: AR-preserving fit-inside at the target's grid density, /16 floor-snapped
  88. and capped at the target's /16 floor (training-matched). Near-matched AR
  89. fills the target grid exactly via a minimal center-crop.
  90. crop: center-crop to the target AR then resize (v1/v1.1 legacy).
  91. """
  92. px_h, px_w = lat_h * 8, lat_w * 8
  93. img = image[..., :3].movedim(-1, 1).float()
  94. ih, iw = img.shape[-2:]
  95. if fit_mode == "fit":
  96. sc = min(px_h / ih, px_w / iw)
  97. if ih * sc >= px_h * (1 - CROP_TOL) and iw * sc >= px_w * (1 - CROP_TOL):
  98. s = max(px_h / ih, px_w / iw)
  99. ch, cw = min(ih, int(round(px_h / s))), min(iw, int(round(px_w / s)))
  100. y0, x0 = (ih - ch) // 2, (iw - cw) // 2
  101. img = img[..., y0:y0 + ch, x0:x0 + cw]
  102. nh, nw = px_h, px_w
  103. else:
  104. nh = min(max(16, int(ih * sc) // 16 * 16), max(16, px_h // 16 * 16))
  105. nw = min(max(16, int(iw * sc) // 16 * 16), max(16, px_w // 16 * 16))
  106. img = F.interpolate(img, size=(nh, nw), mode="bicubic", antialias=True)
  107. else:
  108. s = max(px_h / ih, px_w / iw)
  109. ch, cw = min(ih, int(round(px_h / s))), min(iw, int(round(px_w / s)))
  110. y0, x0 = (ih - ch) // 2, (iw - cw) // 2
  111. img = img[..., y0:y0 + ch, x0:x0 + cw]
  112. img = F.interpolate(img, size=(px_h, px_w), mode="bicubic", antialias=True)
  113. return img.movedim(1, -1).clamp(0, 1)
  114.  
  115.  
  116. def _ref_attn_bias(boosts, boost_mask, txtlen, slens, tgtlen, mask_hw, device, dtype):
  117. """Additive attention-logit bias on the [text | refs... | target] sequence.
  118.  
  119. boosts: per-ref factor on target->ref attention (last entry = last ref = subject).
  120. boost_mask (ComfyUI MASK) restricts the LAST ref's boost to a region (e.g. face).
  121. """
  122. nsrc = len(slens)
  123. offs = [txtlen]
  124. for sl in slens:
  125. offs.append(offs[-1] + sl)
  126. rows0 = offs[-1]
  127. L = rows0 + tgtlen
  128. bias = torch.zeros(1, 1, L, L, device=device, dtype=dtype)
  129. for i, b in enumerate(boosts):
  130. if b == 1.0:
  131. continue
  132. off, sl = offs[i], slens[i]
  133. if boost_mask is not None and i == nsrc - 1 and mask_hw is not None and len(mask_hw) > i:
  134. mask = boost_mask[:1]
  135. if mask.ndim == 2:
  136. mask = mask[None]
  137. mask = F.interpolate(mask[None].float(), size=mask_hw[i], mode="area")[0, 0]
  138. cols = off + torch.nonzero(mask.reshape(-1) > 0.5, as_tuple=True)[0].to(device)
  139. else:
  140. cols = torch.arange(off, off + sl, device=device)
  141. bias[:, :, rows0:, cols] = math.log(max(b, 1e-4))
  142. return bias
  143.  
  144.  
  145. def krea2_edit_forward(m, x, timesteps, context, src_latent, transformer_options,
  146. ref_boost=1.0, ref_boost_a=1.0, ref_boost_mask=None,
  147. ref_native=False, pos_mode="anchor"):
  148. """Krea2 SingleStreamDiT._forward with source block(s) prepended.
  149.  
  150. m : the SingleStreamDiT (LoRA-patched at sample time)
  151. x : (B,C,H,W) or (B,C,T,H,W) noisy TARGET latent
  152. src_latent : clean SOURCE latent (VAE-encoded), 4D/5D — or a LIST of them
  153. (multi-ref: [scene, subject], frames 1..N)
  154. context : (B, seq, txtlayers*txtdim) — the 12-layer Qwen3-VL stack
  155. """
  156. patch = m.patch
  157.  
  158. temporal = x.ndim == 5
  159. if temporal:
  160. b5, c5, t5, h5, w5 = x.shape
  161. x = _to_4d(x)
  162. bs, c, H_orig, W_orig = x.shape
  163.  
  164. x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch, patch), padding_mode="replicate")
  165. H, W = x.shape[-2], x.shape[-1]
  166. h_, w_ = H // patch, W // patch
  167.  
  168. # source(s) -> list of (bs, C, H', W')
  169. src_list = src_latent if isinstance(src_latent, (list, tuple)) else [src_latent]
  170. srcs = []
  171. for sl in src_list:
  172. src = _to_4d(sl).to(x.device, x.dtype)
  173. if src.shape[0] != bs:
  174. src = src[:1].expand(bs, *src.shape[1:])
  175. # Native-fit refs that already fit inside the target keep their own grid
  176. # (stride-1 offset positions). Everything else is fitted in latent space.
  177. native = ref_native and src.shape[-2] <= H and src.shape[-1] <= W
  178. if src.shape[-2:] != (H, W) and not native:
  179. src = _fit_latent(src, H, W).to(x.dtype)
  180. srcs.append(comfy.ldm.common_dit.pad_to_patch_size(src, (patch, patch), padding_mode="replicate"))
  181. src_grids = [(s_.shape[-2] // patch, s_.shape[-1] // patch) for s_ in srcs]
  182.  
  183. context = m._unpack_context(context)
  184.  
  185. tgt_img = m.first(rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch))
  186. src_imgs = [m.first(rearrange(s_, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch))
  187. for s_ in srcs]
  188.  
  189. t = m.tmlp(timestep_embedding(timesteps, m.tdim).unsqueeze(1).to(tgt_img.dtype))
  190. tvec = m.tproj(t)
  191.  
  192. context = m.txtfusion(context, mask=None, transformer_options=transformer_options)
  193. context = m.txtmlp(context)
  194.  
  195. txtlen, tgtlen = context.shape[1], tgt_img.shape[1]
  196. srclen = sum(si.shape[1] for si in src_imgs)
  197. combined = torch.cat([context] + src_imgs + [tgt_img], dim=1)
  198.  
  199. device = combined.device
  200. if pos_mode == "stride1" and ref_native:
  201. ref_ids = [_imgids_offset(bs, i + 1, gh, gw, h_, w_, device)
  202. for i, (gh, gw) in enumerate(src_grids)]
  203. else:
  204. ref_ids = [_imgids(bs, i + 1, gh, gw, device) for i, (gh, gw) in enumerate(src_grids)]
  205.  
  206. pos = torch.cat([
  207. torch.zeros(bs, txtlen, 3, device=device, dtype=torch.float32)]
  208. + ref_ids
  209. + [_imgids(bs, 0, h_, w_, device)],
  210. dim=1)
  211. freqs = m.pe_embedder(pos)
  212.  
  213. attn_bias = None
  214. if ref_boost != 1.0 or ref_boost_a != 1.0:
  215. # last ref = subject; earlier refs (scene) get ref_boost_a
  216. boosts = [ref_boost_a] * (len(src_imgs) - 1) + [ref_boost]
  217. attn_bias = _ref_attn_bias(boosts, ref_boost_mask, txtlen,
  218. [si.shape[1] for si in src_imgs], tgtlen,
  219. src_grids, combined.device, combined.dtype)
  220.  
  221. for block in m.blocks:
  222. combined = block(combined, tvec, freqs, attn_bias, transformer_options=transformer_options)
  223.  
  224. final = m.last(combined, t)
  225. out = final[:, txtlen + srclen: txtlen + srclen + tgtlen, :]
  226. out = rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)",
  227. h=h_, w=w_, ph=patch, pw=patch, c=m.channels)
  228. out = out[:, :, :H_orig, :W_orig]
  229. if temporal:
  230. out = out.reshape(b5, t5, m.channels, H_orig, W_orig).movedim(1, 2)
  231. return out
  232.  
  233.  
  234. class Krea2EditModelPatch:
  235. @classmethod
  236. def INPUT_TYPES(cls):
  237. return {"required": {
  238. "model": ("MODEL",),
  239. "source_latent": ("LATENT",),
  240. }, "optional": {
  241. "source_latent_b": ("LATENT", {
  242. "tooltip": "2nd reference (subject photo) for multi-ref LoRAs → RoPE frame=2. "
  243. "Training-matched order: scene first, subject second."}),
  244. "ref_boost": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1000.0, "step": 0.01, "round": 0.001,
  245. "tooltip": "reference-fidelity dial: multiplies target→reference attention. "
  246. "Applies to the LAST ref (= subject in two-ref workflows). "
  247. "1.0 = off; >1 pulls harder (v1.2 author suggests 2–6)."}),
  248. "ref_boost_a": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1000.0, "step": 0.01, "round": 0.001,
  249. "tooltip": "same dial for the FIRST ref (= scene in two-ref workflows). "
  250. "No effect in single-ref. 1.0 = off."}),
  251. "ref_boost_mask": ("MASK", {
  252. "tooltip": "optional region on the (last) reference to boost, e.g. the face; "
  253. "empty = whole reference. Only affects ref_boost."}),
  254. "fit_mode": (["fit", "crop (legacy)"], {"default": "fit",
  255. "tooltip": "how an image source fits a mismatched output AR "
  256. "(needs vae + source_image): fit = v1.2 pixel-space geometry "
  257. "(blur-proof); crop = v1/v1.1 legacy center-crop + resize."}),
  258. "vae": ("VAE", {
  259. "tooltip": "RECOMMENDED with source_image: enables the blur-proof pixel-space path "
  260. "(crop/resize in pixels, encode internally)."}),
  261. "source_image": ("IMAGE", {
  262. "tooltip": "source as IMAGE (with vae): overrides source_latent with exact "
  263. "pixel-space fitting — fixes blurry results from resolution mismatch."}),
  264. "source_image_b": ("IMAGE", {
  265. "tooltip": "2nd reference as IMAGE (with vae)."}),
  266. }}
  267.  
  268. RETURN_TYPES = ("MODEL",)
  269. FUNCTION = "patch"
  270. CATEGORY = "krea2edit"
  271. DESCRIPTION = ("[MODDED] Adds the krea2_edit in-context source-preservation path "
  272. "(source latent as frame=1+ tokens) to a Krea2 model. "
  273. "Upgraded with v1.2 pixel-space fit geometry and regional boost mask.")
  274.  
  275. def patch(self, model, source_latent, source_latent_b=None, ref_boost=1.0, ref_boost_a=1.0,
  276. ref_boost_mask=None, vae=None, source_image=None,
  277. source_image_b=None, fit_mode="fit", **_future):
  278. if _future:
  279. print(f"[krea2edit] WARNING: workflow provides inputs this node version does not "
  280. f"know ({', '.join(sorted(_future))}). Update the node pack and restart "
  281. f"ComfyUI. Continuing without them.", flush=True)
  282.  
  283. m = model.clone()
  284. # Target latent reaches the diffusion model already scaled (process_latent_in);
  285. # scale the source(s) the same way so all share one latent space.
  286. src_samples = model.model.process_latent_in(source_latent["samples"])
  287. if source_latent_b is not None:
  288. src_samples = [src_samples, model.model.process_latent_in(source_latent_b["samples"])]
  289.  
  290. px_cache = {} # pixel-path encoded sources, keyed per target resolution
  291. mm = model.model
  292.  
  293. if fit_mode == "fit" and (vae is None or source_image is None):
  294. print(f"[krea2edit] WARNING: fit_mode='fit' has limited effect without both "
  295. f"'vae' and 'source_image' connected (pixel path). Falling back to latent "
  296. f"crop path where needed.", flush=True)
  297.  
  298. def wrapper(executor, x, timesteps, context, *wargs, **kwargs):
  299. # Signature-drift guard (ComfyUI has added/removed positional args over time).
  300. transformer_options = kwargs.pop("transformer_options", None)
  301. if transformer_options is None:
  302. transformer_options = {}
  303. for a in reversed(wargs):
  304. if isinstance(a, dict):
  305. transformer_options = a
  306. break
  307.  
  308. dm = executor.class_obj # the SingleStreamDiT instance
  309. src = src_samples
  310. use_pixel = vae is not None and source_image is not None
  311.  
  312. if use_pixel:
  313. if not px_cache:
  314. print(f"[krea2edit] pixel path ACTIVE (fit_mode={fit_mode})", flush=True)
  315. xx = _to_4d(x)
  316. Hh, Ww = xx.shape[-2], xx.shape[-1]
  317. mode = "crop" if fit_mode.startswith("crop") else "fit"
  318. key_a = ("a", Hh, Ww, mode)
  319. if key_a not in px_cache:
  320. fitted = fit_image_to_latent(source_image[:1], Hh, Ww, mode)
  321. px_cache[key_a] = mm.process_latent_in(vae.encode(fitted))
  322. lat = px_cache[key_a]
  323. if source_image_b is not None:
  324. key_b = ("b", Hh, Ww, mode)
  325. if key_b not in px_cache:
  326. fitted_b = fit_image_to_latent(source_image_b[:1], Hh, Ww, mode)
  327. px_cache[key_b] = mm.process_latent_in(vae.encode(fitted_b))
  328. lat = [lat, px_cache[key_b]]
  329. src = lat
  330.  
  331. v = krea2_edit_forward(
  332. dm, x, timesteps, context, src, transformer_options,
  333. ref_boost=ref_boost, ref_boost_a=ref_boost_a,
  334. ref_boost_mask=ref_boost_mask,
  335. ref_native=use_pixel and not fit_mode.startswith("crop"),
  336. pos_mode=("stride1" if (use_pixel and not fit_mode.startswith("crop")) else "anchor"),
  337. )
  338. return v
  339.  
  340. to = m.model_options.setdefault("transformer_options", {})
  341. comfy.patcher_extension.add_wrapper_with_key(
  342. comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, "krea2_edit", wrapper, to
  343. )
  344. return (m,)
  345.  
  346.  
  347. class Krea2EditGroundedEncode:
  348. """Image-grounded instruction encode — the SEMANTIC path of krea2_edit.
  349.  
  350. Training always encodes the instruction TOGETHER with the source image through
  351. Qwen3-VL (user turn = <vision tokens: source> + instruction) and taps 12 layers.
  352. Stock CLIPTextEncode is text-only, so inference was running with the grounding
  353. half of the recipe missing (the VAE source tokens carry appearance; THIS carries
  354. scene semantics: "the man on the left", "the sign in the back").
  355.  
  356. Requires a qwen3vl TE checkpoint WITH the vision tower.
  357. grounding_px caps the longest side fed to the VLM — the identity LoRA trained
  358. with 384-768px jitter, so 640-768 is in-distribution; 0 = native res.
  359. For CFG, ground the NEGATIVE too: second node, empty prompt, same image.
  360. """
  361. DEFAULT_SYSTEM = (
  362. "Describe the image by detailing the color, shape, size, "
  363. "texture, quantity, text, spatial relationships of the objects and background:"
  364. )
  365.  
  366. @classmethod
  367. def _template(cls, nimg, system_prompt=""):
  368. sp = system_prompt.strip() or cls.DEFAULT_SYSTEM
  369. vis = "<|vision_start|><|image_pad|><|vision_end|>" * nimg
  370. return ("<|im_start|>system\n" + sp + "<|im_end|>\n<|im_start|>user\n"
  371. + vis + "{}<|im_end|>\n<|im_start|>assistant\n")
  372.  
  373. @classmethod
  374. def INPUT_TYPES(cls):
  375. return {
  376. "required": {
  377. "clip": ("CLIP",),
  378. "prompt": ("STRING", {"multiline": True, "default": ""}),
  379. },
  380. "optional": {
  381. "image": ("IMAGE",),
  382. "image_b": ("IMAGE", {
  383. "tooltip": "2nd reference (subject) for multi-ref LoRAs; "
  384. "vision blocks in training order: scene, subject"}),
  385. "grounding_px": ("INT", {"default": 768, "min": 0, "max": 4096, "step": 64,
  386. "tooltip": "cap longest side fed to Qwen3-VL; 0 = native"}),
  387. "system_prompt": ("STRING", {"multiline": True, "default": "",
  388. "tooltip": "advanced: override the grounding system prompt "
  389. "(empty = training default)."}),
  390. },
  391. }
  392.  
  393. RETURN_TYPES = ("CONDITIONING",)
  394. FUNCTION = "encode"
  395. CATEGORY = "krea2edit"
  396. DESCRIPTION = "[MODDED] Encodes the edit instruction grounded on the source image (training-matched semantic path)."
  397.  
  398. def _prep(self, image, grounding_px):
  399. samples = image.movedim(-1, 1)
  400. h, w = samples.shape[2], samples.shape[3]
  401. if grounding_px and max(h, w) > grounding_px:
  402. s = grounding_px / max(h, w)
  403. samples = comfy.utils.common_upscale(samples, round(w * s), round(h * s), "area", "disabled")
  404. return samples.movedim(1, -1)[:, :, :, :3]
  405.  
  406. def encode(self, clip, prompt, image=None, image_b=None, grounding_px=768, system_prompt="", **_future):
  407. if _future:
  408. print(f"[krea2edit] WARNING: workflow provides inputs this node version does not "
  409. f"know ({', '.join(sorted(_future))}). Update and restart ComfyUI.", flush=True)
  410. if image is None:
  411. tokens = clip.tokenize(prompt)
  412. return (clip.encode_from_tokens_scheduled(tokens),)
  413. imgs = [self._prep(image, grounding_px)]
  414. if image_b is not None:
  415. imgs.append(self._prep(image_b, grounding_px))
  416. template = self._template(len(imgs), system_prompt)
  417. tokens = clip.tokenize(prompt, images=imgs, llama_template=template)
  418. return (clip.encode_from_tokens_scheduled(tokens),)
  419.  
  420.  
  421. NODE_CLASS_MAPPINGS = {
  422. "Krea2EditModelPatch": Krea2EditModelPatch,
  423. "Krea2EditGroundedEncode": Krea2EditGroundedEncode,
  424. }
  425. NODE_DISPLAY_NAME_MAPPINGS = {
  426. "Krea2EditModelPatch": "Krea2 Edit [MODDED] (source patch)",
  427. "Krea2EditGroundedEncode": "Krea2 Edit [MODDED] (grounded encode)",
  428. }
  429.  
  430.  
  431. def _pack_version():
  432. try:
  433. import os, re
  434. p = os.path.join(os.path.dirname(__file__), "pyproject.toml")
  435. with open(p) as f:
  436. m = re.search(r'^version\s*=\s*"([^"]+)"', f.read(), re.M)
  437. return m.group(1) if m else "2.0-upgraded"
  438. except Exception:
  439. return "2.0-upgraded"
  440.  
  441.  
  442. print(f"[krea2edit] nodes v{_pack_version()} loaded (ModelPatch + v1.2 latent handling + mask)", flush=True)
  443.  
Advertisement
Comments
  • User was banned
Add Comment
Please, Sign In to add comment
Public Pastes
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy.  OK, I Understand
Not a member of Pastebin yet?
Sign Up, it unlocks many cool features!