Upload folder using huggingface_hub

model_pong.pt

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

predict.py

@@ -1,4 +1,4 @@
-"""Ensemble hybrid: AR+direct ensemble for Sonic, AR for Pong, direct for PP."""
+"""Compact ensemble: small Pong AR, Sonic AR+direct ensemble, PP direct. All 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: AR model (3 outputs)
+    # Pong: compact AR model (3 outputs, smaller architecture)
     pong = UNet(in_channels=24, out_channels=3,
-                enc_channels=(32, 64, 128), bottleneck_channels=128,
+                enc_channels=(24, 48, 96), bottleneck_channels=96,
                 upsample_mode="bilinear").to(DEVICE)
     sd = torch.load(os.path.join(model_dir, "model_pong.pt"),
                     map_location=DEVICE, weights_only=True)
@@ -86,14 +86,14 @@ def load_model(model_dir: str):
 
 
 def _predict_8frames_direct(model, context_tensor, last_tensor):
-    output = model(context_tensor)  # (1, 24, 64, 64)
+    output = model(context_tensor)
     residuals = output.reshape(1, PRED_FRAMES, 3, 64, 64)
     last_expanded = last_tensor.unsqueeze(1).expand_as(residuals)
     return torch.clamp(last_expanded + residuals, 0, 1)
 
 
 def _predict_ar_frame(model, context_tensor, last_tensor):
-    residual = model(context_tensor)  # (1, 3, 64, 64)
+    residual = model(context_tensor)
     return torch.clamp(last_tensor + residual, 0, 1)
 
 
@@ -115,7 +115,6 @@ 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":
-        # AR prediction for Pong
         with torch.no_grad():
             context_tensor = torch.from_numpy(context).to(DEVICE)
             last_tensor = torch.from_numpy(last_frame_t).to(DEVICE)
@@ -127,7 +126,6 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         return predicted_np
 
     elif game == "sonic":
-        # Ensemble: AR + direct for Sonic with 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
@@ -146,7 +144,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
             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 for each step
             ar_preds = []
@@ -160,7 +158,6 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
                 ar_flip_back = torch.flip(ar_flip, dims=[3])
                 ar_frame = (ar_orig + ar_flip_back) / 2.0
                 ar_preds.append(ar_frame)
-                # Shift context for next AR step
                 ctx_frames = ctx.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
                 ctx_frames = torch.cat([ctx_frames[:, 1:], ar_orig.unsqueeze(1)], dim=1)
                 ctx = ctx_frames.reshape(1, -1, 64, 64)
@@ -170,9 +167,7 @@ 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)
-
-            # Ensemble: average AR and direct
+            ar_pred = torch.stack(ar_preds, dim=1)
             predicted = (ar_pred + direct_pred) / 2.0
 
         predicted_np = predicted[0].cpu().numpy()
@@ -187,7 +182,6 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         return result
 
     else:
-        # Direct 8-frame for PP 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