diff --git a/core/brownian.py b/core/brownian.py index b31d2ca2..6e1c7ced 100644 --- a/core/brownian.py +++ b/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 diff --git a/core/tests/test_brownian.py b/core/tests/test_brownian.py index 6d39b45a..fc2f2135 100644 --- a/core/tests/test_brownian.py +++ b/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 diff --git a/modules/sampling/attention_bandit/__init__.py b/modules/sampling/attention_bandit/__init__.py index d95ca874..4929a555 100644 --- a/modules/sampling/attention_bandit/__init__.py +++ b/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): diff --git a/modules/sampling/attention_bandit/tests/test_attention_bandit.py b/modules/sampling/attention_bandit/tests/test_attention_bandit.py index cb7be91c..603dc4a1 100644 --- a/modules/sampling/attention_bandit/tests/test_attention_bandit.py +++ b/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 diff --git a/modules/sampling/block_map/__init__.py b/modules/sampling/block_map/__init__.py index fb1fe425..0fc298f4 100644 --- a/modules/sampling/block_map/__init__.py +++ b/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) diff --git a/modules/sampling/block_map/measure.py b/modules/sampling/block_map/measure.py index 593f3a90..a71d8c78 100644 --- a/modules/sampling/block_map/measure.py +++ b/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 diff --git a/modules/sampling/block_map/tests/test_block_map.py b/modules/sampling/block_map/tests/test_block_map.py index 33316eee..b91aafce 100644 --- a/modules/sampling/block_map/tests/test_block_map.py +++ b/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)