Upload folder using huggingface_hub

predict.py

@@ -202,10 +202,10 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
             direct_flipped = torch.flip(direct_flipped, dims=[4])
             direct_pred = (direct_orig + direct_flipped) / 2.0
 
-            # 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, 2.0/255.0]:
-                ar_preds_run = []
+            # Multi-run AR with noise diversity - collect all 6 predictions per step
+            # (3 noise levels x 2 TTA directions) and take per-pixel median
+            all_step_preds = [[] for _ in range(PRED_FRAMES)]
+            for noise_std in [0.0, 1.0/255.0, 2.0/255.0]:
                 ctx = context_tensor.clone()
                 ctx_flip = context_flipped.clone()
                 last_t = last_tensor.clone()
@@ -216,8 +216,8 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
                     ar_orig = _predict_ar_frame(ens.sonic_ar, ctx_in, last_t)
                     ar_flip = _predict_ar_frame(ens.sonic_ar, ctx_flip_in, last_f)
                     ar_flip_back = torch.flip(ar_flip, dims=[3])
-                    ar_frame = (ar_orig + ar_flip_back) / 2.0
-                    ar_preds_run.append(ar_frame)
+                    all_step_preds[step].append(ar_orig)
+                    all_step_preds[step].append(ar_flip_back)
                     ctx_frames = ctx.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
                     ctx_frames = torch.cat([ctx_frames[:, 1:], ar_orig.unsqueeze(1)], dim=1)
                     ctx = ctx_frames.reshape(1, -1, 64, 64)
@@ -226,9 +226,13 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
                     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
-                all_ar_runs.append(torch.stack(ar_preds_run, dim=1))
 
-            ar_pred = sum(all_ar_runs) / len(all_ar_runs)
+            ar_pred_list = []
+            for step in range(PRED_FRAMES):
+                stacked = torch.stack(all_step_preds[step], dim=0)  # [6, 1, 3, 64, 64]
+                median_val = torch.median(stacked, dim=0).values
+                ar_pred_list.append(median_val)
+            ar_pred = torch.stack(ar_pred_list, dim=1)  # [1, 8, 3, 64, 64]
 
             predicted = torch.zeros_like(direct_pred)
             for step in range(PRED_FRAMES):