ddevin2002/tiny-self-forcing-doom

1"""Generation for the playable demo, with no UI framework in sight.2 3Kept separate from app.py so the part that can be wrong -- context bookkeeping, action4history, the rolling window -- is testable without installing Gradio. The UI is a thin5shell over `act`.6 7Design note for ZeroGPU: the GPU is released between calls, so nothing GPU-resident may8live in session state. A session carries twelve uint8 frames and twelve action ids on the9CPU and each step rebuilds the KV cache from them, which costs a few milliseconds and10lets the same code run unchanged on a laptop.11"""12 13from __future__ import annotations14 15import sys16import tempfile17import time18from pathlib import Path19 20import numpy as np21import torch22from PIL import Image23 24sys.path.insert(0, str(Path(__file__).resolve().parent))25 26from model.dit import DiTConfig, VideoDiT27from model.masks import chunk_causal28from sample.rollout import CausalSampler, to_uint829 30# ZeroGPU when deployed, a no-op when run anywhere else, so the file that gets tested is31# the file that gets shipped.32try:33    import spaces34    # 5, not 25: ZeroGPU charges the *declared* duration against the visitor's daily35    # quota, and one chunk takes tens of milliseconds. Declaring 25 s burned an36    # anonymous visitor's entire allowance in three presses. Shorter also means higher37    # queue priority.38    gpu = spaces.GPU(duration=5)39    ON_ZEROGPU = True40except Exception:                                    # noqa: BLE00141    def gpu(fn):42        return fn43    ON_ZEROGPU = False44 45# ZeroGPU wants the model placed on cuda at module level, not moved inside the decorated46# function: outside @spaces.GPU a CUDA emulation mode accepts the placement, and doing it47# per call is "significantly less efficient" because transfers are optimised for startup.48DEVICE = ("cuda" if ON_ZEROGPU or torch.cuda.is_available()49          else "mps" if torch.backends.mps.is_available() else "cpu")50 51HERE = Path(__file__).resolve().parent52CHUNK, STEPS, CTX = 4, 4, 1253DECISION_FPS = 8.75                                  # the rate the data was collected at54NOOP, FORWARD, LEFT, RIGHT, FWD_LEFT, FWD_RIGHT = range(6)55MAX_RECORD = 480                                     # ~55 s of game time, bounded memory56 57 58def load_model(path: Path | None = None) -> VideoDiT:59    ck = torch.load(path or HERE / "student_ema_fp16.pt", map_location="cpu",60                    weights_only=False)61    model = VideoDiT(DiTConfig(**ck["model_cfg"]))62    # stored fp16 to halve the download, run fp32 so the demo matches the evaluation,63    # which was measured without autocast64    model.load_state_dict({k: v.float() for k, v in ck["ema"].items()})65    model.set_mask(chunk_causal(model.cfg.frames, chunk=CHUNK))66    model.eval().requires_grad_(False)67    model.meta = {"step": ck.get("step"), "fingerprint": str(ck.get("fingerprint"))[:12]}68    return model69 70 71MODEL: VideoDiT | None = None72STARTS: np.ndarray | None = None73 74 75def init() -> None:76    """Load once, onto DEVICE. Called at app.py's module level so the placement happens77    at startup; kept out of import here so tests can stub the model without the weights."""78    global MODEL, STARTS79    if MODEL is None:80        MODEL = load_model().to(DEVICE)81    if STARTS is None:82        STARTS = np.load(HERE / "starts.npz")["frames"]83 84 85def new_session(which: int | None = None) -> dict:86    init()87    i = int(np.random.randint(len(STARTS))) if which is None else which % len(STARTS)88    frames = STARTS[i].copy()89    return {"frames": frames, "actions": [NOOP] * CTX, "steps": 0,90            "record": list(frames), "start": i, "ms": 0.0}91 92 93def _to_tensor(frames: np.ndarray, device) -> torch.Tensor:94    x = torch.from_numpy(np.ascontiguousarray(frames)).to(device)95    return (x.permute(0, 3, 1, 2).float() / 127.5 - 1.0)[None]96 97 98def gif(frames, scale: int = 6, ms_per_frame: int = 160) -> str:99    """160 ms per frame: ~1.4x slower than the data's 115 ms decision rate, on request --100    the demo reads better unhurried, and game time is reported in the status line."""101    frames = list(frames)102    h, w, _ = frames[0].shape103    imgs = [Image.fromarray(f).resize((w * scale, h * scale), Image.NEAREST)104            for f in frames]105    path = tempfile.NamedTemporaryFile(suffix=".gif", delete=False).name106    # No loop extension: the clip plays once and holds on its final frame. loop=0107    # meant "repeat forever", so every press replayed its 0.9 s of motion in a loop --108    # a turn snapped back and re-turned endlessly, which read as the world shaking.109    imgs[0].save(path, save_all=True, append_images=imgs[1:], duration=ms_per_frame)110    return path111 112 113@gpu114def generate(frames: np.ndarray, actions: list[int], plan: list[int]):115    """Generate len(plan) frames, frame i conditioned on plan[i].116    Returns (uint8 [len(plan),H,W,3], milliseconds).117 118    The loop is inside the decorated function on purpose: each call out to ZeroGPU is119    charged separately, so generating two chunks in one call costs half the quota of120    two calls. The visitor's daily allowance is the scarce resource here, not compute.121 122    A per-frame plan rather than one held action, because the model conditions each123    frame on its own action id. That is what lets a turn be a *tap* -- a chunk of124    turning followed by a chunk of standing -- instead of 0.9 s of key-held spinning,125    at identical quota cost.126    """127    init()128    if len(plan) % CHUNK:129        raise ValueError(f"plan length {len(plan)} is not a whole number of chunks")130    device = DEVICE131    sampler = CausalSampler(MODEL, chunk=CHUNK, steps_per_chunk=STEPS,132                            context_chunks=CTX // CHUNK)133    sampler.reset(batch=1, init_frames=_to_tensor(frames, device), device=device)134    hist = list(actions)135    out = []136    t0 = time.perf_counter()137    for c in range(len(plan) // CHUNK):138        chunk_plan = plan[c * CHUNK:(c + 1) * CHUNK]139        a = torch.tensor(chunk_plan, dtype=torch.long, device=device)[None]140        # the actions that actually produced the context, not the current one repeated:141        # the model conditions each context frame on its own action142        ctx_a = torch.tensor(hist[-CTX:], dtype=torch.long, device=device)[None]143        x = sampler.next_chunk(a, ctx_a)144        sampler.push(x, a)145        hist.extend(chunk_plan)146        out.append(x)147    if device == "cuda":148        torch.cuda.synchronize()149    ms = (time.perf_counter() - t0) * 1000.0150    return to_uint8(torch.cat(out, dim=1))[0], ms151 152 153# One decision step of turning pans the view ~6 px of the 64-px frame -- roughly 8154# degrees (90 degree FOV) -- and one step of walking covers a matching stride. Devin,155# playing: presses should be *small* -- a short step, a slight turn -- with the world156# settling in between, like nudging through the space rather than lunging. So every157# motion press is a short tap followed by stillness. TAP_FRAMES started at 2, but a158# measured session decayed from 0.036 to 0.028 Laplacian sharpness: during stillness159# the model regenerates a nearly-static scene over and over, low-passing it a little160# each round. Four frames of motion re-anchor the image with fresh geometry -- the same161# session at 4+4 *gained* sharpness (0.039 -> 0.057). So: ~33 degrees or a stride per162# press, then four frames of rest.163TAP_FRAMES = 4164 165 166def plan_for(action_id: int, n_chunks: int) -> list[int]:167    """Per-frame actions for one press: a short tap of motion, then stillness."""168    total = CHUNK * max(1, n_chunks)169    if action_id == NOOP:170        return [NOOP] * total171    return [action_id] * TAP_FRAMES + [NOOP] * (total - TAP_FRAMES)172 173 174def act(state: dict | None, action_id: int, n_chunks: int = 2) -> tuple[dict, str, str]:175    """Advance n_chunks. Turns are taps (turn, then settle); movement is held."""176    if state is None:177        state = new_session()178    n_chunks = max(1, int(n_chunks))179    plan = plan_for(action_id, n_chunks)180    new, ms = generate(state["frames"], state["actions"], plan)181    state = dict(state)182    keep = CTX183    state["frames"] = np.concatenate([state["frames"], new], axis=0)[-keep:]184    state["actions"] = (state["actions"] + plan)[-keep:]185    state["steps"] += n_chunks186    state["ms"] = ms187    state["chunks"] = n_chunks188    state["record"] = (state["record"] + list(new))[-MAX_RECORD:]189    return state, gif(new), status(state)190 191 192def status(state: dict) -> str:193    secs = state["steps"] * CHUNK / DECISION_FPS194    made = CHUNK * state.get("chunks", 1)195    fps = made / (state["ms"] / 1000.0) if state["ms"] else 0.0196    return (f"**{state['steps'] * CHUNK} frames generated** ({secs:.1f} s of game time)"197            f" · last press {state['ms']:.0f} ms, {fps:.1f} fps"198            f" · start #{state['start']}")199