Upload folder using huggingface_hub

predict.py

@@ -230,13 +230,15 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
 
             ar_pred = sum(all_ar_runs) / len(all_ar_runs)
 
+            # Blend in residual space
+            last_ctx = last_tensor.unsqueeze(1).expand_as(direct_pred)
             predicted = torch.zeros_like(direct_pred)
-            eps = 1e-6
             for step in range(PRED_FRAMES):
                 w = 0.65 - (step / (PRED_FRAMES - 1)) * 0.3
-                ar_safe = torch.clamp(ar_pred[:, step], eps, 1.0)
-                direct_safe = torch.clamp(direct_pred[:, step], eps, 1.0)
-                predicted[:, step] = torch.clamp(ar_safe.pow(w) * direct_safe.pow(1.0 - w), 0, 1)
+                ar_residual = ar_pred[:, step] - last_ctx[:, step]
+                direct_residual = direct_pred[:, step] - last_ctx[:, step]
+                blended_residual = w * ar_residual + (1.0 - w) * direct_residual
+                predicted[:, step] = torch.clamp(last_ctx[:, step] + blended_residual, 0, 1)
 
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []