Upload folder using huggingface_hub

predict.py

@@ -106,11 +106,11 @@ def load_model(model_dir: str):
     return ens
 
 
-def _predict_8frames_direct(model, context_tensor, last_tensor):
+def _predict_8frames_direct(model, context_tensor, last_tensor, residual_scale=1.0):
     output = model(context_tensor)
     residuals = output.reshape(1, PRED_FRAMES, 3, 64, 64)
     last_expanded = last_tensor.unsqueeze(1).expand_as(residuals)
-    return torch.clamp(last_expanded + residuals, 0, 1)
+    return torch.clamp(last_expanded + residual_scale * residuals, 0, 1)
 
 
 def _predict_ar_frame(model, context_tensor, last_tensor, residual_scale=1.0):
@@ -166,7 +166,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
 
             predicted = torch.zeros_like(direct_pred)
             for step in range(PRED_FRAMES):
-                ar_weight = 0.90 - (step / (PRED_FRAMES - 1)) * 0.3
+                ar_weight = 0.85 - (step / (PRED_FRAMES - 1)) * 0.3
                 direct_weight = 1.0 - ar_weight
                 predicted[:, step] = ar_weight * ar_pred[:, step] + direct_weight * direct_pred[:, step]
 
@@ -232,7 +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):
-                ar_weight = 0.70 - (step / (PRED_FRAMES - 1)) * 0.3
+                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]
 
@@ -261,10 +261,10 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
             context_tensor = torch.from_numpy(context).to(DEVICE)
             last_tensor = torch.from_numpy(last_frame_t).to(DEVICE)
 
-            predicted_orig = _predict_8frames_direct(ens.models["pole_position"], context_tensor, last_tensor)
+            predicted_orig = _predict_8frames_direct(ens.models["pole_position"], context_tensor, last_tensor, residual_scale=1.03)
             context_flipped = torch.flip(context_tensor, dims=[3])
             last_flipped = torch.flip(last_tensor, dims=[3])
-            predicted_flipped = _predict_8frames_direct(ens.models["pole_position"], context_flipped, last_flipped)
+            predicted_flipped = _predict_8frames_direct(ens.models["pole_position"], context_flipped, last_flipped, residual_scale=1.03)
             predicted_flipped = torch.flip(predicted_flipped, dims=[4])
             predicted = (predicted_orig + predicted_flipped) / 2.0