commit 13b79b3fc0bb1f571a1773717d90038743558bb0 Author: Ethanfel Date: Thu Jun 25 03:09:21 2026 +0200 Initial release: Text Encode (Krea2) node Vision-aware text conditioning for the Krea2/K2 model (Qwen3-VL-4B text encoder). Forces the Krea2 descriptor template even with images, omits the no-op VAE input (Krea2 has no reference-latent pathway), and accepts an unbounded, auto-growing set of image+mask reference pairs. Optional per-image mask crops the image to the masked bounding box before the vision encoder. Co-Authored-By: Claude Opus 4.8 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f27eeed --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +__pycache__/ +*.py[cod] +*.egg-info/ +.DS_Store +.ipynb_checkpoints/ diff --git a/README.md b/README.md new file mode 100644 index 0000000..d02e4b8 --- /dev/null +++ b/README.md @@ -0,0 +1,58 @@ +# ComfyUI-Krea2TexTEncoder + +**Text Encode (Krea2)** — vision-aware text conditioning for the Krea2 / K2 model +(`kreaturbo.safetensors`), whose text encoder is **Qwen3-VL-4B** (12-layer tap). + +## Why this node exists + +Krea2's text encoder is a vision-language model, so a reference image can be pushed +through its **vision path** to make the conditioning visually aware of that image — a +"prompt from a picture" effect. People were reaching for the core +`TextEncodeQwenImageEdit` node to do this, which has two problems for Krea2: + +1. **The VAE input does nothing.** Krea2's DiT (`comfy/ldm/krea2/model.py`, + `SingleStreamDiT`) builds its token sequence as `[text_tokens, noisy_image_patches]` + — there is **no slot for a reference latent**, and `Krea2.extra_conds` + (`comfy/model_base.py`) never reads `reference_latents`. So a connected VAE produces a + `reference_latents` entry the model silently discards. (Real pixel-faithful editing + would require training a reference-latent pathway into the DiT — it can't be done from + a node alone.) +2. **Wrong template with images.** With an image attached, the core node falls back to + Qwen3-VL's *plain* image template instead of the Krea2 *descriptor* template the model + was conditioned with. + +This node fixes both: it **forces the Krea2 descriptor template** even with images, and +it **omits the VAE input** entirely. It also accepts an **unbounded, auto-growing set of +image+mask pairs** (connect `image1`, a fresh `image2`/`mask2` pair appears, and so on). + +## Masks + +Each reference image has an optional companion mask. When 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 — **not inpainting**: Krea2 has no +concat/inpaint pathway to regenerate a masked region of the output. (To spatially restrict +where the *prompt* applies in the output, use ComfyUI's standard `ConditioningSetMask` +downstream of this node — that works generically at the sampler level.) + +## Inputs + +| Input | Type | Notes | +|-------|------|-------| +| `clip` | CLIP | Load with **CLIPLoader → type `krea2`**. | +| `prompt` | STRING | Your text prompt. | +| `image1…N` | IMAGE | Optional reference images; slots grow as you connect them. | +| `mask1…N` | MASK | Optional per-image mask; crops `imageN` to the masked bounding box. | +| `vision_megapixels` | FLOAT | Area each image is resized to before the vision encoder (default `1.0`). | + +Output is a standard `CONDITIONING` for the Krea2 sampler. With no images connected it +works as a plain Krea2 text encoder. + +## Install + +Symlink (or copy) this repo into ComfyUI's `custom_nodes/`, matching the existing setup: + +```bash +ln -s /media/p5/ComfyUI-Krea2TexTEncoder /media/p5/Comfyui/custom_nodes/ComfyUI-Krea2TexTEncoder +``` + +Then restart ComfyUI. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..8ce4cd7 --- /dev/null +++ b/__init__.py @@ -0,0 +1,8 @@ +"""ComfyUI-Krea2TexTEncoder — vision-aware text conditioning for the Krea2 / K2 model.""" + +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +# Serves web/ as a frontend extension (auto-growing image inputs). +WEB_DIRECTORY = "./web" + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..03968f9 --- /dev/null +++ b/nodes.py @@ -0,0 +1,138 @@ +"""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" + ) + + +class TextEncodeKrea2: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "clip": ("CLIP",), + "prompt": ("STRING", {"multiline": True, "dynamicPrompts": True}), + }, + "optional": { + # 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": "Target size (in megapixels) each reference image is area-resized to " + "before the Qwen3-VL vision encoder. Higher = more visual detail, more VRAM.", + }), + }, + } + + 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): + """Crop image (B,H,W,C) to the bounding box of mask. 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]) + return image[:, y0:y1 + 1, x0:x1 + 1, :] + + def encode(self, clip, prompt, vision_megapixels=1.0, **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)) + + samples = image.movedim(-1, 1) + scale_by = 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|>" + + tokens = clip.tokenize(image_prompt + prompt, images=images_vl, llama_template=KREA2_TEMPLATE) + conditioning = clip.encode_from_tokens_scheduled(tokens) + return (conditioning,) + + +NODE_CLASS_MAPPINGS = { + "TextEncodeKrea2": TextEncodeKrea2, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "TextEncodeKrea2": "Text Encode (Krea2)", +} diff --git a/web/krea2_dynamic_images.js b/web/krea2_dynamic_images.js new file mode 100644 index 0000000..1123eee --- /dev/null +++ b/web/krea2_dynamic_images.js @@ -0,0 +1,101 @@ +import { app } from "../../scripts/app.js"; + +// Auto-growing reference slots for "Text Encode (Krea2)": each reference is an +// (imageN, maskN) pair. Keep exactly one empty trailing pair, appending a fresh +// pair whenever the last one gets connected. +const NODE_NAME = "TextEncodeKrea2"; +const IMAGE_RE = /^image(\d+)$/; + +function pairNumbers(node) { + const nums = []; + for (const inp of node.inputs || []) { + const m = IMAGE_RE.exec(inp.name); + if (m) nums.push(parseInt(m[1], 10)); + } + nums.sort((a, b) => a - b); + return nums; +} + +function inputIndex(node, name) { + return (node.inputs || []).findIndex((i) => i.name === name); +} + +function linked(node, name) { + const idx = inputIndex(node, name); + return idx >= 0 && node.inputs[idx].link != null; +} + +function addPair(node, n) { + node.addInput(`image${n}`, "IMAGE"); + node.addInput(`mask${n}`, "MASK"); +} + +function removePair(node, n) { + // Remove mask first (it sits after the image, so the image index stays valid). + let idx = inputIndex(node, `mask${n}`); + if (idx >= 0) node.removeInput(idx); + idx = inputIndex(node, `image${n}`); + if (idx >= 0) node.removeInput(idx); +} + +function pairEmpty(node, n) { + return !linked(node, `image${n}`) && !linked(node, `mask${n}`); +} + +function syncPairs(node) { + if (pairNumbers(node).length === 0) addPair(node, 1); + + // Collapse trailing fully-empty pairs down to a single spare. + for (;;) { + const nums = pairNumbers(node); + if (nums.length <= 1) break; + const last = nums[nums.length - 1]; + const prev = nums[nums.length - 2]; + if (pairEmpty(node, last) && pairEmpty(node, prev)) { + removePair(node, last); + } else { + break; + } + } + + // If the last pair is in use, append a fresh spare pair. + const nums = pairNumbers(node); + const last = nums[nums.length - 1]; + if (!pairEmpty(node, last)) { + addPair(node, last + 1); + } + + node.setDirtyCanvas(true, true); +} + +app.registerExtension({ + name: "krea2.textencode.dynamicimages", + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData.name !== NODE_NAME) return; + + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined; + syncPairs(this); + return r; + }; + + const onConnectionsChange = nodeType.prototype.onConnectionsChange; + nodeType.prototype.onConnectionsChange = function (slotType) { + const r = onConnectionsChange ? onConnectionsChange.apply(this, arguments) : undefined; + // LiteGraph.INPUT === 1; only react to input-side connection changes. + if (slotType === 1) { + syncPairs(this); + } + return r; + }; + + // Restore the spare pair after a saved graph is loaded. + const onConfigure = nodeType.prototype.onConfigure; + nodeType.prototype.onConfigure = function () { + const r = onConfigure ? onConfigure.apply(this, arguments) : undefined; + syncPairs(this); + return r; + }; + }, +});