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

@@ -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,
         ]