Review: Brownian noise is the stock noise again; block map profile survives audio+video blocks and cancel

Queries are clamped into torchsde's span instead of moving the span. A cancelled profile saves nothing.
Attention explorer temperature stays above zero under the map; unmeasured groups are not given 3x.

core/brownian.py

@@ -3,29 +3,25 @@
 torchsde rounds a BrownianTree's start to 6 decimals (its default tolerance) and then compares the first query
 against the ROUNDED start. A schedule whose smallest sigma has a 7th decimal of 5 or more (0.4465445876 rounds to
 0.446545) is therefore queried "before" the tree starts and torchsde prints
-`Should have ta>=t0 but got ta=... and t0=...`. The query is clamped by 4e-7, so the picture is not hurt, but the
-log fills with warnings on every SDE-type sampler, and Two stages / Turbo hand such sigmas to the sampler.
+`Should have ta>=t0 but got ta=... and t0=...`, then clamps the query to the start itself. The picture is not hurt,
+but the log fills with warnings on every SDE-type sampler, and Two stages / Turbo hand such sigmas to the sampler.
 
-ComfyUI builds its trees from `sigmas[sigmas > 0].min()` and `sigmas.max()`. The class below widens that span to
-the next 6-decimal value outside it (plus one more step of margin) before the tree rounds, so the start is never
-above the first query and the end never below the last. Noise inside the span is unchanged.
+The class below builds the tree exactly as ComfyUI does (so the noise of a seed is the stock noise) and clamps a
+query to the tree's own span BEFORE torchsde sees it, which is the very same clamp torchsde would apply, silently.
 """
 
-import math
-
 import torch
 
-_SCALE = 1e6            # torchsde's default tol=1e-6 rounds to 6 decimals
 _installed = False
 
 
-def widened(lo, hi):
-    """(lo, hi) as tensors of the same device and dtype as the inputs, pushed outwards past torchsde's rounding."""
-    lo, hi = torch.as_tensor(lo), torch.as_tensor(hi)
-    new_lo = math.floor(float(lo) * _SCALE) / _SCALE - 1.0 / _SCALE
-    new_hi = math.ceil(float(hi) * _SCALE) / _SCALE + 1.0 / _SCALE
-    return (torch.as_tensor(new_lo).to(device=lo.device, dtype=lo.dtype),
-            torch.as_tensor(new_hi).to(device=hi.device, dtype=hi.dtype))
+def clamped(t, lo, hi):
+    """`t` limited to [lo, hi] (python floats) on t's own device; a value already inside comes back untouched.
+    The bound is held in float64 so a bf16 query is not rounded back out of the span."""
+    v = float(t)
+    if lo <= v <= hi:
+        return t
+    return torch.tensor(min(max(v, lo), hi), dtype=torch.float64, device=t.device)
 
 
 def install():
@@ -40,12 +36,12 @@ def install():
     class BrownianTreeNoiseSampler(original):
         _funui_widened = True
 
-        def __init__(self, x, sigma_min, sigma_max, seed=None, transform=lambda v: v, cpu=False):
-            super().__init__(x, sigma_min, sigma_max, seed=seed, transform=transform, cpu=cpu)
-            t0, t1 = self.transform(torch.as_tensor(sigma_min)), self.transform(torch.as_tensor(sigma_max))
-            lo, hi = (t0, t1) if t0 < t1 else (t1, t0)
-            lo, hi = widened(lo, hi)
-            self.tree = ksampling.BatchedBrownianTree(x, lo, hi, seed, cpu=cpu)
+        def __call__(self, sigma, sigma_next):
+            t0, t1 = self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next))
+            tree = self.tree.trees[0]._interval
+            lo, hi = float(tree._start), float(tree._end)
+            noise = self.tree(clamped(t0, lo, hi), clamped(t1, lo, hi)).to(dtype=tree.dtype)      # a float64 bound must not widen the noise
+            return noise / (t1 - t0).abs().sqrt()
 
     ksampling.BrownianTreeNoiseSampler = BrownianTreeNoiseSampler
     _installed = True

core/tests/test_brownian.py

@@ -56,9 +56,26 @@ def test_install_twice_does_not_stack(patched):
     assert ksampling.BrownianTreeNoiseSampler is before
 
 
-def test_widened_follows_dtype_and_device():
+def test_the_noise_of_a_seed_is_the_stock_noise(patched):
+    """The warning is silenced without moving the tree: same seed, same noise as ComfyUI's own sampler."""
+    stock = patched.__mro__[1]
+    for schedule in (SIGMAS, [1.0, 0.5, 0.25, 0.0], [0.9, 0.7, 0.5, 0.3, 0.1, 0.0]):
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            want = _run(stock, torch.float32, schedule)
+        got = _run(patched, torch.float32, schedule)
+        assert all(torch.equal(a, b) for a, b in zip(want, got))
+
+
+def test_a_query_is_clamped_on_its_own_device_and_never_rounded_out_of_the_span():
     mod = _brownian()
-    for dtype in (torch.float32, torch.float64, torch.bfloat16):
-        lo, hi = mod.widened(torch.tensor(0.4465445876, dtype=dtype), torch.tensor(0.99, dtype=dtype))
-        assert lo.dtype == hi.dtype == dtype
-        assert float(lo) < 0.4465445876 - 5e-7 and float(hi) > 0.99 or dtype == torch.bfloat16
+    for dtype in (torch.float32, torch.bfloat16):
+        for device in ("cpu", "meta"):
+            t = torch.tensor(0.4465445876, dtype=dtype, device="cpu")
+            if device == "meta":
+                t = t.to("meta")
+                continue                                      # float() needs data: the meta stand-in cannot be read
+            out = mod.clamped(t, 0.4475, 0.99)               # bf16 holds 0.4473 for it: below the bound, so it is lifted to the bound itself
+            assert out.device == t.device and float(out) == 0.4475
+    inside = torch.tensor(0.6)
+    assert mod.clamped(inside, 0.446545, 0.99) is inside

modules/sampling/attention_bandit/__init__.py

@@ -289,7 +289,7 @@ def followed(factors, options, index, blocks):
     from 1 is spread by how much that block can take at this point of the run and softly limited."""
     if (options or {}).get("funui_block_map") is None:
         return factors
-    return tuple(1.0 + dit_hooks.block_map_value(options, "attention", index, v - 1.0, blocks) for v in factors)
+    return tuple(max(0.05, 1.0 + dit_hooks.block_map_value(options, "attention", index, v - 1.0, blocks)) for v in factors)       # never zero or negative: that would flip the query
 
 
 def _prompt_of(guider):

modules/sampling/attention_bandit/tests/test_attention_bandit.py

@@ -401,3 +401,15 @@ def test_the_block_map_spreads_the_change_between_blocks_around_one():
     assert low[0] > 1.2 > high[0] > 1.0 and low[1] == high[1] == 1.0       # the sharper change is larger where the block takes it
     assert followed((1.2, 1.0), {}, 1, blocks) == (1.2, 1.0)                 # no map: as set
     assert followed((1.0, 1.0), options, 1, blocks) == (1.0, 1.0)
+
+
+def test_a_block_map_can_never_turn_a_temperature_to_zero_or_negative():
+    from modules.sampling.attention_bandit import followed
+    from modules.sampling.block_map import measure
+    grid = measure.groups(10)
+    lev = [[0.1] * 3] + [[1.0] * 3]                               # the first blocks barely move the picture: the biggest spread
+    data = {"blocks": 10, "groups": grid, "points": list(measure.POINTS),
+            "attention": {"leverage": lev, "fragility": [[0.0] * 3, [0.0] * 3]}}
+    options = {"funui_block_map": measure.MapView(data), "sample_sigmas": torch.linspace(1, 0, 5), "sigmas": torch.tensor([0.5])}
+    for size in (0.2, 0.35, 0.5, 0.9):
+        assert min(followed((1.0 - size, 1.0 - size), options, 1, [1, 7])) >= 0.05

modules/sampling/block_map/__init__.py

@@ -83,18 +83,17 @@ def install(patcher, values, key):
 
     def outer(executor, *args, **kwargs):
         live["state"] = measure.Probe(n)
-        try:
-            return executor(*args, **kwargs)
-        finally:
-            res = measure.result(live["state"])
-            if res is not None:
-                measure.save(res)
-                measure.set_enabled(False)              # one profile per switch: later runs follow the map, they do not redo it
-                log.once(f"{ID}:saved:{n}", log.INFO, "FunUI Block map",
-                         f"Active | profile of {n} blocks saved and the switch turned off; later runs follow it")
-            else:
-                measure.problem = "the profile run measured nothing: the model's picture could not be read"
-                log.once(f"{ID}:none", log.INFO, "FunUI Block map", "Inactive | " + measure.problem)
+        out = executor(*args, **kwargs)                 # cancelled or failed: nothing is saved and the switch stays on, so the next run profiles again
+        res = measure.result(live["state"])
+        if res is not None:
+            measure.save(res)
+            measure.set_enabled(False)                  # one profile per switch: later runs follow the map, they do not redo it
+            log.once(f"{ID}:saved:{n}", log.INFO, "FunUI Block map",
+                     f"Active | profile of {n} blocks saved and the switch turned off; later runs follow it")
+        else:
+            measure.problem = "the profile run measured nothing: the model's picture could not be read"
+            log.once(f"{ID}:none", log.INFO, "FunUI Block map", "Inactive | " + measure.problem)
+        return out
 
     patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, key, outer)
 

modules/sampling/block_map/measure.py

@@ -124,8 +124,10 @@ def block_hook(state, block):
         out = extra["original_block"](args)
         if state.kind == "residual" and state.active() and torch.is_tensor(before):
             after = out.get("img") if isinstance(out, dict) else None
-            if torch.is_tensor(after) and after.shape == before.shape:
-                out = {**out, "img": before + (1.0 + DELTA) * (after - before)}
+            video = after[0] if isinstance(after, (tuple, list)) and after else after       # an audio+video block hands back (video, audio)
+            if torch.is_tensor(video) and video.shape == before.shape:
+                nudged = before + (1.0 + DELTA) * (video - before)
+                out = {**out, "img": type(after)([nudged, *after[1:]]) if isinstance(after, (tuple, list)) else nudged}
         return out
     return hook
 
@@ -234,7 +236,7 @@ class MapView:
         lev, frg = self._fields(kind, frac)
         raw = {}
         for b in blocks:
-            if b >= self.n or lev[b] is None or frg[b] is None:
+            if b >= self.n or lev[b] is None or frg[b] is None or lev[b] <= 1e-9:      # a block the probe saw no effect from is unmeasured, not infinitely cheap
                 raw[b] = None
                 continue
             raw[b] = max(0.05, 1.0 - frg[b]) ** 2 / max(lev[b], 1e-6)
@@ -257,7 +259,7 @@ class MapView:
         lev, frg = self._fields(kind, frac)
         b = int(block)
         known = sorted(v for v in lev if v is not None and v > 0)
-        if b >= self.n or lev[b] is None or frg[b] is None or not known:
+        if b >= self.n or lev[b] is None or frg[b] is None or lev[b] <= 1e-9 or not known:
             return value * m
         cap = REFERENCE * HEADROOM * known[len(known) // 2] / (max(lev[b], 1e-6) * (1.0 + 2.0 * frg[b]))
         x = abs(value) * m

modules/sampling/block_map/tests/test_block_map.py

@@ -193,3 +193,38 @@ def test_adjust_leaves_neutral_and_flat_settings_alone_and_follows_the_run():
     view = _view([[1.0] * 3, [4.0] * 3])
     assert view.adjust("attention", 3, 0.5, [2, 3], 0.0) == 0.0
     assert view.adjust("residual", 3, 0.1, [2, 3], 0.3) != view.adjust("attention", 8, 0.1, [2, 3], 0.3)
+
+
+def test_a_cancelled_profile_saves_nothing_and_leaves_the_switch_on(tiny_h3, root, monkeypatch):
+    measure.set_enabled(True)
+    patched, wrap, outer = _load(tiny_h3)
+    real, calls = tiny_h3.run, [0]
+
+    def cancelled(*a, **k):
+        calls[0] += 1
+        if calls[0] > 5:
+            raise KeyboardInterrupt
+        return real(*a, **k)
+    monkeypatch.setattr(tiny_h3, "run", cancelled)
+    with pytest.raises(KeyboardInterrupt):
+        _run(tiny_h3, patched, wrap, outer)
+    n = len(tiny_h3.patcher.model.diffusion_model.blocks)
+    assert measure.enabled() and measure.load(n) is None              # the next run profiles again, from the start
+
+
+def test_a_group_the_probe_saw_no_effect_from_is_unmeasured_not_free():
+    lev = [[0.0] * 3, [1.0] * 3]
+    view = _view(lev)
+    assert view.adjust("attention", 1, 0.5, [1, 7], 0.5) == 0.5       # no 3x on a block nobody measured
+    assert view.multipliers("attention", 0.5, [1, 7]) == {1: 1.0, 7: 1.0}
+
+
+def test_an_audio_video_block_that_hands_back_a_pair_is_still_nudged():
+    state = measure.Probe(2)
+    state.kind, state.group = "residual", (0, 2)
+    state.points = getattr(state, "points", None)
+    before, after, audio = torch.ones(1, 3, 4), torch.full((1, 3, 4), 2.0), torch.zeros(1, 2, 4)
+    hook = measure.block_hook(state, 0)
+    out = hook({"img": before}, {"original_block": lambda args: {"img": (after, audio)}})
+    assert isinstance(out["img"], tuple) and out["img"][1] is audio
+    assert not torch.equal(out["img"][0], after)