---
license: apache-2.0
library_name: pytorch
pipeline_tag: other
tags:
  - world-model
  - diffusion
  - pong
  - video
  - causal
  - game
---

# 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.

![demo](demo.gif)

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

```bash
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:

```bash
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.