Upload folder using huggingface_hub

predict.py

@@ -135,13 +135,10 @@ 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]
 
-    # Compute per-pixel variance across context frames for motion mask
-    # frames_t: [8, 3, 64, 64]
-    pixel_var = np.var(frames_t, axis=0)  # [3, 64, 64]
-    pixel_var_mean = pixel_var.mean(axis=0)  # [64, 64] - average across channels
-    # Static mask: 1.0 for static pixels (low variance), 0.0 for dynamic
-    var_thresh = 5.0 / (255.0 * 255.0)  # variance in [0,1] scale (5/255^2)
-    static_mask = (pixel_var_mean < var_thresh).astype(np.float32)  # [64, 64]
+    # 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
@@ -181,10 +178,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []
         for i in range(PRED_FRAMES):
-            frame = np.transpose(predicted_np[i], (1, 2, 0))  # [64, 64, 3]
-            # Apply motion mask: static pixels blend 80% context / 20% prediction
-            mask_3d = static_mask[:, :, np.newaxis]  # [64, 64, 1]
-            frame = frame * (1.0 - 0.8 * mask_3d) + last_frame * 0.8 * mask_3d
+            frame = np.transpose(predicted_np[i], (1, 2, 0))
             frame = (frame * 255).clip(0, 255).astype(np.uint8)
             ens.direct_cache.append(frame)
 
@@ -250,10 +244,9 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []
         for i in range(PRED_FRAMES):
-            frame = np.transpose(predicted_np[i], (1, 2, 0))  # [64, 64, 3]
-            # Apply motion mask: static pixels blend 80% context / 20% prediction
-            mask_3d = static_mask[:, :, np.newaxis]  # [64, 64, 1]
-            frame = frame * (1.0 - 0.8 * mask_3d) + last_frame * 0.8 * mask_3d
+            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)
 
@@ -285,10 +278,9 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []
         for i in range(PRED_FRAMES):
-            frame = np.transpose(predicted_np[i], (1, 2, 0))  # [64, 64, 3]
-            # Apply motion mask: static pixels blend 80% context / 20% prediction
-            mask_3d = static_mask[:, :, np.newaxis]  # [64, 64, 1]
-            frame = frame * (1.0 - 0.8 * mask_3d) + last_frame * 0.8 * mask_3d
+            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)