nodes/olivv-cs-ComfyUI-FunPack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: forget fires only once the first run is queued, one owner per Generate; velocity bias banks likes only, mixed sizes import, says when no structure step, refuses windows/second-order/negative-only calls

Browse Source 
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
2026-10-03 16:34:03 +03:00
co-authored by Claude Sonnet 5.5
parent 59c2153137
commit b6319c51c4
6 changed files with 153 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/legacy/store.js
+11 -6
View File
@@ -4019,10 +4019,6 @@
return false;
}
_genRunActive = true;
if (_tasteGenPending) { // the first run of a Generate: what the last one left unrated is forgotten
_tasteGenPending = false;
await _newTasteGeneration();
}
// A run started through the normal click path owns its own clock
// lifecycle (the _clockedX wrapper's `finally`, spanning the whole
// batch) — any stale flag left over from a PRIOR adopted run must not
@@ -4074,6 +4070,13 @@
// clobber it. Setting it last means an unwired reference is always
// what the person sees, never lost to a race with an unrelated note.
if (unwiredNotice) { state.notice = unwiredNotice; notify(); }
if (queued && _tasteGenPending) {
// The first run of a Generate is really queued: what the last Generate left unrated is
// forgotten now (a refused or failed click above never gets here). Sampling takes seconds;
// this takes milliseconds, so it lands before this run's own capture.
_tasteGenPending = false;
await _newTasteGeneration();
}
if (!queued) {
// Refused before it ever reached run.start() (an incomplete pipeline,
// or the queue itself saying no) -- the bridge's own "say" hook
@@ -4338,14 +4341,16 @@
}
// Runs `fn` as the owner of one Generate -- unless one is already in flight, in which case this click
// is refused later anyway and must not make the running Generate forget its unrated clips.
let _generateOwned = false; // an entry point holds a Generate (set before its first await)
async function _ownsGenerate(fn) {
const owner = !_genRunActive;
const owner = !_genRunActive && !_generateOwned;
if (owner) {
_generateOwned = true;
_tasteGenPending = true;
// The settings of this Generate, taken once: every run of it is queued with these.
_frozenSettings = window.PipelineState && window.PipelineState.frozenInputs ? window.PipelineState.frozenInputs() : null;
}
try { return await fn(); } finally { if (owner) { _tasteGenPending = false; _frozenSettings = null; } }
try { return await fn(); } finally { if (owner) { _tasteGenPending = false; _frozenSettings = null; _generateOwned = false; } }
}
async function _clockedGenerate(onlyScene) {
_genClockStart(onlyScene ? "scene" : "all");
modules/sampling/velocity_bias/__init__.py
+36 -10
View File
@@ -18,6 +18,7 @@ import torch.nn.functional as F
from comfy.patcher_extension import WrappersMP
from ..._core import dit_hooks, log, registry, streams
from ..score_slider import cond_row
ID = "velocity_bias"
TITLE = "Velocity bias"
@@ -123,46 +124,71 @@ def install(patcher, values, key):
return None
strength = max(0.0, min(3.0, float(values.get("strength", 0.15))))
nearest = bool(values.get("nearest"))
live = {"rows": []}
live = {"rows": [], "hit": False, "seen": {}, "multi": False}
def fresh():
live["hit"], live["multi"] = False, False
live["seen"].clear()
live["rows"] = taste.rows(KIND)
liked = sum(1 for r in live["rows"] if r["reward"] > 0)
log.once(f"{ID}:state", log.INFO, "FunPack Velocity bias",
f"key {taste.key!r}: {liked} liked clip(s) banked" if liked else
f"key {taste.key!r}: learning, nothing liked banked yet")
captured = taste.collect(patcher, key, KIND, keep=MAX_ROWS, fresh=fresh)
# Only liked clips are banked (nothing here learns from a disliked one); clips of different sizes
# are meant to coexist, so a key holding two resolutions still exports and imports.
captured = taste.collect(patcher, key, KIND, keep=MAX_ROWS, fresh=fresh, only="liked", mixed=True)
def apply_model(executor, x, t, *args, **kwargs):
named = streams.model_args(args, kwargs)
to = named.get("transformer_options")
if dit_hooks.probing(to) or not structure_step(to):
if dit_hooks.probing(to):
return executor(x, t, *args, **kwargs)
if (to or {}).get("context_window") is not None:
_say("off this run: it cannot follow clips cut into context windows (each window is a different size)")
return executor(x, t, *args, **kwargs)
where = dit_hooks.current_step(to)
if where is not None and where[0] >= where[1] - 1 and not live["hit"] and not live["multi"]:
_say("did nothing this run: no step of this schedule is near 90% of the starting noise, "
"so there was no structure step to bias or learn from")
if not structure_step(to):
return executor(x, t, *args, **kwargs)
kinds = tuple(int(v) for v in ((to or {}).get("cond_or_uncond") or ()))
c = named.get("c_crossattn")
row = cond_row(c, to) if torch.is_tensor(c) and c.dim() == 3 else 0
if row is None: # only the negative prompt in this call
return executor(x, t, *args, **kwargs)
if live["seen"].get(kinds) == where[0] or live["multi"]:
if not live["multi"]:
live["multi"] = True
_say("off for the rest of this run: the model is called more than once per step "
"(a second-order sampler); use euler-style sampling")
return executor(x, t, *args, **kwargs)
live["seen"][kinds] = where[0]
sig = None
if torch.is_tensor(c) and c.dim() == 3:
words = registry.current().ask("text_rows", named, int(c.shape[1]))
row = c[0] if words is None else c[0][words.to(c.device)]
sig = row.float().mean(0).detach().cpu()
r = c[row] if words is None else c[row][words.to(c.device)]
sig = r.float().mean(0).detach().cpu()
split_in = streams.video_of(x, named)
if split_in is None:
_say("off this run: could not find the picture in this model's latent")
return executor(x, t, *args, **kwargs)
live["hit"] = True
video_in, rebuild = split_in
sigma = float(t.max())
ratio = sigma / float(to["sample_sigmas"][0])
direction, how = reference(live["rows"], video_in[0].shape, sig, nearest) if strength > 0 else (None, "")
direction, how = reference(live["rows"], video_in[row].shape, sig, nearest) if strength > 0 else (None, "")
run_x = x
if direction is not None:
run_x = rebuild(rotate(video_in, direction.unsqueeze(0), strength, ratio))
log.once(f"{ID}:applied", log.INFO, "FunPack Velocity bias", f"applied: {how}")
elif strength > 0 and live["rows"]:
_say("not applied this run: " + reference(live["rows"], video_in[0].shape, sig, nearest)[1])
log.once(f"{ID}:applied:{how}", log.INFO, "FunPack Velocity bias", f"applied: {how}")
elif strength > 0:
_say("not applied this run: " + reference(live["rows"], video_in[row].shape, sig, nearest)[1])
out = executor(run_x, t, *args, **kwargs)
split_out, split_run = streams.video_of(out, named), streams.video_of(run_x, named)
if split_out is not None and split_run is not None:
d = ((split_run[0].float() - split_out[0].float()) / max(sigma, 1e-6))[0] # where the model moves the picture
d = ((split_run[0].float() - split_out[0].float()) / max(sigma, 1e-6))[row] # where the model moves the picture
captured["v"] = d.detach().half().cpu()
if sig is not None:
captured["sig"] = sig
modules/sampling/velocity_bias/tests/test_velocity_bias.py
+46
View File
@@ -111,3 +111,49 @@ def test_strength_zero_only_learns(tiny_h3):
seen = []
wrap(lambda x, *a, **k: seen.append(x) or x0, x0, sched[1:2], None, None, None, _to(1, sched), latent_shapes=shapes)
assert torch.equal(seen[0], x0)
def _packed(shapes_seed=0):
from conftest import packed_av
torch.manual_seed(shapes_seed)
return packed_av(torch.randn(1, 4, 3, 8, 8), torch.randn(1, 8, 5))
def test_a_schedule_with_no_structure_step_says_so_on_its_last_step(tiny_h3):
from core import log
wrap, _ = _load(tiny_h3, strength=1.0)
x0, shapes = _packed()
sched = torch.tensor([1.0, 0.75, 0.5, 0.25, 0.0])
log.new_run()
before = len(log.history())
for i in range(4):
wrap(lambda x, *a, **k: x0, x0, sched[i:i + 1], None, None, None, _to(i, sched), latent_shapes=shapes)
assert any("no step of this schedule is near 90%" in e["message"] for e in log.history()[before:])
def test_context_windows_and_second_order_calls_are_refused_in_words(tiny_h3):
from core import log
x0, shapes = _packed()
sched = torch.tensor([1.0, 0.909375, 0.725, 0.421875, 0.0])
wrap, _ = _load(tiny_h3, strength=1.0)
log.new_run()
wrap(lambda x, *a, **k: x0, x0, sched[1:2], None, None, None, {**_to(1, sched), "context_window": object()},
latent_shapes=shapes)
assert any("context windows" in e["message"] for e in log.history())
wrap2, outer = _load(tiny_h3, strength=1.0)
log.new_run()
for o in outer[:1]:
o(lambda: None) # a run starts: the state is fresh
for _ in range(2): # the same structure step twice: a corrector call
wrap2(lambda x, *a, **k: x0, x0, sched[1:2], None, None, None, _to(1, sched), latent_shapes=shapes)
assert any("more than once per step" in e["message"] for e in log.history())
def test_a_call_that_carries_only_the_negative_prompt_neither_biases_nor_learns(tiny_h3):
wrap, _ = _load(tiny_h3, strength=1.0)
x0, shapes = _packed()
sched = torch.tensor([1.0, 0.909375, 0.725, 0.421875, 0.0])
seen = []
to = {**_to(1, sched), "cond_or_uncond": [1]}
wrap(lambda x, *a, **k: seen.append(x) or x0, x0, sched[1:2], None, torch.ones(1, 4, 8), None, to, latent_shapes=shapes)
assert torch.equal(seen[0], x0)
modules/system/taste/__init__.py
+4 -4
View File
@@ -38,8 +38,8 @@ class Handle:
def __init__(self, key):
self.key = key
def capture(self, kind, rows, keep=store.MAX_ROWS):
store.capture(self.key, kind, rows, keep=keep)
def capture(self, kind, rows, keep=store.MAX_ROWS, only=None, mixed=False):
store.capture(self.key, kind, rows, keep=keep, only=only, mixed=mixed)
def rows(self, kind, blind_to=None):
"""Every rated row of `kind`, oldest first: {"reward", "rows", "prompt_id"}.
@@ -47,7 +47,7 @@ class Handle:
neutral -- for a learner that cannot tell whether that half was the problem."""
return store.blind(store.load(self.key, kind)["rows"], blind_to)
def collect(self, patcher, key, kind, keep=store.MAX_ROWS, fresh=None):
def collect(self, patcher, key, kind, keep=store.MAX_ROWS, fresh=None, only=None, mixed=False):
"""A dict to fill during sampling; kept as `kind`'s capture when the
sampling call finishes. An interrupted or failed run keeps nothing, so
a half-finished clip can never be what a rating teaches.
@@ -64,7 +64,7 @@ class Handle:
if fresh is not None:
fresh()
out = executor(*args, **kwargs)
self.capture(kind, rows, keep=keep)
self.capture(kind, rows, keep=keep, only=only, mixed=mixed)
return out
patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, key, outer)
modules/system/taste/store.py
+26 -11
View File
@@ -134,12 +134,20 @@ def _pending_path(key, kind, prompt_id=None):
return _dir(key) / (f"{kind}.pending.pt" if prompt_id is None else f"{kind}.{prompt_id}.pending.pt")
def capture(key, kind, rows, prompt_id=None, keep=MAX_ROWS):
def _prune_aside(key, kind):
aside = sorted(_dir(key).glob(f"{kind}.*.pending.pt"), key=lambda p: p.stat().st_mtime_ns)
for old in aside[:-MAX_PENDING]:
old.unlink(missing_ok=True)
def capture(key, kind, rows, prompt_id=None, keep=MAX_ROWS, only=None, mixed=False):
"""This run's capture for `kind`, waiting for a rating until the next Generate starts.
The previous run's capture of the same Generate is kept aside, not overwritten. Kept in the
dtype given: a banked latent stays half precision. `keep` caps how many rated rows this kind
holds, oldest dropped first.
holds, oldest dropped first. `only="liked"`: a kind that learns from liked clips alone does not
bank a disliked one (it could only push the liked ones out under `keep`). `mixed`: its rows are
legitimately of different shapes (a clip bank is matched by size), so import does not demand one.
"""
if not rows:
return
@@ -152,10 +160,8 @@ def capture(key, kind, rows, prompt_id=None, keep=MAX_ROWS):
if earlier is not None and earlier != prompt_id:
if _PID.match(str(earlier)):
os.replace(latest, _pending_path(key, kind, earlier))
_write(latest, {"prompt_id": prompt_id, "rows": clean, "keep": int(keep)})
aside = sorted(_dir(key).glob(f"{kind}.*.pending.pt"), key=lambda p: p.stat().st_mtime_ns)
for old in aside[:-MAX_PENDING]:
old.unlink(missing_ok=True)
_write(latest, {"prompt_id": prompt_id, "rows": clean, "keep": int(keep), "only": only, "mixed": bool(mixed)})
_prune_aside(key, kind)
state = _read_latest()
runs = state.get("runs", {})
kinds = runs.setdefault(prompt_id, {}).setdefault(key, [])
@@ -195,7 +201,7 @@ def blind(rows, to):
return [dict(r, reward=0.0) if r.get("axis") == to else r for r in rows]
def _restore_pending(key, kind, prompt_id, rows, keep):
def _restore_pending(key, kind, prompt_id, rows, keep, only=None, mixed=False):
"""A cleared rating puts the capture back to waiting (until the next Generate), so the clip can be
rated again. The latest slot if it is free or already this run's, else aside under the prompt id."""
latest = _pending_path(key, kind)
@@ -204,7 +210,8 @@ def _restore_pending(key, kind, prompt_id, rows, keep):
if not _PID.match(str(prompt_id)):
return
latest = _pending_path(key, kind, prompt_id)
_write(latest, {"prompt_id": prompt_id, "rows": rows, "keep": int(keep)})
_write(latest, {"prompt_id": prompt_id, "rows": rows, "keep": int(keep), "only": only, "mixed": bool(mixed)})
_prune_aside(key, kind)
def rate(prompt_id, rating, axis=None):
@@ -237,7 +244,8 @@ def rate(prompt_id, rating, axis=None):
if rating is None:
data["rows"] = [r for r in data["rows"] if r.get("prompt_id") != prompt_id]
if path.stem in entry.get(key, []):
_restore_pending(key, path.stem, prompt_id, hit[-1]["rows"], hit[-1].get("keep", MAX_ROWS))
_restore_pending(key, path.stem, prompt_id, hit[-1]["rows"], hit[-1].get("keep", MAX_ROWS),
hit[-1].get("only"), hit[-1].get("mixed", False))
else:
for r in hit:
r["reward"] = REWARD[rating]
@@ -269,7 +277,12 @@ def rate(prompt_id, rating, axis=None):
continue
data = load(key, kind)
keep = int(pending.get("keep", MAX_ROWS))
row = {"prompt_id": prompt_id, "reward": REWARD[rating], "rows": pending["rows"], "keep": keep}
if pending.get("only") == "liked" and rating != "liked":
pending_path.unlink(missing_ok=True) # nothing here learns from it: not banked
out.setdefault("skipped", []).append(kind)
continue
row = {"prompt_id": prompt_id, "reward": REWARD[rating], "rows": pending["rows"], "keep": keep,
"only": pending.get("only"), "mixed": bool(pending.get("mixed"))}
if axis:
row["axis"] = axis
data["rows"].append(row)
@@ -277,7 +290,7 @@ def rate(prompt_id, rating, axis=None):
_write(_dir(key) / f"{kind}.pt", data)
pending_path.unlink(missing_ok=True)
out["recorded"].append(kind)
if not out["recorded"]:
if not out["recorded"] and not out.get("skipped"):
out["why"] = ("this clip's capture was not kept (more runs waited than are kept, or it was "
"cleared); nothing was learned")
return out
@@ -384,6 +397,8 @@ def _sound_rows(data):
return False
# What is stacked later must agree: one name, one shape (a key trained across two model
# families would otherwise abort every sampling that reads it).
if row.get("mixed") is True:
continue # a clip bank: rows of different sizes are the point
for name, v in captured.items():
if shapes.setdefault(name, tuple(v.shape)) != tuple(v.shape):
return False
modules/system/taste/tests/test_taste.py
+30
View File
@@ -210,3 +210,33 @@ def test_a_run_waiting_under_two_keys_teaches_both():
store.capture("wolf", "reins", {25: _v(2)}, prompt_id="p1")
assert store.rate("p1", "liked")["recorded"] == ["reins", "reins"]
assert store.counts("fox", "reins") == (1, 0) and store.counts("wolf", "reins") == (1, 0)
def test_a_liked_only_kind_banks_likes_and_never_lets_dislikes_push_them_out():
for i in range(3):
store.capture("fox", "vel", {"v": _v(i)}, prompt_id=f"l{i}", keep=4, only="liked")
assert store.rate(f"l{i}", "liked")["recorded"] == ["vel"]
for i in range(10):
store.capture("fox", "vel", {"v": _v(9)}, prompt_id=f"d{i}", keep=4, only="liked")
out = store.rate(f"d{i}", "disliked")
assert out["recorded"] == [] and out["skipped"] == ["vel"] and out["why"] is None
assert store.counts("fox", "vel") == (3, 0)
def test_a_kind_of_mixed_sizes_exports_and_imports(tmp_path):
store.capture("fox", "vel", {"v": torch.ones(2, 2)}, prompt_id="a", mixed=True)
store.rate("a", "liked")
store.capture("fox", "vel", {"v": torch.ones(3, 3)}, prompt_id="b", mixed=True)
store.rate("b", "liked")
out = tmp_path / "k.zip"
store.export_key("fox", out)
assert store.import_key("wolf", out) == 1
assert store.counts("wolf", "vel") == (2, 0)
# a kind that did not say so still must hold one shape per name
store.capture("bear", "reins", {"v": torch.ones(2, 2)}, prompt_id="c")
store.rate("c", "liked")
store.capture("bear", "reins", {"v": torch.ones(3, 3)}, prompt_id="d")
store.rate("d", "liked")
store.export_key("bear", tmp_path / "b.zip")
with pytest.raises(ValueError):
store.import_key("cub", tmp_path / "b.zip")