| # BeatFormer |
|
|
| Generate Beat Saber maps and lighting from music. Qwen3-0.6B with a learned audio projection, frozen MERT-v1-95M layer-6 features, per-beat interleaved audio/map tokens and text conditioning. The released joint90k model predicts notes and lighting; lights_v2 is the dedicated lighter. |
| |
| ## Install |
| |
| Python 3.12. Install `ffmpeg` (including `ffprobe`) and `curl` through your system package manager. Install a CUDA-enabled PyTorch build appropriate to your machine, then the remaining Python dependencies: |
| |
| ```sh |
| uv venv |
| uv pip install -r requirements.txt |
| ``` |
| |
| The static-KV studio integration is pinned to `transformers==5.16.1`. `requirements-optional.txt` lists W&B logging, Beat This! automatic tempo detection and experimental torchao FP8 support. These are not needed for ordinary training on cached features; enter BPM manually if Beat This! is not installed. No API key is required except for the optional `data_pipeline/caption_rewrite.py` OpenRouter call, which reads `OPENROUTER_API_KEY` from the environment. |
| |
| ## Prepare data |
| |
| Run from this repository. `meta/beatsaver_kept.jsonl` contains the original pinned download URLs; `meta/manifests/gx10-{2,3,4}.jsonl` are the exact original joint90k selections with historical paths and an accompanying manifest-hash audit. To rebuild that selection, use the IDs and versions in those files, rather than assuming every downloaded map was trained on. The original three manifests contain 201,406 rows / 57,007 map IDs. |
|
|
| ```sh |
| python data_pipeline/beatsaver_download.py meta/beatsaver_kept.jsonl --out outputs/zips --workers 6 |
| python data_pipeline/dataset_pass.py --zips outputs/zips --extract outputs/extracted --out outputs/dataset --workers 4 |
| python data_pipeline/profile_pass.py --extract outputs/extracted --out outputs/profiles --workers 4 |
| python data_pipeline/tagger.py --profiles outputs/profiles/profiles.jsonl --out outputs/tags.jsonl |
| python data_pipeline/tokenize_pass.py --extract outputs/extracted --out outputs/tokens --lattice 96 --mirror --flags outputs/dataset/maps.tsv --workers 4 |
| python data_pipeline/lights_pass.py --manifest outputs/tokens/manifest_96.jsonl --host 0 --out outputs/lights_audit.jsonl |
| python data_pipeline/features_pass.py --manifest outputs/tokens/manifest_96.jsonl --out outputs/features --device cuda |
| ``` |
|
|
| `data_pipeline/tagger_reference.json` is the frozen per-difficulty percentile reference, extracted without recalculation from the original `profiles_ref.json`. The original tags/phrases and full BeatSaver conditioning metadata are in `meta/`. Use the generated tag file when training on new maps. Optional alignment corrections can be supplied to tokenization with `--align`; these are not inferred merely by running the dataset report. |
|
|
| For replaying historical manifests after recreating their token files: |
|
|
| ```sh |
| python data_pipeline/remap_manifest.py --manifest meta/manifests/gx10-2.jsonl --maps outputs/extracted --tokens outputs/tokens/96 --out outputs/shards/gx10-2.jsonl |
| ``` |
|
|
| Repeat for GX3/GX4. The old run used per-rank sampling weights; concatenating the shards changes that sampling distribution. The archival handoff has per-file token/MERT hashes for comparing rebuilt caches. Audio decoding, numerical backend and preprocessing versions can affect regenerated bytes. |
|
|
| The original autolighter census inputs are included in `data_pipeline/autolighter_filter/census/`. `clean_manifest.py --census PATH --out flags.jsonl` rebuilds flags from a supplied census; its defaults use the included census and write beside the script. For new maps, run the five `census_*.py` scripts with explicit `--manifest`, `--host` and `--out` arguments first. The filter's mapper-population and fitted-discriminator behavior depends on the census population; adding maps is not a frozen-classifier-only operation. The preserved `meta/lights_drop.json` is the original training exclusion list. |
|
|
| ## Train |
|
|
| ```sh |
| python train.py --config configs/joint90k.json |
| ``` |
|
|
| This recipe starts from `checkpoints/beatformer_year_metadata_70000.pt`, expands the lighting vocabulary and trains to step 90,000 with a fresh optimizer. Paths point to the preparation outputs above. Config values are defaults; explicit flags override them. For example, add `--style_lines outputs/tags.jsonl` for new maps. `--lights_only` trains only the lighting targets. Header fields have independent conditioning dropout in `caption.py`. |
|
|
| Continue a joint checkpoint with matching vocabulary: |
|
|
| ```sh |
| python train.py --config configs/joint90k.json --resume checkpoints/beatformer_joint_90000.pt --steps 100000 --lr_change_step 90000 --lr_schedule_steps 10000 --out outputs/runs/continue |
| ``` |
|
|
| Checkpoints are local: `<out>/checkpoints/ckpt_STEP.pt` contains weights; `resume.pt` additionally contains AdamW state. Add `--resume_optimizer` for an exact optimizer restart from that file. W&B is disabled by default; enable with `--wandb_mode offline` or `online`. GPU allocation defaults to a 0.7 cap and is configurable with `--gpu_memory_fraction`. Use `torchrun` with the normal distributed launch arguments for DDP; `--batch` is per rank and `--rank_weights` supplies relative rank weights. |
|
|
| ## Generate and view |
|
|
| ```sh |
| python studio/fetch_viewer.py |
| BEATFORMER_BACKEND=local BEATFORMER_CHECKPOINT=checkpoints/beatformer_joint_90000.pt BEATFORMER_LIGHTS_CHECKPOINT=checkpoints/beatformer_lights_v2_90000.pt python studio/server.py |
| ``` |
|
|
| Open `http://127.0.0.1:8790`. Upload audio, enter BPM/NPS/style and generate; the embedded viewer starts muted. Studio inference uses CUDA, keeps the model loaded and supports compiled decoding through `BEATFORMER_COMPILE=1`. Qwen and MERT are downloaded by Transformers on first use unless already cached. Optional Beat This! lives in `BEATFORMER_BT_PYTHON`; manual BPM does not require it. `BEATFORMER_OUTPUT` and `BEATFORMER_GPU_FRACTION` configure output storage and the CUDA memory cap. ArcViewer assets are fetched from the pinned manifest; no workstation-specific viewer directory is required. |
|
|
| `generate.py --help` exposes the direct map-directory CLI. `quality/note_quality.py`, `quality/light_quality.py` and `mapcheck.py` provide the retained map diagnostics. The `studio/runtime/` implementation is the released inference variant; the top-level modules are the training implementation. They have distinct lighting-loss and decode behavior, so they are deliberately retained separately. |
|
|
| ## Released weights |
|
|
| | Checkpoint | Training stage | Held CE | |
| |---|---|---| |
| | `beatformer_year_metadata_70000.pt` | full-pool notes, metadata conditioning | notes 0.676 | |
| | `beatformer_joint_87000.pt` | joint notes + lighting | notes 0.588 / lights 0.736 | |
| | `beatformer_joint_90000.pt` | joint endpoint | notes 0.589 / lights 0.736 | |
| | `beatformer_lights_v2_90000.pt` | lighting-only endpoint | lights 0.571 | |
|
|
| `checkpoints/SHA256SUMS` identifies the unchanged released weights. Cached upstream revisions observed on the training boxes: MERT `12af15fef9d0ac838c3f475bfbbf26d2060dd4f5`, Qwen `c1899de289a04d12100db370d81485cdf75e47ca`. |
|
|
