Upload folder using huggingface_hub

predict.py

@@ -106,18 +106,6 @@ def load_model(model_dir: str):
     return ens
 
 
-def _unsharp_mask(frames_tensor, alpha=0.15):
-    """Apply unsharp mask sharpening. frames_tensor: [1, N, 3, 64, 64]"""
-    import torch.nn.functional as F
-    kernel = torch.tensor([[1, 2, 1], [2, 4, 2], [1, 2, 1]], dtype=frames_tensor.dtype, device=frames_tensor.device).float() / 16.0
-    kernel = kernel.unsqueeze(0).unsqueeze(0).expand(3, 1, 3, 3)
-    b, n, c, h, w = frames_tensor.shape
-    x = frames_tensor.reshape(b * n, c, h, w)
-    blurred = F.conv2d(x, kernel, padding=1, groups=3)
-    sharpened = x + alpha * (x - blurred)
-    return torch.clamp(sharpened.reshape(b, n, c, h, w), 0, 1)
-
-
 def _predict_8frames_direct(model, context_tensor, last_tensor):
     output = model(context_tensor)
     residuals = output.reshape(1, PRED_FRAMES, 3, 64, 64)
@@ -216,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, 1.0/255.0, 2.0/255.0]:
+            for noise_std in [0.0, 0.5/255.0, 1.0/255.0, 1.5/255.0]:
                 ar_preds_run = []
                 ctx = context_tensor.clone()
                 ctx_flip = context_flipped.clone()
@@ -248,8 +236,6 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
                 direct_weight = 1.0 - ar_weight
                 predicted[:, step] = ar_weight * ar_pred[:, step] + direct_weight * direct_pred[:, step]
 
-            predicted = _unsharp_mask(predicted, alpha=0.15)
-
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []
         for i in range(PRED_FRAMES):
@@ -281,7 +267,6 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
             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
-            predicted = _unsharp_mask(predicted, alpha=0.15)
 
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []