Upload folder using huggingface_hub

predict.py

@@ -202,17 +202,17 @@ 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 noise diversity
+            # Multi-run AR with brightness diversity
             all_ar_runs = []
-            for noise_std in [0.0, 1.0/255.0, 2.0/255.0]:
+            for brightness in [1.0, 1.02, 0.98]:
                 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 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)
+                    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
                     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])
@@ -230,17 +230,11 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
 
             ar_pred = sum(all_ar_runs) / len(all_ar_runs)
 
-            # Channel-specific AR/direct blend: blue gets more direct weight
-            # ar_start, ar_end per channel: R(0.65,0.35), G(0.60,0.30), B(0.50,0.20)
-            ch_ar_start = [0.65, 0.60, 0.50]
-            ch_ar_end = [0.35, 0.30, 0.20]
             predicted = torch.zeros_like(direct_pred)
             for step in range(PRED_FRAMES):
-                t = step / (PRED_FRAMES - 1)
-                for c in range(3):
-                    ar_w = ch_ar_start[c] - t * (ch_ar_start[c] - ch_ar_end[c])
-                    di_w = 1.0 - ar_w
-                    predicted[:, step, c] = ar_w * ar_pred[:, step, c] + di_w * direct_pred[:, step, c]
+                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]
 
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []