Upload folder using huggingface_hub

predict.py

@@ -211,11 +211,12 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
                 ctx_flip = context_flipped.clone()
                 last_t = last_tensor.clone()
                 last_f = last_flipped.clone()
+                sonic_scales = [1.04, 1.04, 1.04, 1.08, 1.08, 1.08, 1.12, 1.12]
                 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)
+                    ar_orig = _predict_ar_frame(ens.sonic_ar, ctx_in, last_t, residual_scale=sonic_scales[step])
+                    ar_flip = _predict_ar_frame(ens.sonic_ar, ctx_flip_in, last_f, residual_scale=sonic_scales[step])
                     ar_flip_back = torch.flip(ar_flip, dims=[3])
                     ar_frame = (ar_orig + ar_flip_back) / 2.0
                     ar_preds_run.append(ar_frame)