Upload folder using huggingface_hub

predict.py

@@ -1,4 +1,4 @@
-"""Inference tricks: Pong AR caching, Sonic step-dependent blending."""
+"""Inference v2: Pong AR+TTA caching, Sonic aggressive blending."""
 import sys
 import os
 import numpy as np
@@ -42,7 +42,6 @@ class EnsembleModels:
 def load_model(model_dir: str):
     ens = EnsembleModels()
 
-    # Pong: full AR model (3 outputs)
     pong = UNet(in_channels=24, out_channels=3,
                 enc_channels=(32, 64, 128), bottleneck_channels=128,
                 upsample_mode="bilinear").to(DEVICE)
@@ -52,7 +51,6 @@ def load_model(model_dir: str):
     pong.eval()
     ens.models["pong"] = pong
 
-    # Sonic AR model (3 outputs)
     sonic_ar = UNet(in_channels=24, out_channels=3,
                     enc_channels=(48, 96, 192), bottleneck_channels=256,
                     upsample_mode="bilinear").to(DEVICE)
@@ -62,7 +60,6 @@ def load_model(model_dir: str):
     sonic_ar.eval()
     ens.sonic_ar = sonic_ar
 
-    # Sonic direct model (24 outputs)
     sonic_direct = UNet(in_channels=24, out_channels=24,
                         enc_channels=(48, 96, 192), bottleneck_channels=256,
                         upsample_mode="bilinear").to(DEVICE)
@@ -72,7 +69,6 @@ def load_model(model_dir: str):
     sonic_direct.eval()
     ens.sonic_direct = sonic_direct
 
-    # PP: compact direct 8-frame model (24 outputs)
     pp = UNet(in_channels=24, out_channels=24,
               enc_channels=(24, 48, 96), bottleneck_channels=128,
               upsample_mode="bilinear").to(DEVICE)
@@ -115,7 +111,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
     last_frame_t = np.transpose(last_frame, (2, 0, 1))[np.newaxis]
 
     if game == "pong":
-        # Pong: AR with internal caching (all 8 steps in float32)
+        # Pong: AR with TTA and internal float32 caching
         if ens.direct_cache is not None and n > CONTEXT_FRAMES and ens.cache_step < PRED_FRAMES:
             result = ens.direct_cache[ens.cache_step]
             ens.cache_step += 1
@@ -129,22 +125,36 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
             context_tensor = torch.from_numpy(context).to(DEVICE)
             last_tensor = torch.from_numpy(last_frame_t).to(DEVICE)
 
-            # Run all 8 AR steps in float32 (no uint8 quantization between steps)
-            preds = []
+            # Original AR
+            preds_orig = []
             ctx = context_tensor.clone()
             last_t = last_tensor.clone()
             for step in range(PRED_FRAMES):
                 predicted = _predict_ar_frame(model, ctx, last_t)
-                preds.append(predicted)
-                # Shift context with float32 prediction (no quantization)
+                preds_orig.append(predicted)
                 ctx_frames = ctx.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
                 ctx_frames = torch.cat([ctx_frames[:, 1:], predicted.unsqueeze(1)], dim=1)
                 ctx = ctx_frames.reshape(1, -1, 64, 64)
                 last_t = predicted
 
+            # TTA: horizontal flip AR
+            context_flipped = torch.flip(context_tensor, dims=[3])
+            last_flipped = torch.flip(last_tensor, dims=[3])
+            preds_flip = []
+            ctx_f = context_flipped.clone()
+            last_f = last_flipped.clone()
+            for step in range(PRED_FRAMES):
+                predicted_f = _predict_ar_frame(model, ctx_f, last_f)
+                preds_flip.append(torch.flip(predicted_f, dims=[3]))  # flip back
+                ctx_frames_f = ctx_f.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
+                ctx_frames_f = torch.cat([ctx_frames_f[:, 1:], predicted_f.unsqueeze(1)], dim=1)
+                ctx_f = ctx_frames_f.reshape(1, -1, 64, 64)
+                last_f = predicted_f
+
         ens.direct_cache = []
         for i in range(PRED_FRAMES):
-            frame = preds[i][0].cpu().numpy()
+            avg = (preds_orig[i] + preds_flip[i]) / 2.0
+            frame = avg[0].cpu().numpy()
             frame = np.transpose(frame, (1, 2, 0))
             frame = (frame * 255).clip(0, 255).astype(np.uint8)
             ens.direct_cache.append(frame)
@@ -154,7 +164,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         return result
 
     elif game == "sonic":
-        # Sonic: step-dependent AR/direct blending with caching
+        # Sonic: step-dependent blending with more aggressive shift to direct
         if ens.direct_cache is not None and n > CONTEXT_FRAMES and ens.cache_step < PRED_FRAMES:
             result = ens.direct_cache[ens.cache_step]
             ens.cache_step += 1
@@ -167,15 +177,13 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
             context_tensor = torch.from_numpy(context).to(DEVICE)
             last_tensor = torch.from_numpy(last_frame_t).to(DEVICE)
 
-            # Direct prediction with TTA
             direct_orig = _predict_8frames_direct(ens.sonic_direct, context_tensor, last_tensor)
             context_flipped = torch.flip(context_tensor, dims=[3])
             last_flipped = torch.flip(last_tensor, dims=[3])
             direct_flipped = _predict_8frames_direct(ens.sonic_direct, context_flipped, last_flipped)
             direct_flipped = torch.flip(direct_flipped, dims=[4])
-            direct_pred = (direct_orig + direct_flipped) / 2.0  # (1, 8, 3, 64, 64)
+            direct_pred = (direct_orig + direct_flipped) / 2.0
 
-            # AR prediction with TTA
             ar_preds = []
             ctx = context_tensor.clone()
             ctx_flip = context_flipped.clone()
@@ -196,14 +204,13 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
                 ctx_flip = ctx_flip_frames.reshape(1, -1, 64, 64)
                 last_f = ar_flip
 
-            ar_pred = torch.stack(ar_preds, dim=1)  # (1, 8, 3, 64, 64)
+            ar_pred = torch.stack(ar_preds, dim=1)
 
-            # Step-dependent blending: AR weight goes from 0.7 (step 0) to 0.3 (step 7)
-            # Direct weight goes from 0.3 (step 0) to 0.7 (step 7)
+            # More aggressive blending: AR weight 0.8 -> 0.2
             predicted = torch.zeros_like(direct_pred)
             for step in range(PRED_FRAMES):
-                ar_weight = 0.7 - (step / (PRED_FRAMES - 1)) * 0.4  # 0.7 -> 0.3
-                direct_weight = 1.0 - ar_weight  # 0.3 -> 0.7
+                ar_weight = 0.8 - (step / (PRED_FRAMES - 1)) * 0.6  # 0.8 -> 0.2
+                direct_weight = 1.0 - ar_weight
                 predicted[:, step] = ar_weight * ar_pred[:, step] + direct_weight * direct_pred[:, step]
 
         predicted_np = predicted[0].cpu().numpy()