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