Upload folder using huggingface_hub

predict.py

@@ -149,31 +149,28 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
             context_tensor = torch.from_numpy(context).to(DEVICE)
             last_tensor = torch.from_numpy(last_frame_t).to(DEVICE)
 
-            # Pong AR-only with hflip TTA
-            context_flipped = torch.flip(context_tensor, dims=[3])
-            last_flipped = torch.flip(last_tensor, dims=[3])
+            direct_pred = _predict_8frames_direct(ens.pong_direct, context_tensor, last_tensor)
 
             ar_preds = []
             ctx = context_tensor.clone()
-            ctx_flip = context_flipped.clone()
             last_t = last_tensor.clone()
-            last_f = last_flipped.clone()
             for step in range(PRED_FRAMES):
-                ar_orig = _predict_ar_frame(ens.models["pong"], ctx, last_t)
-                ar_flip = _predict_ar_frame(ens.models["pong"], ctx_flip, last_f)
-                ar_flip_back = torch.flip(ar_flip, dims=[3])
-                ar_frame = (ar_orig + ar_flip_back) / 2.0
-                ar_preds.append(ar_frame)
+                predicted = _predict_ar_frame(ens.models["pong"], ctx, last_t)
+                ar_preds.append(predicted)
                 ctx_frames = ctx.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
-                ctx_frames = torch.cat([ctx_frames[:, 1:], ar_orig.unsqueeze(1)], dim=1)
+                ctx_frames = torch.cat([ctx_frames[:, 1:], predicted.unsqueeze(1)], dim=1)
                 ctx = ctx_frames.reshape(1, -1, 64, 64)
-                last_t = ar_orig
-                ctx_flip_frames = ctx_flip.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
-                ctx_flip_frames = torch.cat([ctx_flip_frames[:, 1:], ar_flip.unsqueeze(1)], dim=1)
-                ctx_flip = ctx_flip_frames.reshape(1, -1, 64, 64)
-                last_f = ar_flip
+                last_t = predicted
 
-            predicted = torch.stack(ar_preds, dim=1)
+            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]
+                direct_weight = 1.0 - ar_weight
+                predicted[:, step] = ar_weight * ar_pred[:, step] + direct_weight * direct_pred[:, step]
 
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []