diff --git a/README.md b/README.md index 869a056..151b074 100644 --- a/README.md +++ b/README.md @@ -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) diff --git a/nodes.py b/nodes.py index 5974552..1d960f6 100644 --- a/nodes.py +++ b/nodes.py @@ -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", } diff --git a/web/krea2_dynamic_images.js b/web/krea2_dynamic_images.js index a9c940f..154a6d3 100644 --- a/web/krea2_dynamic_images.js +++ b/web/krea2_dynamic_images.js @@ -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 () {