Upload folder using huggingface_hub

predict.py

@@ -113,9 +113,9 @@ def _predict_8frames_direct(model, context_tensor, last_tensor):
     return torch.clamp(last_expanded + residuals, 0, 1)
 
 
-def _predict_ar_frame(model, context_tensor, last_tensor):
+def _predict_ar_frame(model, context_tensor, last_tensor, residual_scale=1.0):
     residual = model(context_tensor)
-    return torch.clamp(last_tensor + residual, 0, 1)
+    return torch.clamp(last_tensor + residual_scale * residual, 0, 1)
 
 
 def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
@@ -202,9 +202,8 @@ 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, track TTA disagreement
+            # Multi-run AR with noise diversity
             all_ar_runs = []
-            tta_disagreements = []  # per-step disagreement from clean run
             for noise_std in [0.0, 1.0/255.0, 2.0/255.0]:
                 ar_preds_run = []
                 ctx = context_tensor.clone()
@@ -214,13 +213,9 @@ 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)
-                    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_orig = _predict_ar_frame(ens.sonic_ar, ctx_in, last_t, residual_scale=0.95)
+                    ar_flip = _predict_ar_frame(ens.sonic_ar, ctx_flip_in, last_f, residual_scale=0.95)
                     ar_flip_back = torch.flip(ar_flip, dims=[3])
-                    # Track TTA disagreement for clean run
-                    if noise_std == 0:
-                        disagreement = torch.abs(ar_orig - ar_flip_back).mean().item()
-                        tta_disagreements.append(disagreement)
                     ar_frame = (ar_orig + ar_flip_back) / 2.0
                     ar_preds_run.append(ar_frame)
                     ctx_frames = ctx.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
@@ -237,14 +232,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
 
             predicted = torch.zeros_like(direct_pred)
             for step in range(PRED_FRAMES):
-                base_ar_weight = 0.65 - (step / (PRED_FRAMES - 1)) * 0.3
-                # Adjust based on TTA disagreement
-                disagreement = tta_disagreements[step]
-                if disagreement < 5.0 / 255.0:
-                    base_ar_weight += 0.1  # AR confident
-                elif disagreement > 15.0 / 255.0:
-                    base_ar_weight -= 0.1  # AR uncertain
-                ar_weight = max(0.15, min(0.85, base_ar_weight))
+                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]