STEVE-1 v1: style-conditioned terrain VAE (ONNX + mappings + base_v1 fine-tuning checkpoint)

README.md

@@ -1,3 +1,108 @@
 ---
 license: mit
+tags:
+- minecraft
+- terrain-generation
+- vae
+- onnx
+- voxel
+- film
+- transfer-learning
+- games
+datasets:
+- ghosteau/minecraft-chunks
+pipeline_tag: other
 ---
+
+# STEVE-1 — Style-conditioned Terrain Encoder VAE
+
+STEVE-1 generates full Minecraft chunks (16 x 384 x 16 voxels, bedrock to build limit) from **latent noise + a biome ID + a style ID** in a **single ONNX forward pass** — fast enough to run inside a live PaperMC server. It is also a **pretrained base for transfer learning**: anyone can fine-tune a new terrain *style* onto it from a folder of their own world's chunks, in minutes, training ~1.8% of the weights, without touching the base.
+
+- **Training pipeline, plugin, and fine-tuning notebook:** [github.com/ghosteau/generative-terrain](https://github.com/ghosteau/generative-terrain)
+- **Training data:** [ghosteau/minecraft-chunks](https://huggingface.co/datasets/ghosteau/minecraft-chunks) (2,040 chunks, ~200M voxels, 22 biomes)
+
+## Architecture
+
+A conditional VAE built around one hard constraint: in-game inference must be a single fast graph call (no diffusion, no autoregression).
+
+| Component | What it does |
+|---|---|
+| Spatial latent `z` `[8, 4, 16, 4]` | A coarse 3D latent grid instead of one broadcast vector, so terrain varies from place to place (hills, valleys, caves) |
+| Heightmap stage | The decoder first predicts a per-column surface height, then feeds a signed-distance-to-surface channel to the voxel head — the model commits to a surface instead of hedging |
+| 3-level 3D U-Net (~4.6M params) | Chunk-scale receptive field: at the bottleneck the chunk is 2 x 48 x 2, so horizontal context is global |
+| FiLM style conditioning | Every U-Net block and the heightmap head is modulated (per-channel scale/shift) by a style embedding. Style 0 = vanilla; 8 slots reserved for custom styles. This is the fine-tuning mechanism |
+| 27 block groups | The model predicts semantic groups (`STONE`, `GRASS`, `WATER`, ores, ...); the plugin expands them to concrete blocks with local context |
+
+Conditioning uses **only what exists before terrain does** (position, biome, noise, style). It deliberately never sees neighbour blocks, light, or surface flags — the data-leakage mistake that sank earlier versions of this project.
+
+Training: EMA weights (exported), mixed precision, warmup + cosine LR, KL annealing with free bits, biome-balanced sampling, and biome dropout so a reserved `UNKNOWN` biome row learns generic terrain (that is what makes fine-tuning work on custom worlds with unrecognised biomes).
+
+## Files
+
+| File | Purpose |
+|---|---|
+| `terrain_vae_decoder.onnx` | The generator: `(z, biome_id, style_id) -> logits [1, 27, 16, 384, 16]` |
+| `block_group_mapping.json` | Class id -> block group name |
+| `biome_id_mapping.json` | Biome name -> id (includes `UNKNOWN`) |
+| `style_mapping.json` | Style name -> id (`BASE` = 0) |
+| `base_v1/` | The versioned base checkpoint for fine-tuning: PyTorch weights + a manifest with the exact architecture and mappings |
+
+## Use it in a Minecraft server
+
+Drop `terrain_vae_decoder.onnx` and the three JSON mappings into `<server>/plugins/GenerativeTerrain/` (plugin jar from the [GitHub repo](https://github.com/ghosteau/generative-terrain)), then in game:
+
+```
+/generateterrain                  base style, current chunk
+/generateterrain style <name>     a fine-tuned custom style
+```
+
+## Use it from Python
+
+```python
+import numpy as np
+import onnxruntime as ort
+from huggingface_hub import hf_hub_download
+
+sess = ort.InferenceSession(hf_hub_download("ghosteau/STEVE-1", "terrain_vae_decoder.onnx"))
+
+logits = sess.run(["logits"], {
+    "z":        np.random.randn(1, 8, 4, 16, 4).astype(np.float32),  # the variety
+    "biome_id": np.array([15], dtype=np.int64),                      # PLAINS (see biome_id_mapping.json)
+    "style_id": np.zeros(1, dtype=np.int64),                         # 0 = BASE
+})[0]
+
+chunk = logits.argmax(1)[0]        # [16, 384, 16] block-group ids, Y index 0 = world Y -64
+```
+
+Sample a fresh `z` per chunk; scale it (`z * 0.8`) for tamer terrain, (`z * 1.2`) for wilder.
+
+## Fine-tune it on your own terrain
+
+This is the headline feature. The `base_v1/` directory is a self-contained transfer-learning checkpoint: weights plus a manifest that the pipeline's `fine_tune()` uses to rebuild the exact architecture, so version mismatches are impossible.
+
+1. Export 50+ chunks of your world with the plugin: `/grabchunkarea 5`
+2. Clone the [GitHub repo](https://github.com/ghosteau/generative-terrain), open `ml/notebooks/FineTune.ipynb`
+3. Point it at `base_v1/`, your CSV folder, and a style name; run it
+
+Fine-tuning trains the new style vector and the FiLM projections (~84k of 4.6M parameters) while everything else stays frozen — a unit test in the repo asserts the base weights come out bit-identical. The exported ONNX then carries both the vanilla style and yours, selectable per command. Up to 8 custom styles fit in one model.
+
+## Training results
+
+Vertical block-distribution profile, real vs generated (the terrain-plausibility check — bedrock floor, deepslate/stone body, thin surface band, air above):
+
+![Real vs generated vertical profiles](assets/profiles.png)
+
+Training curves (EMA-validated loss):
+
+![Training curves](assets/vae_curves.png)
+
+## Limitations
+
+- **Chunks generate independently** — edges between adjacent generated chunks do not line up yet. Cross-chunk edge conditioning is the next planned lever.
+- **Rare biomes are undertrained**: the dataset has 1-5 chunks of TAIGA, FLOWER_FOREST, and SWAMP; those styles lean on the balanced sampler and will look generic.
+- Trees and vegetation are coarse (single log/leaf groups expanded with heuristics by the plugin).
+- Ores are placed procedurally after generation (as vanilla does), not modelled.
+
+## License and attribution
+
+MIT. Not affiliated with or endorsed by Mojang or Microsoft; "Minecraft" is a trademark of Mojang Synergies AB. The name STEVE is, of course, for Steve.

base_v1/manifest.json

@@ -0,0 +1,169 @@
+{
+  "format": 1,
+  "version": "v1",
+  "created": "2026-07-02T20:08:08",
+  "num_classes": 27,
+  "num_biomes": 23,
+  "arch": {
+    "base_channels": 32,
+    "biome_embed_dim": 16,
+    "latent_channels": 8,
+    "latent_grid": [
+      4,
+      16,
+      4
+    ],
+    "style_dim": 64,
+    "max_custom_styles": 8
+  },
+  "geometry": {
+    "chunk_width": 16,
+    "chunk_height": 384,
+    "chunk_depth": 16,
+    "min_y": -64
+  },
+  "biome_mapping": {
+    "BADLANDS": 0,
+    "BAMBOO_JUNGLE": 1,
+    "BEACH": 2,
+    "BIRCH_FOREST": 3,
+    "COLD_OCEAN": 4,
+    "DARK_FOREST": 5,
+    "DESERT": 6,
+    "ERODED_BADLANDS": 7,
+    "FLOWER_FOREST": 8,
+    "FOREST": 9,
+    "JUNGLE": 10,
+    "LUKEWARM_OCEAN": 11,
+    "MANGROVE_SWAMP": 12,
+    "OCEAN": 13,
+    "OLD_GROWTH_BIRCH_FOREST": 14,
+    "PLAINS": 15,
+    "RIVER": 16,
+    "SAVANNA": 17,
+    "SPARSE_JUNGLE": 18,
+    "SUNFLOWER_PLAINS": 19,
+    "SWAMP": 20,
+    "TAIGA": 21,
+    "UNKNOWN": 22
+  },
+  "block_groups": {
+    "AIR": [
+      "AIR",
+      "CAVE_AIR",
+      "VOID_AIR"
+    ],
+    "STONE": [
+      "STONE",
+      "ANDESITE",
+      "DIORITE",
+      "GRANITE",
+      "TUFF"
+    ],
+    "DEEPSLATE": [
+      "DEEPSLATE"
+    ],
+    "DIRT": [
+      "DIRT",
+      "COARSE_DIRT",
+      "ROOTED_DIRT"
+    ],
+    "GRASS": [
+      "GRASS_BLOCK",
+      "PODZOL"
+    ],
+    "SAND": [
+      "SAND",
+      "RED_SAND"
+    ],
+    "GRAVEL": [
+      "GRAVEL"
+    ],
+    "CLAY": [
+      "CLAY"
+    ],
+    "WATER": [
+      "WATER"
+    ],
+    "LAVA": [
+      "LAVA"
+    ],
+    "BEDROCK": [
+      "BEDROCK"
+    ],
+    "COAL_ORE": [
+      "COAL_ORE",
+      "DEEPSLATE_COAL_ORE"
+    ],
+    "IRON_ORE": [
+      "IRON_ORE",
+      "DEEPSLATE_IRON_ORE"
+    ],
+    "COPPER_ORE": [
+      "COPPER_ORE",
+      "DEEPSLATE_COPPER_ORE"
+    ],
+    "GOLD_ORE": [
+      "GOLD_ORE",
+      "DEEPSLATE_GOLD_ORE"
+    ],
+    "REDSTONE_ORE": [
+      "REDSTONE_ORE",
+      "DEEPSLATE_REDSTONE_ORE"
+    ],
+    "LAPIS_ORE": [
+      "LAPIS_ORE",
+      "DEEPSLATE_LAPIS_ORE"
+    ],
+    "DIAMOND_ORE": [
+      "DIAMOND_ORE",
+      "DEEPSLATE_DIAMOND_ORE"
+    ],
+    "EMERALD_ORE": [
+      "EMERALD_ORE",
+      "DEEPSLATE_EMERALD_ORE"
+    ],
+    "OAK_WOOD": [
+      "OAK_LOG",
+      "OAK_LEAVES"
+    ],
+    "SPRUCE_WOOD": [
+      "SPRUCE_LOG",
+      "SPRUCE_LEAVES"
+    ],
+    "BIRCH_WOOD": [
+      "BIRCH_LOG",
+      "BIRCH_LEAVES"
+    ],
+    "JUNGLE_WOOD": [
+      "JUNGLE_LOG",
+      "JUNGLE_LEAVES"
+    ],
+    "ACACIA_WOOD": [
+      "ACACIA_LOG",
+      "ACACIA_LEAVES"
+    ],
+    "DARK_OAK_WOOD": [
+      "DARK_OAK_LOG",
+      "DARK_OAK_LEAVES"
+    ],
+    "VEGETATION": [
+      "SHORT_GRASS",
+      "TALL_GRASS",
+      "SEAGRASS",
+      "TALL_SEAGRASS",
+      "KELP",
+      "KELP_PLANT",
+      "LILY_PAD"
+    ],
+    "MISC": [
+      "GLOW_LICHEN",
+      "MOSS_BLOCK",
+      "MOSS_CARPET",
+      "SCULK",
+      "SCULK_VEIN",
+      "POINTED_DRIPSTONE"
+    ]
+  },
+  "dataset_chunks": 2040
+}

base_v1/style_mapping.json

@@ -0,0 +1,3 @@
+{
+  "BASE": 0
+}

base_v1/vae_base.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0228c462212366303bc74c1900c8a3f847691f860fa4d76f59ce03434879e3f7
+size 18532484

biome_id_mapping.json

@@ -0,0 +1,25 @@
+{
+  "BADLANDS": 0,
+  "BAMBOO_JUNGLE": 1,
+  "BEACH": 2,
+  "BIRCH_FOREST": 3,
+  "COLD_OCEAN": 4,
+  "DARK_FOREST": 5,
+  "DESERT": 6,
+  "ERODED_BADLANDS": 7,
+  "FLOWER_FOREST": 8,
+  "FOREST": 9,
+  "JUNGLE": 10,
+  "LUKEWARM_OCEAN": 11,
+  "MANGROVE_SWAMP": 12,
+  "OCEAN": 13,
+  "OLD_GROWTH_BIRCH_FOREST": 14,
+  "PLAINS": 15,
+  "RIVER": 16,
+  "SAVANNA": 17,
+  "SPARSE_JUNGLE": 18,
+  "SUNFLOWER_PLAINS": 19,
+  "SWAMP": 20,
+  "TAIGA": 21,
+  "UNKNOWN": 22
+}

block_group_mapping.json

@@ -0,0 +1,29 @@
+{
+  "0": "AIR",
+  "1": "STONE",
+  "2": "DEEPSLATE",
+  "3": "DIRT",
+  "4": "GRASS",
+  "5": "SAND",
+  "6": "GRAVEL",
+  "7": "CLAY",
+  "8": "WATER",
+  "9": "LAVA",
+  "10": "BEDROCK",
+  "11": "COAL_ORE",
+  "12": "IRON_ORE",
+  "13": "COPPER_ORE",
+  "14": "GOLD_ORE",
+  "15": "REDSTONE_ORE",
+  "16": "LAPIS_ORE",
+  "17": "DIAMOND_ORE",
+  "18": "EMERALD_ORE",
+  "19": "OAK_WOOD",
+  "20": "SPRUCE_WOOD",
+  "21": "BIRCH_WOOD",
+  "22": "JUNGLE_WOOD",
+  "23": "ACACIA_WOOD",
+  "24": "DARK_OAK_WOOD",
+  "25": "VEGETATION",
+  "26": "MISC"
+}

style_mapping.json

@@ -0,0 +1,3 @@
+{
+  "BASE": 0
+}

terrain_vae_decoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8311a665929d33ad3ebcf4c92b57ec1913f3ff123f037a791c3890215a794559
+size 16751514