From 734fd92b14266cb92b7290388cd81886c9caa664 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Mon, 13 Jul 2026 17:55:10 +0200 Subject: [PATCH 1/9] Add spatial fusion functionality for models using Qwen3-VL encoders, which allows for fusing of image inputs for style, structure and concept blending directly from images. New node added: TextEncodeQwenImageEditFusion --- comfy_extras/nodes_qwen.py | 161 +++++++++++++++++++++++++++++++++++++ 1 file changed, 161 insertions(+) diff --git a/comfy_extras/nodes_qwen.py b/comfy_extras/nodes_qwen.py index 4960774db..4e8ac97b2 100644 --- a/comfy_extras/nodes_qwen.py +++ b/comfy_extras/nodes_qwen.py @@ -7,6 +7,80 @@ import comfy.model_management import torch import nodes + +def _spatial_fusion_mask(height, width, num_sources, method, block_size, dither_ratio, device): + rows = torch.arange(height, device=device).unsqueeze(1) + columns = torch.arange(width, device=device).unsqueeze(0) + + if method == "spatial-checkerboard": + return ((rows + columns) % num_sources).flatten() + if method == "spatial-block-interleave": + return ((rows // block_size + columns // block_size) % num_sources).flatten() + if method == "spatial-dither-random": + generator = torch.Generator(device=device).manual_seed(42) + random = torch.rand((height, width), generator=generator, device=device) + if num_sources == 2: + return torch.where(random < dither_ratio, 0, 1).flatten() + return (random * num_sources).long().flatten() + raise ValueError(f"Unsupported visual fusion method: {method}") + + +def _visual_token_span(tokens, cond_length, visual_tokens): + if len(tokens) != 1: + raise ValueError("Visual fusion requires a Qwen3-VL or Krea2 text encoder.") + + token_pairs = next(iter(tokens.values()))[0] + image_positions = [i for i, pair in enumerate(token_pairs) if isinstance(pair[0], dict) and pair[0].get("type") == "image"] + if len(image_positions) != 1: + raise ValueError("Visual fusion requires exactly one visual token block per encoding pass.") + + image_position = image_positions[0] + if any(not isinstance(pair[0], (int, float)) for pair in token_pairs[image_position + 1:]): + raise ValueError("Visual fusion does not support embeddings after the image token block.") + + end = cond_length - (len(token_pairs) - image_position - 1) + start = end - visual_tokens + if start < 0 or end > cond_length: + raise ValueError("Could not locate the visual token block in the encoded conditioning.") + return start, end + + +def _fuse_conditionings(conditionings, tokens, visual_height, visual_width, method, block_size, dither_ratio): + schedule_count = len(conditionings[0]) + if any(len(source) != schedule_count for source in conditionings): + raise ValueError("All visual fusion sources must use the same CLIP schedule.") + + visual_tokens = visual_height * visual_width + fused = [] + for schedule in range(schedule_count): + source_conds = [source[schedule][0] for source in conditionings] + spans = [_visual_token_span(source_tokens, cond.shape[1], visual_tokens) for source_tokens, cond in zip(tokens, source_conds)] + if any(span != spans[0] for span in spans[1:]): + raise ValueError("Visual fusion sources produced different token layouts.") + + start, end = spans[0] + visuals = torch.stack([cond[:, start:end] for cond in source_conds], dim=2) + mask = _spatial_fusion_mask(visual_height, visual_width, len(source_conds), method, block_size, dither_ratio, visuals.device) + blended_visual = torch.take_along_dim(visuals, mask[None, :, None, None], dim=2).squeeze(2) + + blended = source_conds[0].clone() + blended[:, start:end] = blended_visual + fused.append([blended, conditionings[0][schedule][1].copy()]) + return fused + + +def _flatten_images(images): + sources = [] + for name in sorted(images, key=lambda value: int(value.rsplit("_", 1)[-1])): + image = images[name] + if image is None: + continue + if image.ndim == 3: + image = image.unsqueeze(0) + sources.extend(image[i:i + 1] for i in range(image.shape[0])) + return sources + + class TextEncodeQwenImageEdit(io.ComfyNode): @classmethod def define_schema(cls): @@ -106,6 +180,92 @@ class TextEncodeQwenImageEditPlus(io.ComfyNode): return io.NodeOutput(conditioning) +class TextEncodeQwenImageEditFusion(io.ComfyNode): + @classmethod + def define_schema(cls): + images = io.Autogrow.TemplateNames( + io.Image.Input("image"), + names=[f"image_{i}" for i in range(1, 17)], + min=2, + ) + return io.Schema( + node_id="TextEncodeQwenImageEditFusion", + display_name="Text Encode Qwen Image Edit (Visual Fusion)", + category="model/conditioning/qwen image", + description="Encodes images separately and spatially interleaves their Qwen3-VL visual conditioning tokens.", + inputs=[ + io.Clip.Input("clip"), + io.String.Input("prompt", multiline=True, dynamic_prompts=True), + io.Autogrow.Input("images", template=images), + io.Combo.Input( + "fusion_method", + options=["spatial-checkerboard", "spatial-block-interleave", "spatial-dither-random"], + default="spatial-checkerboard", + ), + io.Int.Input("block_size", default=2, min=1, max=8, step=1, advanced=True), + io.Float.Input( + "dither_ratio", + default=0.5, + min=0.0, + max=1.0, + step=0.01, + advanced=True, + tooltip="For two sources, the probability of selecting the first source. Three or more sources are selected uniformly.", + ), + io.Vae.Input("vae", optional=True), + ], + outputs=[io.Conditioning.Output()], + ) + + @classmethod + def execute(cls, clip, prompt, images: io.Autogrow.Type, fusion_method, block_size=2, dither_ratio=0.5, vae=None) -> io.NodeOutput: + sources = _flatten_images(images) + if len(sources) < 2: + raise ValueError("Visual fusion requires at least two images.") + + first = sources[0].movedim(-1, 1) + total = 384 * 384 + scale_by = math.sqrt(total / (first.shape[3] * first.shape[2])) + width = max(32, round(first.shape[3] * scale_by)) + height = max(32, round(first.shape[2] * scale_by)) + + processed = [] + for source in sources: + samples = source[:, :, :, :3].movedim(-1, 1) + resized = comfy.utils.common_upscale(samples, width, height, "area", "center") + processed.append(resized.movedim(1, -1)) + + factor = 32 + visual_height = max(factor, round(height / factor) * factor) // factor + visual_width = max(factor, round(width / factor) * factor) // factor + + full_prompt = ( + "<|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<|vision_start|><|image_pad|><|vision_end|>" + prompt + "<|im_end|>\n" + "<|im_start|>assistant\n" + ) + tokens = [clip.tokenize(full_prompt, images=[image]) for image in processed] + token_key = next(iter(tokens[0]), None) + if token_key not in ("qwen3vl_4b", "qwen3vl_8b") or any(next(iter(source_tokens), None) != token_key for source_tokens in tokens): + raise ValueError("Visual fusion requires a Qwen3-VL or Krea2 text encoder.") + + conditionings = [clip.encode_from_tokens_scheduled(source_tokens) for source_tokens in tokens] + conditioning = _fuse_conditionings(conditionings, tokens, visual_height, visual_width, fusion_method, block_size, dither_ratio) + + if vae is not None: + ref_latents = [] + for source in sources: + samples = source[:, :, :, :3].movedim(-1, 1) + scale_by = math.sqrt((1024 * 1024) / (samples.shape[3] * samples.shape[2])) + latent_width = max(8, round(samples.shape[3] * scale_by / 8.0) * 8) + latent_height = max(8, round(samples.shape[2] * scale_by / 8.0) * 8) + resized = comfy.utils.common_upscale(samples, latent_width, latent_height, "area", "disabled") + ref_latents.append(vae.encode(resized.movedim(1, -1))) + conditioning = node_helpers.conditioning_set_values(conditioning, {"reference_latents": ref_latents}, append=True) + + return io.NodeOutput(conditioning) + + class EmptyQwenImageLayeredLatentImage(io.ComfyNode): @classmethod def define_schema(cls): @@ -136,6 +296,7 @@ class QwenExtension(ComfyExtension): return [ TextEncodeQwenImageEdit, TextEncodeQwenImageEditPlus, + TextEncodeQwenImageEditFusion, EmptyQwenImageLayeredLatentImage, ] From 658e71135660546773b58e46d1a26f063522183c Mon Sep 17 00:00:00 2001 From: silveroxides Date: Mon, 13 Jul 2026 18:11:10 +0200 Subject: [PATCH 2/9] test: add visual fusion coverage --- .../test_qwen_visual_fusion.py | 102 ++++++++++++++++++ 1 file changed, 102 insertions(+) create mode 100644 tests-unit/comfy_extras_test/test_qwen_visual_fusion.py diff --git a/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py b/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py new file mode 100644 index 000000000..bed7fe9f4 --- /dev/null +++ b/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py @@ -0,0 +1,102 @@ +import torch + +from comfy_extras.nodes_qwen import TextEncodeQwenImageEditFusion, _flatten_images, _fuse_conditionings, _spatial_fusion_mask, _visual_token_span + + +def _tokens(image_position=1, suffix=1): + pairs = [(1, 1.0)] * image_position + pairs.append(({"type": "image", "data": torch.zeros(1, 32, 32, 3)}, 1.0)) + pairs.extend([(2, 1.0)] * suffix) + return {"qwen3vl_4b": [pairs]} + + +def test_checkerboard_mask_multiple_sources(): + mask = _spatial_fusion_mask(2, 3, 3, "spatial-checkerboard", 2, 0.5, "cpu") + assert mask.tolist() == [0, 1, 2, 1, 2, 0] + + +def test_block_interleave_mask(): + mask = _spatial_fusion_mask(4, 4, 2, "spatial-block-interleave", 2, 0.5, "cpu") + assert mask.reshape(4, 4).tolist() == [ + [0, 0, 1, 1], + [0, 0, 1, 1], + [1, 1, 0, 0], + [1, 1, 0, 0], + ] + + +def test_dither_mask_is_deterministic_and_honors_two_source_ratio(): + first = _spatial_fusion_mask(4, 4, 2, "spatial-dither-random", 2, 0.5, "cpu") + second = _spatial_fusion_mask(4, 4, 2, "spatial-dither-random", 2, 0.5, "cpu") + assert torch.equal(first, second) + assert _spatial_fusion_mask(2, 2, 2, "spatial-dither-random", 2, 1.0, "cpu").tolist() == [0, 0, 0, 0] + assert _spatial_fusion_mask(2, 2, 2, "spatial-dither-random", 2, 0.0, "cpu").tolist() == [1, 1, 1, 1] + + +def test_visual_span_accounts_for_stripped_prefix(): + tokens = _tokens(image_position=3, suffix=4) + assert _visual_token_span(tokens, cond_length=9, visual_tokens=4) == (1, 5) + + +def test_fusion_replaces_only_visual_tokens_and_preserves_dtype_and_metadata(): + tokens = [_tokens(), _tokens()] + first = torch.tensor([[[10], [10], [10], [10], [10], [20]]], dtype=torch.float16) + second = torch.tensor([[[30], [30], [30], [30], [30], [40]]], dtype=torch.float16) + metadata = {"pooled_output": torch.tensor([1.0]), "marker": "first"} + conditionings = [ + [[first, metadata]], + [[second, {"pooled_output": torch.tensor([2.0])}]], + ] + + fused = _fuse_conditionings(conditionings, tokens, 2, 2, "spatial-checkerboard", 2, 0.5) + output, output_metadata = fused[0] + + assert output.dtype == torch.float16 + assert output.flatten().tolist() == [10, 10, 30, 30, 10, 20] + assert output_metadata == metadata + assert output_metadata is not metadata + + +def test_flatten_images_uses_numeric_input_order_and_splits_batches(): + images = { + "image_10": torch.full((1, 2, 2, 3), 10.0), + "image_2": torch.stack([torch.full((2, 2, 3), 2.0), torch.full((2, 2, 3), 3.0)]), + "image_1": torch.full((1, 2, 2, 3), 1.0), + } + + sources = _flatten_images(images) + assert [source[0, 0, 0, 0].item() for source in sources] == [1.0, 2.0, 3.0, 10.0] + + +def test_node_uses_custom_krea_prompt_and_returns_fused_conditioning(): + class FakeClip: + def tokenize(self, text, images): + assert text.startswith("<|im_start|>system\nDescribe the image by detailing") + assert "Picture 1:" not in text + pairs = [ + (1, 1.0), + ({"type": "image", "data": images[0]}, 1.0), + (2, 1.0), + ] + return {"qwen3vl_4b": [pairs]} + + def encode_from_tokens_scheduled(self, tokens): + image = next(pair[0]["data"] for pair in tokens["qwen3vl_4b"][0] if isinstance(pair[0], dict)) + value = image.mean() + return [[torch.full((1, 146, 1), value, dtype=torch.float16), {"source": float(value)}]] + + result = TextEncodeQwenImageEditFusion.execute( + FakeClip(), + "test prompt", + {"image_1": torch.zeros(1, 32, 32, 3), "image_2": torch.ones(1, 32, 32, 3)}, + "spatial-checkerboard", + ) + conditioning = result.args[0] + output, metadata = conditioning[0] + + assert output.dtype == torch.float16 + assert output.shape == (1, 146, 1) + assert output[:, 0].item() == 0.0 + assert output[:, -1].item() == 0.0 + assert set(output[:, 1:-1].flatten().tolist()) == {0.0, 1.0} + assert metadata == {"source": 0.0} From ba9bcce65f98d94d2499e26ded2b2fdf897688db Mon Sep 17 00:00:00 2001 From: silveroxides Date: Mon, 13 Jul 2026 19:09:22 +0200 Subject: [PATCH 3/9] test: configure CPU mode for fusion tests --- tests-unit/comfy_extras_test/test_qwen_visual_fusion.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py b/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py index bed7fe9f4..a03739eba 100644 --- a/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py +++ b/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py @@ -1,5 +1,10 @@ import torch +from comfy.cli_args import args as cli_args + +if not torch.cuda.is_available(): + cli_args.cpu = True + from comfy_extras.nodes_qwen import TextEncodeQwenImageEditFusion, _flatten_images, _fuse_conditionings, _spatial_fusion_mask, _visual_token_span From 1e7e97713cd9aff632dc342d8bf6eb0004c242b7 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Mon, 13 Jul 2026 19:17:54 +0200 Subject: [PATCH 4/9] test: restore CLI state after node import --- tests-unit/comfy_extras_test/test_qwen_visual_fusion.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py b/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py index a03739eba..0e8b70b6d 100644 --- a/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py +++ b/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py @@ -2,10 +2,14 @@ import torch from comfy.cli_args import args as cli_args +prior_cpu = cli_args.cpu if not torch.cuda.is_available(): cli_args.cpu = True -from comfy_extras.nodes_qwen import TextEncodeQwenImageEditFusion, _flatten_images, _fuse_conditionings, _spatial_fusion_mask, _visual_token_span +try: + from comfy_extras.nodes_qwen import TextEncodeQwenImageEditFusion, _flatten_images, _fuse_conditionings, _spatial_fusion_mask, _visual_token_span +finally: + cli_args.cpu = prior_cpu def _tokens(image_position=1, suffix=1): From 1f04c502352461460d4dec191fcede03ccc91ffa Mon Sep 17 00:00:00 2001 From: Silver <65376327+silveroxides@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:20:51 +0200 Subject: [PATCH 5/9] Copy tensor slice in _flatten_images function Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> --- comfy_extras/nodes_qwen.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy_extras/nodes_qwen.py b/comfy_extras/nodes_qwen.py index 4e8ac97b2..280eedfe7 100644 --- a/comfy_extras/nodes_qwen.py +++ b/comfy_extras/nodes_qwen.py @@ -77,7 +77,7 @@ def _flatten_images(images): continue if image.ndim == 3: image = image.unsqueeze(0) - sources.extend(image[i:i + 1] for i in range(image.shape[0])) + sources.extend(image[i:i + 1].clone() for i in range(image.shape[0])) return sources From 6aed818effcf366498a23959a5256f19c610070f Mon Sep 17 00:00:00 2001 From: silveroxides Date: Mon, 13 Jul 2026 22:32:43 +0200 Subject: [PATCH 6/9] feat: add seeded multi-source spatial dither --- comfy_extras/nodes_qwen.py | 20 ++++---- .../test_qwen_visual_fusion.py | 51 +++++++++++++++++-- 2 files changed, 57 insertions(+), 14 deletions(-) diff --git a/comfy_extras/nodes_qwen.py b/comfy_extras/nodes_qwen.py index 280eedfe7..66dd97939 100644 --- a/comfy_extras/nodes_qwen.py +++ b/comfy_extras/nodes_qwen.py @@ -8,7 +8,7 @@ import torch import nodes -def _spatial_fusion_mask(height, width, num_sources, method, block_size, dither_ratio, device): +def _spatial_fusion_mask(height, width, num_sources, method, block_size, dither_ratio, device, seed=0): rows = torch.arange(height, device=device).unsqueeze(1) columns = torch.arange(width, device=device).unsqueeze(0) @@ -17,11 +17,10 @@ def _spatial_fusion_mask(height, width, num_sources, method, block_size, dither_ if method == "spatial-block-interleave": return ((rows // block_size + columns // block_size) % num_sources).flatten() if method == "spatial-dither-random": - generator = torch.Generator(device=device).manual_seed(42) + generator = torch.Generator(device=device).manual_seed(seed) random = torch.rand((height, width), generator=generator, device=device) - if num_sources == 2: - return torch.where(random < dither_ratio, 0, 1).flatten() - return (random * num_sources).long().flatten() + other_sources = 1 + ((rows + columns) % (num_sources - 1)) + return torch.where(random < dither_ratio, 0, other_sources).flatten() raise ValueError(f"Unsupported visual fusion method: {method}") @@ -45,7 +44,7 @@ def _visual_token_span(tokens, cond_length, visual_tokens): return start, end -def _fuse_conditionings(conditionings, tokens, visual_height, visual_width, method, block_size, dither_ratio): +def _fuse_conditionings(conditionings, tokens, visual_height, visual_width, method, block_size, dither_ratio, seed=0): schedule_count = len(conditionings[0]) if any(len(source) != schedule_count for source in conditionings): raise ValueError("All visual fusion sources must use the same CLIP schedule.") @@ -60,7 +59,7 @@ def _fuse_conditionings(conditionings, tokens, visual_height, visual_width, meth start, end = spans[0] visuals = torch.stack([cond[:, start:end] for cond in source_conds], dim=2) - mask = _spatial_fusion_mask(visual_height, visual_width, len(source_conds), method, block_size, dither_ratio, visuals.device) + mask = _spatial_fusion_mask(visual_height, visual_width, len(source_conds), method, block_size, dither_ratio, visuals.device, seed) blended_visual = torch.take_along_dim(visuals, mask[None, :, None, None], dim=2).squeeze(2) blended = source_conds[0].clone() @@ -210,15 +209,16 @@ class TextEncodeQwenImageEditFusion(io.ComfyNode): max=1.0, step=0.01, advanced=True, - tooltip="For two sources, the probability of selecting the first source. Three or more sources are selected uniformly.", + tooltip="Probability of selecting the first source. Remaining sources are selected with a checkerboard pattern.", ), + io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff, control_after_generate=True, advanced=True, tooltip="Seed for the spatial-dither-random pattern."), io.Vae.Input("vae", optional=True), ], outputs=[io.Conditioning.Output()], ) @classmethod - def execute(cls, clip, prompt, images: io.Autogrow.Type, fusion_method, block_size=2, dither_ratio=0.5, vae=None) -> io.NodeOutput: + def execute(cls, clip, prompt, images: io.Autogrow.Type, fusion_method, block_size=2, dither_ratio=0.5, vae=None, seed=0) -> io.NodeOutput: sources = _flatten_images(images) if len(sources) < 2: raise ValueError("Visual fusion requires at least two images.") @@ -250,7 +250,7 @@ class TextEncodeQwenImageEditFusion(io.ComfyNode): raise ValueError("Visual fusion requires a Qwen3-VL or Krea2 text encoder.") conditionings = [clip.encode_from_tokens_scheduled(source_tokens) for source_tokens in tokens] - conditioning = _fuse_conditionings(conditionings, tokens, visual_height, visual_width, fusion_method, block_size, dither_ratio) + conditioning = _fuse_conditionings(conditionings, tokens, visual_height, visual_width, fusion_method, block_size, dither_ratio, seed) if vae is not None: ref_latents = [] diff --git a/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py b/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py index 0e8b70b6d..e918cbecd 100644 --- a/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py +++ b/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py @@ -1,3 +1,4 @@ +import pytest import torch from comfy.cli_args import args as cli_args @@ -34,14 +35,28 @@ def test_block_interleave_mask(): ] -def test_dither_mask_is_deterministic_and_honors_two_source_ratio(): - first = _spatial_fusion_mask(4, 4, 2, "spatial-dither-random", 2, 0.5, "cpu") - second = _spatial_fusion_mask(4, 4, 2, "spatial-dither-random", 2, 0.5, "cpu") +def test_dither_mask_honors_seed_and_two_source_ratio(): + first = _spatial_fusion_mask(4, 4, 2, "spatial-dither-random", 2, 0.5, "cpu", 7) + second = _spatial_fusion_mask(4, 4, 2, "spatial-dither-random", 2, 0.5, "cpu", 7) + changed = _spatial_fusion_mask(4, 4, 2, "spatial-dither-random", 2, 0.5, "cpu", 8) assert torch.equal(first, second) + assert not torch.equal(first, changed) assert _spatial_fusion_mask(2, 2, 2, "spatial-dither-random", 2, 1.0, "cpu").tolist() == [0, 0, 0, 0] assert _spatial_fusion_mask(2, 2, 2, "spatial-dither-random", 2, 0.0, "cpu").tolist() == [1, 1, 1, 1] +def test_dither_ratio_selects_first_source_or_remaining_checkerboard(): + assert _spatial_fusion_mask(2, 3, 4, "spatial-dither-random", 2, 1.0, "cpu").tolist() == [0, 0, 0, 0, 0, 0] + assert _spatial_fusion_mask(2, 3, 4, "spatial-dither-random", 2, 0.0, "cpu").tolist() == [1, 2, 3, 2, 3, 1] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") +def test_dither_mask_is_seeded_on_cuda(): + first = _spatial_fusion_mask(4, 4, 3, "spatial-dither-random", 2, 0.5, "cuda", 7) + second = _spatial_fusion_mask(4, 4, 3, "spatial-dither-random", 2, 0.5, "cuda", 7) + assert torch.equal(first, second) + + def test_visual_span_accounts_for_stripped_prefix(): tokens = _tokens(image_position=3, suffix=4) assert _visual_token_span(tokens, cond_length=9, visual_tokens=4) == (1, 5) @@ -66,6 +81,19 @@ def test_fusion_replaces_only_visual_tokens_and_preserves_dtype_and_metadata(): assert output_metadata is not metadata +def test_dither_seed_changes_fused_conditioning(): + tokens = [_tokens(), _tokens()] + conditionings = [ + [[torch.zeros((1, 6, 1)), {}]], + [[torch.ones((1, 6, 1)), {}]], + ] + + first = _fuse_conditionings(conditionings, tokens, 2, 2, "spatial-dither-random", 2, 0.5, 7)[0][0] + second = _fuse_conditionings(conditionings, tokens, 2, 2, "spatial-dither-random", 2, 0.5, 8)[0][0] + + assert not torch.equal(first, second) + + def test_flatten_images_uses_numeric_input_order_and_splits_batches(): images = { "image_10": torch.full((1, 2, 2, 3), 10.0), @@ -98,7 +126,15 @@ def test_node_uses_custom_krea_prompt_and_returns_fused_conditioning(): FakeClip(), "test prompt", {"image_1": torch.zeros(1, 32, 32, 3), "image_2": torch.ones(1, 32, 32, 3)}, - "spatial-checkerboard", + "spatial-dither-random", + seed=7, + ) + changed_seed = TextEncodeQwenImageEditFusion.execute( + FakeClip(), + "test prompt", + {"image_1": torch.zeros(1, 32, 32, 3), "image_2": torch.ones(1, 32, 32, 3)}, + "spatial-dither-random", + seed=8, ) conditioning = result.args[0] output, metadata = conditioning[0] @@ -109,3 +145,10 @@ def test_node_uses_custom_krea_prompt_and_returns_fused_conditioning(): assert output[:, -1].item() == 0.0 assert set(output[:, 1:-1].flatten().tolist()) == {0.0, 1.0} assert metadata == {"source": 0.0} + assert not torch.equal(output, changed_seed.args[0][0][0]) + + +def test_node_exposes_seed_control(): + inputs = {value.id: value for value in TextEncodeQwenImageEditFusion.define_schema().inputs} + assert inputs["seed"].default == 0 + assert inputs["seed"].control_after_generate is True From af91a7c44a50acf96e09b667341bcf86b396c552 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Wed, 15 Jul 2026 20:46:11 +0200 Subject: [PATCH 7/9] refactor: generalize image conditioning fusion Move fusion out of Qwen Image conditioning, preserve tokenizer-owned templates, and align visual grids without resizing source images. --- comfy/text_encoders/qwen_vl.py | 31 ++-- comfy_extras/nodes_cond.py | 156 +++++++++++++++++ comfy_extras/nodes_qwen.py | 161 ------------------ .../test_qwen_visual_fusion.py | 72 ++++++-- 4 files changed, 231 insertions(+), 189 deletions(-) diff --git a/comfy/text_encoders/qwen_vl.py b/comfy/text_encoders/qwen_vl.py index 924eb6ad8..22043f5d6 100644 --- a/comfy/text_encoders/qwen_vl.py +++ b/comfy/text_encoders/qwen_vl.py @@ -6,6 +6,23 @@ import math from comfy.ldm.modules.attention import optimized_attention_for_device +def qwen2vl_image_size(height, width, min_pixels=3136, max_pixels=12845056, patch_size=14, merge_size=2): + factor = patch_size * merge_size + resized_height = round(height / factor) * factor + resized_width = round(width / factor) * factor + + if resized_height * resized_width > max_pixels: + beta = math.sqrt((height * width) / max_pixels) + resized_height = max(factor, math.floor(height / beta / factor) * factor) + resized_width = max(factor, math.floor(width / beta / factor) * factor) + elif resized_height * resized_width < min_pixels: + beta = math.sqrt(min_pixels / (height * width)) + resized_height = math.ceil(height * beta / factor) * factor + resized_width = math.ceil(width * beta / factor) * factor + + return resized_height, resized_width + + def process_qwen2vl_images( images: torch.Tensor, min_pixels: int = 3136, @@ -30,19 +47,7 @@ def process_qwen2vl_images( grid_thw_list = [] img = images[0] - factor = patch_size * merge_size - - h_bar = round(height / factor) * factor - w_bar = round(width / factor) * factor - - if h_bar * w_bar > max_pixels: - beta = math.sqrt((height * width) / max_pixels) - h_bar = max(factor, math.floor(height / beta / factor) * factor) - w_bar = max(factor, math.floor(width / beta / factor) * factor) - elif h_bar * w_bar < min_pixels: - beta = math.sqrt(min_pixels / (height * width)) - h_bar = math.ceil(height * beta / factor) * factor - w_bar = math.ceil(width * beta / factor) * factor + h_bar, w_bar = qwen2vl_image_size(height, width, min_pixels, max_pixels, patch_size, merge_size) img_resized = F.interpolate( img.unsqueeze(0), diff --git a/comfy_extras/nodes_cond.py b/comfy_extras/nodes_cond.py index c8091b7a4..0a2dc5ea2 100644 --- a/comfy_extras/nodes_cond.py +++ b/comfy_extras/nodes_cond.py @@ -1,8 +1,163 @@ +import torch +import torch.nn.functional as F from typing_extensions import override +import comfy.text_encoders.qwen_vl from comfy_api.latest import ComfyExtension, io +def _spatial_fusion_mask(height, width, num_sources, method, block_size, dither_ratio, device, seed=0): + rows = torch.arange(height).unsqueeze(1) + columns = torch.arange(width).unsqueeze(0) + + if method == "spatial-checkerboard": + mask = (rows + columns) % num_sources + elif method == "spatial-block-interleave": + mask = (rows // block_size + columns // block_size) % num_sources + elif method == "spatial-dither-random": + generator = torch.Generator().manual_seed(seed) + random = torch.rand((height, width), generator=generator) + other_sources = 1 + ((rows + columns) % (num_sources - 1)) + mask = torch.where(random < dither_ratio, 0, other_sources) + else: + raise ValueError(f"Unsupported visual fusion method: {method}") + return mask.flatten().to(device) + + +def _visual_token_span(tokens, cond_length, visual_tokens): + if len(tokens) != 1: + raise ValueError("Image fusion requires a compatible multimodal CLIP encoder with one token stream.") + + token_pairs = next(iter(tokens.values()))[0] + image_positions = [i for i, pair in enumerate(token_pairs) if isinstance(pair[0], dict) and pair[0].get("type") == "image"] + if len(image_positions) != 1: + raise ValueError("Image fusion requires exactly one visual token block per encoding pass.") + + image_position = image_positions[0] + if any(not isinstance(pair[0], (int, float)) for pair in token_pairs[image_position + 1:]): + raise ValueError("Image fusion does not support embeddings after the image token block.") + + end = cond_length - (len(token_pairs) - image_position - 1) + start = end - visual_tokens + if start < 0 or end > cond_length: + raise ValueError("Could not locate the visual token block in the encoded conditioning.") + return start, end + + +def _visual_grid(image): + height, width = image.shape[1:3] + height, width = comfy.text_encoders.qwen_vl.qwen2vl_image_size(height, width, patch_size=16, merge_size=2) + return height // 32, width // 32 + + +def _resize_visual_tokens(visual, source_grid, target_grid): + if source_grid == target_grid: + return visual + + dtype = visual.dtype + batch, _, dimensions = visual.shape + height, width = source_grid + target_height, target_width = target_grid + visual = visual.reshape(batch, height, width, dimensions).permute(0, 3, 1, 2).float() + visual = F.interpolate(visual, size=(target_height, target_width), mode="bilinear", align_corners=False) + return visual.permute(0, 2, 3, 1).reshape(batch, target_height * target_width, dimensions).to(dtype=dtype) + + +def _fuse_conditionings(conditionings, tokens, visual_grids, method, block_size, dither_ratio, seed=0): + schedule_count = len(conditionings[0]) + if any(len(source) != schedule_count for source in conditionings): + raise ValueError("All image fusion sources must use the same CLIP schedule.") + + target_grid = visual_grids[0] + fused = [] + for schedule in range(schedule_count): + source_conds = [source[schedule][0] for source in conditionings] + spans = [_visual_token_span(source_tokens, cond.shape[1], height * width) for source_tokens, cond, (height, width) in zip(tokens, source_conds, visual_grids)] + prefix_length = spans[0][0] + suffix_length = source_conds[0].shape[1] - spans[0][1] + if any(start != prefix_length or cond.shape[1] - end != suffix_length for cond, (start, end) in zip(source_conds[1:], spans[1:])): + raise ValueError("Image fusion sources produced different text token layouts.") + + visuals = [] + for cond, (start, end), grid in zip(source_conds, spans, visual_grids): + visual = _resize_visual_tokens(cond[:, start:end], grid, target_grid) + visuals.append(visual.to(dtype=source_conds[0].dtype, device=source_conds[0].device)) + + visuals = torch.stack(visuals, dim=2) + target_height, target_width = target_grid + mask = _spatial_fusion_mask(target_height, target_width, len(source_conds), method, block_size, dither_ratio, visuals.device, seed) + blended_visual = torch.take_along_dim(visuals, mask[None, :, None, None], dim=2).squeeze(2) + + start, end = spans[0] + blended = source_conds[0].clone() + blended[:, start:end] = blended_visual + fused.append([blended, conditionings[0][schedule][1].copy()]) + return fused + + +def _flatten_images(images): + sources = [] + for name in sorted(images, key=lambda value: int(value.rsplit("_", 1)[-1])): + image = images[name] + if image is None: + continue + if image.ndim == 3: + image = image.unsqueeze(0) + sources.extend(image[i:i + 1].clone() for i in range(image.shape[0])) + return sources + + +class CLIPTextEncodeImageFusion(io.ComfyNode): + @classmethod + def define_schema(cls): + images = io.Autogrow.TemplateNames( + io.Image.Input("image"), + names=[f"image_{i}" for i in range(1, 17)], + min=2, + ) + return io.Schema( + node_id="CLIPTextEncodeImageFusion", + display_name="CLIP Text Encode (Image Fusion)", + category="model/conditioning", + description="Encodes images separately and spatially interleaves their visual conditioning tokens.", + inputs=[ + io.Clip.Input("clip"), + io.String.Input("text", multiline=True, dynamic_prompts=True), + io.Autogrow.Input("images", template=images), + io.Combo.Input( + "fusion_method", + options=["spatial-checkerboard", "spatial-block-interleave", "spatial-dither-random"], + default="spatial-checkerboard", + ), + io.Int.Input("block_size", default=2, min=1, max=8, step=1, advanced=True), + io.Float.Input( + "dither_ratio", + default=0.5, + min=0.0, + max=1.0, + step=0.01, + advanced=True, + tooltip="Probability of selecting the first source. Remaining sources are selected with a checkerboard pattern.", + ), + io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff, control_after_generate=True, advanced=True, tooltip="Seed for the spatial-dither-random pattern."), + ], + outputs=[io.Conditioning.Output()], + ) + + @classmethod + def execute(cls, clip, text, images: io.Autogrow.Type, fusion_method, block_size=2, dither_ratio=0.5, seed=0) -> io.NodeOutput: + sources = _flatten_images(images) + if len(sources) < 2: + raise ValueError("Image fusion requires at least two images.") + + sources = [source[:, :, :, :3] for source in sources] + visual_grids = [_visual_grid(source) for source in sources] + tokens = [clip.tokenize(text, images=[source]) for source in sources] + conditionings = [clip.encode_from_tokens_scheduled(source_tokens) for source_tokens in tokens] + conditioning = _fuse_conditionings(conditionings, tokens, visual_grids, fusion_method, block_size, dither_ratio, seed) + return io.NodeOutput(conditioning) + + class CLIPTextEncodeControlnet(io.ComfyNode): @classmethod def define_schema(cls) -> io.Schema: @@ -61,6 +216,7 @@ class CondExtension(ComfyExtension): @override async def get_node_list(self) -> list[type[io.ComfyNode]]: return [ + CLIPTextEncodeImageFusion, CLIPTextEncodeControlnet, T5TokenizerOptions, ] diff --git a/comfy_extras/nodes_qwen.py b/comfy_extras/nodes_qwen.py index 66dd97939..4960774db 100644 --- a/comfy_extras/nodes_qwen.py +++ b/comfy_extras/nodes_qwen.py @@ -7,79 +7,6 @@ import comfy.model_management import torch import nodes - -def _spatial_fusion_mask(height, width, num_sources, method, block_size, dither_ratio, device, seed=0): - rows = torch.arange(height, device=device).unsqueeze(1) - columns = torch.arange(width, device=device).unsqueeze(0) - - if method == "spatial-checkerboard": - return ((rows + columns) % num_sources).flatten() - if method == "spatial-block-interleave": - return ((rows // block_size + columns // block_size) % num_sources).flatten() - if method == "spatial-dither-random": - generator = torch.Generator(device=device).manual_seed(seed) - random = torch.rand((height, width), generator=generator, device=device) - other_sources = 1 + ((rows + columns) % (num_sources - 1)) - return torch.where(random < dither_ratio, 0, other_sources).flatten() - raise ValueError(f"Unsupported visual fusion method: {method}") - - -def _visual_token_span(tokens, cond_length, visual_tokens): - if len(tokens) != 1: - raise ValueError("Visual fusion requires a Qwen3-VL or Krea2 text encoder.") - - token_pairs = next(iter(tokens.values()))[0] - image_positions = [i for i, pair in enumerate(token_pairs) if isinstance(pair[0], dict) and pair[0].get("type") == "image"] - if len(image_positions) != 1: - raise ValueError("Visual fusion requires exactly one visual token block per encoding pass.") - - image_position = image_positions[0] - if any(not isinstance(pair[0], (int, float)) for pair in token_pairs[image_position + 1:]): - raise ValueError("Visual fusion does not support embeddings after the image token block.") - - end = cond_length - (len(token_pairs) - image_position - 1) - start = end - visual_tokens - if start < 0 or end > cond_length: - raise ValueError("Could not locate the visual token block in the encoded conditioning.") - return start, end - - -def _fuse_conditionings(conditionings, tokens, visual_height, visual_width, method, block_size, dither_ratio, seed=0): - schedule_count = len(conditionings[0]) - if any(len(source) != schedule_count for source in conditionings): - raise ValueError("All visual fusion sources must use the same CLIP schedule.") - - visual_tokens = visual_height * visual_width - fused = [] - for schedule in range(schedule_count): - source_conds = [source[schedule][0] for source in conditionings] - spans = [_visual_token_span(source_tokens, cond.shape[1], visual_tokens) for source_tokens, cond in zip(tokens, source_conds)] - if any(span != spans[0] for span in spans[1:]): - raise ValueError("Visual fusion sources produced different token layouts.") - - start, end = spans[0] - visuals = torch.stack([cond[:, start:end] for cond in source_conds], dim=2) - mask = _spatial_fusion_mask(visual_height, visual_width, len(source_conds), method, block_size, dither_ratio, visuals.device, seed) - blended_visual = torch.take_along_dim(visuals, mask[None, :, None, None], dim=2).squeeze(2) - - blended = source_conds[0].clone() - blended[:, start:end] = blended_visual - fused.append([blended, conditionings[0][schedule][1].copy()]) - return fused - - -def _flatten_images(images): - sources = [] - for name in sorted(images, key=lambda value: int(value.rsplit("_", 1)[-1])): - image = images[name] - if image is None: - continue - if image.ndim == 3: - image = image.unsqueeze(0) - sources.extend(image[i:i + 1].clone() for i in range(image.shape[0])) - return sources - - class TextEncodeQwenImageEdit(io.ComfyNode): @classmethod def define_schema(cls): @@ -179,93 +106,6 @@ class TextEncodeQwenImageEditPlus(io.ComfyNode): return io.NodeOutput(conditioning) -class TextEncodeQwenImageEditFusion(io.ComfyNode): - @classmethod - def define_schema(cls): - images = io.Autogrow.TemplateNames( - io.Image.Input("image"), - names=[f"image_{i}" for i in range(1, 17)], - min=2, - ) - return io.Schema( - node_id="TextEncodeQwenImageEditFusion", - display_name="Text Encode Qwen Image Edit (Visual Fusion)", - category="model/conditioning/qwen image", - description="Encodes images separately and spatially interleaves their Qwen3-VL visual conditioning tokens.", - inputs=[ - io.Clip.Input("clip"), - io.String.Input("prompt", multiline=True, dynamic_prompts=True), - io.Autogrow.Input("images", template=images), - io.Combo.Input( - "fusion_method", - options=["spatial-checkerboard", "spatial-block-interleave", "spatial-dither-random"], - default="spatial-checkerboard", - ), - io.Int.Input("block_size", default=2, min=1, max=8, step=1, advanced=True), - io.Float.Input( - "dither_ratio", - default=0.5, - min=0.0, - max=1.0, - step=0.01, - advanced=True, - tooltip="Probability of selecting the first source. Remaining sources are selected with a checkerboard pattern.", - ), - io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff, control_after_generate=True, advanced=True, tooltip="Seed for the spatial-dither-random pattern."), - io.Vae.Input("vae", optional=True), - ], - outputs=[io.Conditioning.Output()], - ) - - @classmethod - def execute(cls, clip, prompt, images: io.Autogrow.Type, fusion_method, block_size=2, dither_ratio=0.5, vae=None, seed=0) -> io.NodeOutput: - sources = _flatten_images(images) - if len(sources) < 2: - raise ValueError("Visual fusion requires at least two images.") - - first = sources[0].movedim(-1, 1) - total = 384 * 384 - scale_by = math.sqrt(total / (first.shape[3] * first.shape[2])) - width = max(32, round(first.shape[3] * scale_by)) - height = max(32, round(first.shape[2] * scale_by)) - - processed = [] - for source in sources: - samples = source[:, :, :, :3].movedim(-1, 1) - resized = comfy.utils.common_upscale(samples, width, height, "area", "center") - processed.append(resized.movedim(1, -1)) - - factor = 32 - visual_height = max(factor, round(height / factor) * factor) // factor - visual_width = max(factor, round(width / factor) * factor) // factor - - full_prompt = ( - "<|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<|vision_start|><|image_pad|><|vision_end|>" + prompt + "<|im_end|>\n" - "<|im_start|>assistant\n" - ) - tokens = [clip.tokenize(full_prompt, images=[image]) for image in processed] - token_key = next(iter(tokens[0]), None) - if token_key not in ("qwen3vl_4b", "qwen3vl_8b") or any(next(iter(source_tokens), None) != token_key for source_tokens in tokens): - raise ValueError("Visual fusion requires a Qwen3-VL or Krea2 text encoder.") - - conditionings = [clip.encode_from_tokens_scheduled(source_tokens) for source_tokens in tokens] - conditioning = _fuse_conditionings(conditionings, tokens, visual_height, visual_width, fusion_method, block_size, dither_ratio, seed) - - if vae is not None: - ref_latents = [] - for source in sources: - samples = source[:, :, :, :3].movedim(-1, 1) - scale_by = math.sqrt((1024 * 1024) / (samples.shape[3] * samples.shape[2])) - latent_width = max(8, round(samples.shape[3] * scale_by / 8.0) * 8) - latent_height = max(8, round(samples.shape[2] * scale_by / 8.0) * 8) - resized = comfy.utils.common_upscale(samples, latent_width, latent_height, "area", "disabled") - ref_latents.append(vae.encode(resized.movedim(1, -1))) - conditioning = node_helpers.conditioning_set_values(conditioning, {"reference_latents": ref_latents}, append=True) - - return io.NodeOutput(conditioning) - - class EmptyQwenImageLayeredLatentImage(io.ComfyNode): @classmethod def define_schema(cls): @@ -296,7 +136,6 @@ class QwenExtension(ComfyExtension): return [ TextEncodeQwenImageEdit, TextEncodeQwenImageEditPlus, - TextEncodeQwenImageEditFusion, EmptyQwenImageLayeredLatentImage, ] diff --git a/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py b/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py index e918cbecd..d878ee0be 100644 --- a/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py +++ b/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py @@ -8,7 +8,7 @@ if not torch.cuda.is_available(): cli_args.cpu = True try: - from comfy_extras.nodes_qwen import TextEncodeQwenImageEditFusion, _flatten_images, _fuse_conditionings, _spatial_fusion_mask, _visual_token_span + from comfy_extras.nodes_cond import CLIPTextEncodeImageFusion, _flatten_images, _fuse_conditionings, _resize_visual_tokens, _spatial_fusion_mask, _visual_grid, _visual_token_span finally: cli_args.cpu = prior_cpu @@ -72,7 +72,7 @@ def test_fusion_replaces_only_visual_tokens_and_preserves_dtype_and_metadata(): [[second, {"pooled_output": torch.tensor([2.0])}]], ] - fused = _fuse_conditionings(conditionings, tokens, 2, 2, "spatial-checkerboard", 2, 0.5) + fused = _fuse_conditionings(conditionings, tokens, [(2, 2), (2, 2)], "spatial-checkerboard", 2, 0.5) output, output_metadata = fused[0] assert output.dtype == torch.float16 @@ -88,8 +88,8 @@ def test_dither_seed_changes_fused_conditioning(): [[torch.ones((1, 6, 1)), {}]], ] - first = _fuse_conditionings(conditionings, tokens, 2, 2, "spatial-dither-random", 2, 0.5, 7)[0][0] - second = _fuse_conditionings(conditionings, tokens, 2, 2, "spatial-dither-random", 2, 0.5, 8)[0][0] + first = _fuse_conditionings(conditionings, tokens, [(2, 2), (2, 2)], "spatial-dither-random", 2, 0.5, 7)[0][0] + second = _fuse_conditionings(conditionings, tokens, [(2, 2), (2, 2)], "spatial-dither-random", 2, 0.5, 8)[0][0] assert not torch.equal(first, second) @@ -103,13 +103,47 @@ def test_flatten_images_uses_numeric_input_order_and_splits_batches(): sources = _flatten_images(images) assert [source[0, 0, 0, 0].item() for source in sources] == [1.0, 2.0, 3.0, 10.0] + images["image_2"][0, 0, 0, 0] = 99.0 + assert sources[1][0, 0, 0, 0].item() == 2.0 -def test_node_uses_custom_krea_prompt_and_returns_fused_conditioning(): +def test_visual_tokens_interpolate_in_two_dimensions_and_restore_dtype(): + visual = torch.tensor([[[0.0], [2.0]]], dtype=torch.float16) + + resized = _resize_visual_tokens(visual, (2, 1), (2, 2)) + + assert resized.dtype == torch.float16 + assert resized.flatten().tolist() == [0.0, 0.0, 2.0, 2.0] + + +def test_fusion_interpolates_to_first_image_grid(): + first = torch.zeros((1, 6, 1), dtype=torch.float16) + second = torch.tensor([[[0.0], [0.0], [2.0], [0.0]]], dtype=torch.float16) + conditionings = [[[first, {}]], [[second, {}]]] + + fused = _fuse_conditionings(conditionings, [_tokens(), _tokens()], [(2, 2), (2, 1)], "spatial-dither-random", 2, 0.0)[0][0] + + assert fused.shape == first.shape + assert fused.flatten().tolist() == [0.0, 0.0, 0.0, 2.0, 2.0, 0.0] + + +def test_fusion_rejects_different_text_layouts(): + conditionings = [ + [[torch.zeros((1, 6, 1)), {}]], + [[torch.zeros((1, 7, 1)), {}]], + ] + + with pytest.raises(ValueError, match="different text token layouts"): + _fuse_conditionings(conditionings, [_tokens(), _tokens(suffix=2)], [(2, 2), (2, 2)], "spatial-checkerboard", 2, 0.5) + + +def test_node_preserves_images_uses_tokenizer_template_and_returns_fused_conditioning(): + seen_shapes = [] + class FakeClip: def tokenize(self, text, images): - assert text.startswith("<|im_start|>system\nDescribe the image by detailing") - assert "Picture 1:" not in text + assert text == "test prompt" + seen_shapes.append(images[0].shape) pairs = [ (1, 1.0), ({"type": "image", "data": images[0]}, 1.0), @@ -120,27 +154,30 @@ def test_node_uses_custom_krea_prompt_and_returns_fused_conditioning(): def encode_from_tokens_scheduled(self, tokens): image = next(pair[0]["data"] for pair in tokens["qwen3vl_4b"][0] if isinstance(pair[0], dict)) value = image.mean() - return [[torch.full((1, 146, 1), value, dtype=torch.float16), {"source": float(value)}]] + height, width = _visual_grid(image) + return [[torch.full((1, height * width + 2, 1), value, dtype=torch.float16), {"source": float(value)}]] - result = TextEncodeQwenImageEditFusion.execute( + images = {"image_1": torch.zeros(1, 32, 64, 4), "image_2": torch.ones(1, 64, 32, 4)} + result = CLIPTextEncodeImageFusion.execute( FakeClip(), "test prompt", - {"image_1": torch.zeros(1, 32, 32, 3), "image_2": torch.ones(1, 32, 32, 3)}, + images, "spatial-dither-random", seed=7, ) - changed_seed = TextEncodeQwenImageEditFusion.execute( + changed_seed = CLIPTextEncodeImageFusion.execute( FakeClip(), "test prompt", - {"image_1": torch.zeros(1, 32, 32, 3), "image_2": torch.ones(1, 32, 32, 3)}, + images, "spatial-dither-random", seed=8, ) conditioning = result.args[0] output, metadata = conditioning[0] + assert seen_shapes == [torch.Size([1, 32, 64, 3]), torch.Size([1, 64, 32, 3])] * 2 assert output.dtype == torch.float16 - assert output.shape == (1, 146, 1) + assert output.shape == (1, 8, 1) assert output[:, 0].item() == 0.0 assert output[:, -1].item() == 0.0 assert set(output[:, 1:-1].flatten().tolist()) == {0.0, 1.0} @@ -148,7 +185,12 @@ def test_node_uses_custom_krea_prompt_and_returns_fused_conditioning(): assert not torch.equal(output, changed_seed.args[0][0][0]) -def test_node_exposes_seed_control(): - inputs = {value.id: value for value in TextEncodeQwenImageEditFusion.define_schema().inputs} +def test_node_exposes_generic_interface_without_vae(): + schema = CLIPTextEncodeImageFusion.define_schema() + inputs = {value.id: value for value in schema.inputs} + assert schema.node_id == "CLIPTextEncodeImageFusion" + assert schema.category == "model/conditioning" + assert "text" in inputs + assert "vae" not in inputs assert inputs["seed"].default == 0 assert inputs["seed"].control_after_generate is True From 20a8d3b9aecdf86487f8a7532364927a30749a4d Mon Sep 17 00:00:00 2001 From: silveroxides Date: Wed, 15 Jul 2026 20:58:59 +0200 Subject: [PATCH 8/9] fix: preserve Krea visual tokens --- comfy/text_encoders/krea2.py | 27 ++++++++++--------- .../test_qwen_visual_fusion.py | 17 ++++++++++++ 2 files changed, 32 insertions(+), 12 deletions(-) diff --git a/comfy/text_encoders/krea2.py b/comfy/text_encoders/krea2.py index 408a03566..0341dd6e3 100644 --- a/comfy/text_encoders/krea2.py +++ b/comfy/text_encoders/krea2.py @@ -20,6 +20,20 @@ KREA2_TAP_LAYERS = [2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35] 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" +def _krea2_template_end(tok_pairs): + image_position = next((i for i, pair in enumerate(tok_pairs) if isinstance(pair[0], dict) and pair[0].get("type") == "image"), len(tok_pairs)) + template_end = -1 + for i in range(image_position): + if i + 2 >= len(tok_pairs): + break + values = [tok_pairs[j][0] for j in range(i, i + 3)] + if all(not torch.is_tensor(value) and isinstance(value, numbers.Integral) for value in values) and values == [151644, 872, 198]: + template_end = i + 3 + if template_end == -1: + raise ValueError("Could not locate the Krea 2 user prompt template.") + return template_end + + class Krea2Tokenizer(comfy.text_encoders.qwen3vl.Qwen3VLTokenizer): def __init__(self, embedding_directory=None, tokenizer_data={}): super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, model_type="qwen3vl_4b") @@ -44,19 +58,8 @@ class Krea2TEModel(sd1_clip.SD1ClipModel): out, pooled, extra = super().encode_token_weights(token_weight_pairs) # out: (B, 12, seq, 2560) tok_pairs = token_weight_pairs["qwen3vl_4b"][0] - # Strip the system + user-opening prefix - count_im_start = 0 if template_end == -1: - for i, v in enumerate(tok_pairs): - elem = v[0] - if not torch.is_tensor(elem) and isinstance(elem, numbers.Integral): - if elem == 151644 and count_im_start < 2: - template_end = i - count_im_start += 1 - if out.shape[2] > (template_end + 3): - if tok_pairs[template_end + 1][0] == 872: # "user" - if tok_pairs[template_end + 2][0] == 198: # "\n" - template_end += 3 + template_end = _krea2_template_end(tok_pairs) out = out[:, :, template_end:] diff --git a/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py b/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py index d878ee0be..ff3c104c7 100644 --- a/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py +++ b/tests-unit/comfy_extras_test/test_qwen_visual_fusion.py @@ -8,6 +8,7 @@ if not torch.cuda.is_available(): cli_args.cpu = True try: + from comfy.text_encoders.krea2 import KREA2_TEMPLATE, Krea2Tokenizer, _krea2_template_end from comfy_extras.nodes_cond import CLIPTextEncodeImageFusion, _flatten_images, _fuse_conditionings, _resize_visual_tokens, _spatial_fusion_mask, _visual_grid, _visual_token_span finally: cli_args.cpu = prior_cpu @@ -62,6 +63,22 @@ def test_visual_span_accounts_for_stripped_prefix(): assert _visual_token_span(tokens, cond_length=9, visual_tokens=4) == (1, 5) +def test_krea_template_stripping_preserves_visual_tokens(): + tokenizer = Krea2Tokenizer() + image = torch.zeros(1, 64, 64, 3) + image_prompt = "<|vision_start|><|image_pad|><|vision_end|>test prompt" + + for text in ("test prompt", KREA2_TEMPLATE.format(image_prompt)): + tokens = tokenizer.tokenize_with_weights(text, images=[image]) + token_pairs = tokens["qwen3vl_4b"][0] + image_position = next(i for i, pair in enumerate(token_pairs) if isinstance(pair[0], dict)) + template_end = _krea2_template_end(token_pairs) + cond_length = len(token_pairs) - 1 + 4 - template_end + + assert template_end <= image_position + assert _visual_token_span(tokens, cond_length, 4) == (image_position - template_end, image_position - template_end + 4) + + def test_fusion_replaces_only_visual_tokens_and_preserves_dtype_and_metadata(): tokens = [_tokens(), _tokens()] first = torch.tensor([[[10], [10], [10], [10], [10], [20]]], dtype=torch.float16) From 411ad961b081a7824b4320a563ef0f6e1d4c70d8 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Wed, 15 Jul 2026 21:06:17 +0200 Subject: [PATCH 9/9] perf: preserve dtype during visual resize --- comfy_extras/nodes_cond.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/comfy_extras/nodes_cond.py b/comfy_extras/nodes_cond.py index 0a2dc5ea2..262e5bd30 100644 --- a/comfy_extras/nodes_cond.py +++ b/comfy_extras/nodes_cond.py @@ -54,13 +54,12 @@ def _resize_visual_tokens(visual, source_grid, target_grid): if source_grid == target_grid: return visual - dtype = visual.dtype batch, _, dimensions = visual.shape height, width = source_grid target_height, target_width = target_grid - visual = visual.reshape(batch, height, width, dimensions).permute(0, 3, 1, 2).float() + visual = visual.reshape(batch, height, width, dimensions).permute(0, 3, 1, 2) visual = F.interpolate(visual, size=(target_height, target_width), mode="bilinear", align_corners=False) - return visual.permute(0, 2, 3, 1).reshape(batch, target_height * target_width, dimensions).to(dtype=dtype) + return visual.permute(0, 2, 3, 1).reshape(batch, target_height * target_width, dimensions) def _fuse_conditionings(conditionings, tokens, visual_grids, method, block_size, dither_ratio, seed=0):