506 lines
17 KiB
Python
506 lines
17 KiB
Python
"""
|
|
Krea2 MultiRef custom nodes for ComfyUI.
|
|
Adapted from Ostris's implementation, with RoPE coordinate fix and
|
|
training-specific adjustments.
|
|
|
|
Training-specific details:
|
|
- Reference tag format: <reference_N><|vision_start|><|image_pad|><|vision_end|></reference_N>
|
|
- VLM max pixels: 512*512 = 262144 (patched from 384*384 during training)
|
|
- VAE ref latent max pixels: 1024*1024, snapped to /16
|
|
"""
|
|
|
|
import math
|
|
|
|
import torch
|
|
from einops import rearrange
|
|
|
|
import comfy.conds
|
|
import comfy.ldm.common_dit
|
|
import comfy.utils
|
|
import node_helpers
|
|
from comfy.ldm.flux.layers import timestep_embedding
|
|
|
|
# Training-patched values
|
|
VLM_MAX_PIXELS = 512 * 512 # patched from 384*384
|
|
REF_LATENT_MAX_PIXELS = 1024 * 1024
|
|
REF_SNAP = 16 # VAE f8 * patch 2
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Image resizing helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _fit_area(samples, max_pixels, snap=1):
|
|
"""
|
|
Downscale (never upscale) a (B, C, H, W) tensor to fit max_pixels,
|
|
keeping aspect ratio, snapping dims to snap.
|
|
"""
|
|
h, w = samples.shape[2], samples.shape[3]
|
|
scale = min(1.0, math.sqrt(max_pixels / (w * h)))
|
|
nw = max(round(w * scale / snap) * snap, snap)
|
|
nh = max(round(h * scale / snap) * snap, snap)
|
|
if (nh, nw) == (h, w):
|
|
return samples
|
|
return comfy.utils.common_upscale(samples, nw, nh, "area", "disabled")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Text + vision encode node
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class Krea2ReferenceTextEncode:
|
|
"""
|
|
Encodes prompt + reference images through the Krea2 Qwen3-VL text encoder.
|
|
|
|
Reference images are injected into the prompt using the tag format
|
|
that was used during training:
|
|
<reference_N><|vision_start|><|image_pad|><|vision_end|></reference_N>
|
|
|
|
The VLM sees a downscaled version (max 512*512 px) for semantic understanding.
|
|
When a VAE is connected, full-resolution reference latents (max 1MP, snapped
|
|
to /16) are also attached for the model patch to use.
|
|
|
|
This node replaces the standard CLIPTextEncode when using Krea2 edit LoRAs.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP",),
|
|
"prompt": ("STRING", {
|
|
"multiline": True,
|
|
"dynamicPrompts": True,
|
|
}),
|
|
},
|
|
"optional": {
|
|
"vae": ("VAE",),
|
|
"image1": ("IMAGE",),
|
|
"image2": ("IMAGE",),
|
|
"image3": ("IMAGE",),
|
|
"image4": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "encode"
|
|
CATEGORY = "krea2/editing"
|
|
DESCRIPTION = (
|
|
"Encode prompt + reference images for a Krea2 edit LoRA. "
|
|
"Images are fed to the Qwen3-VL text encoder via training-matching "
|
|
"reference tags. Connect a VAE to also attach reference latents for "
|
|
"the Krea2ModelPatcher node."
|
|
)
|
|
|
|
def encode(self, clip, prompt,
|
|
vae=None, image1=None, image2=None, image3=None, image4=None):
|
|
images_vl = []
|
|
ref_latents = []
|
|
image_prompt = ""
|
|
|
|
for image in (image1, image2, image3, image4):
|
|
if image is None:
|
|
continue
|
|
|
|
n = len(images_vl) + 1
|
|
|
|
# Downscale for VLM (semantic understanding)
|
|
samples_hw = image.movedim(-1, 1) # (B, C, H, W)
|
|
vl_img = _fit_area(samples_hw, VLM_MAX_PIXELS).movedim(1, -1)
|
|
images_vl.append(vl_img)
|
|
|
|
# Build reference tag matching training format
|
|
image_prompt += (
|
|
f"<reference_{n}>"
|
|
f"<|vision_start|><|image_pad|><|vision_end|>"
|
|
f"</reference_{n}>"
|
|
)
|
|
|
|
# Encode reference latent for model patcher
|
|
if vae is not None:
|
|
s = _fit_area(samples_hw, REF_LATENT_MAX_PIXELS, snap=REF_SNAP)
|
|
ref_latents.append(
|
|
vae.encode(s.movedim(1, -1)[:, :, :, :3])
|
|
)
|
|
|
|
# Krea2 uses its own template (KREA2_TEMPLATE) — try to import it,
|
|
# fall back to passing None which uses the model's default.
|
|
try:
|
|
from comfy.text_encoders.krea2 import KREA2_TEMPLATE
|
|
llama_template = KREA2_TEMPLATE
|
|
except ImportError:
|
|
llama_template = None
|
|
|
|
tokens = clip.tokenize(
|
|
image_prompt + prompt,
|
|
images=images_vl if images_vl else None,
|
|
llama_template=llama_template,
|
|
)
|
|
conditioning = clip.encode_from_tokens_scheduled(tokens)
|
|
|
|
if ref_latents:
|
|
conditioning = node_helpers.conditioning_set_values(
|
|
conditioning,
|
|
{"reference_latents": ref_latents},
|
|
append=True,
|
|
)
|
|
|
|
return (conditioning,)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Conditioning helper (for the simple latent-only encoder)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _append_ref_to_cond(conditioning, ref_tensor, strength):
|
|
out = []
|
|
scaled = ref_tensor * strength
|
|
for embedding, meta in conditioning:
|
|
meta = dict(meta)
|
|
refs = list(meta.get("reference_latents") or [])
|
|
refs.append(scaled)
|
|
meta["reference_latents"] = refs
|
|
out.append((embedding, meta))
|
|
return out
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Reference packing — Kontext-style RoPE layout
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _pack_refs(dit, ref_latents, bs, device, dtype, target_w):
|
|
"""
|
|
Patchify reference latents into tokens + RoPE positions.
|
|
|
|
RoPE layout (same axis-0=0 as target, offset in x):
|
|
target: axis-0=0, y in [0, h_), x in [0, target_w)
|
|
ref 0: axis-0=0, y in [0, rh0), x in [target_w, target_w+rw0)
|
|
ref 1: axis-0=0, y in [0, rh1), x in [target_w+rw0, ...)
|
|
"""
|
|
patch = dit.patch
|
|
ref_tokens = []
|
|
ref_pos = []
|
|
x_offset = target_w
|
|
|
|
for i, ref in enumerate(ref_latents):
|
|
if ref.ndim == 5:
|
|
rb, rc, rt, rh5, rw5 = ref.shape
|
|
ref = ref.reshape(rb * rt, rc, rh5, rw5)
|
|
|
|
ref = comfy.ldm.common_dit.pad_to_patch_size(
|
|
ref.to(device, dtype), (patch, patch)
|
|
)
|
|
ref = comfy.utils.repeat_to_batch_size(ref, bs)
|
|
|
|
rh = ref.shape[-2] // patch
|
|
rw = ref.shape[-1] // patch
|
|
|
|
ref_tokens.append(
|
|
rearrange(ref, "b c (h ph) (w pw) -> b (h w) (c ph pw)",
|
|
ph=patch, pw=patch)
|
|
)
|
|
|
|
rid = torch.zeros(rh, rw, 3, device=device, dtype=torch.float32)
|
|
rid[..., 0] = 0.0
|
|
rid[..., 1] = torch.arange(rh, device=device,
|
|
dtype=torch.float32)[:, None]
|
|
rid[..., 2] = (
|
|
torch.arange(rw, device=device, dtype=torch.float32)[None, :]
|
|
+ x_offset
|
|
)
|
|
ref_pos.append(rid.reshape(1, rh * rw, 3).repeat(bs, 1, 1))
|
|
x_offset += rw
|
|
|
|
return torch.cat(ref_tokens, dim=1), torch.cat(ref_pos, dim=1)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Per-block forward with split modulation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _block_ref_forward(block, x, vec, refvec, split, freqs, transformer_options):
|
|
"""
|
|
SingleStreamBlock forward with per-span AdaLN:
|
|
[:split] text + noisy image -> vec (real timestep)
|
|
[split:] clean references -> refvec (timestep=0)
|
|
"""
|
|
m = block.mod(vec)
|
|
r = block.mod(refvec)
|
|
|
|
def mod(h, scale, shift):
|
|
return torch.cat(
|
|
(
|
|
(1 + m[scale]) * h[:, :split] + m[shift],
|
|
(1 + r[scale]) * h[:, split:] + r[shift],
|
|
),
|
|
dim=1,
|
|
)
|
|
|
|
def gate(h, g):
|
|
return torch.cat(
|
|
(m[g] * h[:, :split], r[g] * h[:, split:]),
|
|
dim=1,
|
|
)
|
|
|
|
x = x + gate(
|
|
block.attn(
|
|
mod(block.prenorm(x), 0, 1),
|
|
freqs, None,
|
|
transformer_options=transformer_options,
|
|
),
|
|
2,
|
|
)
|
|
x = x + gate(block.mlp(mod(block.postnorm(x), 3, 4)), 5)
|
|
return x
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Full forward with reference latents
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _forward_with_refs(dit, x, timesteps, context, ref_latents,
|
|
transformer_options):
|
|
temporal = x.ndim == 5
|
|
if temporal:
|
|
b5, c5, t5, h5, w5 = x.shape
|
|
x = x.reshape(b5 * t5, c5, h5, w5)
|
|
|
|
bs, c, H_orig, W_orig = x.shape
|
|
patch = dit.patch
|
|
device = x.device
|
|
|
|
x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch, patch))
|
|
H, W = x.shape[-2], x.shape[-1]
|
|
h_, w_ = H // patch, W // patch
|
|
|
|
context = dit._unpack_context(context)
|
|
|
|
img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)",
|
|
ph=patch, pw=patch)
|
|
|
|
reftok, refpos = _pack_refs(dit, ref_latents, bs, device, x.dtype,
|
|
target_w=w_)
|
|
reflen = reftok.shape[1]
|
|
|
|
img = dit.first(torch.cat((img, reftok), dim=1))
|
|
|
|
t = dit.tmlp(timestep_embedding(timesteps, dit.tdim)
|
|
.unsqueeze(1).to(img.dtype))
|
|
tvec = dit.tproj(t)
|
|
t0 = dit.tmlp(timestep_embedding(torch.zeros_like(timesteps), dit.tdim)
|
|
.unsqueeze(1).to(img.dtype))
|
|
tvec0 = dit.tproj(t0)
|
|
|
|
context = dit.txtfusion(context, mask=None,
|
|
transformer_options=transformer_options)
|
|
context = dit.txtmlp(context)
|
|
|
|
txtlen = context.shape[1]
|
|
imglen = img.shape[1]
|
|
|
|
combined = torch.cat((context, img), dim=1)
|
|
split = txtlen + imglen - reflen
|
|
|
|
txtpos = torch.zeros(bs, txtlen, 3, device=device, dtype=torch.float32)
|
|
imgids = torch.zeros(h_, w_, 3, device=device, dtype=torch.float32)
|
|
imgids[..., 1] = torch.arange(h_, device=device,
|
|
dtype=torch.float32)[:, None]
|
|
imgids[..., 2] = torch.arange(w_, device=device,
|
|
dtype=torch.float32)[None, :]
|
|
imgpos = imgids.reshape(1, h_ * w_, 3).repeat(bs, 1, 1)
|
|
|
|
pos = torch.cat((txtpos, imgpos, refpos), dim=1)
|
|
freqs = dit.pe_embedder(pos)
|
|
|
|
for block in dit.blocks:
|
|
combined = _block_ref_forward(
|
|
block, combined, tvec, tvec0, split, freqs, transformer_options
|
|
)
|
|
|
|
final = dit.last(combined, t)
|
|
out = final[:, txtlen:split, :]
|
|
out = rearrange(
|
|
out,
|
|
"b (h w) (c ph pw) -> b c (h ph) (w pw)",
|
|
h=h_, w=w_, ph=patch, pw=patch, c=dit.channels,
|
|
)
|
|
out = out[:, :, :H_orig, :W_orig]
|
|
|
|
if temporal:
|
|
out = out.reshape(b5, t5, dit.channels, H_orig, W_orig).movedim(1, 2)
|
|
return out
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Model patcher
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class Krea2ModelPatcher:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"model": ("MODEL",)}}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
CATEGORY = "krea2/editing"
|
|
FUNCTION = "patch"
|
|
|
|
def patch(self, model):
|
|
patched = model.clone()
|
|
base_model = patched.model
|
|
dit = patched.get_model_object("diffusion_model")
|
|
|
|
orig_forward = dit.forward
|
|
orig_extra_conds = base_model.extra_conds
|
|
orig_extra_conds_shapes = base_model.extra_conds_shapes
|
|
|
|
def extra_conds(**kwargs):
|
|
out = orig_extra_conds(**kwargs)
|
|
refs = kwargs.get("reference_latents")
|
|
if refs is not None:
|
|
out["ref_latents"] = comfy.conds.CONDList(
|
|
[base_model.process_latent_in(lat) for lat in refs]
|
|
)
|
|
return out
|
|
|
|
def extra_conds_shapes(**kwargs):
|
|
out = orig_extra_conds_shapes(**kwargs)
|
|
refs = kwargs.get("reference_latents")
|
|
if refs is not None:
|
|
out["ref_latents"] = [
|
|
1, 16,
|
|
sum(math.prod(r.size()) for r in refs) // 16,
|
|
]
|
|
return out
|
|
|
|
def forward(
|
|
x,
|
|
timesteps,
|
|
context,
|
|
attention_mask=None,
|
|
transformer_options={},
|
|
ref_latents=None,
|
|
**kwargs,
|
|
):
|
|
if ref_latents is None or len(ref_latents) == 0:
|
|
return orig_forward(
|
|
x, timesteps, context,
|
|
attention_mask=attention_mask,
|
|
transformer_options=transformer_options,
|
|
**kwargs,
|
|
)
|
|
return _forward_with_refs(
|
|
dit, x, timesteps, context, ref_latents, transformer_options,
|
|
)
|
|
|
|
patched.add_object_patch("extra_conds", extra_conds)
|
|
patched.add_object_patch("extra_conds_shapes", extra_conds_shapes)
|
|
patched.add_object_patch("diffusion_model.forward", forward)
|
|
return (patched,)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Simple latent-only reference encoder (no VLM — use when you already have
|
|
# a text conditioning from a normal CLIPTextEncode and just want to attach
|
|
# latent refs without re-encoding text)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class Krea2ReferenceEncoder:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"positive": ("CONDITIONING",),
|
|
"base_latent": ("LATENT", {
|
|
"tooltip": "VAE-encoded image to edit. Reference slot 1."
|
|
}),
|
|
"strength": ("FLOAT", {
|
|
"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05,
|
|
}),
|
|
},
|
|
"optional": {
|
|
"reference_latent_2": ("LATENT",),
|
|
"reference_latent_3": ("LATENT",),
|
|
"reference_latent_4": ("LATENT",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
CATEGORY = "krea2/editing"
|
|
FUNCTION = "encode"
|
|
|
|
def encode(
|
|
self,
|
|
positive,
|
|
base_latent,
|
|
strength = 1.0,
|
|
reference_latent_2 = None,
|
|
reference_latent_3 = None,
|
|
reference_latent_4 = None,
|
|
):
|
|
slots = [base_latent, reference_latent_2,
|
|
reference_latent_3, reference_latent_4]
|
|
tensors = [s["samples"] for s in slots if s is not None]
|
|
|
|
new_positive = positive
|
|
for ref in tensors:
|
|
new_positive = _append_ref_to_cond(new_positive, ref, strength)
|
|
return (new_positive,)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Diagnostic
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class Krea2DiagnosticNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"model": ("MODEL",)}}
|
|
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "krea2/editing"
|
|
FUNCTION = "run"
|
|
|
|
def run(self, model):
|
|
import inspect
|
|
dit = model.get_model_object("diffusion_model")
|
|
print("\n=== DIT TYPE ===")
|
|
print(type(dit).__name__, type(dit).__module__)
|
|
print("\n=== DIT NAMED CHILDREN ===")
|
|
for name, mod in dit.named_children():
|
|
print(f" {name}: {type(mod).__name__}")
|
|
print("\n=== DIT NON-MODULE ATTRS ===")
|
|
for k, v in vars(dit).items():
|
|
if not k.startswith("_") and not isinstance(v, torch.nn.Module):
|
|
print(f" {k}: {type(v).__name__} = {repr(v)[:120]}")
|
|
print("\n=== DIT FORWARD SIGNATURE ===")
|
|
print(inspect.signature(dit.forward))
|
|
print("\n=== FIRST BLOCK ===")
|
|
if hasattr(dit, "blocks"):
|
|
block = dit.blocks[0]
|
|
print(f" {type(block).__name__}")
|
|
for n, m in block.named_children():
|
|
print(f" {n}: {type(m).__name__}")
|
|
print(f" sig: {inspect.signature(block.forward)}")
|
|
print(f"\n has _unpack_context: {hasattr(dit, '_unpack_context')}")
|
|
return ()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Registration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Krea2ModelPatcher": Krea2ModelPatcher,
|
|
"Krea2ReferenceEncoder": Krea2ReferenceEncoder,
|
|
"Krea2ReferenceTextEncode": Krea2ReferenceTextEncode,
|
|
"Krea2DiagnosticNode": Krea2DiagnosticNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Krea2ModelPatcher": "Krea2 Model Patcher",
|
|
"Krea2ReferenceEncoder": "Krea2 Reference Encoder (latent only)",
|
|
"Krea2ReferenceTextEncode": "Krea2 Reference Text Encode",
|
|
"Krea2DiagnosticNode": "Krea2 Diagnostic",
|
|
} |