Upload folder using huggingface_hub

predict.py

@@ -135,11 +135,6 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
     last_frame = frames_norm[-1]
     last_frame_t = np.transpose(last_frame, (2, 0, 1))[np.newaxis]
 
-    # Static background mask (for Sonic/PP only)
-    pixel_var = np.var(frames_t, axis=0).mean(axis=0)  # [64, 64]
-    var_thresh = 2.0 / (255.0 * 255.0)
-    static_mask = (pixel_var < var_thresh).astype(np.float32)[:, :, np.newaxis]  # [64, 64, 1]
-
     if game == "pong":
         # Pong: AR+direct ensemble, float32 caching, no TTA
         if ens.direct_cache is not None and n > CONTEXT_FRAMES and ens.cache_step < PRED_FRAMES:
@@ -169,9 +164,11 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
 
             ar_pred = torch.stack(ar_preds, dim=1)
 
+            # U-shaped AR weights: less AR early/late, more in middle
+            pong_ar_weights = [0.70, 0.70, 0.85, 0.85, 0.85, 0.60, 0.60, 0.60]
             predicted = torch.zeros_like(direct_pred)
             for step in range(PRED_FRAMES):
-                ar_weight = 0.85 - (step / (PRED_FRAMES - 1)) * 0.3
+                ar_weight = pong_ar_weights[step]
                 direct_weight = 1.0 - ar_weight
                 predicted[:, step] = ar_weight * ar_pred[:, step] + direct_weight * direct_pred[:, step]
 
@@ -245,8 +242,6 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         ens.direct_cache = []
         for i in range(PRED_FRAMES):
             frame = np.transpose(predicted_np[i], (1, 2, 0))
-            # Mild static mask: 10% context blend on static pixels
-            frame = frame * (1.0 - 0.1 * static_mask) + last_frame * 0.1 * static_mask
             frame = (frame * 255).clip(0, 255).astype(np.uint8)
             ens.direct_cache.append(frame)
 
@@ -279,8 +274,6 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         ens.direct_cache = []
         for i in range(PRED_FRAMES):
             frame = np.transpose(predicted_np[i], (1, 2, 0))
-            # Mild static mask: 10% context blend on static pixels
-            frame = frame * (1.0 - 0.1 * static_mask) + last_frame * 0.1 * static_mask
             frame = (frame * 255).clip(0, 255).astype(np.uint8)
             ens.direct_cache.append(frame)