From 69f3019d9cdc26028ec812ada894cee2da819b70 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sat, 3 Oct 2026 00:07:09 +0300 Subject: [PATCH] fix: track modules/system/taste (an unanchored .gitignore rule had kept the whole module out of git); key import validates row shape, clears a stale latest, reports damaged zips; wizard Esc yields to layers above it and keeps a resumable run; icons view keeps its actions Co-Authored-By: Claude Sonnet 5.5 --- .gitignore | 2 +- app/legacy/api.js | 4 +- app/legacy/onboarding.js | 9 +- app/legacy/styles.css | 3 +- modules/system/taste/__init__.py | 227 +++++++++++ modules/system/taste/store.py | 351 ++++++++++++++++++ modules/system/taste/tests/test_taste.py | 117 ++++++ .../system/taste/tests/test_taste_routes.py | 126 +++++++ modules/system/taste/tests/test_value.py | 35 ++ modules/system/taste/value.py | 126 +++++++ 10 files changed, 995 insertions(+), 5 deletions(-) create mode 100644 modules/system/taste/__init__.py create mode 100644 modules/system/taste/store.py create mode 100644 modules/system/taste/tests/test_taste.py create mode 100644 modules/system/taste/tests/test_taste_routes.py create mode 100644 modules/system/taste/tests/test_value.py create mode 100644 modules/system/taste/value.py diff --git a/.gitignore b/.gitignore index 29983fab..b9607fc3 100644 --- a/.gitignore +++ b/.gitignore @@ -40,7 +40,7 @@ app/playwright-report/ /shortcut_revolver.json # Taste keys: what your ratings taught, per rental -taste/ +/taste/ /block_influence.enabled /block_influence.latest.json /phrase_probe.enabled diff --git a/app/legacy/api.js b/app/legacy/api.js index 374bda06..ba6d27ed 100644 --- a/app/legacy/api.js +++ b/app/legacy/api.js @@ -398,7 +398,9 @@ refinementKeys: () => j("GET", API("/api/m/taste/keys")), // A key travels as a zip (the exported file), sent as the raw body: it can be tens of MB. importRefinementKey: async (file, opts) => { - const name = String(file.name || "key").replace(/\.zip$/i, "").trim(); + // A downloaded copy may be "fox (1).zip": reduce it to a name a key can have. + const name = String(file.name || "key").replace(/\.zip$/i, "").replace(/[^A-Za-z0-9 _.-]+/g, "_") + .replace(/^[^A-Za-z0-9]+/, "").replace(/[. ]+$/, "").slice(0, 64) || "key"; const q = new URLSearchParams({ name }); if (opts && opts.overwrite) q.set("overwrite", "1"); const res = await fetch(API("/api/m/taste/keys/import?" + q), { method: "POST", body: file }); diff --git a/app/legacy/onboarding.js b/app/legacy/onboarding.js index 481b0ef7..481f9b9e 100644 --- a/app/legacy/onboarding.js +++ b/app/legacy/onboarding.js @@ -698,6 +698,7 @@ const back = () => go(-1); function renderStep() { + if (!root) return; // closed while a step was still working const name = steps[idx]; const st = STEPS[name]; paint(st.render(), name); @@ -705,6 +706,7 @@ } function paint(node, name) { + if (!root || !stage) return; if (stage?._child?._cleanup) stage._child._cleanup(); clear(stage); root.querySelector(".oo-actions")?.remove(); @@ -785,7 +787,9 @@ // The wizard must always be leavable: a modal that can only be completed is a trap. Esc and the // ✕ close it; an unfinished run is remembered (save()) the way a refresh would have left it. function onEscape(e) { - if (e.key !== "Escape" || !root) return; + if (e.key !== "Escape" || !root || e.defaultPrevented) return; + // A layer open above the wizard (Settings, a floating window, the restart overlay) owns Esc. + if (document.querySelector(".restart-overlay, .sw-overlay, .fw:not([hidden])")) return; const t = document.activeElement; if (t && (t.tagName === "INPUT" || t.tagName === "TEXTAREA" || t.tagName === "SELECT")) return; e.preventDefault(); @@ -793,7 +797,8 @@ } function leave() { if (panel) { closePane(); return; } - try { if (ctx) save(false); } catch (_) {} + // A run waiting out the wizard's own ComfyUI restart stays resumable: only an ordinary run is rewritten. + try { if (ctx && !(restore() || {}).resumable) save(false); } catch (_) {} close(); } diff --git a/app/legacy/styles.css b/app/legacy/styles.css index 5148bda7..e5699299 100644 --- a/app/legacy/styles.css +++ b/app/legacy/styles.css @@ -721,7 +721,8 @@ input[type=number] { font-family: var(--font-ui); } #media-zone .media-grid[data-view="list"] .media-name-actions { opacity: 1; } #media-zone .media-grid[data-view="icons"][data-cols] { grid-template-columns: repeat(auto-fill, minmax(46px, 1fr)); gap: 4px; --mb-cell: 46px; } #media-zone .media-grid[data-view="icons"] .media-thumb { height: 46px; } -#media-zone .media-grid[data-view="icons"] .media-name-row { display: none; } +#media-zone .media-grid[data-view="icons"] .media-name { display: none; } +#media-zone .media-grid[data-view="icons"] .media-name-row { position: absolute; top: 0; right: 0; left: 0; min-height: 0; padding: 1px; justify-content: flex-end; flex-wrap: wrap; background: color-mix(in srgb, var(--ink-void) 70%, transparent); } .media-card { position: relative; border: 1px solid var(--line); border-radius: 8px; overflow: hidden; cursor: pointer; background: var(--ink-0); } .media-card:hover { border-color: var(--muted); } .media-card.previewing { border-color: var(--teal); box-shadow: 0 0 0 1px color-mix(in srgb, var(--teal) 35%, transparent); } diff --git a/modules/system/taste/__init__.py b/modules/system/taste/__init__.py new file mode 100644 index 00000000..fab8c5c0 --- /dev/null +++ b/modules/system/taste/__init__.py @@ -0,0 +1,227 @@ +"""Taste keys: named stores of what you liked and disliked, for features that learn. + +This module learns nothing itself. It names the key a run teaches, keeps each +learning feature's capture of the latest run until you rate it, and answers +"which way is liked?" per feature. Features reach it through the `taste_store` +capability, so turning this module off turns every learning feature off with a +message instead of breaking it. +""" + +import asyncio +import os +import tempfile + +from . import store, value + +ID = "taste" +TITLE = "Taste key" +MOUNT = "generation.sampling" +STAGE = "load" +CATEGORY = "system" +STATUS = "proven" + +SETTINGS = { + "key": { + "type": "text", "default": "", + "label": "Taste key", + "hint": "Name what your ratings teach, e.g. a style or a character. Empty = learning features stay off.", + }, +} + +_OPTION = "funpack_taste_key" +_judges = {} + + +class Handle: + """One key, as a learning feature sees it: capture now, ask for a direction later.""" + + 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 rows(self, kind, blind_to=None): + """Every rated row of `kind`, oldest first: {"reward", "rows", "prompt_id"}. + `blind_to` ("image"/"composition"): dislikes blamed on that axis alone read as + 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): + """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. + + `fresh()` runs as each sampling starts. What was learned must be read + THERE, not at install: ComfyUI caches the node that installs modifiers + (only the seed changes between runs), so install-time reads would ignore + every rating after the first run.""" + from comfy.patcher_extension import WrappersMP + rows = {} + + def outer(executor, *args, **kwargs): + rows.clear() + if fresh is not None: + fresh() + out = executor(*args, **kwargs) + self.capture(kind, rows, keep=keep) + return out + + patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, key, outer) + return rows + + def direction(self, kind, name): + return store.direction(self.key, kind, name) + + def directions(self, kind, names, strength): + """-> ({name: unit vector}, one-line summary) for the names that have + enough ratings, with the ones still waiting named in the summary.""" + found, steering, waiting = {}, [], [] + for name in names: + d, liked, disliked = self.direction(kind, name) + if d is None: + waiting.append(f"{name} ({liked}/{disliked})") + else: + found[name] = d + steering.append(f"{name} ({liked}/{disliked})") + parts = [] + if steering: + parts.append(f"steering {', '.join(steering)} at {strength:g}" if strength > 0 + else f"strength 0: learning only at {', '.join(steering)}") + if waiting: + parts.append(f"learning, needs 2 liked + 2 disliked at {', '.join(waiting)}") + return found, f"key {self.key!r}: " + "; ".join(parts or ["learning"]) + + def describe(self, video): + return value.describe(video) + + def judge(self, kind, name): + """A Judge trained on this key's `kind` rows holding `name`, or None + while there are too few. Retrained only when the rows change.""" + rows = [r for r in store.load(self.key, kind)["rows"] if str(name) in r["rows"]] + slot = (self.key, kind, str(name)) + stamp = tuple((r.get("prompt_id"), r["reward"]) for r in rows) + hit = _judges.get(slot) + if hit is None or hit[0] != stamp: + hit = _judges[slot] = (stamp, value.Judge.train_on( + [r["rows"][str(name)] for r in rows], [r["reward"] for r in rows])) + return hit[1] + + def counts(self, kind, name=None): + return store.counts(self.key, kind, name) + + +def install(patcher, values, key): + name = str(values.get("key") or "").strip() + if not name: + return None + if not store.valid(name): + raise ValueError(f"{name!r} can't be a key name: letters, digits, space, _ . - only") + patcher.model_options[_OPTION] = name + return f"teaching {name!r}" + + +def taste_store(patcher): + """The key this model's run teaches, or None when no key is set.""" + name = (getattr(patcher, "model_options", None) or {}).get(_OPTION) + return Handle(name) if name else None + + +class Kind: + """One kind of one key, for a panel that reads or clears it with no run going.""" + + def __init__(self, key, kind): + if not store.valid(key): + raise ValueError(f"{key!r} is not a usable taste key name") + self.key, self.kind = key, kind + + def rows(self): + return store.load(self.key, self.kind)["rows"] + + def path(self): + return store.kind_path(self.key, self.kind) + + def clear(self): + store.clear_kind(self.key, self.kind) + + +def taste_latest_key(): + return store.latest_key() + + +def taste_kind(key, kind): + return Kind(key, kind) + + +def routes(table, base, web): + @table.get(base + "/keys") + async def _keys(_req): + return web.json_response({"keys": store.keys()}) + + @table.post(base + "/rate") + async def _rate(req): + try: + body = await req.json() + out = store.rate(body.get("prompt_id"), body.get("rating"), body.get("axis")) + except (ValueError, AttributeError) as exc: + return web.json_response({"why": str(exc)}, status=400) + return web.json_response(out) + + @table.delete(base + "/keys/{name}") + async def _delete(req): + name = req.match_info["name"] + try: + removed = len(list(store._dir(name).glob("*"))) if store._dir(name).is_dir() else 0 + store.delete(name) + except ValueError as exc: + return web.json_response({"why": str(exc)}, status=400) + return web.json_response({"keys": store.keys(), "deleted": name, "removed": removed}) + + @table.get(base + "/keys/{name}/export") + async def _export(req): + """A key as a zip, to carry to another machine. Written to a temp file off the loop.""" + name = req.match_info["name"] + fd, tmp = tempfile.mkstemp(suffix=".zip") + os.close(fd) + try: + await asyncio.to_thread(store.export_key, name, tmp) + except ValueError as exc: + os.remove(tmp) + return web.json_response({"why": str(exc)}, status=400) + try: + with open(tmp, "rb") as fh: + body = fh.read() + finally: + os.remove(tmp) + return web.Response(body=body, content_type="application/zip", + headers={"Content-Disposition": f'attachment; filename="{name}.zip"'}) + + @table.post(base + "/keys/import") + async def _import(req): + """Take an exported key (the zip, as the raw request body) under `?name=`. Streamed to a + temp file: a key can be tens of MB, past what reading a body whole allows.""" + name = req.query.get("name", "") + if not store.valid(name): + return web.json_response({"why": f"{name!r} is not a usable taste key name"}, status=400) + fd, tmp = tempfile.mkstemp(suffix=".zip") + size = 0 + try: + with os.fdopen(fd, "wb") as out: + async for chunk in req.content.iter_chunked(1 << 20): + size += len(chunk) + if size > store.MAX_BYTES: + return web.json_response({"why": "that file is larger than 4 GB"}, status=413) + out.write(chunk) + try: + count = await asyncio.to_thread(store.import_key, name, tmp, req.query.get("overwrite") == "1") + except FileExistsError: + return web.json_response({"why": f'A key named "{name}" already exists.', "exists": True, "key": name}, status=409) + except ValueError as exc: + return web.json_response({"why": str(exc)}, status=400) + finally: + if os.path.exists(tmp): + os.remove(tmp) + return web.json_response({"imported": name, "kinds": count, "keys": store.keys()}) + + +PROVIDES = {"modifier": install, "taste_store": taste_store, "taste_kind": taste_kind, "taste_latest_key": taste_latest_key, "routes": routes} diff --git a/modules/system/taste/store.py b/modules/system/taste/store.py new file mode 100644 index 00000000..074f4f7b --- /dev/null +++ b/modules/system/taste/store.py @@ -0,0 +1,351 @@ +"""Named taste keys on disk, and pairing a run's capture with its rating. + +Layout: `/taste//.pt` holds `{"rows": [{"prompt_id", "reward", +"rows": {name: tensor}}]}`; `.pending.pt` holds the latest run's capture. +A key is a folder, so deleting one is one folder and a kind is one file. + +"Latest only" (user's choice, 2026-09-27): each capture overwrites the pending +one, so a rating can only teach from the most recent run -- but it is TAGGED with +the ComfyUI prompt id it came from, so rating any other clip is refused and said, +never paired with the wrong run. Changing your mind on a clip already recorded +updates its row instead of adding a second one. + +Keys are disposable (retrained per rental), so there is no migration from v4's +`refinements/` files. +""" + +import json +import os +import re +import shutil +import tempfile +import threading +import zipfile +import zlib +from pathlib import Path + +import torch + +from ..._core import config + +ROOT = config.ROOT / "taste" +_KEY = re.compile(r"\A[A-Za-z0-9][A-Za-z0-9 _.-]{0,63}\Z") +_LOCK = threading.RLock() +REWARD = {"liked": 1.0, "disliked": -1.0} +# Each rating's vote shrinks by this per later rating: half the vote sits in the +# last ~7 (v4's RECENCY_DECAY; without it early ratings of one style outvoted a +# new style until out-rated one-for-one). +RECENCY_DECAY = 0.9 +# Two liked AND two disliked before a direction exists: one of each is two +# points, not a direction. +MIN_PER_GROUP = 2 +# A dislike can say WHICH half went wrong: "image" (the picture was bad) or "composition" +# (the picture was fine, the shot plan was not). Plain disliked blames both. +AXES = ("image", "composition") +MAX_ROWS = 200 + + +def valid(key) -> bool: + return isinstance(key, str) and bool(_KEY.match(key)) and not key.endswith((".", " ")) + + +def _dir(key): + if not valid(key): + raise ValueError(f"{key!r} is not a usable taste key name") + return ROOT / key + + +def keys(): + if not ROOT.is_dir(): + return [] + return sorted(p.name for p in ROOT.iterdir() if p.is_dir() and valid(p.name)) + + +def _read(path, default): + try: + return torch.load(path, map_location="cpu", weights_only=True) + except FileNotFoundError: + return default + + +def _write(path, data): + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(path.suffix + ".tmp") + torch.save(data, tmp) + os.replace(tmp, path) + + +_loaded = {} + + +def load(key, kind): + """Rated rows of a kind. Re-read only when the file changed: features ask at + the start of every run, and a REINS log is tens of MB.""" + path = _dir(key) / f"{kind}.pt" + try: + stamp = path.stat().st_mtime_ns + except FileNotFoundError: + return {"rows": []} + hit = _loaded.get(path) + if hit is None or hit[0] != stamp: + hit = _loaded[path] = (stamp, _read(path, {"rows": []})) + return hit[1] + + +def latest_key(): + """The key the most recent capture went to, or None.""" + return _read_latest().get("key") + + +def kind_path(key, kind): + """Where a kind's rated rows live on disk (it may not exist yet).""" + return _dir(key) / f"{kind}.pt" + + +def clear_kind(key, kind): + """Forget every rated row and the waiting capture of one kind.""" + for path in (kind_path(key, kind), _dir(key) / f"{kind}.pending.pt"): + path.unlink(missing_ok=True) + + +def current_prompt_id(): + """The ComfyUI prompt running now -- set before any node executes.""" + try: + from server import PromptServer + return getattr(PromptServer.instance, "last_prompt_id", None) + except Exception: # noqa: BLE001 -- tests, headless + return None + + +def _latest_path(): + return ROOT / "latest.json" + + +def capture(key, kind, rows, prompt_id=None, keep=MAX_ROWS): + """This run's capture for `kind`, waiting for a rating. Overwrites the last. + + Kept in the dtype given: a banked latent stays half precision. `keep` caps + how many rated rows this kind holds, oldest dropped first. + """ + if not rows: + return + prompt_id = prompt_id or current_prompt_id() + clean = {str(k): v.detach().cpu() for k, v in rows.items()} + with _LOCK: + _write(_dir(key) / f"{kind}.pending.pt", + {"prompt_id": prompt_id, "rows": clean, "keep": int(keep)}) + latest = _read_latest() + kinds = latest.get("kinds", []) if latest.get("prompt_id") == prompt_id \ + and latest.get("key") == key else [] + if kind not in kinds: + kinds.append(kind) + ROOT.mkdir(parents=True, exist_ok=True) + _latest_path().write_text(json.dumps({"key": key, "prompt_id": prompt_id, + "kinds": kinds})) + + +def _read_latest(): + try: + return json.loads(_latest_path().read_text()) + except (FileNotFoundError, ValueError): + return {} + + +def blind(rows, to): + """`rows` as a learner that cannot judge the `to` axis sees them: a dislike blamed on + that axis alone is neutral (reward 0, so it counts as neither liked nor disliked).""" + if not to: + return rows + return [dict(r, reward=0.0) if r.get("axis") == to else r for r in rows] + + +def rate(prompt_id, rating, axis=None): + """-> {"recorded": [kinds], "updated": [kinds], "why": str|None}. + + `rating` None clears: the row is removed, nothing is learned from that clip. + `axis` only goes with "disliked" (see AXES); any other rating drops it. + """ + if rating is not None and rating not in REWARD: + raise ValueError(f"{rating!r} is not a rating") + if axis is not None and axis not in AXES: + raise ValueError(f"{axis!r} is not a rating axis") + if axis is not None and rating != "disliked": + raise ValueError("only a dislike can name what went wrong") + out = {"recorded": [], "updated": [], "why": None} + with _LOCK: + # A clip already recorded: change or remove its row, in every key/kind. + for key in keys(): + for path in _dir(key).glob("*.pt"): + if path.name.endswith(".pending.pt"): + continue + data = _read(path, {"rows": []}) + hit = [r for r in data["rows"] if r.get("prompt_id") == prompt_id] + if not hit: + continue + if rating is None: + data["rows"] = [r for r in data["rows"] if r.get("prompt_id") != prompt_id] + else: + for r in hit: + r["reward"] = REWARD[rating] + r.pop("axis", None) + if axis: + r["axis"] = axis + _write(path, data) + out["updated"].append(path.stem) + if out["updated"]: + return out + + latest = _read_latest() + if not prompt_id or latest.get("prompt_id") != prompt_id: + out["why"] = ("only the most recent generation can teach the taste key -- " + "this clip's capture has been replaced by a newer run") + return out + if rating is None: + return out + key = latest["key"] + for kind in latest.get("kinds", []): + pending_path = _dir(key) / f"{kind}.pending.pt" + pending = _read(pending_path, None) + if not pending or pending.get("prompt_id") != prompt_id: + continue + data = load(key, kind) + row = {"prompt_id": prompt_id, "reward": REWARD[rating], "rows": pending["rows"]} + if axis: + row["axis"] = axis + data["rows"].append(row) + data["rows"] = data["rows"][-int(pending.get("keep", MAX_ROWS)):] + _write(_dir(key) / f"{kind}.pt", data) + pending_path.unlink(missing_ok=True) + out["recorded"].append(kind) + return out + + +def counts(key, kind, name=None): + """(liked, disliked) rows for a kind, optionally only rows holding `name`.""" + rows = [r for r in load(key, kind)["rows"] if name is None or str(name) in r["rows"]] + return (sum(r["reward"] > 0 for r in rows), sum(r["reward"] < 0 for r in rows)) + + +def direction(key, kind, name): + """-> (unit vector | None, n_liked, n_disliked). + + Recency-weighted mean(liked) - mean(disliked) over the rows holding `name`, + split by the SIGN of the rating only. Raw rows, not normalised: a hidden + state's magnitude carries content (v4's trajectory-probe lesson). + """ + name = str(name) + rows = [r for r in load(key, kind)["rows"] if name in r["rows"]] + aged = [(r["rows"][name], r["reward"], RECENCY_DECAY ** (len(rows) - 1 - i)) + for i, r in enumerate(rows)] + liked = [(d, a) for d, w, a in aged if w > 0] + disliked = [(d, a) for d, w, a in aged if w < 0] + if len(liked) < MIN_PER_GROUP or len(disliked) < MIN_PER_GROUP: + return None, len(liked), len(disliked) + + def wmean(group): + ds = torch.stack([d for d, _ in group]).float() + a = torch.tensor([a for _, a in group], dtype=ds.dtype) + return (ds * a.view(-1, *([1] * (ds.dim() - 1)))).sum(0) / a.sum() + + diff = wmean(liked) - wmean(disliked) + norm = diff.norm() + if not torch.isfinite(norm) or norm <= 1e-8: + return None, len(liked), len(disliked) + return diff / norm, len(liked), len(disliked) + + +def delete(key): + with _LOCK: + shutil.rmtree(_dir(key), ignore_errors=True) + if _read_latest().get("key") == key: + _latest_path().unlink(missing_ok=True) + + +# --- moving a key between machines ------------------------------------------- +# +# A key is a folder of `.pt` files, so exporting is a zip of them. Importing reads a zip a +# stranger may have made: only plainly named `.pt` files, bounded in count and size, each one +# loadable with weights_only (no pickled code) before anything touches the real folder. + +_KIND_FILE = re.compile(r"\A[A-Za-z0-9_]{1,48}\.pt\Z") +MAX_FILES = 64 +MAX_BYTES = 4 * 1024 ** 3 + + +def export_key(key, out_path): + """Zip `key`'s rated rows into `out_path` (waiting captures are not part of what was learned).""" + folder = _dir(key) + if not folder.is_dir(): + raise ValueError(f"there is no taste key called {key!r}") + with _LOCK: + found = [p for p in sorted(folder.glob("*.pt")) if _KIND_FILE.match(p.name)] + if not found: + raise ValueError(f"{key!r} has not learned anything yet: there is nothing to export") + with zipfile.ZipFile(out_path, "w", zipfile.ZIP_STORED) as zf: + for path in found: + zf.write(path, path.name) + return out_path + + +def _sound_rows(data): + """True when `data` is what the learners consume: a dict of rated rows, each with a number + for a reward and tensors for what was captured. Loading it safely is not enough -- a row of + the wrong shape would break every later rating.""" + if not isinstance(data, dict) or not isinstance(data.get("rows"), list): + return False + for row in data["rows"]: + if not isinstance(row, dict) or isinstance(row.get("reward"), bool) or not isinstance(row.get("reward"), (int, float)): + return False + if row.get("axis") not in (None, *AXES): + return False + captured = row.get("rows") + if not isinstance(captured, dict) or not all(isinstance(v, torch.Tensor) for v in captured.values()): + return False + return True + + +def import_key(key, zip_path, overwrite=False): + """-> number of kinds imported. FileExistsError if `key` exists and `overwrite` is off.""" + dest = _dir(key) + try: + zf = zipfile.ZipFile(zip_path) + except zipfile.BadZipFile as exc: + raise ValueError("that is not an exported taste key (not a zip file)") from exc + with zf: + infos = zf.infolist() + if not infos or len(infos) > MAX_FILES: + raise ValueError(f"an exported key holds 1 to {MAX_FILES} files, this has {len(infos)}") + for info in infos: + if not _KIND_FILE.match(info.filename): + raise ValueError(f"{info.filename!r} is not part of a taste key") + if sum(i.file_size for i in infos) > MAX_BYTES: + raise ValueError("that key is larger than 4 GB") + ROOT.mkdir(parents=True, exist_ok=True) + stage = Path(tempfile.mkdtemp(prefix=".import-", dir=ROOT)) + try: + for info in infos: + target = stage / info.filename + try: + with zf.open(info) as src, open(target, "wb") as out: + shutil.copyfileobj(src, out) + except (zipfile.BadZipFile, zlib.error, RuntimeError, NotImplementedError, EOFError) as exc: + raise ValueError(f"{info.filename} cannot be read from the zip ({type(exc).__name__}): " + f"it is damaged, encrypted or compressed in an unsupported way") from exc + try: + data = torch.load(target, map_location="cpu", weights_only=True) + except Exception as exc: # noqa: BLE001 + raise ValueError(f"{info.filename} is not a readable taste file ({type(exc).__name__})") from exc + if not _sound_rows(data): + raise ValueError(f"{info.filename} does not hold rated rows of the shape FunPack reads") + with _LOCK: + if dest.exists(): + if not overwrite: + raise FileExistsError(key) + shutil.rmtree(dest) + os.replace(stage, dest) + # Anything waiting to be rated under this name died with the old folder. + if _read_latest().get("key") == key: + _latest_path().unlink(missing_ok=True) + return len(infos) + finally: + shutil.rmtree(stage, ignore_errors=True) diff --git a/modules/system/taste/tests/test_taste.py b/modules/system/taste/tests/test_taste.py new file mode 100644 index 00000000..b8d0fddb --- /dev/null +++ b/modules/system/taste/tests/test_taste.py @@ -0,0 +1,117 @@ +"""Taste keys: capture -> rate pairing, re-rating, and the direction math.""" + +import pytest +import torch + +from modules.system import taste +from modules.system.taste import store + + +@pytest.fixture(autouse=True) +def root(tmp_path, monkeypatch): + monkeypatch.setattr(store, "ROOT", tmp_path / "taste") + return tmp_path / "taste" + + +def _v(x): + return torch.tensor([float(x), 0.0]) + + +def test_latest_run_is_recorded_once_and_a_replaced_one_is_refused(): + store.capture("fox", "reins", {25: _v(1)}, prompt_id="p1") + store.capture("fox", "reins", {25: _v(2)}, prompt_id="p2") + out = store.rate("p1", "liked") + assert out["recorded"] == [] and "most recent" in out["why"] + assert store.rate("p2", "liked")["recorded"] == ["reins"] + assert store.counts("fox", "reins") == (1, 0) + # The pending capture is spent: rating p2 again changes it, never adds a row. + assert store.rate("p2", "disliked")["updated"] == ["reins"] + assert store.counts("fox", "reins") == (0, 1) + store.rate("p2", None) + assert store.counts("fox", "reins") == (0, 0) + + +def test_every_kind_of_one_run_is_recorded_together(): + store.capture("fox", "reins", {1: _v(1)}, prompt_id="p1") + store.capture("fox", "dynashift", {1: _v(1)}, prompt_id="p1") + assert sorted(store.rate("p1", "liked")["recorded"]) == ["dynashift", "reins"] + + +def test_direction_needs_two_of_each_and_points_liked_minus_disliked(): + for i, (x, r) in enumerate([(3, "liked"), (-3, "disliked"), (5, "liked")]): + store.capture("fox", "k", {7: _v(x)}, prompt_id=f"p{i}") + store.rate(f"p{i}", r) + assert store.direction("fox", "k", 7)[0] is None + store.capture("fox", "k", {7: _v(-1)}, prompt_id="p9") + store.rate("p9", "disliked") + d, liked, disliked = store.direction("fox", "k", 7) + assert (liked, disliked) == (2, 2) + assert torch.allclose(d, torch.tensor([1.0, 0.0])) + + +def test_recent_ratings_outvote_old_ones(): + rows = [(10, "liked"), (-10, "disliked"), (10, "liked"), (-10, "disliked")] + rows += [(0, "liked")] * 30 + [(1, "disliked")] * 2 + for i, (x, r) in enumerate(rows): + store.capture("fox", "k", {0: torch.tensor([float(x), float(i % 2)])}, prompt_id=f"p{i}") + store.rate(f"p{i}", r) + d, *_ = store.direction("fox", "k", 0) + assert d[0] < 0 # the recent disliked +1 dominates the old +-10 + + +def test_bad_names_are_refused(): + for bad in ("", "../x", "a/b", ".hidden", "x."): + assert not store.valid(bad) + with pytest.raises(ValueError): + store.capture("../x", "k", {0: _v(1)}, prompt_id="p") + + +class _Patcher: + def __init__(self): + self.model_options = {} + + +def test_modifier_names_the_key_the_run_teaches(): + p = _Patcher() + assert taste.install(p, {"key": ""}, "k") is None + assert taste.taste_store(p) is None + assert "fox" in taste.install(p, {"key": "fox"}, "k") + handle = taste.taste_store(p) + handle.capture("reins", {3: _v(1)}) + assert (store.ROOT / "fox" / "reins.pending.pt").exists() + with pytest.raises(ValueError): + taste.install(p, {"key": "a/b"}, "k") + + +def test_a_dislike_can_name_the_axis_and_a_blind_learner_reads_it_as_neutral(): + for i, axis in enumerate([None, "image", "composition"]): + store.capture("fox", "k", {1: _v(i)}, prompt_id=f"p{i}") + store.rate(f"p{i}", "disliked", axis) + rows = store.load("fox", "k")["rows"] + assert [r.get("axis") for r in rows] == [None, "image", "composition"] + # composition learners cannot judge a bad PICTURE; the picture learner cannot judge a bad plan + assert [r["reward"] for r in store.blind(rows, "image")] == [-1.0, 0.0, -1.0] + assert [r["reward"] for r in store.blind(rows, "composition")] == [-1.0, -1.0, 0.0] + assert [r["reward"] for r in store.blind(rows, None)] == [-1.0] * 3 + assert rows[1]["reward"] == -1.0 # the stored rows are never edited + + +def test_changing_the_rating_drops_or_sets_the_axis_and_bad_axes_are_refused(): + store.capture("fox", "k", {1: _v(1)}, prompt_id="p1") + store.rate("p1", "disliked", "image") + store.rate("p1", "liked") + assert "axis" not in store.load("fox", "k")["rows"][0] + store.rate("p1", "disliked", "composition") + assert store.load("fox", "k")["rows"][0]["axis"] == "composition" + for bad in [("liked", "image"), ("disliked", "colour"), (None, "image")]: + with pytest.raises(ValueError): + store.rate("p1", *bad) + + +def test_the_learners_read_through_the_blind_view(): + for i, axis in enumerate(["image", "composition"]): + store.capture("fox", "k", {1: _v(i)}, prompt_id=f"p{i}") + store.rate(f"p{i}", "disliked", axis) + h = taste.Handle("fox") + assert [r["reward"] for r in h.rows("k", blind_to="image")] == [0.0, -1.0] + assert [r["reward"] for r in h.rows("k")] == [-1.0, -1.0] diff --git a/modules/system/taste/tests/test_taste_routes.py b/modules/system/taste/tests/test_taste_routes.py new file mode 100644 index 00000000..67409261 --- /dev/null +++ b/modules/system/taste/tests/test_taste_routes.py @@ -0,0 +1,126 @@ +"""The rate route over real HTTP, mounted by core's module-routes hook.""" + +import pytest +import torch + +pytest.importorskip("aiohttp") + +from core.tests.test_routes_shortcuts import _request, server # noqa: F401,E402 +from modules.system.taste import store # noqa: E402 + + +@pytest.fixture(autouse=True) +def root(tmp_path, monkeypatch): + monkeypatch.setattr(store, "ROOT", tmp_path / "taste") + + +def test_rate_and_list(server): # noqa: F811 + store.capture("fox", "reins", {1: torch.ones(2)}, prompt_id="p1") + status, out = _request(server, "POST", "/funpack/api/m/taste/rate", + {"prompt_id": "p1", "rating": "liked"}) + assert status == 200 and out["recorded"] == ["reins"], out + status, out = _request(server, "POST", "/funpack/api/m/taste/rate", + {"prompt_id": "old", "rating": "liked"}) + assert status == 200 and out["why"], out + status, out = _request(server, "POST", "/funpack/api/m/taste/rate", + {"prompt_id": "p1", "rating": "meh"}) + assert status == 400 + assert _request(server, "GET", "/funpack/api/m/taste/keys")[1] == {"keys": ["fox"]} + assert _request(server, "DELETE", "/funpack/api/m/taste/keys/fox")[1] == {"keys": [], "deleted": "fox", "removed": 1} + + +def test_a_dislike_over_http_carries_its_axis(server): # noqa: F811 + store.capture("fox", "reins", {1: torch.ones(2)}, prompt_id="p1") + status, out = _request(server, "POST", "/funpack/api/m/taste/rate", + {"prompt_id": "p1", "rating": "disliked", "axis": "image"}) + assert status == 200 and out["recorded"] == ["reins"] + assert store.load("fox", "reins")["rows"][0]["axis"] == "image" + status, _ = _request(server, "POST", "/funpack/api/m/taste/rate", + {"prompt_id": "p1", "rating": "liked", "axis": "image"}) + assert status == 400 + + +def _raw(server, method, path, body=b""): + import http.client + conn = http.client.HTTPConnection("127.0.0.1", server) + conn.request(method, path, body=body) + resp = conn.getresponse() + data = resp.read() + return resp.status, data + + +def test_a_key_exports_and_imports_back_under_another_name(server): # noqa: F811 + store.capture("fox", "reins", {1: torch.ones(2)}, prompt_id="p1") + _request(server, "POST", "/funpack/api/m/taste/rate", {"prompt_id": "p1", "rating": "liked"}) + status, zipped = _raw(server, "GET", "/funpack/api/m/taste/keys/fox/export") + assert status == 200 and zipped[:2] == b"PK" + status, out = _raw(server, "POST", "/funpack/api/m/taste/keys/import?name=wolf", zipped) + assert status == 200 and b'"wolf"' in out + assert store.load("wolf", "reins")["rows"][0]["reward"] == 1.0 + status, out = _raw(server, "POST", "/funpack/api/m/taste/keys/import?name=wolf", zipped) + assert status == 409 and b'"exists": true' in out + assert _raw(server, "POST", "/funpack/api/m/taste/keys/import?name=wolf&overwrite=1", zipped)[0] == 200 + + +def test_an_import_refuses_what_is_not_a_key(server): # noqa: F811 + import io + import zipfile + assert _raw(server, "POST", "/funpack/api/m/taste/keys/import?name=x", b"not a zip")[0] == 400 + for name in ("../evil.pt", "sub/dir.pt", "notes.txt", "reins.pending.pt"): + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as zf: + zf.writestr(name, b"x") + assert _raw(server, "POST", "/funpack/api/m/taste/keys/import?name=x", buf.getvalue())[0] == 400, name + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as zf: + zf.writestr("reins.pt", b"this is not a torch file") + status, out = _raw(server, "POST", "/funpack/api/m/taste/keys/import?name=x", buf.getvalue()) + assert status == 400 and b"readable" in out + assert store.keys() == [] # nothing half-imported + assert _raw(server, "POST", "/funpack/api/m/taste/keys/import?name=..%2Fx", b"")[0] == 400 + + +def _zip_of(files): + import io + import zipfile + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as zf: + for name, data in files.items(): + zf.writestr(name, data) + return buf.getvalue() + + +def _torch_bytes(obj): + import io + buf = io.BytesIO() + torch.save(obj, buf) + return buf.getvalue() + + +def test_an_import_with_rows_the_learners_cannot_read_is_refused(server): # noqa: F811 + for bad in ({"rows": [1, 2]}, {"rows": [{"reward": "high", "rows": {}}]}, + {"rows": [{"reward": 1.0, "rows": {"1": 5}}]}, + {"rows": [{"reward": 1.0, "axis": "mood", "rows": {"1": torch.ones(1)}}]}): + status, out = _raw(server, "POST", "/funpack/api/m/taste/keys/import?name=x", + _zip_of({"reins.pt": _torch_bytes(bad)})) + assert status == 400 and b"shape FunPack reads" in out, bad + good = {"rows": [{"prompt_id": "p", "reward": 1.0, "rows": {"1": torch.ones(2)}}]} + assert _raw(server, "POST", "/funpack/api/m/taste/keys/import?name=x", _zip_of({"reins.pt": _torch_bytes(good)}))[0] == 200 + + +def test_overwriting_a_key_drops_the_waiting_capture_so_the_next_rating_says_so(server): # noqa: F811 + good = _zip_of({"reins.pt": _torch_bytes({"rows": []})}) + store.capture("fox", "reins", {1: torch.ones(2)}, prompt_id="p1") + assert _raw(server, "POST", "/funpack/api/m/taste/keys/import?name=fox&overwrite=1", good)[0] == 200 + assert store.latest_key() is None + + +def test_a_damaged_zip_and_an_empty_key_are_sentences_not_500s(server): # noqa: F811 + data = bytearray(_zip_of({"reins.pt": b"x" * 4000})) + data[40:60] = b"\xff" * 20 # inside the entry's bytes + status, out = _raw(server, "POST", "/funpack/api/m/taste/keys/import?name=x", bytes(data)) + assert status == 400, out + store.capture("fox", "reins", {1: torch.ones(2)}, prompt_id="p1") # waiting capture only + status, out = _raw(server, "GET", "/funpack/api/m/taste/keys/fox/export") + assert status == 400 and b"nothing to export" in out + assert not store.valid("a ") diff --git a/modules/system/taste/tests/test_value.py b/modules/system/taste/tests/test_value.py new file mode 100644 index 00000000..341860da --- /dev/null +++ b/modules/system/taste/tests/test_value.py @@ -0,0 +1,35 @@ +"""The judge: needs enough ratings, learns their order, and its push helps.""" + +import torch + +from modules.system.taste import value + + +def _rated(n=12): + torch.manual_seed(0) + liked_dir = torch.randn(value.DIM) + descs, rewards = [], [] + for i in range(n): + sign = 1.0 if i % 2 else -1.0 + descs.append(sign * liked_dir + 0.3 * torch.randn(value.DIM)) + rewards.append(sign) + return descs, rewards, liked_dir + + +def test_too_few_or_all_alike_is_no_judge(): + descs, rewards, _ = _rated(9) + assert value.Judge.train_on(descs, rewards) is None + descs, _r, _ = _rated(12) + assert value.Judge.train_on(descs, [1.0] * 12) is None + + +def test_its_push_raises_the_score_and_points_at_liked(): + descs, rewards, liked_dir = _rated() + judge = value.Judge.train_on(descs, rewards) + video = torch.randn(1, 4, 2, 8, 8) + moved = judge.nudge(video, 0.05) + score = lambda v: float(sum(net(value.describe(v).unsqueeze(0)) for net in judge.nets)) + assert score(moved) > score(video) + shift = value.describe(moved) - value.describe(video) + assert torch.nn.functional.cosine_similarity(shift, liked_dir, dim=0) > 0.3 + assert torch.allclose((moved - video).norm(), 0.05 * video.norm(), rtol=1e-3) diff --git a/modules/system/taste/value.py b/modules/system/taste/value.py new file mode 100644 index 00000000..a7370e00 --- /dev/null +++ b/modules/system/taste/value.py @@ -0,0 +1,126 @@ +"""A tiny learned judge of rated clips, for features that steer by its gradient. + +Three small MLPs, each trained on its own random pairs of rated descriptors to +rank the liked one above the disliked one (Bradley-Terry: only the ORDER must +be right, which is what a handful of ratings can support). Their gradients are +averaged and scaled by how much they agree on the direction: where they +disagree the judge is guessing, and the push shrinks toward zero. A single +judge in v4's first version walked confidently into a "static character" +collapse; the disagreement damping is the fix for that, kept here. + +Rebuilt from the rated rows whenever they change, never updated in place: the +rows are the record, a judge is a pure function of them. +""" + +import random + +import torch +import torch.nn as nn +import torch.nn.functional as F + +DIM = 512 +MIN_SAMPLES = 10 +ENSEMBLE = 3 +BATCH = 16 +# ponytail: one fixed training budget over the whole buffer instead of v4's +# replay (20 steps per rating as it arrived); same objective, bounded cost. +STEPS = 400 +BUFFER = 100 + + +def describe(video): + """Any-shaped picture latent -> [DIM], differentiably (adaptive pooling, so + clips of any length or size share one judge).""" + c = video.float().reshape(1, 1, -1) + return F.adaptive_avg_pool1d(c, DIM).reshape(-1) + + +class Judge(nn.Module): + def __init__(self, seed=0): + super().__init__() + gen = torch.Generator().manual_seed(seed) + + def net(): + layers = nn.Sequential(nn.Linear(DIM, 128), nn.SiLU(), nn.Linear(128, 64), + nn.SiLU(), nn.Linear(64, 1)) + for p in layers.parameters(): + with torch.no_grad(): + p.copy_(torch.randn(p.shape, generator=gen) * (1.0 / max(1, p.shape[-1])) ** 0.5) + return layers + + self.nets = nn.ModuleList(net() for _ in range(ENSEMBLE)) + + @classmethod + def train_on(cls, descriptors, rewards, seed=0): + """-> a trained Judge, or None under MIN_SAMPLES or with no two ratings + that differ (nothing to rank).""" + pairs = list(zip(descriptors, rewards))[-BUFFER:] + if len(pairs) < MIN_SAMPLES or len({r for _d, r in pairs}) < 2: + return None + judge = cls(seed) + pick = random.Random(seed) + device = "cuda" if torch.cuda.is_available() else "cpu" + xs = torch.stack([d.float().reshape(-1) for d, _r in pairs]).to(device) + rs = [float(r) for _d, r in pairs] + n = len(rs) + with torch.inference_mode(False), torch.enable_grad(): + judge.to(device) + opt = torch.optim.Adam(judge.parameters(), lr=2e-3) + for _ in range(STEPS): + opt.zero_grad() + loss = None + for net in judge.nets: + win, lose = [], [] + for _ in range(BATCH): + i, j = pick.sample(range(n), 2) + if rs[i] == rs[j]: + continue + if rs[i] < rs[j]: + i, j = j, i + win.append(i) + lose.append(j) + if win: + term = -F.logsigmoid(net(xs[win]) - net(xs[lose])).mean() + loss = term if loss is None else loss + term + if loss is not None: + loss.backward() + opt.step() + judge.cpu() + judge.requires_grad_(False) + return judge + + def gradient(self, video): + """d(score)/d(video), same shape, scaled by the members' agreement.""" + device = video.device + if next(self.parameters()).device != device: + self.to(device) + with torch.inference_mode(False), torch.enable_grad(): + x = torch.empty(video.shape, dtype=torch.float32, device=device) + x.copy_(video) + x.requires_grad_(True) + d = describe(x) + grads = [torch.autograd.grad(net(d.unsqueeze(0)).sum(), x, retain_graph=True)[0] + for net in self.nets] + unit = torch.stack([F.normalize(g.flatten(), dim=0) for g in grads]) + n = len(grads) + agreement = ((unit @ unit.T).sum() - n) / (n * (n - 1)) if n > 1 else torch.ones(()) + return (torch.stack(grads).mean(0) * agreement.clamp(0.0, 1.0)).to(video.dtype) + + def nudge(self, video, amount): + """video + this judge's gradient, sized to `amount` x the video's own + norm: the raw gradient reaches each value through a pooling window, so + used as-is it would be far too small to matter.""" + grad = self.gradient(video) + gn = grad.float().norm() + if not bool(gn > 0): + return video + return video + grad * (amount * video.float().norm() / gn).to(video.dtype) + + def score(self, video): + """The members' mean score for a picture: higher = more like what was liked.""" + device = video.device + if next(self.parameters()).device != device: + self.to(device) + with torch.no_grad(): + d = describe(video.float()).unsqueeze(0) + return float(sum(net(d) for net in self.nets)) / len(self.nets)