refresh context

app.py

@@ -131,6 +131,8 @@ def generate_loop_continuation_with_mrt(
 
     return out, loud_stats
 
+
+
 # ----------------------------
 # FastAPI app with lazy, thread-safe model init
 # ----------------------------
@@ -298,10 +300,18 @@ def jam_start(
     target_sr = int(target_sample_rate or input_sr)
 
     params = JamParams(
-        bpm=bpm, beats_per_bar=beats_per_bar, bars_per_chunk=bars_per_chunk,
-        target_sr=target_sr, loudness_mode=loudness_mode, headroom_db=loudness_headroom_db,
-        style_vec=style_vec, ref_loop=loop_tail,
-        guidance_weight=guidance_weight, temperature=temperature, topk=topk
+        bpm=bpm, 
+        beats_per_bar=beats_per_bar, 
+        bars_per_chunk=bars_per_chunk,
+        target_sr=target_sr, 
+        loudness_mode=loudness_mode, 
+        headroom_db=loudness_headroom_db,
+        style_vec=style_vec, 
+        ref_loop=loop_tail,                    # For loudness matching
+        combined_loop=loop,                    # NEW: Full loop for context setup
+        guidance_weight=guidance_weight, 
+        temperature=temperature, 
+        topk=topk
     )
 
     worker = JamWorker(mrt, params)

jam_worker.py

@@ -17,12 +17,13 @@ from math import gcd
 class JamParams:
     bpm: float
     beats_per_bar: int
-    bars_per_chunk: int           # 4 or 8
+    bars_per_chunk: int
     target_sr: int
     loudness_mode: str = "auto"
     headroom_db: float = 1.0
-    style_vec: np.ndarray | None = None   # combined_style vector
+    style_vec: np.ndarray | None = None
     ref_loop: any = None                  # au.Waveform at model SR for 1st-chunk loudness
+    combined_loop: any = None             # NEW: Full combined audio for context setup
     guidance_weight: float = 1.1
     temperature: float = 1.1
     topk: int = 40
@@ -38,7 +39,12 @@ class JamWorker(threading.Thread):
         super().__init__(daemon=True)
         self.mrt = mrt
         self.params = params
+        # Initialize fresh state
         self.state = mrt.init_state()
+        
+        # CRITICAL: Set up fresh context from the new combined audio
+        if params.combined_loop is not None:
+            self._setup_context_from_combined_loop()
         self.idx = 0
         self.outbox: list[JamChunk] = []
         self._stop_event = threading.Event()
@@ -46,6 +52,47 @@ class JamWorker(threading.Thread):
         self.last_chunk_completed_at = None
         self._lock = threading.Lock()
 
+    def _setup_context_from_combined_loop(self):
+        """Set up MRT context tokens from the combined loop audio"""
+        try:
+            # Import the utility functions (same as used in main generation)
+            from utils import make_bar_aligned_context, take_bar_aligned_tail
+            
+            # Extract context from combined loop (same logic as generate_loop_continuation_with_mrt)
+            codec_fps = float(self.mrt.codec.frame_rate)
+            ctx_seconds = float(self.mrt.config.context_length_frames) / codec_fps
+            
+            # Take tail portion for context (matches main generation)
+            loop_for_context = take_bar_aligned_tail(
+                self.params.combined_loop, 
+                self.params.bpm, 
+                self.params.beats_per_bar, 
+                ctx_seconds
+            )
+            
+            # Encode to tokens
+            tokens_full = self.mrt.codec.encode(loop_for_context).astype(np.int32)
+            tokens = tokens_full[:, :self.mrt.config.decoder_codec_rvq_depth]
+            
+            # Create bar-aligned context
+            context_tokens = make_bar_aligned_context(
+                tokens, 
+                bpm=self.params.bpm, 
+                fps=int(self.mrt.codec.frame_rate),
+                ctx_frames=self.mrt.config.context_length_frames, 
+                beats_per_bar=self.params.beats_per_bar
+            )
+            
+            # Set context on state - this is the key fix!
+            self.state.context_tokens = context_tokens
+            
+            print(f"✅ JamWorker: Set up fresh context from combined loop")
+            print(f"   Context shape: {context_tokens.shape if context_tokens is not None else None}")
+            
+        except Exception as e:
+            print(f"❌ Failed to setup context from combined loop: {e}")
+            # Continue without context rather than crashing
+
     def stop(self):
         self._stop_event.set()
 
@@ -137,3 +184,5 @@ class JamWorker(threading.Thread):
                 self.last_chunk_completed_at = time.time()
 
         # optional: cleanup here if needed
+
+