kerzgrr/diffusionbreakout

Diffusion Breakout

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

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.

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:

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