EthanfelandClaude Opus 4.8
9bede2cbc5
Editable system_prompt input + Krea2 System Prompt helper node, plus mask self-heal/padding, downscale-only vision cap, and logo/banner. Also moves the system_prompt socket directly above image1. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
222 lines
10 KiB
Python
222 lines
10 KiB
Python
"""Text Encode (Krea2) — vision-aware conditioning for the Krea2 / K2 DiT.
|
|
|
|
Krea2 conditions on a 12-layer Qwen3-VL-4B tap (see ``comfy/text_encoders/krea2.py``).
|
|
Because that text encoder is a vision-language model, a reference image can be fed through
|
|
its *vision* path so the conditioning becomes visually informed by the image — without any
|
|
VAE / reference-latent. The Krea2 DiT (``comfy/ldm/krea2/model.py``) is pure text-to-image:
|
|
its sequence is ``[text_tokens, noisy_image_patches]`` with no slot for a reference latent,
|
|
so a VAE input would be a no-op here and is deliberately omitted.
|
|
|
|
Each reference image has an optional companion mask. When a mask is connected the image is
|
|
cropped to the mask's bounding box before the vision encoder, so the VLM only "sees" the
|
|
masked region. (This is reference-image masking; it is not inpainting — Krea2 has no
|
|
inpaint/concat pathway to regenerate a masked output region.)
|
|
|
|
This node differs from ``TextEncodeQwenImageEdit`` in two ways:
|
|
* it forces the Krea2 *descriptor* conditioning template even when images are attached
|
|
(the core Qwen-Edit node falls back to Qwen3-VL's plain image template), and
|
|
* it has no VAE input, and it accepts an unbounded, auto-growing set of image+mask slots.
|
|
"""
|
|
|
|
import math
|
|
import re
|
|
|
|
import torch
|
|
|
|
import comfy.utils
|
|
|
|
# Keep in sync with the model's own template; fall back to a literal copy on non-Krea2 builds.
|
|
try:
|
|
from comfy.text_encoders.krea2 import KREA2_TEMPLATE
|
|
except Exception: # pragma: no cover - portability shim
|
|
KREA2_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{}<|im_end|>\n<|im_start|>assistant\n"
|
|
)
|
|
|
|
# The user-facing system_prompt field holds just the system *message text*; the node wraps it in
|
|
# the chat-template scaffolding. Pull the default (Krea2's trained descriptor) out of the template
|
|
# so it stays in sync with whatever comfy ships.
|
|
_sys = re.search(r"<\|im_start\|>system\n(.*?)<\|im_end\|>", KREA2_TEMPLATE, re.S)
|
|
KREA2_SYSTEM_DEFAULT = _sys.group(1) if _sys else (
|
|
"Describe the image by detailing the color, shape, size, texture, quantity, text, "
|
|
"spatial relationships of the objects and background:"
|
|
)
|
|
|
|
# Instruct/edit-style framing (à la TextEncodeQwenImageEditPlus): paste this into system_prompt to
|
|
# make the VLM fuse the user's text WITH the reference image instead of just describing it.
|
|
# Out-of-distribution for Krea2's trained descriptor — experimental.
|
|
KREA2_INSTRUCT_SYSTEM = (
|
|
"Describe the key features of the reference image (color, shape, size, texture, objects, "
|
|
"background), then explain how the user's instruction should combine with or alter it, and "
|
|
"generate a new image meeting the instruction while staying consistent with the reference "
|
|
"where appropriate:"
|
|
)
|
|
|
|
|
|
class TextEncodeKrea2:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP",),
|
|
"prompt": ("STRING", {"multiline": True, "dynamicPrompts": True}),
|
|
},
|
|
"optional": {
|
|
# system_prompt sits just above the image slots.
|
|
"system_prompt": ("STRING", {
|
|
"forceInput": True,
|
|
"tooltip": "Optional system-instruction input. Wire a text node to override how the "
|
|
"VLM frames the reference + your prompt; leave unconnected to use Krea2's "
|
|
"trained descriptor (in-distribution). Use an instruct/edit-style "
|
|
"instruction (see README) to fuse the prompt with the image. The node "
|
|
"adds the chat-template scaffolding; provide just the instruction text.",
|
|
}),
|
|
# image1/mask1 are the seed pair; the web extension grows image2/mask2, ... on connect.
|
|
"image1": ("IMAGE",),
|
|
"mask1": ("MASK",),
|
|
"vision_megapixels": ("FLOAT", {
|
|
"default": 1.0, "min": 0.1, "max": 8.0, "step": 0.1,
|
|
"tooltip": "Maximum size (in megapixels) for each reference before the Qwen3-VL "
|
|
"vision encoder. References larger than this are downscaled; smaller "
|
|
"ones (e.g. a tight mask crop) are kept at native size, never upscaled.",
|
|
}),
|
|
"mask_padding": ("FLOAT", {
|
|
"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.02,
|
|
"tooltip": "Context kept around the mask before cropping, as a fraction of the "
|
|
"image size added on EACH side. 0 = tight crop to the mask; 0.1 = ~10% "
|
|
"margin of surroundings. Only applies when a mask is connected.",
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "encode"
|
|
CATEGORY = "model/conditioning/krea2"
|
|
DESCRIPTION = ("Krea2 (K2) text conditioning with optional vision prompting. Reference images are "
|
|
"fed through the Qwen3-VL vision path; an optional per-image mask crops the image to "
|
|
"the masked region. No VAE is used (Krea2 has no reference-latent pathway).")
|
|
|
|
@staticmethod
|
|
def _collect_indexed(kwargs, prefix):
|
|
pattern = re.compile(r"^{}(\d+)$".format(prefix))
|
|
out = {}
|
|
for key, value in kwargs.items():
|
|
match = pattern.match(key)
|
|
if match is not None and value is not None:
|
|
out[int(match.group(1))] = value
|
|
return out
|
|
|
|
@staticmethod
|
|
def _crop_to_mask(image, mask, padding=0.0):
|
|
"""Crop image (B,H,W,C) to the mask bounding box, expanded by `padding` (a
|
|
fraction of the image size) on each side. No-op if mask empty/None."""
|
|
if mask is None:
|
|
return image
|
|
|
|
if mask.dim() == 2:
|
|
mask = mask.unsqueeze(0)
|
|
elif mask.dim() == 4: # (B,1,H,W) or similar -> (B,H,W)
|
|
mask = mask.reshape(-1, mask.shape[-2], mask.shape[-1])
|
|
|
|
h, w = image.shape[1], image.shape[2]
|
|
if mask.shape[-2:] != (h, w):
|
|
resized = comfy.utils.common_upscale(mask.unsqueeze(1), w, h, "bilinear", "disabled")
|
|
mask = resized[:, 0]
|
|
|
|
presence = (mask > 0.5).any(dim=0) # collapse batch -> (H,W)
|
|
if not bool(presence.any()):
|
|
return image # nothing selected: keep the whole image
|
|
|
|
rows = torch.where(torch.any(presence, dim=1))[0]
|
|
cols = torch.where(torch.any(presence, dim=0))[0]
|
|
y0, y1 = int(rows[0]), int(rows[-1])
|
|
x0, x1 = int(cols[0]), int(cols[-1])
|
|
|
|
if padding > 0.0: # grow the box outward for surrounding context, clamped to the image
|
|
pad_x = round(padding * w)
|
|
pad_y = round(padding * h)
|
|
x0 = max(0, x0 - pad_x)
|
|
x1 = min(w - 1, x1 + pad_x)
|
|
y0 = max(0, y0 - pad_y)
|
|
y1 = min(h - 1, y1 + pad_y)
|
|
|
|
return image[:, y0:y1 + 1, x0:x1 + 1, :]
|
|
|
|
def encode(self, clip, prompt, vision_megapixels=1.0, mask_padding=0.0,
|
|
system_prompt=KREA2_SYSTEM_DEFAULT, **kwargs):
|
|
images = self._collect_indexed(kwargs, "image")
|
|
masks = self._collect_indexed(kwargs, "mask")
|
|
ordered = sorted(images.keys())
|
|
|
|
images_vl = []
|
|
image_prompt = ""
|
|
total = int(vision_megapixels * 1024 * 1024)
|
|
|
|
for slot, n in enumerate(ordered):
|
|
image = self._crop_to_mask(images[n], masks.get(n), padding=mask_padding)
|
|
|
|
samples = image.movedim(-1, 1)
|
|
# vision_megapixels is an upper CAP, not a fixed target: only downscale
|
|
# oversized references, never upscale. Otherwise a small mask crop gets
|
|
# magnified to fill the VLM frame and the subject reads as huge/zoomed.
|
|
scale_by = min(1.0, math.sqrt(total / (samples.shape[3] * samples.shape[2])))
|
|
width = round(samples.shape[3] * scale_by)
|
|
height = round(samples.shape[2] * scale_by)
|
|
|
|
s = comfy.utils.common_upscale(samples, width, height, "area", "disabled")
|
|
images_vl.append(s.movedim(1, -1)[:, :, :, :3])
|
|
|
|
if len(ordered) > 1:
|
|
image_prompt += "Picture {}: <|vision_start|><|image_pad|><|vision_end|>".format(slot + 1)
|
|
else:
|
|
image_prompt += "<|vision_start|><|image_pad|><|vision_end|>"
|
|
|
|
system = system_prompt.strip() or KREA2_SYSTEM_DEFAULT
|
|
template = ("<|im_start|>system\n" + system + "<|im_end|>\n"
|
|
"<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n")
|
|
tokens = clip.tokenize(image_prompt + prompt, images=images_vl, llama_template=template)
|
|
conditioning = clip.encode_from_tokens_scheduled(tokens)
|
|
return (conditioning,)
|
|
|
|
|
|
class Krea2SystemPrompt:
|
|
"""Generic text node preloaded with the instruct/edit-style system prompt. Wire its
|
|
output into TextEncodeKrea2's `system_prompt` input to make the prompt fuse with the
|
|
reference image (experimental / out-of-distribution). Edit the text freely."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"text": ("STRING", {
|
|
"multiline": True, "default": KREA2_INSTRUCT_SYSTEM,
|
|
"tooltip": "System instruction for Krea2's VLM. Defaults to an instruct/edit-style "
|
|
"framing that fuses your prompt with the reference image. Edit as needed; "
|
|
"paste the plain descriptor to fall back to default behavior.",
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("system_prompt",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "model/conditioning/krea2"
|
|
DESCRIPTION = ("Text node preloaded with an instruct-style system prompt for Text Encode (Krea2). "
|
|
"Wire its output into the encoder's system_prompt input.")
|
|
|
|
def run(self, text):
|
|
return (text,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"TextEncodeKrea2": TextEncodeKrea2,
|
|
"Krea2SystemPrompt": Krea2SystemPrompt,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"TextEncodeKrea2": "Text Encode (Krea2)",
|
|
"Krea2SystemPrompt": "Krea2 System Prompt",
|
|
}
|