Upload folder using huggingface_hub

predict.py

@@ -155,8 +155,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
             ctx = context_tensor.clone()
             last_t = last_tensor.clone()
             for step in range(PRED_FRAMES):
-                pong_scale = 1.06 if step >= 4 else 1.0
-                predicted = _predict_ar_frame(ens.models["pong"], ctx, last_t, residual_scale=pong_scale)
+                predicted = _predict_ar_frame(ens.models["pong"], ctx, last_t, residual_scale=1.03)
                 ar_preds.append(predicted)
                 ctx_frames = ctx.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
                 ctx_frames = torch.cat([ctx_frames[:, 1:], predicted.unsqueeze(1)], dim=1)
@@ -214,9 +213,8 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
                 for step in range(PRED_FRAMES):
                     ctx_in = ctx if noise_std == 0 else torch.clamp(ctx + torch.randn_like(ctx) * noise_std, 0, 1)
                     ctx_flip_in = ctx_flip if noise_std == 0 else torch.clamp(ctx_flip + torch.randn_like(ctx_flip) * noise_std, 0, 1)
-                    sonic_scale = 1.12 if step >= 3 else 1.0
-                    ar_orig = _predict_ar_frame(ens.sonic_ar, ctx_in, last_t, residual_scale=sonic_scale)
-                    ar_flip = _predict_ar_frame(ens.sonic_ar, ctx_flip_in, last_f, residual_scale=sonic_scale)
+                    ar_orig = _predict_ar_frame(ens.sonic_ar, ctx_in, last_t, residual_scale=1.08)
+                    ar_flip = _predict_ar_frame(ens.sonic_ar, ctx_flip_in, last_f, residual_scale=1.08)
                     ar_flip_back = torch.flip(ar_flip, dims=[3])
                     ar_frame = (ar_orig + ar_flip_back) / 2.0
                     ar_preds_run.append(ar_frame)
@@ -232,6 +230,22 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
 
             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)
+
             predicted = torch.zeros_like(direct_pred)
             for step in range(PRED_FRAMES):
                 ar_weight = 0.65 - (step / (PRED_FRAMES - 1)) * 0.3