feat: add seeded multi-source spatial dither

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 = []

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