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 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
co-authored by
Claude Sonnet 5.5
Claude Sonnet 5.5
parent
0edb26f637
commit
69f3019d9c
+1
-1
@@ -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
|
||||
|
||||
+3
-1
@@ -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 });
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
|
||||
@@ -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); }
|
||||
|
||||
@@ -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}
|
||||
@@ -0,0 +1,351 @@
|
||||
"""Named taste keys on disk, and pairing a run's capture with its rating.
|
||||
|
||||
Layout: `<pack>/taste/<key>/<kind>.pt` holds `{"rows": [{"prompt_id", "reward",
|
||||
"rows": {name: tensor}}]}`; `<kind>.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 `<kind>.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)
|
||||
@@ -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]
|
||||
@@ -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 ")
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user