From e1801fdb3ecb382c86c0ae964513ed1cd5a054dc Mon Sep 17 00:00:00 2001
From: digital-garbage <133010426+digital-garbage@users.noreply.github.com>
Date: Mon, 1 Jun 2026 18:56:06 +0300
Subject: [PATCH] Clear velocity-bias/rescue memory on Studio session reset
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

VELOCITY_BIAS_MEMORY is an in-process dict, so reset_session (which only
removed on-disk enhancement files via clear_refinement_data) left stale
good-trajectory memory behind — rescue and velocity bias would keep steering
toward the previous session's content. Add clear_velocity_bias_memory(key)
in samplers.py (drops global averages + prompt clusters for that key) and
call it from clear_refinement_data. Key matches: refine_v2 and the velocity
key default both resolve to the same effective refinement key.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 ltx_enhancements.py | 11 +++++++++++
 samplers.py         | 15 +++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/ltx_enhancements.py b/ltx_enhancements.py
index 6c51a122..4051bef2 100644
--- a/ltx_enhancements.py
+++ b/ltx_enhancements.py
@@ -217,6 +217,17 @@ def clear_refinement_data(refinement_key):
                 os.remove(path)
         except Exception as e:
             print(f"[FunPackEnhancements] Cleanup failed for {path}: {e}")
+    # Velocity-bias / rescue trajectory memory is in-process (not a file), so clear
+    # it here too — otherwise a session reset leaves stale good-trajectory memory that
+    # rescue and velocity bias would keep steering toward.
+    try:
+        try:
+            from .samplers import clear_velocity_bias_memory
+        except ImportError:
+            from samplers import clear_velocity_bias_memory
+        clear_velocity_bias_memory(refinement_key)
+    except Exception as e:
+        print(f"[FunPackEnhancements] Velocity memory cleanup failed: {e}")
 
 
 def bless_attention_maps(refinement_key):
diff --git a/samplers.py b/samplers.py
index 8746a140..1ba2f2b2 100644
--- a/samplers.py
+++ b/samplers.py
@@ -148,6 +148,21 @@ def _apply_motion_pulse(x, sigma, sigma_next, pulse_noise, noise_sampler):
     return x + noise_sampler(sigma, sigma_next) * (pulse_noise * sigma_delta)
 
 
+def clear_velocity_bias_memory(refinement_key):
+    """Drop all stored velocity-bias / rescue trajectory memory (global averages and
+    prompt clusters) for a key. Called on Studio session reset so rescue and velocity
+    bias start from a clean slate instead of steering toward last session's content."""
+    norm = str(refinement_key or "default").strip() or "default"
+    removed = [k for k in list(VELOCITY_BIAS_MEMORY)
+               if isinstance(k, tuple) and k and k[0] == norm]
+    for k in removed:
+        VELOCITY_BIAS_MEMORY.pop(k, None)
+    if removed:
+        print(f"[FunPack rescue] cleared {len(removed)} velocity-bias memory bucket(s) "
+              f"for key '{norm}' (session reset)")
+    return len(removed)
+
+
 def _velocity_bias_enabled(mode, action):
     mode = (mode or "off").lower()
     if mode == "capture_and_apply":