Upload folder using huggingface_hub

predict.py

@@ -202,45 +202,43 @@ 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 combined noise+brightness diversity (6 runs)
-            # (noise_std, brightness_mult)
-            perturbations = [
-                (0.0, 1.0),           # normal
-                (1.0/255.0, 1.0),     # noise only
-                (2.0/255.0, 1.0),     # more noise
-                (0.0, 1.02),          # brighter
-                (0.0, 0.98),          # darker
-                (1.0/255.0, 1.01),    # noise + slight bright
-            ]
+            # Multi-run AR with noise diversity, using direct step 1 as AR init
+            # Get direct model's step 1 prediction for seeding AR context
+            direct_step1 = direct_pred[:, 0]  # [1, 3, 64, 64]
+            direct_step1_flip = torch.flip(direct_step1, dims=[3])
+
             all_ar_runs = []
-            for noise_std, bright in perturbations:
+            for noise_std in [0.0, 1.0/255.0, 2.0/255.0]:
                 ar_preds_run = []
                 ctx = context_tensor.clone()
                 ctx_flip = context_flipped.clone()
                 last_t = last_tensor.clone()
                 last_f = last_flipped.clone()
                 for step in range(PRED_FRAMES):
-                    ctx_in = ctx
-                    ctx_flip_in = ctx_flip
-                    if noise_std > 0:
-                        ctx_in = torch.clamp(ctx_in + torch.randn_like(ctx_in) * noise_std, 0, 1)
-                        ctx_flip_in = torch.clamp(ctx_flip_in + torch.randn_like(ctx_flip_in) * noise_std, 0, 1)
-                    if bright != 1.0:
-                        ctx_in = torch.clamp(ctx_in * bright, 0, 1)
-                        ctx_flip_in = torch.clamp(ctx_flip_in * bright, 0, 1)
+                    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_flip_back = torch.flip(ar_flip, dims=[3])
                     ar_frame = (ar_orig + ar_flip_back) / 2.0
                     ar_preds_run.append(ar_frame)
+
+                    # For step 0: use direct model's prediction as context seed
+                    if step == 0:
+                        feed_orig = direct_step1
+                        feed_flip = direct_step1_flip
+                    else:
+                        feed_orig = ar_orig
+                        feed_flip = ar_flip
+
                     ctx_frames = ctx.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
-                    ctx_frames = torch.cat([ctx_frames[:, 1:], ar_orig.unsqueeze(1)], dim=1)
+                    ctx_frames = torch.cat([ctx_frames[:, 1:], feed_orig.unsqueeze(1)], dim=1)
                     ctx = ctx_frames.reshape(1, -1, 64, 64)
-                    last_t = ar_orig
+                    last_t = feed_orig
                     ctx_flip_frames = ctx_flip.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
-                    ctx_flip_frames = torch.cat([ctx_flip_frames[:, 1:], ar_flip.unsqueeze(1)], dim=1)
+                    ctx_flip_frames = torch.cat([ctx_flip_frames[:, 1:], feed_flip.unsqueeze(1)], dim=1)
                     ctx_flip = ctx_flip_frames.reshape(1, -1, 64, 64)
-                    last_f = ar_flip
+                    last_f = feed_flip
                 all_ar_runs.append(torch.stack(ar_preds_run, dim=1))
 
             ar_pred = sum(all_ar_runs) / len(all_ar_runs)