From f32aaac55072eb894061803ea66dc92dd0e073fc Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 1 Nov 2026 11:59:10 +0000 Subject: [PATCH] Block map: a profile of how much each block can take; REINS, Q steer, Attention temperature and Attention explorer nudge their strength by it, a big setting is kept up to what the block takes Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01Tw9fjW75aYB8UyVRZo6fCr --- core/dit_hooks.py | 15 + modules/sampling/attention_bandit/__init__.py | 10 +- .../tests/test_attention_bandit.py | 14 + .../attention_temperature/__init__.py | 13 +- modules/sampling/block_map/__init__.py | 174 ++++++++++++ modules/sampling/block_map/measure.py | 267 ++++++++++++++++++ .../block_map/tests/test_block_map.py | 194 +++++++++++++ modules/sampling/q_steer/__init__.py | 3 +- modules/sampling/reins/__init__.py | 4 +- 9 files changed, 686 insertions(+), 8 deletions(-) create mode 100644 modules/sampling/block_map/__init__.py create mode 100644 modules/sampling/block_map/measure.py create mode 100644 modules/sampling/block_map/tests/test_block_map.py diff --git a/core/dit_hooks.py b/core/dit_hooks.py index 38aae8e0..88780a3d 100644 --- a/core/dit_hooks.py +++ b/core/dit_hooks.py @@ -189,6 +189,21 @@ def current_step(transformer_options): return int(hit[0]), int(sched.shape[0]) - 1 +def block_map_value(transformer_options, kind, block, value, blocks) -> float: + """The feature's own `value` (its change from neutral: a temperature minus 1, a push strength) for `block`, nudged by the + Block map: spread between `blocks` by how much each can take, and softly limited so a very large setting is respected + but cannot push a block past what its profile says it handles. `value` as given with no map, a flat one, or an unmeasured kind.""" + view = (transformer_options or {}).get("funui_block_map") + if view is None: + return float(value) + where = current_step(transformer_options) + frac = 0.5 if where is None else where[0] / max(1, where[1] - 1) + try: + return float(view.adjust(kind, block, frac, blocks, value)) + except Exception: # noqa: BLE001 -- a bad map never costs the picture + return float(value) + + # Set on the guider by Two stages: 1 while the small first stage runs, 2 while the full-size one does. A # SAMPLER_SAMPLE wrapper runs once per stage and reads it (its `model_wrap` is the guider) to do its # once-per-run work in stage 1 only. diff --git a/modules/sampling/attention_bandit/__init__.py b/modules/sampling/attention_bandit/__init__.py index ff7658ea..d95ca874 100644 --- a/modules/sampling/attention_bandit/__init__.py +++ b/modules/sampling/attention_bandit/__init__.py @@ -284,6 +284,14 @@ def make_plan(rows, cond, signature, n_groups, focus, manual_spread=None, region "unknown": sum(1 for s in sims if s is None)} +def followed(factors, options, index, blocks): + """(temperature, prompt gain) for one block, nudged by the block map when this model has one: each change + 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) + + def _prompt_of(guider): try: p = guider.conds["positive"][0]["cross_attn"] @@ -403,7 +411,7 @@ def install(patcher, values, key): if f is None or (f == (1.0, 1.0) and not live["region_wanted"]): return extra["original_block"](args) to = dict(args.get("transformer_options") or {}) - to[_SCOPE] = f + to[_SCOPE] = followed(f, to, index, blocks) return extra["original_block"]({**args, "transformer_options": to}) return scoped diff --git a/modules/sampling/attention_bandit/tests/test_attention_bandit.py b/modules/sampling/attention_bandit/tests/test_attention_bandit.py index 8af72d31..cb7be91c 100644 --- a/modules/sampling/attention_bandit/tests/test_attention_bandit.py +++ b/modules/sampling/attention_bandit/tests/test_attention_bandit.py @@ -387,3 +387,17 @@ def test_with_the_sound_kept_the_sound_is_left_alone_and_the_log_says_so(tiny_h3 assert any("sound left alone in" in x for x in messages) drawn = [x for x in messages if "clip regions looser" in x] assert drawn and not any("end of the clip, sound" in x for x in drawn) # the sound is never among the changed regions + + +def test_the_block_map_spreads_the_change_between_blocks_around_one(): + from modules.sampling.attention_bandit import followed + from modules.sampling.block_map import measure + grid = measure.groups(10) + data = {"blocks": 10, "groups": grid, "points": list(measure.POINTS), + "attention": {"leverage": [[1.0] * 3, [4.0] * 3], "fragility": [[0.1] * 3, [0.1] * 3]}} + options = {"funui_block_map": measure.MapView(data), "sample_sigmas": torch.linspace(1, 0, 5), "sigmas": torch.tensor([0.5])} + blocks = [1, 2, 7, 8] + low, high = followed((1.2, 1.0), options, 1, blocks), followed((1.2, 1.0), options, 7, blocks) + 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) diff --git a/modules/sampling/attention_temperature/__init__.py b/modules/sampling/attention_temperature/__init__.py index b172df2a..4ebba735 100644 --- a/modules/sampling/attention_temperature/__init__.py +++ b/modules/sampling/attention_temperature/__init__.py @@ -64,13 +64,16 @@ def install(patcher, values, key): # The widget floor keeps this above zero; the clamp is for API callers. temperature = max(1.0 + amount, 0.05) - def scoped(args, extra): - to = dict(args.get("transformer_options") or {}) - to[_SCOPE] = temperature - return extra["original_block"]({**args, "transformer_options": to}) + def scoped_for(index): + def scoped(args, extra): + to = dict(args.get("transformer_options") or {}) + # The Block map (when this model has a profile) spreads the change between the blocks, 1 on average. + to[_SCOPE] = max(1.0 + dit_hooks.block_map_value(to, "attention", index, temperature - 1.0, blocks), 0.05) + return extra["original_block"]({**args, "transformer_options": to}) + return scoped for index in blocks: - dit_hooks.add_block_hook(patcher, key, index, scoped) + dit_hooks.add_block_hook(patcher, key, index, scoped_for(index)) def override(func, q, k, v, *args, **kwargs): t = (kwargs.get("transformer_options") or {}).get(_SCOPE) diff --git a/modules/sampling/block_map/__init__.py b/modules/sampling/block_map/__init__.py new file mode 100644 index 00000000..65e60252 --- /dev/null +++ b/modules/sampling/block_map/__init__.py @@ -0,0 +1,174 @@ +"""Block map: where in the model, and how hard, a per-block feature should act. + +A profile run (switched on from Settings > Refinement & Taste, off by default) nudges one group of +blocks at a time and measures how far the picture guess moves and how much of that move is fine-detail +noise (see measure.py). It costs extra model calls on three steps of one run, then writes a map for this +model. Features that act on chosen blocks (Attention explorer, REINS, ...) read the map through the +model's transformer_options and scale their strength block by block, around 1 on average, so a strength +setting keeps its meaning. With no map, or a flat one, every multiplier is 1 and nothing changes. +""" + +from comfy.patcher_extension import WrappersMP + +from ..._core import dit_hooks, log, streams +from . import measure + +ID = "block_map" +TITLE = "Block map" +MOUNT = "generation.sampling" +STAGE = "post" +CATEGORY = "sampling" +STATUS = "experimental" +REQUIRES = ["dit_block_hooks"] +# Innermost, like the other measurement modules: only real forwards are nudged. +AFTER = ["video_detail", "late_branch", "score_slider", "sharpen"] + +PUBLISHED = "funui_block_map" + +SETTINGS = { + "follow": { + "type": "bool", "default": True, + "label": "Follow the block map", + "hint": "Per-block features spread their strength by how much each block can take, from a profile of this model. " + "Nothing changes until a profile has been made.", + }, +} + + +def _say(message): + log.once(f"{ID}:{message}", log.ALERT, "FunUI Block map", message) + + +def install(patcher, values, key): + n = dit_hooks.block_count(patcher) + if n == 0: + return None + view = None + data = measure.load(n) if values.get("follow", True) else None + if data: + try: + view = measure.MapView(data) + except (KeyError, TypeError, ValueError): + _say("the saved map for this model is unreadable and is not used") + dit_hooks.publish(patcher, PUBLISHED, view) + note = [] + if view is not None: + for kind in measure.KINDS: + if view.has(kind): + note.append(f"{kind}: " + ("flat, no effect" if view.flat(kind) else "in use")) + + if not measure.enabled(): + measure.problem = None + return ("block map " + ", ".join(note)) if note else None + + measure.problem = None + live = {"state": measure.Probe(n)} + + def hook_for(block): + def hook(args, extra): + return measure.block_hook(live["state"], block)(args, extra) + return hook + + for b in range(n): + dit_hooks.add_block_hook(patcher, key, b, hook_for(b)) + + def attention(func, *a, **kw): + return measure.attention_override(live["state"])(func, *a, **kw) + + dit_hooks.add_attention_override(patcher, key, attention) + dit_hooks.without_compiler(patcher, 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) + log.once(f"{ID}:saved:{n}", log.INFO, "FunUI Block map", + f"Active | profile of {n} blocks saved; 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) + + patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, key, outer) + + def apply_model(executor, x, t, *args, **kwargs): + out = executor(x, t, *args, **kwargs) + named = streams.model_args(args, kwargs) + to = named.get("transformer_options") or {} + st = live["state"] + where = dit_hooks.current_step(to) + if where is None or dit_hooks.probing(to) or st.kind is not None: + return out + point = measure.point_of(*where) + if point is None or point in st.done: + return out + st.done.add(point) + base = streams.video_of(out, named) + if base is None: + return out + video = base[0] + st.size = [int(video.shape[-2]), int(video.shape[-1])] + if point == min(st.done): + log.once(f"{ID}:cost", log.INFO, "FunUI Block map", + f"profiling: {2 * len(measure.groups(n))} extra model calls on each of {len(measure.POINTS)} steps") + probe_args, probe_kwargs = streams.with_options(args, kwargs, {**to, dit_hooks.PROBE: True}) + for kind in measure.KINDS: + for g, span in enumerate(measure.groups(n)): + st.kind, st.group, st.block = kind, span, 0 + try: + probed = streams.video_of(executor(x, t, *probe_args, **probe_kwargs), named) + if probed is not None: + st.cells[(kind, g, point)] = measure.response(video, probed[0]) + except Exception as exc: # noqa: BLE001 -- a probe never costs the step + _say(f"a probe pass failed ({type(exc).__name__}: {exc}); that cell is not measured") + finally: + st.kind = None + return out + + patcher.add_wrapper_with_key(WrappersMP.APPLY_MODEL, key, apply_model) + return "profiling blocks on this run (switched on in Settings > Refinement & Taste)" + + +def routes(table, base, web): + def view(n=None): + data = measure.load(n) if n else None + return {"enabled": measure.enabled(), "problem": measure.problem, + "maps": sorted(p.name for p in measure.FOLDER.glob("*_blocks.json")) if measure.FOLDER.exists() else []} + + @table.get(base + "/status") + async def _status(_req): + return web.json_response(view()) + + @table.post(base + "/enabled") + async def _enabled(req): + try: + body = await req.json() + except Exception: # noqa: BLE001 + body = {} + measure.set_enabled((body or {}).get("enabled") is True) + return web.json_response(view()) + + @table.post(base + "/clear") + async def _clear(_req): + measure.clear() + return web.json_response(view()) + + @table.get(base + "/export/{blocks}") + async def _export(req): + try: + data = measure.load(int(req.match_info["blocks"])) + except ValueError: + data = None + if data is None: + return web.json_response({"why": "no map for that block count"}, status=404) + return web.json_response(data, headers={"Content-Disposition": 'attachment; filename="block_map.json"'}) + + +def cache_key(): + return "on" if measure.enabled() else "off" + + +PROVIDES = {"modifier": install, "routes": routes, "cache_key": cache_key} diff --git a/modules/sampling/block_map/measure.py b/modules/sampling/block_map/measure.py new file mode 100644 index 00000000..593f3a90 --- /dev/null +++ b/modules/sampling/block_map/measure.py @@ -0,0 +1,267 @@ +"""The block map's maths and its file: what a profile run measures, and how a feature reads it. + +A profile run nudges one group of blocks at a time, in one of two ways, and measures how far the model's +picture guess moves: + +* leverage: the size of that move, relative to the picture and to the size of the nudge. A block with a + big leverage changes the picture a lot for a small nudge. +* fragility: the share of the move that sits in the finest detail (what is left after a 3x3 blur). That + is the share that looks like noise or damaged edges rather than a change of content. + +Two kinds of nudge, because blocks do not react alike to both: "attention" (the attention query scaled, +what Attention explorer and the temperature features do) and "residual" (the block's own push scaled, +what REINS, Block repeat and steering do in effect). + +A feature asks `multiplier(kind, block, fraction_of_run, blocks)` and gets a number around 1: how much +more or less than its set strength to use in that block at that point of the run. The effect share a +block gets is its quality, (1 - fragility)^2, and the nudge that produces that share is the share over +the leverage, so the picture change goes where it is clean and the nudge is small where leverage is big. +Averaged over the blocks the feature uses it is 1, so a strength slider keeps its meaning. +""" + +import json +import math +import time + +import torch +import torch.nn.functional as F + +from ..._core import config + +SWITCH = config.ROOT / "block_map.enabled" +FOLDER = config.ROOT / "block_map" +KINDS = ("attention", "residual") +DELTA = 0.1 # the size of a probe nudge +GROUP_SIZE = 5 # blocks nudged together +POINTS = (0.1, 0.45, 0.8) # where in the run the profile is taken (fraction of the steps) +LOW, HIGH = 0.25, 3.0 # a multiplier stays inside these +REFERENCE = 0.5 # a nudge the typical block is taken to handle without harm ... +HEADROOM = 3.0 # ... times this is the most picture change any block is let produce +FLAT = 0.10 # leverage spread under this, and fragility spread under 0.05: nothing to aim at + +problem = None # why nothing was measured, for the panel + + +def enabled() -> bool: + try: + return SWITCH.read_text().strip() == "1" + except OSError: + return False + + +def set_enabled(on: bool) -> bool: + SWITCH.parent.mkdir(parents=True, exist_ok=True) + SWITCH.write_text("1" if on else "0") + return bool(on) + + +def groups(n): + """[(first, stop)] blocks per probe group.""" + return [(a, min(n, a + GROUP_SIZE)) for a in range(0, n, GROUP_SIZE)] + + +def point_of(step, steps): + """Which profile point this step is, or None. The first step at or after each point's place in the run.""" + if steps < 2: + return None + for i, p in enumerate(POINTS): + if step == min(steps - 1, int(round(p * (steps - 1)))): + return i + return None + + +def path_of(n): + return FOLDER / f"{int(n)}_blocks.json" + + +def response(base, probed, delta=DELTA): + """-> (leverage, fragility) of one probe. Both 5-D pictures [B, C, T, H, W].""" + d = (probed.float() - base.float()) + ref = base.float() + size = float(d.norm()) + lev = size / (max(float(ref.norm()), 1e-8) * delta) + if size <= 1e-12 or d.dim() != 5: + return lev, 0.0 + b, c, t, h, w = d.shape + x = d.reshape(b * c * t, 1, h, w) + low = F.avg_pool2d(x, 3, 1, 1, count_include_pad=False) + frag = min(1.0, max(0.0, float((x - low).norm() / x.norm().clamp(min=1e-12)) ** 2)) + return lev, frag + + +class Probe: + """What the hooks read: which nudge is on right now, and which block is running.""" + + def __init__(self, n): + self.n = n + self.kind = None # "attention" / "residual" while a probe pass runs + self.group = None # (first, stop) + self.block = 0 + self.done = set() + self.cells = {} # (kind, group index, point) -> (leverage, fragility) + self.size = None # (H, W) of the picture latent, for the record + + def active(self): + return self.kind is not None and self.group is not None and self.group[0] <= self.block < self.group[1] + + +def attention_override(state): + def override(func, q, k, v, heads, *args, **kw): + if state.kind == "attention" and state.active() and torch.is_tensor(q): + q = q * (1.0 + DELTA) + return func(q, k, v, heads, *args, **kw) + return override + + +def block_hook(state, block): + def hook(args, extra): + if state.kind is None: + return extra["original_block"](args) + state.block = block + before = args.get("img") + if state.kind == "residual" and state.active() and torch.is_tensor(before): + before = before.clone() # a block may add to its input in place: keep the input as it was + 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)} + return out + return hook + + +def result(state): + """The finished profile as a plain dict, or None when no cell was measured.""" + if not state.cells: + return None + grid = groups(state.n) + out = {"when": time.strftime("%Y-%m-%d %H:%M:%S"), "blocks": state.n, "groups": grid, + "points": list(POINTS), "size": state.size, "delta": DELTA} + for kind in KINDS: + cells = [[state.cells.get((kind, g, p)) for p in range(len(POINTS))] for g in range(len(grid))] + if any(c is not None for row in cells for c in row): + out[kind] = {"leverage": [[None if c is None else c[0] for c in row] for row in cells], + "fragility": [[None if c is None else c[1] for c in row] for row in cells]} + return out + + +def save(res): + FOLDER.mkdir(parents=True, exist_ok=True) + path = path_of(res["blocks"]) + tmp = path.with_suffix(".tmp") + tmp.write_text(json.dumps(res)) + tmp.replace(path) + + +def load(n): + try: + return json.loads(path_of(n).read_text()) + except (OSError, ValueError): + return None + + +def clear(n=None): + paths = [path_of(n)] if n is not None else list(FOLDER.glob("*_blocks.json")) if FOLDER.exists() else [] + for p in paths: + try: + p.unlink() + except OSError: + pass + + +def _interp(xs, ys, x): + """Piecewise linear through the points that have a value; flat outside them.""" + pts = [(a, b) for a, b in zip(xs, ys) if b is not None] + if not pts: + return None + if x <= pts[0][0]: + return pts[0][1] + for (x0, y0), (x1, y1) in zip(pts, pts[1:]): + if x <= x1: + return y0 + (y1 - y0) * (x - x0) / max(1e-9, x1 - x0) + return pts[-1][1] + + +class MapView: + """A saved profile, asked for multipliers.""" + + def __init__(self, data): + self.data = data + self.n = int(data["blocks"]) + self.centres = [(a + b - 1) / 2.0 for a, b in data["groups"]] + self._cache = {} + self._memo = {} + + def has(self, kind): + return kind in self.data + + def _fields(self, kind, frac): + """(leverage, fragility) for every block at this fraction of the run.""" + key = (kind, round(float(frac) * 20)) + if key in self._cache: + return self._cache[key] + d, pts = self.data[kind], self.data["points"] + lev_g, frg_g = [], [] + for row_l, row_f in zip(d["leverage"], d["fragility"]): + lev_g.append(_interp(pts, row_l, frac)) + frg_g.append(_interp(pts, row_f, frac)) + lev = [_interp(self.centres, lev_g, float(b)) for b in range(self.n)] + frg = [_interp(self.centres, frg_g, float(b)) for b in range(self.n)] + self._cache[key] = (lev, frg) + return lev, frg + + def flat(self, kind): + lev, frg = self._fields(kind, 0.5) + lev = [v for v in lev if v is not None] + frg = [v for v in frg if v is not None] + if not lev or not frg: + return True + mean = sum(lev) / len(lev) + return (max(lev) - min(lev)) < FLAT * max(mean, 1e-9) and (max(frg) - min(frg)) < 0.05 + + def multipliers(self, kind, frac, blocks): + """{block: multiplier} for the blocks a feature uses, mean 1; all 1 when the profile is flat or missing.""" + blocks = [int(b) for b in blocks] + memo = (kind, round(float(frac) * 20), tuple(blocks)) + if memo in self._memo: + return self._memo[memo] + self._memo[memo] = self._multipliers(kind, frac, blocks) + return self._memo[memo] + + def _multipliers(self, kind, frac, blocks): + if not blocks or not self.has(kind) or len(blocks) == 1 or self.flat(kind): + return {b: 1.0 for b in blocks} + 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: + raw[b] = None + continue + raw[b] = max(0.05, 1.0 - frg[b]) ** 2 / max(lev[b], 1e-6) + known = [v for v in raw.values() if v is not None] + if not known: + return {b: 1.0 for b in blocks} + mean = sum(known) / len(known) + return {b: (1.0 if v is None else min(HIGH, max(LOW, v / mean))) for b, v in raw.items()} + + def adjust(self, kind, block, frac, blocks, value): + """The feature's own `value` (its change from neutral: a temperature minus 1, a strength) for this block: + spread by `multipliers`, then softly limited so no block is pushed far past what its profile says it can + take. The limit is the picture change a typical block makes from REFERENCE times HEADROOM, divided by this + block's leverage and its fragility, so a setting of 10 is respected up to that limit and not beyond it. + A flat or missing profile returns `value` as it came.""" + value = float(value) + if value == 0.0 or not self.has(kind) or self.flat(kind): + return value + m = self.multipliers(kind, frac, blocks).get(int(block), 1.0) + 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: + 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 + return math.copysign(cap * math.tanh(x / cap), value) + + def multiplier(self, kind, block, frac, blocks): + return self.multipliers(kind, frac, blocks).get(int(block), 1.0) diff --git a/modules/sampling/block_map/tests/test_block_map.py b/modules/sampling/block_map/tests/test_block_map.py new file mode 100644 index 00000000..85e7508b --- /dev/null +++ b/modules/sampling/block_map/tests/test_block_map.py @@ -0,0 +1,194 @@ +"""The block map: the maths on constructed profiles, and a profile run on the real tiny H3.""" + +import pytest +import torch + +from modules.sampling.block_map import measure + + +def _data(n=10, lev=None, frg=None): + grid = measure.groups(n) + g = len(grid) + lev = lev or [[1.0] * 3 for _ in range(g)] + frg = frg or [[0.1] * 3 for _ in range(g)] + return {"blocks": n, "groups": grid, "points": list(measure.POINTS), + "attention": {"leverage": lev, "fragility": frg}} + + +def test_a_flat_profile_changes_nothing(): + view = measure.MapView(_data()) + assert view.flat("attention") + assert view.multipliers("attention", 0.5, [2, 3, 7]) == {2: 1.0, 3: 1.0, 7: 1.0} + + +def test_a_missing_kind_or_a_single_block_changes_nothing(): + view = measure.MapView(_data()) + assert view.multipliers("residual", 0.5, [2, 3]) == {2: 1.0, 3: 1.0} + assert measure.MapView(_data(lev=[[1.0] * 3, [3.0] * 3])).multipliers("attention", 0.5, [4]) == {4: 1.0} + + +def test_nudges_are_smaller_where_leverage_is_big_and_smaller_where_detail_is_fragile(): + lev = [[1.0] * 3, [4.0] * 3] # blocks 5-9 move the picture four times as much + frg = [[0.1] * 3, [0.1] * 3] + m = measure.MapView(_data(lev=lev, frg=frg)).multipliers("attention", 0.5, [1, 2, 7, 8]) + assert m[1] > m[7] and m[2] > m[8] + assert sum(m.values()) / 4 == pytest.approx(1.0, abs=0.35) # about 1 on average (the clamp may bend it) + lev = [[1.0] * 3, [1.0] * 3] + frg = [[0.0] * 3, [0.8] * 3] # blocks 5-9 turn their move into fine noise + m = measure.MapView(_data(lev=lev, frg=frg)).multipliers("attention", 0.5, [1, 2, 7, 8]) + assert m[1] > m[7] + + +def test_multipliers_stay_inside_their_limits_and_average_one_before_clamping(): + lev = [[1.0] * 3, [1000.0] * 3] + m = measure.MapView(_data(lev=lev)).multipliers("attention", 0.5, [0, 1, 8, 9]) + assert all(measure.LOW <= v <= measure.HIGH for v in m.values()) + + +def test_response_reads_the_size_of_the_move_and_how_much_of_it_is_fine_detail(): + base = torch.randn(1, 4, 2, 8, 8) + smooth = torch.ones_like(base) * 0.01 + noisy = torch.randn_like(base) * 0.01 + lev_s, frag_s = measure.response(base, base + smooth) + lev_n, frag_n = measure.response(base, base + noisy) + assert lev_s == pytest.approx(lev_n, rel=0.5) and frag_s < 0.05 < 0.5 < frag_n + assert measure.response(base, base) == (0.0, 0.0) + + +def test_response_follows_the_picture_dtype_and_device(): + base = torch.randn(1, 2, 2, 6, 6, dtype=torch.bfloat16) + lev, frag = measure.response(base, base * 1.01) + assert lev > 0 and 0.0 <= frag <= 1.0 + + +def test_points_are_hit_once_each_and_a_one_step_run_has_none(): + hits = [measure.point_of(i, 8) for i in range(8)] + assert [h for h in hits if h is not None] == [0, 1, 2] + assert all(measure.point_of(i, 1) is None for i in range(1)) + + +# --- the profile run, on the real tiny H3 ------------------------------------------------------------- + +@pytest.fixture +def root(tmp_path, monkeypatch): + monkeypatch.setattr(measure, "SWITCH", tmp_path / "switch") + monkeypatch.setattr(measure, "FOLDER", tmp_path / "maps") + + +def _load(tiny, settings=None): + from comfy.patcher_extension import WrappersMP + from modules.sampling.modifiers.nodes import FunUILoadModifiers + patched, status = FunUILoadModifiers.execute(tiny.patcher, settings or {}).result + wraps = [w for ws in patched.wrappers.get(WrappersMP.APPLY_MODEL, {}).values() for w in ws] + outer = [w for ws in patched.wrappers.get(WrappersMP.OUTER_SAMPLE, {}).values() for w in ws] + return patched, (wraps[-1] if wraps else None), outer + + +def _run(tiny, patched, wrap, outer, steps=4): + from conftest import packed_av + x0, shapes = packed_av(tiny.video, tiny.audio) + seen = [] + + def executor(x, t, c_concat=None, c_crossattn=None, control=None, transformer_options=None, **kw): + to = dict(transformer_options or {}) + seen.append(to) + video, audio = tiny.run(patched, sigma=float(t.max()), + options={k: v for k, v in to.items() if k == "funui_probe"}) + return packed_av(video, audio)[0] + + s = torch.linspace(1.0, 0.0, steps + 1) + call = lambda: [wrap(executor, x0, s[i:i + 1].reshape(1, 1), None, None, None, + {"sample_sigmas": s, "sigmas": s[i:i + 1], "cond_or_uncond": [0]}, + latent_shapes=shapes) for i in range(steps)] + for o in reversed(outer): + call = (lambda w, inner: (lambda: w(lambda: inner())))(o, call) + return call(), seen + + +def test_off_it_installs_no_hooks_and_publishes_no_map(tiny_h3, root): + patched, wrap, outer = _load(tiny_h3) + assert wrap is None and outer == [] + assert patched.model_options["transformer_options"].get("funui_block_map") is None + + +def test_a_profile_run_writes_a_map_and_returns_the_untouched_answers(tiny_h3, root): + patched, _w, outer = _load(tiny_h3) + outs0, _ = _run(tiny_h3, patched, lambda ex, *a, **k: ex(*a, **k), outer) + measure.set_enabled(True) + patched, wrap, outer = _load(tiny_h3) + outs, seen = _run(tiny_h3, patched, wrap, outer) + assert all(torch.equal(a, b) for a, b in zip(outs, outs0)) + n = len(tiny_h3.patcher.model.diffusion_model.blocks) + per_point = 2 * len(measure.groups(n)) + assert len(seen) == 4 + 3 * per_point + assert [bool(t.get("funui_probe")) for t in seen][:1 + per_point] == [False] + [True] * per_point + data = measure.load(n) + assert data["blocks"] == n and set(measure.KINDS) <= set(data) + for kind in measure.KINDS: + lev = data[kind]["leverage"] + assert len(lev) == len(measure.groups(n)) and all(v is not None and v > 0 for row in lev for v in row) + + +def test_a_later_run_publishes_the_map_and_follow_off_hides_it(tiny_h3, root): + measure.set_enabled(True) + patched, wrap, outer = _load(tiny_h3) + _run(tiny_h3, patched, wrap, outer) + measure.set_enabled(False) + patched, _w, _o = _load(tiny_h3) + view = patched.model_options["transformer_options"]["funui_block_map"] + assert view is not None and view.n == len(tiny_h3.patcher.model.diffusion_model.blocks) + patched, _w, _o = _load(tiny_h3, {"block_map": {"follow": False}}) + assert patched.model_options["transformer_options"].get("funui_block_map") is None + + +def test_a_failing_probe_pass_costs_nothing(tiny_h3, root): + from conftest import packed_av + measure.set_enabled(True) + patched, wrap, outer = _load(tiny_h3) + x0, shapes = packed_av(tiny_h3.video, tiny_h3.audio) + + def executor(x, t, *a, **kw): + to = kw.get("transformer_options") or (a[3] if len(a) > 3 else {}) + if (to or {}).get("funui_probe"): + raise RuntimeError("probe broke") + return x0 + s = torch.linspace(1.0, 0.0, 5) + out = wrap(executor, x0, torch.tensor([[1.0]]), None, None, None, + {"sample_sigmas": s, "sigmas": s[0:1], "cond_or_uncond": [0]}, latent_shapes=shapes) + assert out.shape == x0.shape + + +def test_the_helper_every_feature_asks_is_the_value_itself_without_a_map_and_follows_the_step_with_one(): + from modules._core import dit_hooks + assert dit_hooks.block_map_value({}, "residual", 3, 0.4, [3, 4]) == 0.4 + assert dit_hooks.block_map_value(None, "residual", 3, 0.4, [3, 4]) == 0.4 + data = _data(lev=[[1.0] * 3, [4.0] * 3]) + data["residual"] = data["attention"] + options = {"funui_block_map": measure.MapView(data), "sample_sigmas": torch.linspace(1, 0, 5), "sigmas": torch.tensor([0.5])} + assert dit_hooks.block_map_value(options, "residual", 1, 0.4, [1, 8]) > dit_hooks.block_map_value(options, "residual", 8, 0.4, [1, 8]) + assert dit_hooks.block_map_value(options, "attention", 1, 0.4, [1]) == pytest.approx(0.4, rel=0.05) # one block: nothing to spread, only the soft limit + broken = {"funui_block_map": object(), "sample_sigmas": None} + assert dit_hooks.block_map_value(broken, "residual", 1, 0.4, [1, 2]) == 0.4 + + +def _view(lev, frg=None): + data = _data(lev=lev, frg=frg) + data["residual"] = data["attention"] + return measure.MapView(data) + + +def test_a_huge_setting_is_kept_but_cannot_push_a_block_past_what_it_takes(): + view = _view([[1.0] * 3, [4.0] * 3]) + blocks = [1, 2, 7, 8] + small = view.adjust("attention", 7, 0.5, blocks, 0.05) + assert 0.0 < small <= 0.05 * measure.HIGH + 1e-9 # an ordinary setting passes (spread, not limited) + big_low, big_high = view.adjust("attention", 1, 0.5, blocks, 10.0), view.adjust("attention", 7, 0.5, blocks, 10.0) + assert 1.0 < big_low < 10.0 and 0.0 < big_high < big_low # 10 is respected as far as the block allows + assert view.adjust("attention", 7, 0.5, blocks, -10.0) == pytest.approx(-big_high) # the sign is kept + + +def test_adjust_leaves_neutral_and_flat_settings_alone_and_follows_the_run(): + assert measure.MapView(_data()).adjust("attention", 3, 0.5, [2, 3], 0.3) == 0.3 + 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) diff --git a/modules/sampling/q_steer/__init__.py b/modules/sampling/q_steer/__init__.py index 50569416..e595bd78 100644 --- a/modules/sampling/q_steer/__init__.py +++ b/modules/sampling/q_steer/__init__.py @@ -120,7 +120,8 @@ def install(patcher, values, key): return func(q, k, v, *args, **kwargs) # Rotation keeps length, so the rotated rows' size is the natural one's. norm = torch.linalg.vector_norm(rows, dim=-1, dtype=torch.float32).mean() - scale = torch.where(torch.isfinite(norm) & (norm > 0), norm * strength, torch.zeros_like(norm)) + push_strength = dit_hooks.block_map_value(to, "attention", block, strength, steer) # nudged by the Block map + scale = torch.where(torch.isfinite(norm) & (norm > 0), norm * push_strength, torch.zeros_like(norm)) push = (direction.to(device=q.device, dtype=q.dtype).view(heads, 1, dim) * scale.to(q.dtype)).expand(heads, count, dim) q = q.clone() diff --git a/modules/sampling/reins/__init__.py b/modules/sampling/reins/__init__.py index 927a6791..c6cbedb9 100644 --- a/modules/sampling/reins/__init__.py +++ b/modules/sampling/reins/__init__.py @@ -107,8 +107,10 @@ def install(patcher, values, key): if d is None: d = cast[(block, rows.dtype)] = direction.to(device=rows.device, dtype=rows.dtype) norm = torch.linalg.vector_norm(rows, dim=-1, dtype=torch.float32).mean() + # The Block map (when this model has a profile) nudges the strength per block; a big setting is kept, up to what the block takes. + push = dit_hooks.block_map_value(args.get("transformer_options"), "residual", block, strength, steer) # No host sync: a bad norm makes the push zero instead of a branch. - scale = torch.where(torch.isfinite(norm) & (norm > 0), norm * strength, + scale = torch.where(torch.isfinite(norm) & (norm > 0), norm * push, torch.zeros_like(norm)) rows += d * scale.to(rows.dtype) if isinstance(out, (tuple, list)) and rows.data_ptr() != out[0].data_ptr():