From 85cfd3e3742bd39df8337ae553b6904d0e507130 Mon Sep 17 00:00:00 2001 From: Makadi Date: Thu, 30 Jul 2026 05:40:26 +0300 Subject: [PATCH 1/2] add image_weight --- nodes.py | 50 +++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 47 insertions(+), 3 deletions(-) diff --git a/nodes.py b/nodes.py index a588855..6fed833 100644 --- a/nodes.py +++ b/nodes.py @@ -18,6 +18,7 @@ * it has no VAE input, and it accepts an unbounded, auto-growing set of image+mask slots. """ +import contextlib import math import re @@ -76,6 +77,14 @@ def INPUT_TYPES(cls): # image1/mask1 are the seed pair; the web extension grows image2/mask2, ... on connect. "image1": ("IMAGE",), "mask1": ("MASK",), + "image_weight": ("FLOAT", { + "default": 1.0, "min": 0.0, "max": 3.0, "step": 0.05, + "tooltip": "How strongly the reference image(s) influence the conditioning. Scales " + "the Qwen3-VL vision embeddings (and their DeepStack features) before the " + "text encoder: 1.0 = unmodified, <1 weakens the image relative to your " + "text, >1 amplifies it. 0 = vision tokens carry no signal. Applies to " + "every connected reference; no effect without an image.", + }), "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 " @@ -206,7 +215,41 @@ def _fp8_hint(exc, images_vl): ) return None - def encode(self, clip, prompt, vision_megapixels=1.0, mask_padding=0.0, + @staticmethod + @contextlib.contextmanager + def _scaled_vision(clip, weight): + """Temporarily scale the vision embeddings produced by Qwen3-VL's vision tower. + + Comfy's tokenizer disables per-token weights for the Qwen path (and one image_pad token + expands into hundreds of vision tokens, so the usual weight-per-token mechanism cannot + line up anyway). Instead we wrap the model's ``preprocess_embed``, which is where the + vision tower's merged patch embeddings — and the DeepStack features injected at those + same positions — are produced, and scale both. No-op at 1.0 or on unexpected builds.""" + transformer = None + if weight != 1.0: + inner = getattr(clip.cond_stage_model, getattr(clip.cond_stage_model, "clip", ""), None) + transformer = getattr(inner, "transformer", None) + if transformer is None or not hasattr(transformer, "preprocess_embed"): + yield + return + + original = transformer.preprocess_embed + + def preprocess_embed(embed, device): + emb, extra = original(embed, device=device) + if emb is not None and embed.get("type") == "image": + emb = emb * weight + if extra is not None and extra.get("deepstack") is not None: + extra = {**extra, "deepstack": [d * weight for d in extra["deepstack"]]} + return emb, extra + + transformer.preprocess_embed = preprocess_embed + try: + yield + finally: + del transformer.preprocess_embed # instance override removed -> class method again + + def encode(self, clip, prompt, image_weight=1.0, 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) @@ -215,12 +258,13 @@ def encode(self, clip, prompt, vision_megapixels=1.0, mask_padding=0.0, if print_prompt: print("\n========== Text Encode (Krea2) -> Qwen3-VL prompt ==========") print(template.replace("{}", text, 1)) # literal replace: brace-safe - print("---- references: {} ----".format(len(images_vl))) + print("---- references: {} (image_weight {}) ----".format(len(images_vl), image_weight)) print("===========================================================\n") tokens = clip.tokenize(text, images=images_vl, llama_template=template) try: - conditioning = clip.encode_from_tokens_scheduled(tokens) + with self._scaled_vision(clip, image_weight): + conditioning = clip.encode_from_tokens_scheduled(tokens) except NotImplementedError as exc: hint = self._fp8_hint(exc, images_vl) if hint is not None: From 7763f6337482f6ea5bf3584da7c7207eefc52bf2 Mon Sep 17 00:00:00 2001 From: Makadi Date: Thu, 30 Jul 2026 05:49:52 +0300 Subject: [PATCH 2/2] add image_end_percent --- nodes.py | 47 +++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 39 insertions(+), 8 deletions(-) diff --git a/nodes.py b/nodes.py index 6fed833..c869a20 100644 --- a/nodes.py +++ b/nodes.py @@ -85,6 +85,14 @@ def INPUT_TYPES(cls): "text, >1 amplifies it. 0 = vision tokens carry no signal. Applies to " "every connected reference; no effect without an image.", }), + "image_end_percent": ("FLOAT", { + "default": 1.0, "min": 0.01, "max": 1.0, "step": 0.01, + "tooltip": "Fraction of sampling during which the image-informed conditioning is " + "used; after it, a text-only conditioning takes over. 1.0 = image for the " + "whole run (single encoder pass). 0.4 = reference sets composition and " + "identity early, then the prompt alone refines. Costs a second text-encoder " + "pass when below 1.0; no effect without an image.", + }), "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 " @@ -224,7 +232,12 @@ def _scaled_vision(clip, weight): expands into hundreds of vision tokens, so the usual weight-per-token mechanism cannot line up anyway). Instead we wrap the model's ``preprocess_embed``, which is where the vision tower's merged patch embeddings — and the DeepStack features injected at those - same positions — are produced, and scale both. No-op at 1.0 or on unexpected builds.""" + same positions — are produced, and scale both. No-op at 1.0 or on unexpected builds. + + The scaling interpolates each vision token toward the *mean* vision token rather than + toward zero, so embedding magnitudes stay plausible: 0.0 yields a flat block that still + says "an image is here" but carries no specific content, instead of zero vectors that + occupy sequence positions far outside anything the encoder saw in training.""" transformer = None if weight != 1.0: inner = getattr(clip.cond_stage_model, getattr(clip.cond_stage_model, "clip", ""), None) @@ -235,12 +248,16 @@ def _scaled_vision(clip, weight): original = transformer.preprocess_embed + def toward_mean(t): + mean = t.mean(dim=0, keepdim=True) + return mean + (t - mean) * weight + def preprocess_embed(embed, device): emb, extra = original(embed, device=device) if emb is not None and embed.get("type") == "image": - emb = emb * weight + emb = toward_mean(emb) if extra is not None and extra.get("deepstack") is not None: - extra = {**extra, "deepstack": [d * weight for d in extra["deepstack"]]} + extra = {**extra, "deepstack": [toward_mean(d) for d in extra["deepstack"]]} return emb, extra transformer.preprocess_embed = preprocess_embed @@ -249,27 +266,41 @@ def preprocess_embed(embed, device): finally: del transformer.preprocess_embed # instance override removed -> class method again - def encode(self, clip, prompt, image_weight=1.0, vision_megapixels=1.0, mask_padding=0.0, - system_prompt=KREA2_SYSTEM_DEFAULT, vision_position="before prompt", - print_prompt=False, **kwargs): + def encode(self, clip, prompt, image_weight=1.0, image_end_percent=1.0, 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) + # Hand off to a text-only conditioning partway through sampling: the reference then only + # shapes the early (composition/identity) steps. Both halves are ordinary encoder outputs, + # so nothing here is off-distribution -- unlike blending conditioning tensors, which cannot + # be aligned anyway (the vision tokens make the two sequences different lengths). + schedule = bool(images_vl) and image_end_percent < 1.0 if print_prompt: print("\n========== Text Encode (Krea2) -> Qwen3-VL prompt ==========") print(template.replace("{}", text, 1)) # literal replace: brace-safe - print("---- references: {} (image_weight {}) ----".format(len(images_vl), image_weight)) + print("---- references: {} (image_weight {}{}) ----".format( + len(images_vl), image_weight, + ", image until {:.0%} of sampling".format(image_end_percent) if schedule else "")) print("===========================================================\n") tokens = clip.tokenize(text, images=images_vl, llama_template=template) + add_dict = {"start_percent": 0.0, "end_percent": image_end_percent} if schedule else {} try: with self._scaled_vision(clip, image_weight): - conditioning = clip.encode_from_tokens_scheduled(tokens) + conditioning = clip.encode_from_tokens_scheduled(tokens, add_dict=add_dict) except NotImplementedError as exc: hint = self._fp8_hint(exc, images_vl) if hint is not None: raise hint from exc raise + + if schedule: + text_only, _ = self._build_text(system_prompt, prompt, "", vision_position) + tokens_text_only = clip.tokenize(text_only, llama_template=template) + conditioning = conditioning + clip.encode_from_tokens_scheduled( + tokens_text_only, add_dict={"start_percent": image_end_percent, "end_percent": 1.0}) return (conditioning,)