Upload folder using huggingface_hub

predict.py

@@ -230,27 +230,25 @@ 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):
-                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]
+                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]
 
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []
-        last_ctx_uint8 = (last_frame * 255).clip(0, 255).astype(np.uint8)
-        catastrophic = False
         for i in range(PRED_FRAMES):
             frame = np.transpose(predicted_np[i], (1, 2, 0))
             frame = (frame * 255).clip(0, 255).astype(np.uint8)
-            diff = np.abs(frame.astype(np.float32) - last_ctx_uint8.astype(np.float32)).mean()
-            if diff > 100:
-                catastrophic = True
             ens.direct_cache.append(frame)
 
-        if catastrophic:
-            ens.direct_cache = [last_ctx_uint8.copy() for _ in range(PRED_FRAMES)]
-
         result = ens.direct_cache[ens.cache_step]
         ens.cache_step += 1
         return result