Upload folder using huggingface_hub

pole_model.pt

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

pong_model.pt

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

predict.py

@@ -1,30 +1,63 @@
-"""Prediction interface for Multi-Scale Flow-Warp-Mask U-Net v10 with TTA."""
+"""Prediction interface for per-game Flow-Warp-Mask models v12 with motion encoding + TTA."""
 import sys
 import os
 import numpy as np
 import torch
 
 sys.path.insert(0, "/home/coder/code")
-from multiscale_flow_model import MultiScaleFlowUNet
+from flowmask_model import FlowWarpMaskUNet
 from flownet_model import differentiable_warp
 
 CONTEXT_LEN = 4
-CHANNELS = [56, 112, 224]
+GAME_CONFIGS = {
+    "pong": {"channels": [32, 64, 128], "file": "pong_model.pt"},
+    "sonic": {"channels": [40, 80, 160], "file": "sonic_model.pt"},
+    "pole_position": {"channels": [24, 48, 96], "file": "pole_model.pt"},
+}
+
+
+def detect_game(context_frames):
+    mean_val = context_frames.mean()
+    if mean_val < 10:
+        return "pong"
+    elif mean_val < 80:
+        return "sonic"
+    else:
+        return "pole_position"
 
 
 def load_model(model_dir: str):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = MultiScaleFlowUNet(in_channels=12, channels=CHANNELS)
-    model_path = os.path.join(model_dir, "model.pt")
-    state_dict = torch.load(model_path, map_location=device, weights_only=True)
-    state_dict = {k: v.float() for k, v in state_dict.items()}
-    model.load_state_dict(state_dict)
-    model.to(device)
-    model.eval()
-    return {"model": model, "device": device}
-
-
-def _prepare_input(context_frames):
+    models = {}
+    for game, cfg in GAME_CONFIGS.items():
+        model = FlowWarpMaskUNet(in_channels=12, channels=cfg["channels"])
+        model_path = os.path.join(model_dir, cfg["file"])
+        state_dict = torch.load(model_path, map_location=device, weights_only=True)
+        state_dict = {k: v.float() for k, v in state_dict.items()}
+        model.load_state_dict(state_dict)
+        model.to(device)
+        model.eval()
+        models[game] = model
+    return {"models": models, "device": device}
+
+
+def _make_motion_input(frames):
+    """Create motion encoding: last frame (3ch) + 3 pairwise diffs (9ch) = 12ch.
+
+    Args:
+        frames: (4, 3, H, W) tensor in [0,1]
+    Returns:
+        (12, H, W) tensor
+    """
+    last = frames[-1]           # (3, H, W)
+    diff1 = frames[-1] - frames[-2]  # most recent motion
+    diff2 = frames[-2] - frames[-3]  # previous motion
+    diff3 = frames[-3] - frames[-4]  # older motion
+    return torch.cat([last, diff1, diff2, diff3], dim=0)  # (12, H, W)
+
+
+def _prepare_context(context_frames):
+    """Prepare 4-frame context from numpy frames."""
     if len(context_frames) >= CONTEXT_LEN:
         frames = context_frames[-CONTEXT_LEN:]
     else:
@@ -33,20 +66,19 @@ def _prepare_input(context_frames):
         frames = np.concatenate([padding, context_frames], axis=0)
 
     frames_t = torch.from_numpy(frames.astype(np.float32) / 255.0)
-    frames_t = frames_t.permute(0, 3, 1, 2)
+    frames_t = frames_t.permute(0, 3, 1, 2)  # (4, 3, 64, 64)
     return frames_t
 
 
 def _run_model(model, frames_t, device):
-    last_frame = frames_t[-1].unsqueeze(0)
-    inp = frames_t.reshape(1, -1, 64, 64)
+    """Run model with motion encoding input."""
+    last_frame = frames_t[-1].unsqueeze(0)  # (1, 3, 64, 64)
+    inp = _make_motion_input(frames_t).unsqueeze(0)  # (1, 12, 64, 64)
 
     inp = inp.to(device)
     last_frame = last_frame.to(device)
 
-    flows, mask, gen_frame = model(inp)
-    # Use finest flow (last element)
-    flow = flows[-1]
+    flow, mask, gen_frame = model(inp)
     warped = differentiable_warp(last_frame, flow)
     pred = mask * warped + (1 - mask) * gen_frame
     pred = torch.clamp(pred, 0, 1)
@@ -54,10 +86,13 @@ def _run_model(model, frames_t, device):
 
 
 def predict_next_frame(model_dict, context_frames: np.ndarray) -> np.ndarray:
-    model = model_dict["model"]
+    models = model_dict["models"]
     device = model_dict["device"]
 
-    frames_t = _prepare_input(context_frames)
+    game = detect_game(context_frames)
+    model = models[game]
+
+    frames_t = _prepare_context(context_frames)
 
     with torch.no_grad():
         # Original prediction
@@ -73,4 +108,9 @@ def predict_next_frame(model_dict, context_frames: np.ndarray) -> np.ndarray:
 
     pred = pred[0].cpu().permute(1, 2, 0).numpy()
     pred = (pred * 255).clip(0, 255).astype(np.uint8)
+
+    # Post-processing for Pong: clamp dark pixels to pure black
+    if game == "pong":
+        pred[pred < 5] = 0
+
     return pred

sonic_model.pt

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

train.log

@@ -1,24 +1,131 @@
-[12:02:26] Device: cuda
-[12:02:27] Loaded v10 weights from /home/coder/experiments/2026-04-14-080000-multiscale-flow-v10
-[12:02:27] Model parameters: 6,169,586, channels=[56, 112, 224]
-[12:02:27] Fine-tune: 10 epochs of 8-step AR with pure SSIM loss
-[12:02:31]   43855 sequences
-[12:15:01] Epoch 1/10 | loss=0.09345 lr=0.0000098
-[12:27:32] Epoch 2/10 | loss=0.09245 lr=0.0000091
-[12:29:02]   Val SSIM=0.8883 | {'pong': 0.8811, 'sonic': 0.8354, 'pole_position': 0.9485}
-[12:29:02]   New best! SSIM=0.8883
-[12:41:38] Epoch 3/10 | loss=0.09166 lr=0.0000081
-[12:54:16] Epoch 4/10 | loss=0.09095 lr=0.0000069
-[12:55:36]   Val SSIM=0.8887 | {'pong': 0.8824, 'sonic': 0.8352, 'pole_position': 0.9486}
-[12:55:36]   New best! SSIM=0.8887
-[13:08:12] Epoch 5/10 | loss=0.09033 lr=0.0000055
-[13:20:56] Epoch 6/10 | loss=0.08982 lr=0.0000041
-[13:22:21]   Val SSIM=0.8885 | {'pong': 0.8824, 'sonic': 0.8347, 'pole_position': 0.9483}
-[13:35:14] Epoch 7/10 | loss=0.08941 lr=0.0000029
-[13:48:14] Epoch 8/10 | loss=0.08911 lr=0.0000019
-[13:49:31]   Val SSIM=0.8883 | {'pong': 0.8822, 'sonic': 0.8344, 'pole_position': 0.9484}
-[14:02:17] Epoch 9/10 | loss=0.08888 lr=0.0000012
-[14:14:47] Epoch 10/10 | loss=0.08874 lr=0.0000010
-[14:16:07]   Val SSIM=0.8881 | {'pong': 0.8815, 'sonic': 0.8343, 'pole_position': 0.9485}
-[14:16:07] Experiment dir: 12.4 MB
-[14:16:07] Training complete. Best val SSIM: 0.8887
+[14:35:51] Device: cuda
+[14:35:51] 
+=== Training pong ([32, 64, 128]) ===
+[14:35:51]   2,018,278 parameters
+[14:35:51]   Phase 1: 10 epochs single-step
+[14:35:51]   8568 sequences
+[14:36:00]   P1 pong Epoch 1/10 | loss=0.14558
+[14:36:08]   P1 pong Epoch 2/10 | loss=0.10721
+[14:36:17]   P1 pong Epoch 3/10 | loss=0.09795
+[14:36:25]   P1 pong Epoch 4/10 | loss=0.08996
+[14:36:33]   P1 pong Epoch 5/10 | loss=0.08384
+[14:36:41]   P1 pong Epoch 6/10 | loss=0.07755
+[14:36:49]   P1 pong Epoch 7/10 | loss=0.06995
+[14:36:57]   P1 pong Epoch 8/10 | loss=0.06272
+[14:37:05]   P1 pong Epoch 9/10 | loss=0.05640
+[14:37:13]   P1 pong Epoch 10/10 | loss=0.05177
+[14:37:13]   Phase 2: 25 epochs graduated AR
+[14:37:37]   P2 pong Epoch 1/25 (steps=2) | loss=0.09787 lr=0.000500
+[14:37:59]   P2 pong Epoch 2/25 (steps=2) | loss=0.08854 lr=0.000500
+[14:38:21]   P2 pong Epoch 3/25 (steps=2) | loss=0.08343 lr=0.000500
+[14:39:15]   P2 pong Epoch 4/25 (steps=4) | loss=0.13928 lr=0.000500
+[14:40:08]   P2 pong Epoch 5/25 (steps=4) | loss=0.12631 lr=0.000500
+[14:41:04]   P2 pong Epoch 6/25 (steps=4) | loss=0.11644 lr=0.000500
+[14:43:21]   P2 pong Epoch 7/25 (steps=8) | loss=0.18012 lr=0.000500
+[14:45:38]   P2 pong Epoch 8/25 (steps=8) | loss=0.17484 lr=0.000500
+[14:47:57]   P2 pong Epoch 9/25 (steps=8) | loss=0.16717 lr=0.000500
+[14:50:15]   P2 pong Epoch 10/25 (steps=8) | loss=0.15650 lr=0.000500
+[14:52:31]   P2 pong Epoch 11/25 (steps=8) | loss=0.14624 lr=0.000500
+[14:54:46]   P2 pong Epoch 12/25 (steps=8) | loss=0.13932 lr=0.000500
+[14:57:01]   P2 pong Epoch 13/25 (steps=8) | loss=0.12899 lr=0.000493
+[14:59:17]   P2 pong Epoch 14/25 (steps=8) | loss=0.11960 lr=0.000471
+[15:01:35]   P2 pong Epoch 15/25 (steps=8) | loss=0.10872 lr=0.000437
+[15:03:52]   P2 pong Epoch 16/25 (steps=8) | loss=0.09965 lr=0.000392
+[15:06:07]   P2 pong Epoch 17/25 (steps=8) | loss=0.08785 lr=0.000339
+[15:08:27]   P2 pong Epoch 18/25 (steps=8) | loss=0.07890 lr=0.000280
+[15:10:44]   P2 pong Epoch 19/25 (steps=8) | loss=0.06718 lr=0.000220
+[15:13:01]   P2 pong Epoch 20/25 (steps=8) | loss=0.06123 lr=0.000161
+[15:15:20]   P2 pong Epoch 21/25 (steps=8) | loss=0.05374 lr=0.000108
+[15:17:40]   P2 pong Epoch 22/25 (steps=8) | loss=0.04863 lr=0.000063
+[15:19:57]   P2 pong Epoch 23/25 (steps=8) | loss=0.04435 lr=0.000029
+[15:22:13]   P2 pong Epoch 24/25 (steps=8) | loss=0.04174 lr=0.000010
+[15:24:31]   P2 pong Epoch 25/25 (steps=8) | loss=0.04022 lr=0.000010
+[15:24:31]   pong training complete.
+[15:24:31] 
+=== Training sonic ([40, 80, 160]) ===
+[15:24:31]   3,150,686 parameters
+[15:24:31]   Phase 1: 10 epochs single-step
+[15:24:34]   32256 sequences
+[15:25:03]   P1 sonic Epoch 1/10 | loss=0.08400
+[15:25:34]   P1 sonic Epoch 2/10 | loss=0.06966
+[15:26:03]   P1 sonic Epoch 3/10 | loss=0.06589
+[15:26:34]   P1 sonic Epoch 4/10 | loss=0.06327
+[15:27:03]   P1 sonic Epoch 5/10 | loss=0.06111
+[15:27:33]   P1 sonic Epoch 6/10 | loss=0.05881
+[15:28:03]   P1 sonic Epoch 7/10 | loss=0.05682
+[15:28:33]   P1 sonic Epoch 8/10 | loss=0.05514
+[15:29:02]   P1 sonic Epoch 9/10 | loss=0.05358
+[15:29:32]   P1 sonic Epoch 10/10 | loss=0.05256
+[15:29:32]   Phase 2: 25 epochs graduated AR
+[15:30:57]   P2 sonic Epoch 1/25 (steps=2) | loss=0.07446 lr=0.000500
+[15:32:15]   P2 sonic Epoch 2/25 (steps=2) | loss=0.07291 lr=0.000500
+[15:33:41]   P2 sonic Epoch 3/25 (steps=2) | loss=0.07128 lr=0.000500
+[15:37:15]   P2 sonic Epoch 4/25 (steps=4) | loss=0.10220 lr=0.000500
+[15:40:50]   P2 sonic Epoch 5/25 (steps=4) | loss=0.09976 lr=0.000500
+[15:44:24]   P2 sonic Epoch 6/25 (steps=4) | loss=0.09779 lr=0.000500
+[15:53:05]   P2 sonic Epoch 7/25 (steps=8) | loss=0.14037 lr=0.000500
+[16:01:41]   P2 sonic Epoch 8/25 (steps=8) | loss=0.13753 lr=0.000500
+[16:10:26]   P2 sonic Epoch 9/25 (steps=8) | loss=0.13476 lr=0.000500
+[16:19:08]   P2 sonic Epoch 10/25 (steps=8) | loss=0.13232 lr=0.000500
+[16:28:05]   P2 sonic Epoch 11/25 (steps=8) | loss=0.13010 lr=0.000500
+[16:37:18]   P2 sonic Epoch 12/25 (steps=8) | loss=0.12790 lr=0.000500
+[16:46:19]   P2 sonic Epoch 13/25 (steps=8) | loss=0.12592 lr=0.000493
+[16:55:21]   P2 sonic Epoch 14/25 (steps=8) | loss=0.12408 lr=0.000471
+[17:04:34]   P2 sonic Epoch 15/25 (steps=8) | loss=0.12210 lr=0.000437
+[17:13:54]   P2 sonic Epoch 16/25 (steps=8) | loss=0.11900 lr=0.000392
+[17:23:04]   P2 sonic Epoch 17/25 (steps=8) | loss=0.11596 lr=0.000339
+[17:32:08]   P2 sonic Epoch 18/25 (steps=8) | loss=0.11287 lr=0.000280
+[17:41:13]   P2 sonic Epoch 19/25 (steps=8) | loss=0.10939 lr=0.000220
+[17:50:18]   P2 sonic Epoch 20/25 (steps=8) | loss=0.10548 lr=0.000161
+[17:59:23]   P2 sonic Epoch 21/25 (steps=8) | loss=0.10183 lr=0.000108
+[18:08:26]   P2 sonic Epoch 22/25 (steps=8) | loss=0.09841 lr=0.000063
+[18:17:35]   P2 sonic Epoch 23/25 (steps=8) | loss=0.09526 lr=0.000029
+[18:26:41]   P2 sonic Epoch 24/25 (steps=8) | loss=0.09337 lr=0.000010
+[18:35:42]   P2 sonic Epoch 25/25 (steps=8) | loss=0.09193 lr=0.000010
+[18:35:42]   sonic training complete.
+[18:35:42] 
+=== Training pole_position ([24, 48, 96]) ===
+[18:35:42]   1,137,006 parameters
+[18:35:42]   Phase 1: 10 epochs single-step
+[18:35:42]   4284 sequences
+[18:35:46]   P1 pole_position Epoch 1/10 | loss=0.05831
+[18:35:50]   P1 pole_position Epoch 2/10 | loss=0.03691
+[18:35:54]   P1 pole_position Epoch 3/10 | loss=0.03064
+[18:35:57]   P1 pole_position Epoch 4/10 | loss=0.02707
+[18:36:00]   P1 pole_position Epoch 5/10 | loss=0.02428
+[18:36:04]   P1 pole_position Epoch 6/10 | loss=0.02271
+[18:36:07]   P1 pole_position Epoch 7/10 | loss=0.02128
+[18:36:11]   P1 pole_position Epoch 8/10 | loss=0.02013
+[18:36:15]   P1 pole_position Epoch 9/10 | loss=0.01936
+[18:36:19]   P1 pole_position Epoch 10/10 | loss=0.01879
+[18:36:19]   Phase 2: 25 epochs graduated AR
+[18:36:31]   P2 pole_position Epoch 1/25 (steps=2) | loss=0.02742 lr=0.000500
+[18:36:42]   P2 pole_position Epoch 2/25 (steps=2) | loss=0.02621 lr=0.000500
+[18:36:54]   P2 pole_position Epoch 3/25 (steps=2) | loss=0.02502 lr=0.000500
+[18:37:22]   P2 pole_position Epoch 4/25 (steps=4) | loss=0.03779 lr=0.000500
+[18:37:51]   P2 pole_position Epoch 5/25 (steps=4) | loss=0.03543 lr=0.000500
+[18:38:19]   P2 pole_position Epoch 6/25 (steps=4) | loss=0.03421 lr=0.000500
+[18:39:31]   P2 pole_position Epoch 7/25 (steps=8) | loss=0.05263 lr=0.000500
+[18:40:42]   P2 pole_position Epoch 8/25 (steps=8) | loss=0.05159 lr=0.000500
+[18:41:53]   P2 pole_position Epoch 9/25 (steps=8) | loss=0.04987 lr=0.000500
+[18:43:05]   P2 pole_position Epoch 10/25 (steps=8) | loss=0.04848 lr=0.000500
+[18:44:17]   P2 pole_position Epoch 11/25 (steps=8) | loss=0.04744 lr=0.000500
+[18:45:30]   P2 pole_position Epoch 12/25 (steps=8) | loss=0.04603 lr=0.000500
+[18:46:42]   P2 pole_position Epoch 13/25 (steps=8) | loss=0.04495 lr=0.000493
+[18:47:54]   P2 pole_position Epoch 14/25 (steps=8) | loss=0.04383 lr=0.000471
+[18:49:05]   P2 pole_position Epoch 15/25 (steps=8) | loss=0.04233 lr=0.000437
+[18:50:18]   P2 pole_position Epoch 16/25 (steps=8) | loss=0.04089 lr=0.000392
+[18:51:30]   P2 pole_position Epoch 17/25 (steps=8) | loss=0.03911 lr=0.000339
+[18:52:43]   P2 pole_position Epoch 18/25 (steps=8) | loss=0.03667 lr=0.000280
+[18:53:55]   P2 pole_position Epoch 19/25 (steps=8) | loss=0.03494 lr=0.000220
+[18:55:06]   P2 pole_position Epoch 20/25 (steps=8) | loss=0.03271 lr=0.000161
+[18:56:18]   P2 pole_position Epoch 21/25 (steps=8) | loss=0.03049 lr=0.000108
+[18:57:31]   P2 pole_position Epoch 22/25 (steps=8) | loss=0.02831 lr=0.000063
+[18:58:44]   P2 pole_position Epoch 23/25 (steps=8) | loss=0.02653 lr=0.000029
+[18:59:58]   P2 pole_position Epoch 24/25 (steps=8) | loss=0.02527 lr=0.000010
+[19:01:11]   P2 pole_position Epoch 25/25 (steps=8) | loss=0.02460 lr=0.000010
+[19:01:11]   pole_position training complete.
+[19:01:11] Evaluating...
+[19:02:25] Val SSIM=0.8626 | {'pong': 0.862, 'sonic': 0.7822, 'pole_position': 0.9435}
+[19:02:25] Experiment dir: 12.7 MB
+[19:02:25] Training complete.