Upload folder using huggingface_hub

model_sonic_ar.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0853b8b0dad0a55f126be9bfd767d2e55fcc2ea9dcb379a79f6389c997e54816
-size 3129452
+oid sha256:413d9fcfa15f30c74cdfda5f7d7c9dba8958fe027dfc09de563e6209c78378f5
+size 6180566

predict.py

@@ -1,4 +1,4 @@
-"""Int8 ensemble: Sonic models quantized to int8, Pong/PP in fp16."""
+"""Selective int8: only Sonic direct quantized to int8, Sonic AR kept in fp16."""
 import sys
 import os
 import numpy as np
@@ -75,12 +75,13 @@ def load_model(model_dir: str):
     pong_direct.eval()
     ens.pong_direct = pong_direct
 
-    # Sonic AR (int8 quantized, 3 outputs)
+    # Sonic AR (fp16, 3 outputs) - kept in fp16 for AR chain quality
     sonic_ar = UNet(in_channels=24, out_channels=3,
                     enc_channels=(48, 96, 192), bottleneck_channels=256,
                     upsample_mode="bilinear").to(DEVICE)
-    sd = load_int8_state_dict(os.path.join(model_dir, "model_sonic_ar.pt"), DEVICE)
-    sonic_ar.load_state_dict(sd)
+    sd = torch.load(os.path.join(model_dir, "model_sonic_ar.pt"),
+                    map_location=DEVICE, weights_only=True)
+    sonic_ar.load_state_dict({k: v.float() for k, v in sd.items()})
     sonic_ar.eval()
     ens.sonic_ar = sonic_ar
 
@@ -182,7 +183,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         return result
 
     elif game == "sonic":
-        # Sonic: AR+direct with step blending and TTA
+        # Sonic: AR(fp16)+direct(int8) 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