Upload folder using huggingface_hub

predict.py

@@ -202,17 +202,32 @@ 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 brightness diversity
+            # 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
+            ]
             all_ar_runs = []
-            for brightness in [1.0, 1.02, 0.98]:
+            for noise_std, bright in perturbations:
                 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 = torch.clamp(ctx * brightness, 0, 1) if brightness != 1.0 else ctx
-                    ctx_flip_in = torch.clamp(ctx_flip * brightness, 0, 1) if brightness != 1.0 else ctx_flip
+                    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)
                     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])