"""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}