Upload folder using huggingface_hub

predict.py

@@ -166,7 +166,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
 
             predicted = torch.zeros_like(direct_pred)
             for step in range(PRED_FRAMES):
-                ar_weight = 0.80
+                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]
 
@@ -261,12 +261,19 @@ 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)
 
-            predicted_orig = _predict_8frames_direct(ens.models["pole_position"], context_tensor, last_tensor)
             context_flipped = torch.flip(context_tensor, dims=[3])
             last_flipped = torch.flip(last_tensor, dims=[3])
-            predicted_flipped = _predict_8frames_direct(ens.models["pole_position"], context_flipped, last_flipped)
-            predicted_flipped = torch.flip(predicted_flipped, dims=[4])
-            predicted = (predicted_orig + predicted_flipped) / 2.0
+
+            # 3-run noise diversity + TTA for PP
+            all_pp_runs = []
+            for noise_std in [0.0, 0.5/255.0, 1.0/255.0]:
+                ctx_in = context_tensor if noise_std == 0 else torch.clamp(context_tensor + torch.randn_like(context_tensor) * noise_std, 0, 1)
+                ctx_flip_in = context_flipped if noise_std == 0 else torch.clamp(context_flipped + torch.randn_like(context_flipped) * noise_std, 0, 1)
+                pred_orig = _predict_8frames_direct(ens.models["pole_position"], ctx_in, last_tensor)
+                pred_flipped = _predict_8frames_direct(ens.models["pole_position"], ctx_flip_in, last_flipped)
+                pred_flipped = torch.flip(pred_flipped, dims=[4])
+                all_pp_runs.append((pred_orig + pred_flipped) / 2.0)
+            predicted = sum(all_pp_runs) / len(all_pp_runs)
 
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []