Upload folder using huggingface_hub

predict.py

@@ -1,4 +1,4 @@
-"""Full PP swap: Pong direct int8, full PP model, Sonic AR fp16 + direct int8."""
+"""4-way TTA: original + hflip + vflip + both flips for Sonic and PP."""
 import sys
 import os
 import numpy as np
@@ -175,9 +175,6 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         for i in range(PRED_FRAMES):
             frame = np.transpose(predicted_np[i], (1, 2, 0))
             frame = (frame * 255).clip(0, 255).astype(np.uint8)
-            # Pong post-processing: threshold dark/bright pixels
-            frame[frame < 8] = 0
-            frame[frame > 247] = 255
             ens.direct_cache.append(frame)
 
         result = ens.direct_cache[ens.cache_step]
@@ -185,7 +182,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         return result
 
     elif game == "sonic":
-        # Sonic: AR(fp16)+direct(int8) with step blending and TTA
+        # Sonic: AR(fp16)+direct(int8) with step blending and 4-way 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
@@ -198,34 +195,49 @@ 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_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
-
-            ar_preds = []
-            ctx = context_tensor.clone()
-            ctx_flip = context_flipped.clone()
-            last_t = last_tensor.clone()
-            last_f = last_flipped.clone()
-            for step in range(PRED_FRAMES):
-                ar_orig = _predict_ar_frame(ens.sonic_ar, ctx, last_t)
-                ar_flip = _predict_ar_frame(ens.sonic_ar, ctx_flip, last_f)
-                ar_flip_back = torch.flip(ar_flip, dims=[3])
-                ar_frame = (ar_orig + ar_flip_back) / 2.0
-                ar_preds.append(ar_frame)
-                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)
-                last_t = ar_orig
-                ctx_flip_frames = ctx_flip.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
-                ctx_flip_frames = torch.cat([ctx_flip_frames[:, 1:], ar_flip.unsqueeze(1)], dim=1)
-                ctx_flip = ctx_flip_frames.reshape(1, -1, 64, 64)
-                last_f = ar_flip
-
-            ar_pred = torch.stack(ar_preds, dim=1)
+            # 4 augmentations: original, hflip, vflip, both
+            ctx_hf = torch.flip(context_tensor, dims=[3])
+            last_hf = torch.flip(last_tensor, dims=[3])
+            ctx_vf = torch.flip(context_tensor, dims=[2])
+            last_vf = torch.flip(last_tensor, dims=[2])
+            ctx_hv = torch.flip(context_tensor, dims=[2, 3])
+            last_hv = torch.flip(last_tensor, dims=[2, 3])
+
+            # Direct: 4-way TTA
+            d0 = _predict_8frames_direct(ens.sonic_direct, context_tensor, last_tensor)
+            d1 = torch.flip(_predict_8frames_direct(ens.sonic_direct, ctx_hf, last_hf), dims=[4])
+            d2 = torch.flip(_predict_8frames_direct(ens.sonic_direct, ctx_vf, last_vf), dims=[3])
+            d3 = torch.flip(_predict_8frames_direct(ens.sonic_direct, ctx_hv, last_hv), dims=[3, 4])
+            direct_pred = (d0 + d1 + d2 + d3) / 4.0
+
+            # AR: 4-way TTA
+            augs = [
+                (context_tensor.clone(), last_tensor.clone()),
+                (ctx_hf.clone(), last_hf.clone()),
+                (ctx_vf.clone(), last_vf.clone()),
+                (ctx_hv.clone(), last_hv.clone()),
+            ]
+            # flip dims for undoing: hflip=[3], vflip=[2], both=[2,3]
+            all_ar_preds = []
+            for aug_ctx, aug_last in augs:
+                ar_chain = []
+                c = aug_ctx
+                l = aug_last
+                for step in range(PRED_FRAMES):
+                    pred = _predict_ar_frame(ens.sonic_ar, c, l)
+                    ar_chain.append(pred)
+                    cf = c.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
+                    cf = torch.cat([cf[:, 1:], pred.unsqueeze(1)], dim=1)
+                    c = cf.reshape(1, -1, 64, 64)
+                    l = pred
+                all_ar_preds.append(torch.stack(ar_chain, dim=1))
+
+            # Undo flips
+            ar0 = all_ar_preds[0]
+            ar1 = torch.flip(all_ar_preds[1], dims=[4])
+            ar2 = torch.flip(all_ar_preds[2], dims=[3])
+            ar3 = torch.flip(all_ar_preds[3], dims=[3, 4])
+            ar_pred = (ar0 + ar1 + ar2 + ar3) / 4.0
 
             predicted = torch.zeros_like(direct_pred)
             for step in range(PRED_FRAMES):
@@ -245,7 +257,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         return result
 
     else:
-        # PP: direct with TTA and caching
+        # PP: direct with 4-way 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
@@ -258,12 +270,18 @@ 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)
 
-            predicted_orig = _predict_8frames_direct(ens.models["pole_position"], context_tensor, last_tensor)
-            context_flipped = torch.flip(context_tensor, dims=[3])
-            last_flipped = torch.flip(last_tensor, dims=[3])
-            predicted_flipped = _predict_8frames_direct(ens.models["pole_position"], context_flipped, last_flipped)
-            predicted_flipped = torch.flip(predicted_flipped, dims=[4])
-            predicted = (predicted_orig + predicted_flipped) / 2.0
+            ctx_hf = torch.flip(context_tensor, dims=[3])
+            last_hf = torch.flip(last_tensor, dims=[3])
+            ctx_vf = torch.flip(context_tensor, dims=[2])
+            last_vf = torch.flip(last_tensor, dims=[2])
+            ctx_hv = torch.flip(context_tensor, dims=[2, 3])
+            last_hv = torch.flip(last_tensor, dims=[2, 3])
+
+            p0 = _predict_8frames_direct(ens.models["pole_position"], context_tensor, last_tensor)
+            p1 = torch.flip(_predict_8frames_direct(ens.models["pole_position"], ctx_hf, last_hf), dims=[4])
+            p2 = torch.flip(_predict_8frames_direct(ens.models["pole_position"], ctx_vf, last_vf), dims=[3])
+            p3 = torch.flip(_predict_8frames_direct(ens.models["pole_position"], ctx_hv, last_hv), dims=[3, 4])
+            predicted = (p0 + p1 + p2 + p3) / 4.0
 
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []