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, residual_scale=1.0):
+def _predict_8frames_direct(model, context_tensor, last_tensor):
     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 + residual_scale * residuals, 0, 1)
+    return torch.clamp(last_expanded + residuals, 0, 1)
 
 
 def _predict_ar_frame(model, context_tensor, last_tensor, residual_scale=1.0):
@@ -149,13 +149,14 @@ 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)
 
-            direct_pred = _predict_8frames_direct(ens.pong_direct, context_tensor, last_tensor, residual_scale=1.02)
+            direct_pred = _predict_8frames_direct(ens.pong_direct, context_tensor, last_tensor)
 
             ar_preds = []
             ctx = context_tensor.clone()
             last_t = last_tensor.clone()
             for step in range(PRED_FRAMES):
-                predicted = _predict_ar_frame(ens.models["pong"], ctx, last_t, residual_scale=1.03)
+                pong_scale = 1.01 + step * (1.05 - 1.01) / (PRED_FRAMES - 1)
+                predicted = _predict_ar_frame(ens.models["pong"], ctx, last_t, residual_scale=pong_scale)
                 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)
@@ -195,10 +196,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)
 
-            direct_orig = _predict_8frames_direct(ens.sonic_direct, context_tensor, last_tensor, residual_scale=1.05)
+            direct_orig = _predict_8frames_direct(ens.sonic_direct, context_tensor, last_tensor)
             context_flipped = torch.flip(context_tensor, dims=[3])
             last_flipped = torch.flip(last_tensor, dims=[3])
-            direct_flipped = _predict_8frames_direct(ens.sonic_direct, context_flipped, last_flipped, residual_scale=1.05)
+            direct_flipped = _predict_8frames_direct(ens.sonic_direct, context_flipped, last_flipped)
             direct_flipped = torch.flip(direct_flipped, dims=[4])
             direct_pred = (direct_orig + direct_flipped) / 2.0
 
@@ -213,8 +214,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, residual_scale=1.08)
-                    ar_flip = _predict_ar_frame(ens.sonic_ar, ctx_flip_in, last_f, residual_scale=1.08)
+                    sonic_scale = 1.02 + step * (1.16 - 1.02) / (PRED_FRAMES - 1)
+                    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_flip_back = torch.flip(ar_flip, dims=[3])
                     ar_frame = (ar_orig + ar_flip_back) / 2.0
                     ar_preds_run.append(ar_frame)