Merge 411ad961b0 into 8edea4a65d

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:]
 

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,
@@ -31,19 +48,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),

comfy_extras/nodes_cond.py

@@ -1,8 +1,162 @@
+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
+
+    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)
+    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)
+
+
+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 +215,7 @@ class CondExtension(ComfyExtension):
     @override
     async def get_node_list(self) -> list[type[io.ComfyNode]]:
         return [
+            CLIPTextEncodeImageFusion,
             CLIPTextEncodeControlnet,
             T5TokenizerOptions,
         ]

tests-unit/comfy_extras_test/test_qwen_visual_fusion.py

@@ -0,0 +1,213 @@
+import pytest
+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
+
+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
+
+
+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_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)
+
+
+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)
+    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), (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_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), (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)
+
+
+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]
+    images["image_2"][0, 0, 0, 0] = 99.0
+    assert sources[1][0, 0, 0, 0].item() == 2.0
+
+
+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 == "test prompt"
+            seen_shapes.append(images[0].shape)
+            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()
+            height, width = _visual_grid(image)
+            return [[torch.full((1, height * width + 2, 1), value, dtype=torch.float16), {"source": float(value)}]]
+
+    images = {"image_1": torch.zeros(1, 32, 64, 4), "image_2": torch.ones(1, 64, 32, 4)}
+    result = CLIPTextEncodeImageFusion.execute(
+        FakeClip(),
+        "test prompt",
+        images,
+        "spatial-dither-random",
+        seed=7,
+    )
+    changed_seed = CLIPTextEncodeImageFusion.execute(
+        FakeClip(),
+        "test prompt",
+        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, 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}
+    assert metadata == {"source": 0.0}
+    assert not torch.equal(output, changed_seed.args[0][0][0])
+
+
+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