Upload folder using huggingface_hub

predict.py

@@ -164,19 +164,23 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
 
             ar_pred = torch.stack(ar_preds, dim=1)
 
-            # Pure AR for steps 1-4, then blend for steps 5-8
-            pong_ar_weights = [1.0, 1.0, 1.0, 1.0, 0.70, 0.65, 0.60, 0.55]
             predicted = torch.zeros_like(direct_pred)
             for step in range(PRED_FRAMES):
-                ar_weight = pong_ar_weights[step]
+                ar_weight = 0.85 - (step / (PRED_FRAMES - 1)) * 0.3
                 direct_weight = 1.0 - ar_weight
                 predicted[:, step] = ar_weight * ar_pred[:, step] + direct_weight * direct_pred[:, step]
 
         predicted_np = predicted[0].cpu().numpy()
+        last_ctx_uint8 = (last_frame * 255).clip(0, 255).astype(np.uint8)  # [64,64,3]
         ens.direct_cache = []
         for i in range(PRED_FRAMES):
             frame = np.transpose(predicted_np[i], (1, 2, 0))
             frame = (frame * 255).clip(0, 255).astype(np.uint8)
+            # Fallback: if prediction is very different from context, blend with context
+            diff = np.abs(frame.astype(np.float32) - last_ctx_uint8.astype(np.float32))
+            mean_diff = diff.mean()
+            if mean_diff > 30:  # very different prediction
+                frame = ((0.5 * frame.astype(np.float32) + 0.5 * last_ctx_uint8.astype(np.float32))).clip(0, 255).astype(np.uint8)
             ens.direct_cache.append(frame)
 
         result = ens.direct_cache[ens.cache_step]