ddevin2002/tiny-self-forcing-doom

1"""Playable neural DOOM -- a Gradio shell over engine.py.2 3Everything that can be wrong about the generation lives in engine.py and is tested4without Gradio installed. This file is layout and key bindings.5"""6 7from __future__ import annotations8 9import gradio as gr10 11import engine12from engine import FORWARD, FWD_LEFT, FWD_RIGHT, LEFT, NOOP, RIGHT13 14# At module level on purpose: ZeroGPU optimises CUDA placement done at startup, and15# moving the model inside the decorated function costs a transfer on every keypress.16engine.init()17 18CSS = """19.gradio-container{max-width:1020px !important}20#screen img{image-rendering:pixelated; border-radius:4px}21#pad button{font-size:15px; min-height:52px}22.small{font-size:13px; opacity:.78; line-height:1.6}23"""24 25# Gradio has no keyboard primitive; bind the physical keys to the buttons that already26# exist so the demo feels like a game rather than a form.27KEYS_JS = """28() => {29  const map = {ArrowUp:'b-fwd', KeyW:'b-fwd', ArrowLeft:'b-left', KeyA:'b-left',30               ArrowRight:'b-right', KeyD:'b-right', Space:'b-noop',31               KeyQ:'b-fl', KeyE:'b-fr', KeyR:'b-restart'};32  if (window.__doomKeys) return;33  window.__doomKeys = true;34  document.addEventListener('keydown', (e) => {35    const id = map[e.code];36    if (!id || e.repeat) return;37    e.preventDefault();38    const el = document.getElementById(id);39    if (el) el.click();40  });41}42"""43 44 45def restart(_state=None):46    state = engine.new_session()47    return state, engine.gif(state["frames"][-engine.CHUNK:]), engine.status(state)48 49 50def save_clip(state):51    if state is None or len(state["record"]) < 8:52        return gr.update(value=None, visible=False)53    return gr.update(value=engine.gif(state["record"]), visible=True)54 55 56# Gradio 6 moved theme and css from the Blocks constructor to launch().57with gr.Blocks(title="tiny-self-forcing · playable neural DOOM") as demo:58    gr.Markdown(59        "# Playable neural DOOM\n"60        "**No game engine is running.** Every frame is generated by a 58M-parameter "61        "diffusion model, conditioned on the key you press and on the last twelve frames "62        "of its own output. Each press is a small move — a short step or a slight turn — "63        "and the world settles for the rest of the clip, four denoising steps per frame."64    )65    state = gr.State(None)66 67    with gr.Row():68        with gr.Column(scale=3):69            screen = gr.Image(elem_id="screen", show_label=False, height=384)70            status = gr.Markdown("", elem_classes="small")71        with gr.Column(scale=2):72            gr.Markdown("### Controls\nArrow keys or WASD work directly. Q and E veer "73                        "while moving, space stands still, R picks a new start. Every "74                        "press is a nudge: a stride or a ~33\u00b0 turn, then the "75                        "world settles.",76                        elem_classes="small")77            with gr.Column(elem_id="pad"):78                b_fwd = gr.Button("↑　forward (hold)", elem_id="b-fwd")79                with gr.Row():80                    b_left = gr.Button("←\u3000left", elem_id="b-left")81                    b_right = gr.Button("→\u3000right", elem_id="b-right")82                with gr.Row():83                    b_fl = gr.Button("↖\u3000fwd + left", elem_id="b-fl")84                    b_fr = gr.Button("↗\u3000fwd + right", elem_id="b-fr")85                b_noop = gr.Button("·\u3000stand still", elem_id="b-noop")86                b_restart = gr.Button("↻\u3000new start", elem_id="b-restart",87                                      variant="secondary")88            b_save = gr.Button("Save this run as a GIF", variant="primary")89            clip = gr.File(label="your recording", visible=False)90 91    gr.Markdown(92        "### What it gets wrong, and why\n"93        "The context window is **16 frames, about 1.8 seconds**. Turn away from a room "94        "and back and the room is **re-imagined** as something else — the model has no "95        "longer memory. That is a boundary of this design rather than a bug, and it is "96        "what long-memory work such as LingBot-World 2.0 addresses.\n\n"97        "The first twelve frames of a session are real ViZDoom renders. **From frame 13 "98        "on, everything you see is generated.** Before distillation each chunk took 48 "99        "denoising steps; after it, 4 — 557 ms to 69 ms on an RTX 4090, **8.1× faster**, "100        "which is the distance between not interactive and interactive.\n\n"101        "Training, the ablation and the full evaluation: "102        "[github.com/devt287/tiny-self-forcing](https://github.com/devt287/tiny-self-forcing)"103        "\n\n"104        "**On quota**: this runs on Hugging Face's free shared GPU pool, and the daily "105        "allowance is per visitor. Anonymous visitors get 2 minutes at the lowest queue "106        "priority and will hit *exceeded your ZeroGPU runs limit* after a while; "107        "**signing in to any Hugging Face account gets you much further**. The first "108        "press waits ~2 s for a cold GPU allocation, then about 0.9 s each.",109        elem_classes="small")110 111    for btn, aid in ((b_fwd, FORWARD), (b_left, LEFT), (b_right, RIGHT),112                     (b_fl, FWD_LEFT), (b_fr, FWD_RIGHT), (b_noop, NOOP)):113        btn.click(engine.act, [state, gr.State(aid)], [state, screen, status],114                  show_progress="hidden")115    b_restart.click(restart, [state], [state, screen, status], show_progress="hidden")116    b_save.click(save_clip, [state], [clip])117    demo.load(restart, [state], [state, screen, status]).then(None, None, None, js=KEYS_JS)118 119if __name__ == "__main__":120    demo.queue(max_size=32).launch(css=CSS, theme=gr.themes.Base())121