Upload folder using huggingface_hub

predict.py

@@ -229,28 +229,18 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
                 all_ar_runs.append(torch.stack(ar_preds_run, dim=1))
 
             ar_pred = sum(all_ar_runs) / len(all_ar_runs)
-
-            # Apply mild Gaussian blur to AR predictions
-            import torch.nn.functional as F
-            # Create 3x3 Gaussian kernel with sigma=0.5
-            kernel_size = 3
-            sigma = 0.5
-            x = torch.arange(kernel_size, dtype=torch.float32, device=DEVICE) - kernel_size // 2
-            gauss = torch.exp(-x**2 / (2 * sigma**2))
-            kernel_2d = gauss[:, None] * gauss[None, :]
-            kernel_2d = kernel_2d / kernel_2d.sum()
-            kernel_2d = kernel_2d.view(1, 1, kernel_size, kernel_size).expand(3, 1, -1, -1)
-            pad = kernel_size // 2
-            for s in range(PRED_FRAMES):
-                frame = ar_pred[:, s]  # [1, 3, 64, 64]
-                frame_padded = F.pad(frame, (pad, pad, pad, pad), mode='reflect')
-                ar_pred[:, s] = F.conv2d(frame_padded, kernel_2d, groups=3)
+            # Compute per-pixel variance across runs for uncertainty
+            ar_stack = torch.stack(all_ar_runs, dim=0)  # [3, 1, 8, 3, 64, 64]
+            ar_var = ar_stack.var(dim=0)  # [1, 8, 3, 64, 64]
 
             predicted = torch.zeros_like(direct_pred)
             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]
+                base_ar_weight = 0.65 - (step / (PRED_FRAMES - 1)) * 0.3
+                # Per-pixel weight adjustment: reduce AR weight where variance is high
+                uncertainty = torch.clamp(ar_var[:, step] * 50, 0, 0.3)
+                ar_weight_map = base_ar_weight * (1 - uncertainty)
+                direct_weight_map = 1.0 - ar_weight_map
+                predicted[:, step] = ar_weight_map * ar_pred[:, step] + direct_weight_map * direct_pred[:, step]
 
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []