diff --git a/h3_repr_steering.py b/h3_repr_steering.py index b4595854..80dc960c 100644 --- a/h3_repr_steering.py +++ b/h3_repr_steering.py @@ -123,15 +123,16 @@ 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, "slots": {}} try: data = torch.load(path, map_location="cpu", weights_only=True) except Exception: # noqa: BLE001 - return {"rows": [], "pending": None} + return {"rows": [], "pending": None, "slots": {}} if not isinstance(data, dict): - return {"rows": [], "pending": None} + return {"rows": [], "pending": None, "slots": {}} data.setdefault("rows", []) data.setdefault("pending", None) + data.setdefault("slots", {}) return data @@ -198,6 +199,68 @@ def commit(refinement_key, reward): return "recorded" +def save_capture_slot(refinement_key, slot_id, descriptors): + """Like save_pending, but keyed by an explicit slot instead of the single overwritable + `pending` -- for a batch of generations that each need to stay independently rateable, in + any order, any subset, instead of each new run silently evicting the last one waiting to + be rated. `pending` is untouched by this (a slotted run does not also become "the" + pending capture), so the ordinary one-at-a-time rating flow keeps working exactly as + before for anything that doesn't pass a slot id.""" + if not refinement_key or not slot_id or not descriptors: + return + data = _load(refinement_key) + data["slots"][str(slot_id)] = {int(b): d.detach().float().cpu() for b, d in descriptors.items()} + try: + _save(refinement_key, data) + except OSError as e: + _log().failed("H3 representation steering", "save capture slot", e, + f"slot {slot_id}'s capture is not kept") + + +def commit_slot(refinement_key, slot_id, reward): + """Slot equivalent of commit(): pairs ONE named slot's capture with a rating and appends + it to the same rows log direction() reads, then frees that slot. Slots are independent of + each other and of `pending` -- rating one does not consume or disturb any other slot, so + a whole batch can be rated in any order, any subset, whenever -- the actual point of + a batch. Same return contract as commit(): "recorded" | "no_pending" | "no_key".""" + if not refinement_key: + return "no_key" + data = _load(refinement_key) + captured = data["slots"].pop(str(slot_id), None) + if captured is None: + try: + _save(refinement_key, data) + except OSError: + pass + return "no_pending" + try: + data["rows"].append({"desc": captured, "weight": float(reward)}) + except (TypeError, ValueError): + try: + _save(refinement_key, data) + except OSError: + pass + return "no_pending" + try: + _save(refinement_key, data) + except OSError as e: + _log().failed("H3 representation steering", "commit slot", e, "rating not recorded") + return "no_pending" + return "recorded" + + +def discard_slot(refinement_key, slot_id): + """Drop a slot's capture without rating it (Skip) -- costs nothing, frees the state.""" + if not refinement_key: + return + data = _load(refinement_key) + if data["slots"].pop(str(slot_id), None) is not None: + try: + _save(refinement_key, data) + except OSError: + pass + + def clear_all(refinement_key): if not refinement_key: return diff --git a/movie_editor/frontend/api.js b/movie_editor/frontend/api.js index dc3438eb..a8e4352a 100644 --- a/movie_editor/frontend/api.js +++ b/movie_editor/frontend/api.js @@ -203,6 +203,10 @@ reinsStatus: (key) => j("GET", API("/h3_repr_steering") + `?key=${encodeURIComponent(key || "default")}`), reinsSweep: (key, trials) => j("POST", API("/h3_repr_steering/sweep"), { key: key || "default", trials: trials || 2000 }), reinsClear: (key) => j("POST", API("/h3_repr_steering/clear"), { key: key || "default" }), + // Batch/sweep results save under their OWN slot (h3_repr_capture_slot), not the single + // overwritable "pending" — so rating one is a direct commit, no extra generation needed. + reinsRateSlot: (key, slot, rating) => j("POST", API("/h3_repr_steering/rate-slot"), { key: key || "default", slot, rating }), + reinsDiscardSlot: (key, slot) => j("POST", API("/h3_repr_steering/discard-slot"), { key: key || "default", slot }), async reinsExport(key) { const k = key || "default"; const res = await fetch(API("/h3_repr_steering/export") + `?key=${encodeURIComponent(k)}&t=${Date.now()}`, { cache: "no-store" }); diff --git a/movie_editor/frontend/block_sweep.js b/movie_editor/frontend/block_sweep.js index c72855bd..61f1d63f 100644 --- a/movie_editor/frontend/block_sweep.js +++ b/movie_editor/frontend/block_sweep.js @@ -39,7 +39,7 @@ // costs nothing and commits nothing, it just stops offering to rate that card. const _skipped = new WeakSet(); - function _card(item, isLast, onRated) { + function _card(item, onRated) { // A real viewing size, not the media-bin's thumbnail height — the whole point of this // gallery is comparing results without popping each one into a separate tab. const card = el("div", "sw-sweep-card"); @@ -77,49 +77,46 @@ }; card.append(saveBtn); - // Rating only ever pairs with the run whose capture is still "pending" on the - // refinement key — that slot holds ONE entry, overwritten the instant the next run - // samples. Only the LAST card in the gallery is ever validly rateable; anything - // earlier has already been evicted by a later run's own capture. Liking/disliking - // fires ANOTHER generation (same config, rating attached) — it is not free; Skip - // fires nothing, it just stops the row from asking. + // Each result with a REINS half saves its capture under its OWN slot (server-side, + // h3_repr_capture_slot) — independent of every other result's, so EVERY card here is + // rateable, in any order, any subset, same as the real batch training this panel + // replaced. Liking/disliking/skipping is a direct commit against data already on + // disk — no extra generation, unlike the old single-"pending" design this used to have. const outcome = _rated.get(item); if (outcome) { card.append(el("div", "es-hint sw-sweep-rate-note", outcome === "liked" ? "✓ Liked" : "✓ Disliked")); } else if (_skipped.has(item)) { card.append(el("div", "es-hint sw-sweep-rate-note", "Skipped")); - } else if (isLast) { + } else if (item.slotId) { const rateRow = el("div", "sw-sweep-rate"); const mkBtn = (cls, txt, title) => { const b = el("button", cls, txt); b.title = title; return b; }; - const likeBtn = mkBtn("btn ghost tiny", "👍 Like", - "Fires one more generation with a liked rating attached — the only way to commit " - + "a rating against this result's still-pending capture."); - const dislikeBtn = mkBtn("btn ghost tiny", "👎 Dislike", - "Fires one more generation with a disliked rating attached."); - const skipBtn = mkBtn("btn ghost tiny", "Skip", "Don't rate this one — costs nothing."); + const likeBtn = mkBtn("btn ghost tiny", "👍 Like", "Commit a liked rating for this result."); + const dislikeBtn = mkBtn("btn ghost tiny", "👎 Dislike", "Commit a disliked rating for this result."); + const skipBtn = mkBtn("btn ghost tiny", "Skip", "Don't rate this one — discards its capture."); + const setBusy = (busy) => { [likeBtn, dislikeBtn, skipBtn].forEach((b) => { b.disabled = busy; }); }; const rate = async (label, outcomeLabel, btn) => { - [likeBtn, dislikeBtn, skipBtn].forEach((b) => { b.disabled = true; }); - btn.textContent = "…"; + setBusy(true); btn.textContent = "…"; try { const ok = await window.Store.rateComboResult(item, label); if (ok) { _rated.set(item, outcomeLabel); onRated && onRated(); } - else { alert("Rating run failed — see the status line."); } + else { alert("Rating failed — the capture may have already been rated or discarded."); } } finally { - [likeBtn, dislikeBtn, skipBtn].forEach((b) => { b.disabled = false; }); + setBusy(false); btn.textContent = btn === likeBtn ? "👍 Like" : "👎 Dislike"; } }; likeBtn.onclick = () => rate(_LIKE_RATING, "liked", likeBtn); dislikeBtn.onclick = () => rate(_DISLIKE_RATING, "disliked", dislikeBtn); - skipBtn.onclick = () => { _skipped.add(item); onRated && onRated(); }; + skipBtn.onclick = async () => { + setBusy(true); + await window.Store.discardComboResult(item); + _skipped.add(item); onRated && onRated(); + }; rateRow.append(likeBtn, dislikeBtn, skipBtn); card.append(rateRow); - } else { - card.append(el("div", "es-hint sw-sweep-rate-note", - "Superseded by a later run — no longer the pending capture, so it can't be rated.")); } } return card; @@ -194,7 +191,7 @@ : (cs.results && cs.results.length ? `${cs.results.length} result(s), ${seedNote}` : ""); clear(grid); const results = cs.results || []; - results.forEach((item, i) => grid.append(_card(item, i === results.length - 1 && !cs.running, render))); + results.forEach((item) => grid.append(_card(item, render))); } runBtn.onclick = () => { diff --git a/movie_editor/frontend/store.js b/movie_editor/frontend/store.js index 522134cf..4914fa7a 100644 --- a/movie_editor/frontend/store.js +++ b/movie_editor/frontend/store.js @@ -3729,9 +3729,19 @@ }).filter(Boolean); } + // The refinement key REINS reads/writes for — same default the rest of the app uses + // (project.refinement_key), so a sweep's slots land in the same sidecar file everything + // else on this project reads from. + function _reinsKey() { + return (state.project?.refinement_key || "default").trim() || "default"; + } + // Overrides for one combo config. `seedValue` is the seed to force for this run (the - // caller picks a fresh one per line unless "same seed" is checked). - function _comboOverrides(c, seedValue) { + // caller picks a fresh one per line unless "same seed" is checked). `slotId`, when the + // config has a `reins` half, saves THIS run's capture under its own slot instead of the + // single overwritable "pending" — so every result in the gallery stays independently + // rateable later, not just whichever one happened to run last. + function _comboOverrides(c, seedValue, slotId) { const ov = [ { node: "sampler", input: "seed", value: seedValue }, { node: "sampler", input: "h3_repr_steering", value: !!c.reins }, @@ -3741,6 +3751,7 @@ if (c.reins) { ov.push({ node: "sampler", input: "h3_repr_steering_strength", value: c.reins.strength }); ov.push({ node: "sampler", input: "h3_repr_steering_block", value: c.reins.block }); + ov.push({ node: "sampler", input: "h3_repr_capture_slot", value: slotId }); } if (c.sweep) { ov.push({ node: "sampler", input: "h3_block_repeat_span_loop", value: c.sweep.spanLoop }); @@ -3764,35 +3775,47 @@ const fixedSeed = sameSeed ? (state.project.sampler_inputs?.seed ?? Math.floor(Math.random() * 0x7fffffff)) : null; + const reinsKey = _reinsKey(); const results = []; set({ comboSweep: { running: true, results, error: null, total: configs.length, current: 0, sameSeed: !!sameSeed, seed: fixedSeed } }); for (let i = 0; i < configs.length; i++) { const c = configs[i]; const seed = sameSeed ? fixedSeed : Math.floor(Math.random() * 0x7fffffff); + // Unique per run, not just per config — rerunning the same line must not reuse (and + // silently clobber) an earlier, still-unrated slot. + const slotId = c.reins ? `sweep_${Date.now()}_${i}_${Math.random().toString(36).slice(2, 8)}` : null; set({ comboSweep: { running: true, results: [...results], error: null, total: configs.length, current: i + 1, label: c.label, sameSeed: !!sameSeed, seed } }); - const overrides = _comboOverrides(c, seed); + const overrides = _comboOverrides(c, seed, slotId); const ok = await _generateRun([sceneId], sceneId, `Sweep ${i + 1}/${configs.length}: ${c.label}`, false, overrides); if (!ok) { set({ comboSweep: { running: false, results, error: `Stopped at "${c.label}": ${state.gen.msg}`, total: configs.length, sameSeed: !!sameSeed, seed } }); return; } const media = state.sceneRenders[sceneId]?.media; - results.push({ label: c.label, media, sceneId, overrides, seed }); + results.push({ label: c.label, media, sceneId, seed, reinsKey, slotId }); } set({ comboSweep: { running: false, results, error: null, total: configs.length, sameSeed: !!sameSeed, seed: fixedSeed } }); } - // Rate the MOST RECENT combo-sweep result. Ratings only pair with whatever generation's - // capture is still "pending" on the refinement key — that slot holds exactly one entry, - // overwritten by the next run that samples, so only the LAST result in the gallery can - // ever be validly rated; anything earlier has already been evicted by a later run. This - // fires one MORE generation, reusing that same result's exact config, with the rating - // attached — rating is not a free/standalone action, it costs another full render, the - // same way it does everywhere else in this codebase (rating rides the NEXT generation). + // Rate (or discard) ANY combo-sweep result, in any order, any subset — each one's REINS + // capture lives in its own slot (see runComboSweep), independent of every other result's, + // so this is a direct commit against data already on disk. No extra generation, unlike the + // ordinary one-at-a-time Studio rating flow this deliberately does NOT reuse. async function rateComboResult(item, ratingLabel) { - if (!state.project || !item || !ratingLabel) return false; - const overrides = [...item.overrides, { node: "studio", input: "rating", value: ratingLabel }]; - return _generateRun([item.sceneId], item.sceneId, `Rating "${item.label}"`, false, overrides); + if (!item || !item.slotId) return false; + try { + const r = await API.reinsRateSlot(item.reinsKey, item.slotId, ratingLabel); + return r.outcome === "recorded"; + } catch (e) { + console.error("rate-slot failed", e); + return false; + } + } + + async function discardComboResult(item) { + if (!item || !item.slotId) return; + try { await API.reinsDiscardSlot(item.reinsKey, item.slotId); } + catch (e) { console.error("discard-slot failed", e); } } async function generate(onlyScene) { @@ -4630,7 +4653,7 @@ resizeScene, setSceneGapAfter, splitScene, autoMontage, hasPlayableRender, snapFrames, snapFramesFloor, snapFramesCeil, sceneEffFrames, sceneEffFps, scenesOverridingProjectFrames, useProjectFramesEverywhere, setSourceTrim, trimSceneLeft, slipScene, applyEnginePreset, ENGINE_PRESETS, undo, redo, refreshPreview, syncFromPreview, applyGlobalPromptQuiet, scheduleGlobalPromptApply, globalPromptApplyPending, buildGlobalPromptFromTimeline, syncGlobalPromptFromTimeline, generate: _clockedGenerate, generateMontage: _clockedMontage, generateSelected: _clockedSelected, genElapsed, selectedSceneCount, renderFinal, exportSelected, saveSelectedToMediaBin, saveClipToMediaBin, clipSaveableToMediaBin, - runComboSweep, parseComboSweepConfig, rateComboResult, interrupt, loadModels, loadImageTargets, setModelInput, setModelBypass, setModelLink, clearNotice, + runComboSweep, parseComboSweepConfig, rateComboResult, discardComboResult, interrupt, loadModels, loadImageTargets, setModelInput, setModelBypass, setModelLink, clearNotice, projectVariables, setProjectVariables, promptTemplates, savePromptTemplate, deletePromptTemplate, applyPromptTemplate, renamePromptTemplate, clearGlobalPrompt, setConditioningSlot, setSamplerSlot, setSamplerInput, setSamplerInputNow, unsetSamplerInput, setStudioInput, setStudioInputNow, diff --git a/movie_editor/server.py b/movie_editor/server.py index dcf5cd0b..ec6fb971 100644 --- a/movie_editor/server.py +++ b/movie_editor/server.py @@ -2171,6 +2171,36 @@ if web is not None and PromptServer is not None: _rs, state = _repr_steer_state(key) return web.json_response(state) + # Batch/sweep results each save their capture under their OWN slot (h3_repr_capture_slot + # on the sampler) instead of the single overwritable "pending" -- so, unlike the ordinary + # one-at-a-time rating flow, rating one here does NOT require another generation. It's a + # pure commit against a capture that's already on disk, same as batch_training.py's + # existing per-item ratings, just for REINS's own sidecar instead of a value function. + @routes.post(UI_PREFIX + "/api/h3_repr_steering/rate-slot") + async def _repr_steer_rate_slot(req): + body = await req.json() + key = str(body.get("key") or "").strip() + slot_id = str(body.get("slot") or "").strip() + if not key or not slot_id: + return web.json_response({"error": "key and slot are required"}, status=400) + try: + from conditioning import normalize_refiner_v2_rating + except ImportError: + from .. import conditioning # type: ignore + normalize_refiner_v2_rating = conditioning.normalize_refiner_v2_rating + profile = normalize_refiner_v2_rating(str(body.get("rating") or "")) + outcome = _repr_steer_module().commit_slot(key, slot_id, float(profile.get("reward", 0.0))) + return web.json_response({"outcome": outcome}) + + @routes.post(UI_PREFIX + "/api/h3_repr_steering/discard-slot") + async def _repr_steer_discard_slot(req): + body = await req.json() + key = str(body.get("key") or "").strip() + slot_id = str(body.get("slot") or "").strip() + if key and slot_id: + _repr_steer_module().discard_slot(key, slot_id) + return web.json_response({"ok": True}) + # ── block-influence probe ──────────────────────────────────────────────── # Research data, so it lives here rather than in the console: readable, downloadable and # clearable from Settings > Refinement & Taste, with its own collection switch. Same diff --git a/samplers.py b/samplers.py index ec7ba07f..cc8f5876 100644 --- a/samplers.py +++ b/samplers.py @@ -2970,6 +2970,10 @@ class FunPackLTXAVSceneChainSampler: "default": "25", "tooltip": "TEST-ONLY: which block(s) actually get steered (the rest of the sweep's candidate blocks are still captured read-only, same as always). A single block, a range, or a comma list -- 25 | 31-40 | 4,5,6 -- same syntax as h3_block_repeat. Each named block steers with its OWN learned direction, computed independently from its OWN liked/disliked history at that block; this does not change what block_sweep ranks, only which block(s) get injected into this run.", }), + "h3_repr_capture_slot": ("STRING", { + "default": "", + "tooltip": "TEST-ONLY, node_overrides use. Blank (default) = the normal single 'pending' capture, overwritten by whatever generates next -- the ordinary one-at-a-time rating flow. A non-blank id saves this run's REINS capture under that SLOT instead, independent of pending and of every other slot, so a whole batch of runs can each be rated later, in any order, any subset -- not just the most recent one.", + }), "h3_block_repeat": ("STRING", { "default": "", "tooltip": "EXPERIMENTAL, unvalidated (H3). Runs the named block(s) on their OWN output before passing the result on -- a second pass that never leaves latent space (no VAE round trip, no re-noising, no extra sampler step). Because a block's contribution is small, running it twice roughly DOUBLES what that block does, with the nonlinear part being a genuine second refinement of its own result. Blank = off. Accepts a block, a range, or a list: 40 | 38-42 | 10,40,44. Costs one extra block forward per repeat (~2% of a step each). The blocks AFTER a repeated one were trained on its normal output, so a twice-processed stream is mildly out of distribution for them -- expect this to be visible, and not always as an improvement.", @@ -3698,7 +3702,7 @@ class FunPackLTXAVSceneChainSampler: alg_anchor_sigma_threshold=0.975, bounded_attention_enabled=False, h3_repr_steering=False, h3_repr_steering_strength=0.05, - h3_repr_steering_block="25", + h3_repr_steering_block="25", h3_repr_capture_slot="", h3_block_repeat="", h3_block_repeat_times=1, h3_block_repeat_video_only=False, h3_block_repeat_span_loop=False, h3_block_repeat_last_steps=0, @@ -3803,7 +3807,10 @@ class FunPackLTXAVSceneChainSampler: from . import h3_repr_steering as _rs except ImportError: import h3_repr_steering as _rs - _rs.save_pending(refinement_key, _repr_capture[0]) + if h3_repr_capture_slot: + _rs.save_capture_slot(refinement_key, h3_repr_capture_slot, _repr_capture[0]) + else: + _rs.save_pending(refinement_key, _repr_capture[0]) if _influence_capture[0]: try: from . import block_influence as _bi @@ -7454,7 +7461,7 @@ class FunPackLTXAVSceneChainSampler: trajectory_guidance=False, trajectory_guidance_strength=0.02, dynashift=False, dynashift_strength=0.3, dynashift_threshold=0.6, h3_repr_steering=False, h3_repr_steering_strength=0.05, - h3_repr_steering_block="25", + h3_repr_steering_block="25", h3_repr_capture_slot="", h3_block_repeat="", h3_block_repeat_times=1, h3_block_repeat_video_only=False, h3_block_repeat_span_loop=False, h3_block_repeat_last_steps=0, @@ -8352,6 +8359,7 @@ class FunPackLTXAVSceneChainSampler: h3_repr_steering=h3_repr_steering, h3_repr_steering_strength=h3_repr_steering_strength, h3_repr_steering_block=h3_repr_steering_block, + h3_repr_capture_slot=h3_repr_capture_slot, h3_block_repeat=h3_block_repeat, h3_block_repeat_times=h3_block_repeat_times, h3_block_repeat_video_only=h3_block_repeat_video_only, diff --git a/tests/test_h3_repr_steering.py b/tests/test_h3_repr_steering.py index 1bcc78e9..07b5c2c7 100644 --- a/tests/test_h3_repr_steering.py +++ b/tests/test_h3_repr_steering.py @@ -160,6 +160,49 @@ def test_pending_overwrites_not_accumulates(): assert torch.allclose(data["rows"][0]["desc"][rs.DEFAULT_BLOCK], torch.tensor([2.0, 0.0])) +# --- capture slots -- independent, out-of-order rating for a batch/sweep -------- + +def test_slots_are_independent_of_pending(): + """A slotted capture must not become 'the' pending capture, and vice versa -- the + ordinary one-at-a-time Studio rating flow keeps working exactly as before regardless of + what a sweep is doing with slots.""" + rs.save_pending("k", _pend(torch.tensor([1.0, 0.0]))) + rs.save_capture_slot("k", "s1", _pend(torch.tensor([9.0, 0.0]))) + assert rs.commit("k", 1.0) == "recorded" # still pairs with the ORIGINAL pending + data = rs._load("k") + assert torch.allclose(data["rows"][0]["desc"][rs.DEFAULT_BLOCK], torch.tensor([1.0, 0.0])) + assert "s1" in data["slots"] # untouched by the pending commit + + +def test_two_slots_are_rateable_independently_in_any_order(): + """The actual point of a batch: every result stays rateable, not just the most recent.""" + rs.save_capture_slot("k", "a", _pend(torch.tensor([1.0, 0.0]))) + rs.save_capture_slot("k", "b", _pend(torch.tensor([2.0, 0.0]))) + rs.save_capture_slot("k", "c", _pend(torch.tensor([3.0, 0.0]))) + # Rate them out of generation order -- c, then a. b is left alone. + assert rs.commit_slot("k", "c", 1.0) == "recorded" + assert rs.commit_slot("k", "a", -1.0) == "recorded" + data = rs._load("k") + assert len(data["rows"]) == 2 + assert "b" in data["slots"] # never rated, still sitting there + assert "a" not in data["slots"] and "c" not in data["slots"] # consumed on commit + + +def test_committing_an_unknown_or_already_rated_slot_reports_no_pending(): + rs.save_capture_slot("k", "a", _pend(torch.tensor([1.0, 0.0]))) + assert rs.commit_slot("k", "a", 1.0) == "recorded" + assert rs.commit_slot("k", "a", 1.0) == "no_pending" # already consumed + assert rs.commit_slot("k", "does-not-exist", 1.0) == "no_pending" + + +def test_discard_slot_drops_it_without_recording_a_row(): + rs.save_capture_slot("k", "a", _pend(torch.tensor([1.0, 0.0]))) + rs.discard_slot("k", "a") + data = rs._load("k") + assert "a" not in data["slots"] + assert len(data["rows"]) == 0 + + def test_clear_all_removes_the_state_file(): rs.save_pending("k", _pend(torch.tensor([1.0, 0.0]))) rs.commit("k", 1.0)