Card: training progress with live plots

56bf4e8 verified 1 day ago

8.61 kB

license: cc-by-sa-4.0
language:
  - en
tags:
  - music-generation
  - text-to-music
  - audio-tokens
  - rvq
  - rust
  - cuda-oxide
  - staleynet
pipeline_tag: text-to-audio

StaleyNET checkpoints

Weights for StaleyNET, a text-prompted music generator trained from random initialization in Rust on NVIDIA edge GPUs. The Transformer, residual vector quantizer, Griffin-Lim vocoder, and GPU kernels are implemented in the project.

The code, paper, and book live in the companion repository: https://github.com/arpanpathak/staleynet.

Release v0.1.2

GitHub release: https://github.com/arpanpathak/staleynet/releases/tag/v0.1.2
staleynet-audio 0.1.2 on crates.io: https://crates.io/crates/staleynet-audio. Its page explains each audio step with intuition and runnable examples.
Since v0.1.1: the GPU kernels are organised into one module per topic, the AGX Thor builds against cuda-oxide's main branch, and the paper has been revised with hardware, benchmark and tokenizer appendices.
New in these repositories: the lyrics-thor, lyrics-thor-long and pq-large runs, and the band-group token caches they train on.
CI: rustfmt, clippy with warnings denied, cargo-deny, the audio coverage gate, and documentation tests for every example on the crate page.

At a glance

Item	Value
Checkpointed runs	9
Parameters	7,099,904 for text-only runs; 7,165,440 for the lyrics runs; 42,858,496 for `pq-large`
Transformer	width 256, 8 heads, 6 layers, context 768 (`pq-large`: width 512, 8 layers, context 1152)
Audio tokens	4 residual codebooks of 1024 entries (`pq-large`: 4 band groups × 4 levels = 16 codebooks of 1024 entries)
Audio features	128 mel bands at 24 kHz, FFT 1024, hop 256
Training data	Song Describer Dataset and JamendoLyrics MultiLang
License	CC BY-SA 4.0, subject to the upstream dataset terms

Checkpoints

Run	Purpose	Parameters	`best` step	`latest` step	Training tracks
`thor-base`	base text-to-music model	7,099,904	0	None	667
`smoke`	pipeline smoke test	7,099,904	0	1,500	667
`overfit-one-take`	single-take memorisation check	7,099,904	200	200	1
`vocal-scratch`	vocal-filtered model from scratch	7,099,904	5,500	6,000	281
`vocal-finetune`	vocal fine-tune	7,099,904	250	None	281
`lyrics-orin`	lyrics-conditioned model	7,165,440	3,750	4,000	543 (73 with lyrics)
`lyrics-thor`	lyrics-conditioned model trained on the AGX Thor	7,165,440	4,750	8,000	543 (73 with lyrics)
`lyrics-thor-long`	the same model trained for 20,000 steps	7,165,440	19,750	20,000	543 (73 with lyrics)
`pq-large`	larger model on the band-group tokenizer (in progress)	42,858,496	7,000	7,000	742 (75 with lyrics)

The step values are exactly what each model.json records. best/ is the weight-only checkpoint selected by validation. latest/ is the resumable checkpoint with AdamW moments. Some early runs stored only best/, and thor-base and smoke record step 0 in their best/model.json. Where a latest/ exists, prefer it for the late training state.

About `pq-large`

pq-large is still training on the Jetson AGX Thor (128 GB). The files here are a snapshot at step 7,000 of 20,000, and they will be replaced as the run advances. Its validation loss (5.135 nats per token at step 7,000) is averaged over 16 codebooks of a different tokenizer, so it cannot be compared directly with the 4-codebook runs above.

The model reads codes from the band-group tokenizer, which splits the 128 mel bands into 4 groups and quantises each group with 4 residual levels. That tokenizer is stored as checkpoints/pq-large/tokenizer.json, and the matching caches are cache-g4l4f1/, cache-pq-jamendolyrics/ and cache-pq-recordings/ in the dataset repository. Pass --cache data/cache-g4l4f1 when generating with this run.

evaluations/<run>/ holds the paired lyrics test of each lyrics run: the per-track loss with aligned, removed and mismatched lyrics.

Training progress

The runs on the Jetson AGX Thor publish their metrics here every two hours (deploy/hf/publish.sh in the code repository): metrics/<run>.jsonl holds every logged event, and plots/<run>.png shows, from left to right, the training and validation loss in nats per code against optimizer step, the latest validation loss of each codebook (coarse levels are far easier to predict than fine ones), and the throughput in positions per second.

Layout

checkpoints/<run>/best/      weights only       model.json + tensors.bin
checkpoints/<run>/latest/    weights + AdamW    model.json + tensors.bin
configs/                     canonical training configurations
run-configs/                 per-run config.json where the run saved one
metrics/                     full metrics.jsonl for every run
logs/                        captured training logs and the tokenizer sweeps
evaluations/<run>/           paired lyrics evaluations

tensors.bin is little-endian f32 in parameter-visitor order. latest/ appends the AdamW first moments and then the second moments. model.json records the tensor manifest, the step, whether optimizer state is present, and the training configuration.

Download and use

hf download arpanpathak/staleynet-checkpoints \
  --repo-type model --local-dir staleynet-checkpoints

# From the code repository, on a Jetson with the cuda-oxide toolchain:
cargo oxide build -- --release
staleynet sing \
  --cache data/cache \
  --run staleynet-checkpoints/checkpoints/thor-base/best \
  --prompt "slow acoustic ballad, warm female voice" \
  --out song.wav

For the lyrics model:

staleynet sing \
  --run staleynet-checkpoints/checkpoints/lyrics-orin/best \
  --prompt "grunge rock, raw male voice" \
  --lyrics-file song.txt \
  --out song.wav

Architecture: width 256, 8 heads, 6 layers, feed-forward multiplier 4, context length 768, 4 residual codebooks of 1024 entries, 128 mel bands at 24 kHz. lyrics-orin adds a 256-byte lyrics field.

Training data and licensing

Song Describer Dataset (SDD), Zenodo record 10072001: 706 Jamendo recordings with human captions. Captions are CC BY-SA 4.0. The recordings carry the per-track Creative Commons licenses in audio_licenses.txt, including CC BY-NC-ND and CC BY-NC-SA.
JamendoLyrics MultiLang, https://github.com/f90/jamendolyrics: code and metadata are MIT. Songs and lyrics carry the per-song CC licenses in JamendoLyrics.csv.

Because the training audio includes NonCommercial and NoDerivatives tracks, these checkpoints are research artifacts, not a commercial product. Do not use them commercially where the upstream recording licenses forbid it, and do not use them to impersonate a singer without consent.

The checkpoint files are offered under CC BY-SA 4.0 (LICENSE-CC-BY-SA-4.0.txt) to the extent the upstream dataset licenses allow. See NOTICE for the full provenance statement and SHA256SUMS for integrity.

Known limitations

Griffin-Lim produces phasey audio compared with a learned vocoder.
The models are small edge models, not production music systems.
The lyric models were trained on a small lyrics subset, so lyric intelligibility is limited.
The checkpoints are research artifacts and have no warranty.

Citation

@misc{pathak2026staleynet,
  author = {Arpan Pathak},
  title  = {StaleyNET: Text-Conditioned Music Generation Trained from Scratch
            on Edge GPUs in Rust},
  year   = {2026},
  url    = {https://github.com/arpanpathak/staleynet}
}