From ab4f78695d80e969b1d08b8f6575acc799bf83bf Mon Sep 17 00:00:00 2001
From: digital-garbage <133010426+digital-garbage@users.noreply.github.com>
Date: Sun, 14 Jun 2026 23:22:12 +0300
Subject: [PATCH] Multi-key: stop training every refinement key on every clip

Per-scene key attribution diverged from the generation split whenever the
advisor rewrote prompt_to_encode away from the raw prompt, so
resolve_scene_refinement_keys hit its all-keys UNION fallback and assigned
every key to every scene. That union got persisted into the run and read
back at rating time, training every key on every clip (value_fn bloat).

- prompt_scene_shortcut_keys now scans both expanded and original text
  (mirrors split_timeline_verbatim), so identical structure -> matching
  scene count -> precise attribution far more often.
- Divergence fallback is now default-key only (empty sets), not the
  all-keys union; a diverged run under-trains instead of cross-training.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 CHANGELOG.md                  | 11 +++++++++++
 conditioning.py               | 24 +++++++++++++++++++-----
 tests/test_refinement_keys.py |  5 +++--
 3 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2b4f1f78..4fd26a50 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,17 @@
 
 ## [Unreleased]
 
+### Fixed
+- Multi-key refinement: stop training every key on every clip. Per-scene key attribution
+  (`prompt_scene_shortcut_keys`) now scans both the expanded and original prompt — matching the
+  generation splitter `split_timeline_verbatim` — so structurally-identical prompts land on the
+  same scene count and attribution stays precise. When the count still diverges (advisor/repair
+  rewrote `prompt_to_encode` away from the raw prompt), `resolve_scene_refinement_keys` now falls
+  back to the project **default key only** (empty sets) instead of the all-keys **union**. The
+  union assigned every key to every scene, which both steered and — once persisted into the run
+  and read back at rating time — trained every key on every clip (cross-contamination + value
+  function bloat, e.g. an ~80 MB jump on a single 3-clip rating pass).
+
 ## [3.0.1] - 2026-06-14
 
 ### Added
diff --git a/conditioning.py b/conditioning.py
index 8f2a3151..b1ba8309 100644
--- a/conditioning.py
+++ b/conditioning.py
@@ -857,6 +857,7 @@ def prompt_scene_shortcut_keys(prompt):
         split_pattern = re.compile(_GENERIC_SCENE_LABEL_PATTERN, re.IGNORECASE)
 
     cuts = []
+    # (1) Scan the EXPANDED text: catches a shortcut whose expansion resolves to a transition.
     for m in split_pattern.finditer(expanded):
         info = custom_map.get(re.sub(r"\s+", " ", m.group(0).strip().lower())) or {}
         if info.get("placement") == "start":
@@ -864,6 +865,14 @@ def prompt_scene_shortcut_keys(prompt):
         else:
             orig = _project_offset(pieces, m.end(), "after")
         cuts.append(orig)
+    # (2) Scan the ORIGINAL text too: a transition trigger that is ALSO a shortcut gets expanded
+    #     away in (1) before detection, so its split would be lost. Matching the verbatim text
+    #     catches it at its true offset. This mirrors split_timeline_verbatim so this attribution
+    #     split lands on the same scene count as the generation split (fewer divergence fallbacks).
+    for m in split_pattern.finditer(text):
+        info = custom_map.get(re.sub(r"\s+", " ", m.group(0).strip().lower())) or {}
+        orig = m.start() if info.get("placement") == "start" else m.end()
+        cuts.append(orig)
     seen, uniq = set(), []
     for off in sorted(cuts):
         if off <= 0 or off in seen:
@@ -891,10 +900,12 @@ def resolve_scene_refinement_keys(raw_positive, scene_count):
     """Per-scene non-default refinement keys, aligned to split_scene_texts.
 
     A key participates in a scene when one of its bound shortcuts fired in that scene's
-    text (anchor keys count for every scene). Falls back to the union across all scenes
-    when attribution can't be trusted — the raw-prompt scene count diverges from the
-    actual split (advisor/repair rewrote the prompt, unusual transition stacking, etc.).
-    Returns a list of sets aligned to scenes; empty sets => default key (today's path).
+    text (anchor keys count for every scene). Falls back to the project DEFAULT key only
+    (empty sets) when attribution can't be trusted — the raw-prompt scene count diverges
+    from the actual split (advisor/repair rewrote the prompt, unusual transition stacking,
+    etc.). The old behaviour here was the all-keys UNION, which trained EVERY key on EVERY
+    scene (cross-contamination + value-function bloat); under-training one diverged run is
+    far safer. Returns a list of sets aligned to scenes; empty sets => default key.
 
     This is the single source of truth shared by generation (FunPackStudio._v2_scene_refinement_keys)
     and the Movie Editor preview, so what the editor shows is exactly what will run."""
@@ -908,7 +919,10 @@ def resolve_scene_refinement_keys(raw_positive, scene_count):
         return [set() for _ in range(scene_count)]
     if n == scene_count:
         return scene_sets
-    return [set(all_keys) for _ in range(scene_count)]
+    # Divergence: attribution cannot be trusted. Steer/train with the project default key only
+    # (empty sets) instead of the all-keys union, so a diverged run never trains every key on
+    # every scene. all_keys stays available to callers that need the full set (e.g. session reset).
+    return [set() for _ in range(scene_count)]
 
 
 def parse_timeline_segments(prompt):
diff --git a/tests/test_refinement_keys.py b/tests/test_refinement_keys.py
index 58670f6a..debb0138 100644
--- a/tests/test_refinement_keys.py
+++ b/tests/test_refinement_keys.py
@@ -176,9 +176,10 @@ def test_studio_scene_refinement_keys_fallback(monkeypatch):
     raw = "wide shot scene 1 alpha here scene 2 beta here"  # attribution finds 2 scenes
     # scene_count matches -> precise per-scene
     assert studio._v2_scene_refinement_keys(raw, 2) == [{"keyA"}, {"keyB"}]
-    # scene_count diverges (advisor rewrote / stacking) -> safe union for every scene
+    # scene_count diverges (advisor rewrote / stacking) -> default key only (empty sets), NOT the
+    # all-keys union; a union here would train every key on every scene (cross-contamination).
     out = studio._v2_scene_refinement_keys(raw, 3)
-    assert out == [{"keyA", "keyB"}, {"keyA", "keyB"}, {"keyA", "keyB"}]
+    assert out == [set(), set(), set()]
 
 
 def test_studio_scene_refinement_keys_no_bindings(monkeypatch):