kerzgrr/diffusionpong

Diffusion Pong

A small action-conditioned world model that learned to keep a Pong match going. You hold W / S, it invents the next frame.

128x128 at 6 FPS with a 12-frame latent history. Codec is a frozen SDXL VAE (madebyollin/sdxl-vae-fp16-fix). Needs a CUDA GPU with BF16.

Play it

pip install torch numpy pillow safetensors huggingface_hub diffusers

hf download kerzgrr/diffusionpong live_infer.py --local-dir .
python live_infer.py --steps 2

Or pull weights into a folder yourself:

hf download kerzgrr/diffusionpong --local-dir checkpoints/diffusionpong --include "ema.safetensors" --include "config.json" --include "live_infer.py"
python checkpoints/diffusionpong/live_infer.py --local-dir checkpoints/diffusionpong --steps 2 --window-scale 7 --seed 42

Controls: W up, S down (A/D also work). Click the window once so it has focus. Clicking the canvas drops a yellow ball into the latent history for a few frames if you want to poke it.

What's in the files

file	what
ema.safetensors	playable weights
config.json	video / model / codec settings used at train time
live_infer.py	live play script
demo.gif	live rollout clip

Action vector is [W, S] as floats in {0,1}. Hold both and it treats that as neutral.

If something breaks, it's probably the VAE download or a GPU that can't do BF16.