From 19ce43bfb894933e4a097c2dadf28d0d47c4dfe0 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Thu, 3 Sep 2026 20:00:01 +0300 Subject: [PATCH] feat: H3 picker goes binary (Liked/Disliked); REINS gets its own capture pause Picker: h3_repr_steering only ever reads the reward's SIGN now (see the binary-direction commit), so a 1-10 scale was granularity nothing downstream uses. Down to two buttons -- reuses the scale's own "10" and "1" endpoints, zero backend change. REINS capture toggle: independent of the sampler's h3_repr_steering widget, which controls whether the learned direction gets APPLIED -- this new per-key flag controls whether it keeps RECORDING at all, so data collection can be paused without touching the steering setting or discarding what's already captured. Settings > Refinement & Taste gets the toggle; save_pending() is the single gate. --- h3_repr_steering.py | 33 +++++++++++++++++++-- movie_editor/frontend/api.js | 1 + movie_editor/frontend/rating_picker.js | 27 ++++++++++------- movie_editor/frontend/rating_picker.test.js | 6 ++++ movie_editor/frontend/settings_window.js | 33 +++++++++++++++++++-- movie_editor/server.py | 9 ++++++ tests/test_h3_repr_steering.py | 30 +++++++++++++++++++ 7 files changed, 123 insertions(+), 16 deletions(-) diff --git a/h3_repr_steering.py b/h3_repr_steering.py index 05b21eeb..b74a93cb 100644 --- a/h3_repr_steering.py +++ b/h3_repr_steering.py @@ -116,18 +116,43 @@ def capture(hidden_state, video_mask): def _load(refinement_key): path = state_path(refinement_key) if not os.path.exists(path): - return {"rows": [], "pending": None} + return {"rows": [], "pending": None, "enabled": True} try: data = torch.load(path, map_location="cpu", weights_only=True) except Exception: # noqa: BLE001 - return {"rows": [], "pending": None} + return {"rows": [], "pending": None, "enabled": True} if not isinstance(data, dict): - return {"rows": [], "pending": None} + return {"rows": [], "pending": None, "enabled": True} data.setdefault("rows", []) data.setdefault("pending", None) + data.setdefault("enabled", True) return data +def capture_enabled(refinement_key): + """Whether this key is recording at all -- independent of the sampler's own + h3_repr_steering toggle, which also controls whether the learned direction gets APPLIED. + A user who wants to pause data collection (to test something else without polluting the + log, or because a sweep came back null and they want to think before adding more rows) + still needs the widget on for `block_sweep`/status to mean anything, so the pause lives + here instead of overloading that widget.""" + if not refinement_key: + return True + return bool(_load(refinement_key).get("enabled", True)) + + +def set_capture_enabled(refinement_key, enabled): + if not refinement_key: + return + data = _load(refinement_key) + data["enabled"] = bool(enabled) + try: + _save(refinement_key, data) + except OSError as e: + _log().failed("H3 representation steering", "set capture enabled", e, + "the pause/resume did not persist") + + def _save(refinement_key, data): path = state_path(refinement_key) os.makedirs(os.path.dirname(path), exist_ok=True) @@ -146,6 +171,8 @@ def save_pending(refinement_key, descriptors): if not refinement_key or not descriptors: return data = _load(refinement_key) + if not data.get("enabled", True): + return data["pending"] = {int(b): d.detach().float().cpu() for b, d in descriptors.items()} try: _save(refinement_key, data) diff --git a/movie_editor/frontend/api.js b/movie_editor/frontend/api.js index 544a28a0..87038a3d 100644 --- a/movie_editor/frontend/api.js +++ b/movie_editor/frontend/api.js @@ -200,6 +200,7 @@ // H3 representation steering (Settings ▸ Learning) — per-KEY, unlike the probe above. reinsStatus: (key) => j("GET", API("/h3_repr_steering") + `?key=${encodeURIComponent(key || "default")}`), + reinsSetEnabled: (key, enabled) => j("POST", API("/h3_repr_steering/enabled"), { key: key || "default", enabled: !!enabled }), reinsClear: (key) => j("POST", API("/h3_repr_steering/clear"), { key: key || "default" }), async reinsExport(key) { const k = key || "default"; diff --git a/movie_editor/frontend/rating_picker.js b/movie_editor/frontend/rating_picker.js index 13337f5a..679a6a99 100644 --- a/movie_editor/frontend/rating_picker.js +++ b/movie_editor/frontend/rating_picker.js @@ -46,21 +46,23 @@ // H3: the category grid above assumes axis-specific effects (phrase emphasis, dynashift, // ...) this project has not actually proven land differently from one another on H3 — see - // project_reward_model_rework.md. One 1-10 scale instead; the backend already knows how to - // turn a bare number into the nearest existing rating profile (normalize_refiner_v2_rating's - // legacy-score path), so nothing downstream needs to change, only what the picker sends. - const H3_SCALE = Array.from({ length: 10 }, (_, i) => { - const n = i + 1; - return { - label: String(n), display: String(n), - reward: n >= 9 ? "great" : n >= 7 ? "good" : n >= 5 ? "mixed" : n >= 3 ? "bad" : "awful", - hint: n === 10 ? "Exactly what you asked for." : n === 1 ? "Everything wrong." : `${n} of 10.`, - }; - }); + // project_reward_model_rework.md. A 1-10 scale was tried first (still what the backend + // accepts -- normalize_refiner_v2_rating's legacy-score path), but h3_repr_steering only + // ever reads the SIGN of the resulting reward, never its size (see the 2026-09-04 session: + // weighting by magnitude just let a mediocre middle rating dilute the direction at the + // small sample sizes this actually runs at). A scale whose granularity nothing downstream + // uses is a UI lying about how precise the input needs to be, so it's binary now -- "10" + // and "1" are just the scale's own clean-positive/clean-negative endpoints, already wired, + // nothing to change on the backend. + const H3_SCALE = [ + { label: "10", display: "Liked", reward: "+", hint: "This generation was good." }, + { label: "1", display: "Disliked", reward: "−", hint: "This generation was not." }, + ]; // canonical label -> display name, for rendering a stored value (scene button, chips). const DISPLAY_NAMES = {}; for (const cat of CATEGORIES) for (const r of cat.ratings) if (r.display) DISPLAY_NAMES[r.label] = r.display; + for (const r of H3_SCALE) if (r.display) DISPLAY_NAMES[r.label] = r.display; const displayName = (label) => DISPLAY_NAMES[label] || label; const NO_LOVED = new Set([ @@ -88,6 +90,9 @@ const v = String(value || "").trim(); if (!v || v === FORGET_LABEL) return ""; if (v.endsWith("|loved")) return formatLabel(v.slice(0, -6)) + " ♥"; + if (DISPLAY_NAMES[v]) return DISPLAY_NAMES[v]; + // A stray digit from before the H3 picker went binary (2..9) has no display name -- + // still render it as a scale value rather than a bare, unlabelled number. if (/^\d+$/.test(v)) return `${v}/10`; return displayName(v); } diff --git a/movie_editor/frontend/rating_picker.test.js b/movie_editor/frontend/rating_picker.test.js index f880ed64..8e6c6f6e 100644 --- a/movie_editor/frontend/rating_picker.test.js +++ b/movie_editor/frontend/rating_picker.test.js @@ -17,6 +17,12 @@ test("a loved scale number keeps the heart", () => { assert.strictEqual(Picker.formatLabel("9|loved"), "9/10 ♥"); }); +test("the H3 picker's binary endpoints render as Liked/Disliked, not N/10", () => { + assert.strictEqual(Picker.formatLabel("10"), "Liked"); + assert.strictEqual(Picker.formatLabel("1"), "Disliked"); + assert.strictEqual(Picker.buttonLabel("10"), "★ Liked"); +}); + test("an existing category label is untouched", () => { assert.strictEqual(Picker.formatLabel("Perfect"), "Perfect"); }); diff --git a/movie_editor/frontend/settings_window.js b/movie_editor/frontend/settings_window.js index a72cb0c3..d88856f4 100644 --- a/movie_editor/frontend/settings_window.js +++ b/movie_editor/frontend/settings_window.js @@ -408,6 +408,26 @@ return tbl; } + function reinsToggleRow(enabled, onChange) { + const row = el("div", "sw-row"); + const main = el("div", "sw-row-main"); + main.append(el("div", "sw-row-title", "Record ratings for this key")); + main.append(el("div", "sw-row-hint", + "On by default. Independent of the sampler's own toggle — that one controls whether " + + "the learned direction gets APPLIED; this one controls whether it keeps learning at " + + "all. Turn off to pause the log without losing what's captured, or while testing " + + "something else you don't want mixed in.")); + row.append(main); + const lbl = el("label", "chk es-toggle"); + const cb = el("input"); + cb.type = "checkbox"; + cb.checked = !!enabled; + cb.onchange = () => onChange(cb.checked); + lbl.append(cb, el("span", null, "")); + row.append(lbl); + return row; + } + function reinsRows(key) { const box = el("div", "sw-stack"); const paint = () => { @@ -416,12 +436,21 @@ const nLiked = reinsState ? reinsState.n_liked : 0; const nDisliked = reinsState ? reinsState.n_disliked : 0; const ready = !!(reinsState && reinsState.ready); + const enabled = !reinsState || reinsState.enabled !== false; + rows.append(reinsToggleRow(enabled, async (on) => { + reinsError = ""; + try { reinsState = await window.MovieEditorAPI.reinsSetEnabled(key, on); } + catch (e) { reinsError = String(e.message || e); } + paint(); + })); rows.append(infoRow("Rated generations (this key)", `${nLiked} liked / ${nDisliked} disliked`, reinsState ? ready : null)); if (reinsState && !ready) { rows.append(el("div", "sw-hint", - `Needs ${reinsState.min_per_group}+ of each before it steers anything — still ` - + "capturing every generation either way.")); + enabled + ? `Needs ${reinsState.min_per_group}+ of each before it steers anything — still ` + + "capturing every generation either way." + : "Recording is off — nothing new is being captured for this key.")); } rows.append(actionRow("Save this key's REINS data", "Downloads every captured run (all candidate blocks, not just the one that steers) " diff --git a/movie_editor/server.py b/movie_editor/server.py index c6c15890..8fc37591 100644 --- a/movie_editor/server.py +++ b/movie_editor/server.py @@ -2110,6 +2110,7 @@ if web is not None and PromptServer is not None: "min_per_group": rs.MIN_PER_GROUP, "default_block": rs.DEFAULT_BLOCK, "sweep": {str(b): r for b, r in sweep.items()}, + "enabled": rs.capture_enabled(key), } @routes.get(UI_PREFIX + "/api/h3_repr_steering") @@ -2139,6 +2140,14 @@ if web is not None and PromptServer is not None: body=body, content_type="application/octet-stream", headers={"Content-Disposition": f'attachment; filename="{key}.repr_steer.pt"'}) + @routes.post(UI_PREFIX + "/api/h3_repr_steering/enabled") + async def _repr_steer_set_enabled(req): + body = await req.json() + key = str(body.get("key") or "default").strip() or "default" + _repr_steer_module().set_capture_enabled(key, bool(body.get("enabled"))) + _rs, state = _repr_steer_state(key) + return web.json_response(state) + @routes.post(UI_PREFIX + "/api/h3_repr_steering/clear") async def _repr_steer_clear(req): body = await req.json() diff --git a/tests/test_h3_repr_steering.py b/tests/test_h3_repr_steering.py index 3ac7d4cb..07a81362 100644 --- a/tests/test_h3_repr_steering.py +++ b/tests/test_h3_repr_steering.py @@ -142,6 +142,36 @@ def test_pending_overwrites_not_accumulates(): assert torch.allclose(data["rows"][0]["desc"][rs.DEFAULT_BLOCK], torch.tensor([2.0, 0.0])) +def test_capture_enabled_defaults_true(): + assert rs.capture_enabled("k") is True # no file yet -- still True, not False + + +def test_disabling_capture_stops_save_pending(): + rs.set_capture_enabled("k", False) + rs.save_pending("k", _pend(torch.tensor([1.0, 0.0]))) + rs.commit("k", 1.0) + direction, n_pos, n_neg = rs.direction("k") + assert n_pos == 0 and n_neg == 0 # nothing was ever captured to commit + + +def test_re_enabling_capture_resumes_it(): + rs.set_capture_enabled("k", False) + rs.set_capture_enabled("k", True) + assert rs.capture_enabled("k") is True + rs.save_pending("k", _pend(torch.tensor([1.0, 0.0]))) + rs.commit("k", 1.0) + data = rs._load("k") + assert len(data["rows"]) == 1 + + +def test_disabling_capture_does_not_discard_existing_rows(): + rs.save_pending("k", _pend(torch.tensor([1.0, 0.0]))) + rs.commit("k", 1.0) + rs.set_capture_enabled("k", False) + data = rs._load("k") + assert len(data["rows"]) == 1 + + def test_clear_all_removes_the_state_file(): rs.save_pending("k", _pend(torch.tensor([1.0, 0.0]))) rs.commit("k", 1.0)