318 lines
16 KiB
Python
318 lines
16 KiB
Python
"""Krea 2 reference-image conditioning for ComfyUI.
|
|
|
|
Krea 2 conditions on 12 tapped hidden states of Qwen3-VL-4B. Qwen3-VL is a
|
|
vision-language model, so reference images can be interleaved into the prompt
|
|
as native vision tokens and the DiT reads character/style/pose information
|
|
from the same hidden states it was trained on. The Krea 2 DiT has no
|
|
reference-latent pathway (unlike Flux Kontext / Qwen-Image-Edit), so no VAE
|
|
input is needed here.
|
|
"""
|
|
|
|
import re
|
|
|
|
import torch
|
|
|
|
import comfy.utils
|
|
import node_helpers
|
|
|
|
# Longest-side pixel budget fed to the vision tower. More pixels = more vision
|
|
# tokens = stronger, more literal influence from the reference.
|
|
DETAIL = {"low": 256, "normal": 512, "high": 1024, "max": 1280}
|
|
DETAIL_TIP = ("Longest side this reference is scaled to before the vision encoder. Higher = more "
|
|
"vision tokens = stronger, more literal reference. Use high for the character/subject "
|
|
"and low/normal for pose or style hints to avoid side-by-side collages.")
|
|
ROLE_TIP = ("content: the image is used literally (character, subject, pose). "
|
|
"style: its patches are shuffled before the vision encoder, destroying the subject's "
|
|
"shape so only palette/linework/texture transfer — prevents the style image's subject "
|
|
"from appearing in the output.")
|
|
SPLIT_TIP = ("Fraction of the denoise steps that see TEXT-ONLY conditioning before the references "
|
|
"kick in. The first steps fix the layout, so 0.25 strongly prevents side-by-side "
|
|
"collages of the references. 0 disables.")
|
|
|
|
# Vision tower consumes 32px merged patches; align shuffle tiles to that grid.
|
|
VISION_PATCH = 32
|
|
|
|
|
|
MASK_TIP = ("Optional subject mask (e.g. from RemoveBackground). The reference is cropped to the "
|
|
"mask's bounding box and the background is whited out before the vision encoder, so "
|
|
"messy references (scenes, watermarks, other characters) only contribute the subject.")
|
|
BG_TIP = ("Optional: connect Load Background Removal Model (birefnet) and every content-role "
|
|
"reference is auto-cleaned — cropped to its subject with the background whited out. "
|
|
"Per-image masks override it; style-role images are left untouched (their background "
|
|
"often carries the style).")
|
|
|
|
|
|
def isolate_subject(s, mask, pad=0.05):
|
|
"""Crop NCHW image to the mask bbox (+pad) and white out the background."""
|
|
if mask.shape[-2:] != s.shape[-2:]:
|
|
mask = comfy.utils.common_upscale(mask.reshape(-1, 1, *mask.shape[-2:]), s.shape[3], s.shape[2], "bilinear", "disabled").reshape(-1, *s.shape[-2:])
|
|
m = mask[0] > 0.5
|
|
if not m.any():
|
|
return s
|
|
ys, xs = torch.where(m)
|
|
y0, y1 = ys.min().item(), ys.max().item() + 1
|
|
x0, x1 = xs.min().item(), xs.max().item() + 1
|
|
ph, pw = round((y1 - y0) * pad), round((x1 - x0) * pad)
|
|
y0, y1 = max(0, y0 - ph), min(s.shape[2], y1 + ph)
|
|
x0, x1 = max(0, x0 - pw), min(s.shape[3], x1 + pw)
|
|
mm = mask[:1].unsqueeze(1).to(s.dtype)
|
|
s = s * mm + (1.0 - mm)
|
|
return s[:, :, y0:y1, x0:x1]
|
|
|
|
|
|
def shuffle_patches(s, tile):
|
|
"""Randomly permute tile positions (NCHW). Deterministic. Destroys global
|
|
shape while keeping patch-local style statistics."""
|
|
n, c, h, w = s.shape
|
|
gh, gw = max(1, h // tile), max(1, w // tile)
|
|
s = s[:, :, :gh * tile, :gw * tile]
|
|
tiles = s.reshape(n, c, gh, tile, gw, tile).permute(0, 2, 4, 1, 3, 5).reshape(n, gh * gw, c, tile, tile)
|
|
perm = torch.randperm(gh * gw, generator=torch.Generator().manual_seed(0))
|
|
tiles = tiles[:, perm]
|
|
return tiles.reshape(n, gh, gw, c, tile, tile).permute(0, 3, 1, 4, 2, 5).reshape(n, c, gh * tile, gw * tile)
|
|
|
|
# Same edit template Qwen-Image-Edit uses; Krea 2 was trained on Qwen3-VL
|
|
# hidden states over this chat format, and the Krea2 text encoder strips the
|
|
# system + user-opening prefix before conditioning.
|
|
EDIT_TEMPLATE = (
|
|
"<|im_start|>system\n"
|
|
"Describe the key features of the input image (color, shape, size, texture, "
|
|
"objects, background), then explain how the user's text instruction should "
|
|
"alter or modify the image. Generate a new image that meets the user's "
|
|
"requirements while maintaining consistency with the original input where "
|
|
"appropriate.<|im_end|>\n"
|
|
"<|im_start|>user\n{}<|im_end|>\n"
|
|
"<|im_start|>assistant\n"
|
|
)
|
|
|
|
# "image1", "image 2", "img3", "imagen 1" -> "Picture N" (what the encoder sees)
|
|
_IMG_WORD = re.compile(r"\b(?:image|img|imagen)[ _]?([1-4])\b", re.IGNORECASE)
|
|
|
|
|
|
# Auto-added when the prompt doesn't mention a connected image. Mirrors the
|
|
# phrasing the reference LoRA trains on, so casual prompts still bind the refs.
|
|
ROLE_CLAUSE = {
|
|
"content": "the character from Picture {}",
|
|
"style": "in the exact art style of Picture {}",
|
|
}
|
|
|
|
|
|
def build_image_prompt(prompt, image_slots, roles=None):
|
|
"""Returns (full_prompt, connected_slot_indices). Slot numbering is
|
|
preserved so 'image2' still means the image2 input even if image1 is empty.
|
|
Connected images never mentioned in the prompt get a role clause appended."""
|
|
image_prompt = ""
|
|
connected = []
|
|
for i, image in enumerate(image_slots):
|
|
if image is None:
|
|
continue
|
|
connected.append(i)
|
|
image_prompt += "Picture {}: <|vision_start|><|image_pad|><|vision_end|>".format(i + 1)
|
|
text = _IMG_WORD.sub(lambda m: "Picture " + m.group(1), prompt)
|
|
if roles is not None:
|
|
missing = [i for i in connected if "Picture {}".format(i + 1) not in text]
|
|
clauses = [ROLE_CLAUSE[roles[i]].format(i + 1) for i in missing]
|
|
if clauses:
|
|
text = (text.rstrip().rstrip(".,") + ", " if text.strip() else "") + ", ".join(clauses)
|
|
return image_prompt + text, connected
|
|
|
|
|
|
class TextEncodeKrea2Reference:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP",),
|
|
"prompt": ("STRING", {
|
|
"multiline": True,
|
|
"dynamicPrompts": True,
|
|
"tooltip": "Refer to the references as image1..image4 or Picture 1..Picture 4, "
|
|
"e.g. 'Redraw the character from image1 in the pose of image2'.",
|
|
}),
|
|
},
|
|
"optional": {
|
|
"image1": ("IMAGE",),
|
|
"image1_detail": (list(DETAIL), {"default": "high", "tooltip": DETAIL_TIP}),
|
|
"image1_role": (["content", "style"], {"default": "content", "tooltip": ROLE_TIP}),
|
|
"image1_mask": ("MASK", {"tooltip": MASK_TIP}),
|
|
"image2": ("IMAGE",),
|
|
"image2_detail": (list(DETAIL), {"default": "high", "tooltip": DETAIL_TIP}),
|
|
"image2_role": (["content", "style"], {"default": "content", "tooltip": ROLE_TIP}),
|
|
"image2_mask": ("MASK", {"tooltip": MASK_TIP}),
|
|
"image3": ("IMAGE",),
|
|
"image3_detail": (list(DETAIL), {"default": "high", "tooltip": DETAIL_TIP}),
|
|
"image3_role": (["content", "style"], {"default": "content", "tooltip": ROLE_TIP}),
|
|
"image3_mask": ("MASK", {"tooltip": MASK_TIP}),
|
|
"image4": ("IMAGE",),
|
|
"image4_detail": (list(DETAIL), {"default": "high", "tooltip": DETAIL_TIP}),
|
|
"image4_role": (["content", "style"], {"default": "content", "tooltip": ROLE_TIP}),
|
|
"image4_mask": ("MASK", {"tooltip": MASK_TIP}),
|
|
"bg_removal": ("BACKGROUND_REMOVAL", {"tooltip": BG_TIP}),
|
|
"layout_split": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.9, "step": 0.05, "tooltip": SPLIT_TIP}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "encode"
|
|
CATEGORY = "conditioning/krea2"
|
|
DESCRIPTION = ("Encodes a prompt plus up to 4 reference images for Krea 2. References enter "
|
|
"as Qwen3-VL vision tokens, enabling character/style/pose transfer "
|
|
"('Redraw the character from image1 in the pose of image2').")
|
|
|
|
def encode(self, clip, prompt, image1=None, image2=None, image3=None, image4=None,
|
|
image1_detail="high", image2_detail="high", image3_detail="high", image4_detail="high",
|
|
image1_role="content", image2_role="content", image3_role="content", image4_role="content",
|
|
image1_mask=None, image2_mask=None, image3_mask=None, image4_mask=None,
|
|
bg_removal=None, layout_split=0.0):
|
|
slots = (image1, image2, image3, image4)
|
|
details = (image1_detail, image2_detail, image3_detail, image4_detail)
|
|
roles = (image1_role, image2_role, image3_role, image4_role)
|
|
masks = (image1_mask, image2_mask, image3_mask, image4_mask)
|
|
full_prompt, connected = build_image_prompt(prompt, slots, roles)
|
|
|
|
images_vl = []
|
|
for i in connected:
|
|
target = DETAIL[details[i]]
|
|
s = slots[i].movedim(-1, 1) # NHWC -> NCHW
|
|
mask = masks[i]
|
|
if mask is None and bg_removal is not None and roles[i] == "content":
|
|
mask = bg_removal.encode_image(slots[i][:, :, :, :3])
|
|
if mask is not None:
|
|
s = isolate_subject(s, mask)
|
|
scale = target / max(s.shape[2], s.shape[3])
|
|
w = max(VISION_PATCH, round(s.shape[3] * scale))
|
|
h = max(VISION_PATCH, round(s.shape[2] * scale))
|
|
s = comfy.utils.common_upscale(s, w, h, "area", "disabled")
|
|
if roles[i] == "style":
|
|
s = shuffle_patches(s, 2 * VISION_PATCH if target >= 512 else VISION_PATCH)
|
|
images_vl.append(s.movedim(1, -1)[:, :, :, :3])
|
|
|
|
tokens = clip.tokenize(
|
|
full_prompt,
|
|
images=images_vl,
|
|
llama_template=EDIT_TEMPLATE if images_vl else None,
|
|
)
|
|
cond = clip.encode_from_tokens_scheduled(tokens)
|
|
|
|
if images_vl and layout_split > 0.0:
|
|
# First steps see text-only conditioning: layout gets fixed before the
|
|
# reference tokens can push it toward a side-by-side collage.
|
|
text_tokens = clip.tokenize(_IMG_WORD.sub(lambda m: "Picture " + m.group(1), prompt))
|
|
text_cond = clip.encode_from_tokens_scheduled(text_tokens)
|
|
cond = (node_helpers.conditioning_set_values(text_cond, {"start_percent": 0.0, "end_percent": layout_split})
|
|
+ node_helpers.conditioning_set_values(cond, {"start_percent": layout_split, "end_percent": 1.0}))
|
|
return (cond,)
|
|
|
|
|
|
class Krea2ApplyStyleTokens:
|
|
"""Redux-style zero-shot style injection: the style image is encoded ALONE
|
|
(separate Qwen3-VL forward pass, so the causal encoder never contaminates the
|
|
prompt tokens), its rows are grid-mean-pooled to a few tokens (destroying the
|
|
subject's spatial structure in hidden-state space), baseline-subtracted and
|
|
scaled by strength — then appended to the conditioning. Chain several for
|
|
additive multi-reference blending. Negative strength pushes AWAY from the style."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"conditioning": ("CONDITIONING",),
|
|
"clip": ("CLIP",),
|
|
"image": ("IMAGE",),
|
|
"strength": ("FLOAT", {"default": 1.0, "min": -2.0, "max": 2.0, "step": 0.05,
|
|
"tooltip": "0 = no effect, 1 = full style tokens, negative = anti-style."}),
|
|
"style_tokens": ("INT", {"default": 4, "min": 1, "max": 64,
|
|
"tooltip": "How many pooled tokens represent the style. Few (1-8) = "
|
|
"global palette/mood, leak-proof. Many = more texture detail "
|
|
"but the subject's layout starts to survive."}),
|
|
"detail": (list(DETAIL), {"default": "normal", "tooltip": DETAIL_TIP}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "apply"
|
|
CATEGORY = "conditioning/krea2"
|
|
DESCRIPTION = "Append pooled style tokens from an image to Krea 2 conditioning, with continuous (and negative) strength."
|
|
|
|
def apply(self, conditioning, clip, image, strength, style_tokens, detail="normal"):
|
|
if strength == 0.0:
|
|
return (conditioning,)
|
|
|
|
target = DETAIL[detail]
|
|
s = image.movedim(-1, 1)
|
|
scale = target / max(s.shape[2], s.shape[3])
|
|
w = max(VISION_PATCH, round(s.shape[3] * scale))
|
|
h = max(VISION_PATCH, round(s.shape[2] * scale))
|
|
s = comfy.utils.common_upscale(s, w, h, "area", "disabled").movedim(1, -1)[:, :, :, :3]
|
|
|
|
tokens = clip.tokenize("Picture 1: <|vision_start|><|image_pad|><|vision_end|>",
|
|
images=[s], llama_template=EDIT_TEMPLATE)
|
|
style_t = clip.encode_from_tokens_scheduled(tokens)[0][0] # (1, S, 12*2560)
|
|
|
|
base_t = clip.encode_from_tokens_scheduled(clip.tokenize(""))[0][0]
|
|
base = base_t.mean(dim=1, keepdim=True) # (1, 1, 12*2560)
|
|
|
|
k = min(style_tokens, style_t.shape[1])
|
|
pooled = torch.stack([c.mean(dim=1) for c in torch.tensor_split(style_t, k, dim=1)], dim=1)
|
|
injected = base + strength * (pooled - base.to(pooled.device))
|
|
|
|
out = []
|
|
for cond_t, extras in conditioning:
|
|
seg = injected.to(cond_t.device, dtype=cond_t.dtype)
|
|
out.append([torch.cat([cond_t, seg.expand(cond_t.shape[0], -1, -1)], dim=1), dict(extras)])
|
|
return (out,)
|
|
|
|
|
|
class Krea2ReferenceLoraLoader:
|
|
"""Load the reference LoRA with independent strengths per layer group.
|
|
|
|
The reference addressing and anti-bleed behavior live in text_fusion/txt_in
|
|
(the conditioning path); the DiT body tensors mostly carry identity drift.
|
|
strength_text 1.0 + strength_dit 0.0 keeps the base model's rendering while
|
|
following references; raise strength_dit for extra scene adherence."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
import folder_paths
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"lora_name": (folder_paths.get_filename_list("loras"),),
|
|
"strength_text": ("FLOAT", {"default": 1.0, "min": -2.0, "max": 2.0, "step": 0.05,
|
|
"tooltip": "text_fusion + txt_in: reference addressing / anti-bleed."}),
|
|
"strength_dit": ("FLOAT", {"default": 0.0, "min": -2.0, "max": 2.0, "step": 0.05,
|
|
"tooltip": "DiT blocks: scene adherence, but drifts identity above ~0.3."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "loaders"
|
|
|
|
def load(self, model, lora_name, strength_text, strength_dit):
|
|
import comfy.sd
|
|
import folder_paths
|
|
lora = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("loras", lora_name), safe_load=True)
|
|
text_keys = ("transformer.text_fusion.", "transformer.txt_in.")
|
|
groups = (
|
|
({k: v for k, v in lora.items() if k.startswith(text_keys)}, strength_text),
|
|
({k: v for k, v in lora.items() if not k.startswith(text_keys)}, strength_dit),
|
|
)
|
|
for sd, strength in groups:
|
|
if sd and strength != 0.0:
|
|
model, _ = comfy.sd.load_lora_for_models(model, None, sd, strength, 0.0)
|
|
return (model,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"TextEncodeKrea2Reference": TextEncodeKrea2Reference,
|
|
"Krea2ApplyStyleTokens": Krea2ApplyStyleTokens,
|
|
"Krea2ReferenceLoraLoader": Krea2ReferenceLoraLoader,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"TextEncodeKrea2Reference": "Krea 2 Reference Encode",
|
|
"Krea2ApplyStyleTokens": "Krea 2 Apply Style Tokens",
|
|
"Krea2ReferenceLoraLoader": "Krea 2 Reference LoRA (layered)",
|
|
}
|