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

# Diffusion Breakout

An action-conditioned Breakout world model. Hold **A / D** (or the arrow keys) and the model generates the next frame.

![demo](demo.gif)

128x128 at 6 FPS with 12 frames of latent history. The model has about 70M parameters and uses the frozen SDXL VAE from `madebyollin/sdxl-vae-fp16-fix`.

## Run it

Requires a CUDA GPU with BF16 support.

```bash
pip install torch numpy pillow safetensors huggingface_hub diffusers

hf download kerzgrr/diffusionbreakout live_infer.py --local-dir .
python live_infer.py --steps 4
```

To keep the weights in a local folder:

```bash
hf download kerzgrr/diffusionbreakout \
  --local-dir checkpoints/diffusionbreakout \
  --include "ema.safetensors" \
  --include "config.json" \
  --include "live_infer.py"

python checkpoints/diffusionbreakout/live_infer.py \
  --local-dir checkpoints/diffusionbreakout \
  --steps 4 \
  --window-scale 6
```

Click the window for focus. **A / Left** moves left; **D / Right** moves right.

## Files

| file | contents |
|---|---|
| `ema.safetensors` | action-conditioned EMA weights |
| `config.json` | model, video, codec, and training settings |
| `live_infer.py` | single-file live inference |
| `demo.gif` | live rollout |

Actions are encoded as `[left, right]`.