Upload folder using huggingface_hub

model_pole_position.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:27d26875071b536cc75cac27a0840b50cd6c9a8e1956c94f1cd08feacc49621f
-size 2971526
+oid sha256:62d218b9859acd4d19cfcfe6b3aa93ae129485a872175632ed32d6441ae9c7f6
+size 1580934

model_pong.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6e5c823ea30b79baec88c337290bb2491c83906ac288b5a0111decb2f2b49792
-size 1387816
+oid sha256:d6c8b9235347bea94e7e5f5f0f225d4c1dbd13a749d5e28920c75c91902ecb11
+size 2435368

predict.py

@@ -1,4 +1,4 @@
-"""Compact ensemble: small Pong AR, Sonic AR+direct ensemble, PP direct. All under 16MB."""
+"""Compact PP ensemble: full Pong AR, Sonic AR+direct ensemble, compact PP direct. Under 16MB."""
 import sys
 import os
 import numpy as np
@@ -42,9 +42,9 @@ class EnsembleModels:
 def load_model(model_dir: str):
     ens = EnsembleModels()
 
-    # Pong: compact AR model (3 outputs, smaller architecture)
+    # Pong: full AR model (3 outputs)
     pong = UNet(in_channels=24, out_channels=3,
-                enc_channels=(24, 48, 96), bottleneck_channels=96,
+                enc_channels=(32, 64, 128), bottleneck_channels=128,
                 upsample_mode="bilinear").to(DEVICE)
     sd = torch.load(os.path.join(model_dir, "model_pong.pt"),
                     map_location=DEVICE, weights_only=True)
@@ -72,9 +72,9 @@ def load_model(model_dir: str):
     sonic_direct.eval()
     ens.sonic_direct = sonic_direct
 
-    # PP: direct 8-frame model (24 outputs)
+    # PP: compact direct 8-frame model (24 outputs, smaller architecture)
     pp = UNet(in_channels=24, out_channels=24,
-              enc_channels=(32, 64, 128), bottleneck_channels=192,
+              enc_channels=(24, 48, 96), bottleneck_channels=128,
               upsample_mode="bilinear").to(DEVICE)
     sd = torch.load(os.path.join(model_dir, "model_pole_position.pt"),
                     map_location=DEVICE, weights_only=True)
@@ -146,7 +146,7 @@ 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
 
-            # AR prediction with TTA for each step
+            # AR prediction with TTA
             ar_preds = []
             ctx = context_tensor.clone()
             ctx_flip = context_flipped.clone()
@@ -182,6 +182,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         return result
 
     else:
+        # PP: direct 8-frame with caching 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