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

CHANGELOG.md

@@ -4,7 +4,7 @@
 
 ### Added (experimental)
 
-- **Noise merging** (Settings > Engine > Sampling): at the steps you name, the step is answered a second time on the same picture with different noise (a second seed), and the two pictures are blended at your strength. Merge order picks which one is the base. The second seed is random each run unless you type one. Costs one extra model pass per merged step. Untested on a real H3.
+- **Noise merging** (Settings > Engine > Sampling): at the steps you name, the step is answered a second time on the same picture with different noise (a second seed), and the two pictures are blended at your strength. Merge order picks which one is the base. The second seed is random each run unless you type one. Picture and sound are both merged. Costs one extra model pass per merged step. Untested on a real H3.
 - **SLA rising sparsity** (diffusion model loader, `sla_sparsity_start`): skip fewer attention blocks on early steps and more on late ones, in a straight line up to `sla_sparsity`. Off at 0. One real H3 render (4 steps, 0.70 to 0.90): quality held, perhaps slightly better. Idea from GSRS (arXiv 2610.21129); only the schedule is taken, not its other parts.
 
 ## [5.0.3] "Classy Chopper" - 2026-10-25

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