Upload folder using huggingface_hub

predict.py

@@ -164,11 +164,9 @@ 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 = 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]
 
@@ -238,6 +236,13 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
                 direct_weight = 1.0 - ar_weight
                 predicted[:, step] = ar_weight * ar_pred[:, step] + direct_weight * direct_pred[:, step]
 
+            # Motion correction from context
+            ctx_reshaped = context_tensor.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
+            motion = ctx_reshaped[:, -1] - ctx_reshaped[:, -2]  # [1, 3, 64, 64]
+            for step in range(PRED_FRAMES):
+                alpha = 0.05 * (step + 1)
+                predicted[:, step] = torch.clamp(predicted[:, step] + alpha * motion, 0, 1)
+
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []
         for i in range(PRED_FRAMES):