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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Tw9fjW75aYB8UyVRZo6fCr
This commit is contained in:
Claude
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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}
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
|
||||
@@ -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():
|
||||
|
||||
Reference in New Issue
Block a user