---
license: mit
tags:
  - diffusion
  - world-model
  - diffusion-transformer
  - game
  - snake
  - webgpu
---

# nanoOasis — weights

Weights for **[nanoOasis](https://github.com/MaruthiV/nanoOasis)**, a from-scratch reference implementation
of a diffusion **world model** — the "nanoGPT" of the Oasis / GameNGen / DIAMOND paradigm. The model
generates the game **Snake** one frame at a time in response to your arrow keys: there is no game engine.

▶ **[Play the live demo](https://nano-oasis.vercel.app/)** · **[Code + docs](https://github.com/MaruthiV/nanoOasis)**

## Files

| file | what it is |
| --- | --- |
| `dit.pt` | the spatiotemporal **DiT** (13.5M params) — `model` + EMA weights (optimizer state stripped) |
| `vae.pt` | the **ViT-VAE** (7.6M params) — encodes a 256×192 frame to a latent and back |
| `onnx/dit.onnx` | the DiT denoiser exported to ONNX, **FP16** (for ONNX Runtime Web / WebGPU) |
| `onnx/vae_dec.onnx` | the VAE decoder exported to ONNX, FP16 |

## How it works

A ViT-VAE compresses each 256×192 frame to 48 latent tokens (the game is an 8×6 grid, so **one cell = one
DiT token**). A 13.5M-parameter spatiotemporal DiT predicts the next latent from the past 8 latents + your
action, trained with **EDM preconditioning + Diffusion Forcing + context-noise augmentation**. Four Euler
sampling steps make it real-time, so there's no distillation. At play time the model is the only thing
generating frames; a ~30-line deterministic "referee" adjudicates wall/self collisions (the discrete events
diffusion models handle unreliably).

Trained end-to-end for **under $50** on ~500k frames of bot-played Snake.

## Usage

```bash
git clone https://github.com/MaruthiV/nanoOasis && cd nanoOasis
pip install -e .

# play it locally (pygame window, arrow keys)
python infer.py --ckpt dit.pt --vae vae.pt --config small
```

The ONNX files are what the in-browser demo runs; see `export.py` and `demo/` in the repo.

## Intended use & limitations

A **reference implementation** for learning and forking the diffusion-world-model recipe — not a product.
It's a 13.5M-parameter model trained for ~$50, and it plays like one: crisp Snake for the first several
apples, then long-body coherence frays (diffusion models fumble long thin structures, and error accumulates
over a rollout). The demo re-seeds a clean context on death so each life starts fresh.

## License

MIT. Built in the lineage of [DIAMOND](https://arxiv.org/abs/2405.12399),
[GameNGen](https://arxiv.org/abs/2408.14837), and [Oasis](https://github.com/etched-ai/open-oasis),
in the spirit of [nanoGPT](https://github.com/karpathy/nanoGPT).