Upload folder using huggingface_hub

predict.py

@@ -135,6 +135,9 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
     last_frame = frames_norm[-1]
     last_frame_t = np.transpose(last_frame, (2, 0, 1))[np.newaxis]
 
+    # Per-channel means from last 4 context frames for color correction
+    ctx_channel_means = frames_norm[-4:].mean(axis=(0, 1, 2))  # [3]
+
     if game == "pong":
         # Pong: AR+direct ensemble, float32 caching, no TTA
         if ens.direct_cache is not None and n > CONTEXT_FRAMES and ens.cache_step < PRED_FRAMES:
@@ -236,17 +239,17 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
                 direct_weight = 1.0 - ar_weight
                 predicted[:, step] = ar_weight * ar_pred[:, step] + direct_weight * direct_pred[:, step]
 
-            # Motion correction from context
-            ctx_reshaped = context_tensor.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
-            motion = ctx_reshaped[:, -1] - ctx_reshaped[:, -2]  # [1, 3, 64, 64]
-            for step in range(PRED_FRAMES):
-                alpha = 0.05 * (step + 1)
-                predicted[:, step] = torch.clamp(predicted[:, step] + alpha * motion, 0, 1)
-
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []
         for i in range(PRED_FRAMES):
-            frame = np.transpose(predicted_np[i], (1, 2, 0))
+            frame = np.transpose(predicted_np[i], (1, 2, 0))  # [64,64,3]
+            # Color correction: nudge channel means toward context
+            pred_means = frame.mean(axis=(0, 1))  # [3]
+            diff = ctx_channel_means - pred_means
+            thresh = 3.0 / 255.0
+            correction = np.where(np.abs(diff) > thresh, 0.1 * diff, 0.0)
+            frame = frame + correction[np.newaxis, np.newaxis, :]
+            frame = frame.clip(0, 1)
             frame = (frame * 255).clip(0, 255).astype(np.uint8)
             ens.direct_cache.append(frame)
 
@@ -279,6 +282,13 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         ens.direct_cache = []
         for i in range(PRED_FRAMES):
             frame = np.transpose(predicted_np[i], (1, 2, 0))
+            # Color correction
+            pred_means = frame.mean(axis=(0, 1))
+            diff = ctx_channel_means - pred_means
+            thresh = 3.0 / 255.0
+            correction = np.where(np.abs(diff) > thresh, 0.1 * diff, 0.0)
+            frame = frame + correction[np.newaxis, np.newaxis, :]
+            frame = frame.clip(0, 1)
             frame = (frame * 255).clip(0, 255).astype(np.uint8)
             ens.direct_cache.append(frame)