Upload folder using huggingface_hub

predict.py

@@ -204,7 +204,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
 
             # Multi-run AR with noise diversity
             all_ar_runs = []
-            for noise_std in [0.0, 0.5/255.0, 1.0/255.0, 1.5/255.0]:
+            for noise_std in [0.0, 1.0/255.0, 2.0/255.0]:
                 ar_preds_run = []
                 ctx = context_tensor.clone()
                 ctx_flip = context_flipped.clone()
@@ -230,11 +230,15 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
 
             ar_pred = sum(all_ar_runs) / len(all_ar_runs)
 
+            # Agreement-based blending: trust AR where models agree
             predicted = torch.zeros_like(direct_pred)
+            low_thresh = 10.0 / 255.0
+            high_thresh = 50.0 / 255.0
             for step in range(PRED_FRAMES):
-                ar_weight = 0.65 - (step / (PRED_FRAMES - 1)) * 0.3
-                direct_weight = 1.0 - ar_weight
-                predicted[:, step] = ar_weight * ar_pred[:, step] + direct_weight * direct_pred[:, step]
+                disagreement = torch.abs(ar_pred[:, step] - direct_pred[:, step])
+                # Where agree: AR weight=1.0, where disagree strongly: AR weight=0.5
+                ar_w = 1.0 - 0.5 * torch.clamp((disagreement - low_thresh) / (high_thresh - low_thresh), 0, 1)
+                predicted[:, step] = ar_w * ar_pred[:, step] + (1.0 - ar_w) * direct_pred[:, step]
 
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []