Noise merging: merge the sound too; test that run 2 and 3 on the same cached model still make the extra pass

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VmbxSsiiz8wmt226dtbFXe

modules/sampling/noise_merge/__init__.py

@@ -3,7 +3,7 @@ different noise and blend the two answers.
 
 Step N is answered twice. Picture A is the run's own answer. Picture B is the answer to
 the same step after A's picture is re-noised to this step's noise level with a second
-seed (the sound part of the input is left alone). The two are then mixed:
+seed (and so is the sound, when the model has it). The two are then mixed:
 
     merged = base + strength * (other - base)
 
@@ -89,6 +89,14 @@ def renoise(x0, sigma, seed, step):
     return (1.0 - sigma) * x0 + sigma * noise
 
 
+def _sound(packed, video):
+    """The part of a packed [B, 1, N] result that follows the picture (H3's sound), or None."""
+    if getattr(packed, "dim", None) is None or packed.dim() != 3:
+        return None
+    n = video[0].numel()
+    return packed[..., n:] if packed.shape[-1] > n else None
+
+
 def install(patcher, values, key):
     if not values.get("enabled"):
         return None
@@ -134,7 +142,12 @@ def install(patcher, values, key):
                 return out
             sigma = float(t.max())
             video_a, rebuild_a = split_a
-            x_b = split_x[1](renoise(video_a.detach(), sigma, live["seed"], where[0]).to(split_x[0].dtype))
+            tail_a = _sound(out, video_a)
+            gen_step = where[0]
+            x_b = split_x[1](renoise(video_a.detach(), sigma, live["seed"], gen_step).to(split_x[0].dtype))
+            if tail_a is not None:
+                n = x_b.shape[-1] - tail_a.shape[-1]
+                x_b = torch.cat([x_b[..., :n], renoise(tail_a.detach(), sigma, live["seed"] + 1, gen_step)], dim=-1)
             args_b, kwargs_b = streams.with_options(args, kwargs, {**to, dit_hooks.PROBE: True})
             began = time.perf_counter()
             out_b = executor(x_b, t, *args_b, **kwargs_b)
@@ -143,14 +156,22 @@ def install(patcher, values, key):
             if split_b is None:
                 return out
             video_b, rebuild_b = split_b
+            tail_b = _sound(out_b, video_b)
             live["merged"] += 1
             diff = float((video_a - video_b).norm() / video_a.norm().clamp(min=1e-8))
             log.info("FunUI Noise merging",
                      f"step {where[0] + 1}: second answer differs from the first by {diff:.1%}, "
-                     f"merged at {strength:.2f}; the second pass took {took:.1f}s")
+                     f"merged at {strength:.2f}{' (picture and sound)' if tail_a is not None else ''}; "
+                     f"the second pass took {took:.1f}s")
             if b_first:
-                return rebuild_a(merge(video_a, video_b, strength).to(video_a.dtype))
-            return rebuild_b(merge(video_b, video_a, strength).to(video_b.dtype))
+                base, other, base_tail, other_tail, rebuild = video_a, video_b, tail_a, tail_b, rebuild_a
+            else:
+                base, other, base_tail, other_tail, rebuild = video_b, video_a, tail_b, tail_a, rebuild_b
+            merged = rebuild(merge(base, other, strength).to(base.dtype))
+            if base_tail is not None and other_tail is not None:
+                n = merged.shape[-1] - base_tail.shape[-1]
+                merged = torch.cat([merged[..., :n], merge(base_tail, other_tail, strength).to(merged.dtype)], dim=-1)
+            return merged
         except Exception as exc:                    # noqa: BLE001
             log.broke("FunUI Noise merging", exc, "merging the second answer")
             return out

modules/sampling/noise_merge/tests/test_noise_merge.py

@@ -69,7 +69,7 @@ def test_a_named_step_runs_the_model_twice_and_others_once(tiny_h3):
     assert len(calls) == 5 and calls.count(True) == 1
 
 
-def test_the_answer_changes_only_in_the_picture_and_order_matters(tiny_h3):
+def test_the_answer_changes_and_order_matters(tiny_h3):
     from conftest import unpacked
     x0, shapes = _packed()
     model = lambda x, *a, **k: x * 0.5 + 0.3 * x0
@@ -81,7 +81,7 @@ def test_the_answer_changes_only_in_the_picture_and_order_matters(tiny_h3):
     a, b = unpacked(results["b_into_a"], shapes), unpacked(results["a_into_b"], shapes)
     p = unpacked(plain, shapes)
     assert not torch.allclose(a[0], p[0]) and not torch.allclose(a[0], b[0])
-    assert torch.allclose(a[1], p[1])                 # sound is the model's own answer for A
+    assert not torch.allclose(a[1], p[1])             # the sound is merged too
 
 
 def test_nothing_named_or_strength_zero_installs_nothing(tiny_h3):
@@ -122,3 +122,14 @@ def test_the_log_says_the_feature_is_on_and_when_nothing_matched(tiny_h3):
     texts = [e["message"] for e in log.history()]
     assert any(t.startswith("on: steps 9") for t in texts)
     assert any("no step of this run matched" in t for t in texts)
+
+
+def test_the_sound_is_merged_too(tiny_h3):
+    from conftest import unpacked
+    x0, shapes = _packed()
+    model = lambda x, *a, **k: x * 0.5 + 0.3 * x0
+    wrap, _ = _load(tiny_h3, steps="1", strength=0.7)
+    out = wrap(model, x0.clone(), _sched()[0:1], None, None, None, _to(0), latent_shapes=shapes)
+    plain = unpacked(model(x0, None), shapes)
+    got = unpacked(out, shapes)
+    assert not torch.allclose(got[0], plain[0]) and not torch.allclose(got[1], plain[1])

modules/sampling/noise_merge/tests/test_through_the_app.py

@@ -96,3 +96,27 @@ def test_what_the_app_saves_reaches_the_h3_sampler_and_adds_one_pass_at_step_one
     texts = [e["message"] for e in log.history()]
     assert any(t.startswith("on: steps 1") for t in texts)
     assert any(t.startswith("step 1: second answer differs") for t in texts)
+
+
+def test_a_second_run_on_the_same_cached_model_still_makes_the_extra_pass(server, model):
+    """ComfyUI hands the SAME patched model to every run whose Load Modifiers inputs did not change."""
+    from comfy.nested_tensor import NestedTensor
+    from core import log
+    from modules.sampling.modifiers.nodes import FunUILoadModifiers, FunUIModifierSettings
+    from modules.sampling.sampler.nodes import FunUISampler
+    saved = _saved_by_the_app(server, APP_VALUES)
+    patched, _ = FunUILoadModifiers.execute(model, FunUIModifierSettings.execute(saved).result[0]).result
+    counts = []
+    for seed in (1, 2, 3):
+        forwards = []
+        hook = patched.model.diffusion_model.register_forward_hook(lambda *a: forwards.append(1))
+        cond = [[torch.randn(1, 5, 48), {}]]
+        latent = {"samples": NestedTensor((torch.zeros(1, 24, 2, 4, 4), torch.zeros(1, 32, 2, 3)))}
+        before = sum("second pass took" in e["message"] for e in log.history())
+        try:
+            FunUISampler.execute(patched, cond, cond, latent, seed=seed, steps=4, cfg=1.0,
+                                 sampler_name="lcm", scheduler="beta", denoise=1.0)
+        finally:
+            hook.remove()
+        counts.append((len(forwards), sum("second pass took" in e["message"] for e in log.history()) - before))
+    assert counts == [(5, 1)] * 3, counts