Upload folder using huggingface_hub

model_pong_direct.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ab8070ddcde00333d7b52c89a0da9a61eece1e67c46163cd011ce4cd3c422f0c
+size 2436712

predict.py

@@ -1,4 +1,4 @@
-"""Inference v2: Pong AR+TTA caching, Sonic aggressive blending."""
+"""FP16 Pong ensemble: AR+direct for Pong, AR+direct for Sonic, direct for PP."""
 import sys
 import os
 import numpy as np
@@ -31,6 +31,7 @@ class EnsembleModels:
         self.models = {}
         self.sonic_ar = None
         self.sonic_direct = None
+        self.pong_direct = None
         self.direct_cache = None
         self.cache_step = 0
 
@@ -42,6 +43,7 @@ class EnsembleModels:
 def load_model(model_dir: str):
     ens = EnsembleModels()
 
+    # Pong AR (3 outputs)
     pong = UNet(in_channels=24, out_channels=3,
                 enc_channels=(32, 64, 128), bottleneck_channels=128,
                 upsample_mode="bilinear").to(DEVICE)
@@ -51,6 +53,17 @@ def load_model(model_dir: str):
     pong.eval()
     ens.models["pong"] = pong
 
+    # Pong direct (24 outputs)
+    pong_direct = UNet(in_channels=24, out_channels=24,
+                       enc_channels=(32, 64, 128), bottleneck_channels=128,
+                       upsample_mode="bilinear").to(DEVICE)
+    sd = torch.load(os.path.join(model_dir, "model_pong_direct.pt"),
+                    map_location=DEVICE, weights_only=True)
+    pong_direct.load_state_dict({k: v.float() for k, v in sd.items()})
+    pong_direct.eval()
+    ens.pong_direct = pong_direct
+
+    # Sonic AR (3 outputs)
     sonic_ar = UNet(in_channels=24, out_channels=3,
                     enc_channels=(48, 96, 192), bottleneck_channels=256,
                     upsample_mode="bilinear").to(DEVICE)
@@ -60,6 +73,7 @@ def load_model(model_dir: str):
     sonic_ar.eval()
     ens.sonic_ar = sonic_ar
 
+    # Sonic direct (24 outputs)
     sonic_direct = UNet(in_channels=24, out_channels=24,
                         enc_channels=(48, 96, 192), bottleneck_channels=256,
                         upsample_mode="bilinear").to(DEVICE)
@@ -69,6 +83,7 @@ def load_model(model_dir: str):
     sonic_direct.eval()
     ens.sonic_direct = sonic_direct
 
+    # PP compact direct (24 outputs)
     pp = UNet(in_channels=24, out_channels=24,
               enc_channels=(24, 48, 96), bottleneck_channels=128,
               upsample_mode="bilinear").to(DEVICE)
@@ -111,7 +126,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 TTA and internal float32 caching
+        # Pong: AR+direct ensemble with float32 caching, no TTA
         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
@@ -120,42 +135,40 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
             return result
 
         ens.reset_cache()
-        model = ens.models["pong"]
+        model_ar = ens.models["pong"]
+        model_direct = ens.pong_direct
         with torch.no_grad():
             context_tensor = torch.from_numpy(context).to(DEVICE)
             last_tensor = torch.from_numpy(last_frame_t).to(DEVICE)
 
-            # Original AR
-            preds_orig = []
+            # Direct prediction
+            direct_pred = _predict_8frames_direct(model_direct, context_tensor, last_tensor)
+
+            # AR prediction in float32
+            ar_preds = []
             ctx = context_tensor.clone()
             last_t = last_tensor.clone()
             for step in range(PRED_FRAMES):
-                predicted = _predict_ar_frame(model, ctx, last_t)
-                preds_orig.append(predicted)
+                predicted = _predict_ar_frame(model_ar, ctx, last_t)
+                ar_preds.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()
+            ar_pred = torch.stack(ar_preds, dim=1)
+
+            # Step-dependent blending: AR 0.7 -> 0.3
+            predicted = torch.zeros_like(direct_pred)
             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
+                ar_weight = 0.7 - (step / (PRED_FRAMES - 1)) * 0.4
+                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 = []
         for i in range(PRED_FRAMES):
-            avg = (preds_orig[i] + preds_flip[i]) / 2.0
-            frame = avg[0].cpu().numpy()
-            frame = np.transpose(frame, (1, 2, 0))
+            frame = np.transpose(predicted_np[i], (1, 2, 0))
             frame = (frame * 255).clip(0, 255).astype(np.uint8)
             ens.direct_cache.append(frame)
 
@@ -164,7 +177,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         return result
 
     elif game == "sonic":
-        # Sonic: step-dependent blending with more aggressive shift to direct
+        # Sonic: AR+direct with step blending and TTA
         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
@@ -206,10 +219,9 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
 
             ar_pred = torch.stack(ar_preds, dim=1)
 
-            # More aggressive blending: AR weight 0.8 -> 0.2
             predicted = torch.zeros_like(direct_pred)
             for step in range(PRED_FRAMES):
-                ar_weight = 0.8 - (step / (PRED_FRAMES - 1)) * 0.6  # 0.8 -> 0.2
+                ar_weight = 0.7 - (step / (PRED_FRAMES - 1)) * 0.4
                 direct_weight = 1.0 - ar_weight
                 predicted[:, step] = ar_weight * ar_pred[:, step] + direct_weight * direct_pred[:, step]
 
@@ -225,7 +237,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         return result
 
     else:
-        # PP: direct 8-frame with caching and TTA
+        # PP: direct with TTA and 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