Add Krea2 VLM Preview node; drop encoder string output

- New Krea2VLMPreview node: runs the Qwen3-VL encoder generatively on the same image+prompt the encoder feeds, outputs the decoded text (proxy for what the VLM sees). Shares image/template prep with the encoder. - TextEncodeKrea2: remove the vlm_prompt STRING output (back to a single conditioning output); print_prompt console toggle retained. - web JS: auto-growing image/mask pairs now also apply to the preview node. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-27 08:45:53 +02:00
parent ebc552627c
commit 2d98b4c336
3 changed files with 115 additions and 34 deletions
@@ -53,8 +53,20 @@ downstream of this node — that works generically at the sampler level.)
 | `print_prompt` | BOOLEAN | Print the assembled Qwen3-VL prompt to the ComfyUI console (debug). |
 | `vision_megapixels` | FLOAT | Max size before the vision encoder; references are downscaled to this cap, never upscaled (default `1.0`). |

-**Outputs:** `conditioning` (for the Krea2 sampler) and `vlm_prompt` (STRING — the full assembled
-prompt as the Qwen3-VL encoder receives it; wire to a text-preview node to inspect it).
+**Output:** `conditioning` for the Krea2 sampler. (`print_prompt` dumps the assembled Qwen3-VL
+prompt to the console for debugging.)
+
+## Inspecting what the VLM sees — `Krea2 VLM Preview`
+
+The conditioning path only extracts hidden states; it never generates text. To actually *read*
+how Qwen3-VL interprets your image + prompt, use the **`Krea2 VLM Preview`** node: it mirrors the
+encoder's image/template prep, then runs the model **generatively** (`clip.generate` → `decode`)
+and outputs the text. Same `vision_megapixels` / `mask_padding` / `vision_position` / `system_prompt`
+controls, plus `max_length` / `temperature` / `seed`.
+
+Notes: it's a *proxy* (the generated text shares the VLM's forward understanding but isn't the
+hidden-state tensor the sampler uses); it needs the encoder's `lm_head` weights; and image inputs
+require a **bf16** encoder (FP8 vision is unsupported — see below).

 ## System prompt (making the prompt interact with the image)

@@ -102,8 +102,7 @@ class TextEncodeKrea2:
            },
        }

-    RETURN_TYPES = ("CONDITIONING", "STRING")
-    RETURN_NAMES = ("conditioning", "vlm_prompt")
+    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "encode"
    CATEGORY = "model/conditioning/krea2"
    DESCRIPTION = ("Krea2 (K2) text conditioning with optional vision prompting. Reference images are "
@@ -156,11 +155,12 @@ class TextEncodeKrea2:

        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, vision_position="before prompt",
-               print_prompt=False, **kwargs):
-        images = self._collect_indexed(kwargs, "image")
-        masks = self._collect_indexed(kwargs, "mask")
+    @classmethod
+    def _prepare_vision(cls, kwargs, vision_megapixels, mask_padding):
+        """Crop+resize each connected reference and build the vision-token string.
+        Shared by the encoder and the VLM preview so both feed the model identically."""
+        images = cls._collect_indexed(kwargs, "image")
+        masks = cls._collect_indexed(kwargs, "mask")
        ordered = sorted(images.keys())

        images_vl = []
@@ -168,34 +168,54 @@ class TextEncodeKrea2:
        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)
-
+            image = cls._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.
+            # vision_megapixels is an upper CAP, not a fixed target: only downscale oversized
+            # references, never upscale (a small mask crop would otherwise be magnified).
            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|>"
+        return images_vl, image_prompt

+    @staticmethod
+    def _build_text(system_prompt, prompt, image_prompt, vision_position):
+        """Assemble the user text (with vision tokens) and the chat template."""
        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")
        text = (prompt + image_prompt) if vision_position == "after prompt" else (image_prompt + prompt)
-        # literal replace (not .format): safe with braces in the text/system prompt.
-        assembled = template.replace("{}", text, 1)
+        return text, template
+
+    @staticmethod
+    def _fp8_hint(exc, images_vl):
+        """Map the cryptic FP8 vision crash to an actionable error; else None.
+
+        ComfyUI's Qwen3-VL vision tower (qwen35.py fast_pos_embed_interpolate) adds the pos-embed
+        weights without casting, so an FP8-loaded text encoder dies on the image path."""
+        if images_vl and isinstance(exc, NotImplementedError) and "Float8" in str(exc):
+            return RuntimeError(
+                "Krea2: the Qwen3-VL text encoder is loaded in FP8, which ComfyUI's vision tower "
+                "cannot run on the image path ('add_stub not implemented for Float8_e4m3fn'). Load "
+                "a bf16/fp16 Qwen3-VL-4B text encoder (e.g. a qwen3vl_4b *bf16* file) via CLIPLoader "
+                "type 'krea2' when using image references. The FP8 encoder works only text-only."
+            )
+        return None
+
+    def encode(self, clip, prompt, vision_megapixels=1.0, mask_padding=0.0,
+               system_prompt=KREA2_SYSTEM_DEFAULT, vision_position="before prompt",
+               print_prompt=False, **kwargs):
+        images_vl, image_prompt = self._prepare_vision(kwargs, vision_megapixels, mask_padding)
+        text, template = self._build_text(system_prompt, prompt, image_prompt, vision_position)

        if print_prompt:
            print("\n========== Text Encode (Krea2) -> Qwen3-VL prompt ==========")
-            print(assembled)
+            print(template.replace("{}", text, 1))  # literal replace: brace-safe
            print("---- references: {} ----".format(len(images_vl)))
            print("===========================================================\n")

@@ -203,20 +223,11 @@ class TextEncodeKrea2:
        try:
            conditioning = clip.encode_from_tokens_scheduled(tokens)
        except NotImplementedError as exc:
-            # ComfyUI's Qwen3-VL vision tower (qwen35.py fast_pos_embed_interpolate) adds the
-            # pos-embed weights without casting, so an FP8-loaded text encoder crashes on the
-            # image path. Turn the cryptic torch error into an actionable one.
-            if images_vl and "Float8" in str(exc):
-                raise RuntimeError(
-                    "Text Encode (Krea2): the Qwen3-VL text encoder is loaded in FP8, which "
-                    "ComfyUI's vision tower cannot run on the image path ('add_stub not "
-                    "implemented for Float8_e4m3fn' in the pos-embed). Load a bf16/fp16 "
-                    "Qwen3-VL-4B text encoder (e.g. a qwen3vl_4b *bf16* file) via CLIPLoader "
-                    "type 'krea2' when using image references. The FP8 encoder works only for "
-                    "text-only prompts."
-                ) from exc
+            hint = self._fp8_hint(exc, images_vl)
+            if hint is not None:
+                raise hint from exc
            raise
-        return (conditioning, assembled)
+        return (conditioning,)


 class Krea2SystemPrompt:
@@ -248,12 +259,70 @@ class Krea2SystemPrompt:
        return (text,)


+class Krea2VLMPreview:
+    """Runs Krea2's Qwen3-VL encoder GENERATIVELY on the same image+prompt the encoder feeds,
+    returning the model's text output — a proxy for what the VLM 'sees'. Needs the encoder's
+    lm_head in the weights; use a bf16/fp16 encoder for image inputs (FP8 vision is unsupported)."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "clip": ("CLIP",),
+                "prompt": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+            },
+            "optional": {
+                "system_prompt": ("STRING", {
+                    "forceInput": True,
+                    "tooltip": "Same as the encoder's system_prompt input; unconnected = Krea2 descriptor.",
+                }),
+                "image1": ("IMAGE",),
+                "mask1": ("MASK",),
+                "vision_megapixels": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 8.0, "step": 0.1}),
+                "mask_padding": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.02}),
+                "vision_position": (["before prompt", "after prompt"], {"default": "before prompt"}),
+                "max_length": ("INT", {"default": 256, "min": 1, "max": 8192,
+                                       "tooltip": "Max number of tokens to generate."}),
+                "temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.01,
+                                          "tooltip": "0 = greedy/deterministic."}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+            },
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("vlm_text",)
+    FUNCTION = "preview"
+    CATEGORY = "model/conditioning/krea2"
+    DESCRIPTION = ("Generate text from Krea2's Qwen3-VL encoder on the same image+prompt the encoder "
+                   "sends, to inspect how the VLM reads your reference. Needs lm_head weights; use a "
+                   "bf16 encoder for image inputs.")
+
+    def preview(self, clip, prompt, vision_megapixels=1.0, mask_padding=0.0,
+                system_prompt=KREA2_SYSTEM_DEFAULT, vision_position="before prompt",
+                max_length=256, temperature=0.7, seed=0, **kwargs):
+        images_vl, image_prompt = TextEncodeKrea2._prepare_vision(kwargs, vision_megapixels, mask_padding)
+        text, template = TextEncodeKrea2._build_text(system_prompt, prompt, image_prompt, vision_position)
+        tokens = clip.tokenize(text, images=images_vl, llama_template=template)
+        try:
+            ids = clip.generate(tokens, do_sample=(temperature > 0.0), max_length=max_length,
+                                temperature=max(temperature, 0.01), seed=seed)
+            out = clip.decode(ids)
+        except NotImplementedError as exc:
+            hint = TextEncodeKrea2._fp8_hint(exc, images_vl)
+            if hint is not None:
+                raise hint from exc
+            raise
+        return (out,)
+
+
 NODE_CLASS_MAPPINGS = {
    "TextEncodeKrea2": TextEncodeKrea2,
    "Krea2SystemPrompt": Krea2SystemPrompt,
+    "Krea2VLMPreview": Krea2VLMPreview,
 }

 NODE_DISPLAY_NAME_MAPPINGS = {
    "TextEncodeKrea2": "Text Encode (Krea2)",
    "Krea2SystemPrompt": "Krea2 System Prompt",
+    "Krea2VLMPreview": "Krea2 VLM Preview",
 }
@@ -3,7 +3,7 @@ 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 NODE_NAMES = new Set(["TextEncodeKrea2", "Krea2VLMPreview"]);
 const IMAGE_RE = /^image(\d+)$/;

 function pairNumbers(node) {
@@ -82,7 +82,7 @@ function syncPairs(node) {
 app.registerExtension({
    name: "krea2.textencode.dynamicimages",
    async beforeRegisterNodeDef(nodeType, nodeData) {
-        if (nodeData.name !== NODE_NAME) return;
+        if (!NODE_NAMES.has(nodeData.name)) return;

        const onNodeCreated = nodeType.prototype.onNodeCreated;
        nodeType.prototype.onNodeCreated = function () {