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 <noreply@anthropic.com>
2026-06-25 03:09:21 +02:00
commit 13b79b3fc0
5 changed files with 310 additions and 0 deletions
@@ -0,0 +1,5 @@
+__pycache__/
+*.py[cod]
+*.egg-info/
+.DS_Store
+.ipynb_checkpoints/
@@ -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.
@@ -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"]
@@ -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)",
+}
@@ -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;
+        };
+    },
+});