osu!dreamer is a generative model for osu! beatmaps based on diffusion

• sample generated mapset
• video of a generated map

colab notebook (no installation required)

• FFmpeg
• python 3.9
• uv

Clone this repo, then run:

uv sync [--group dev]

This will install osu-dreamer's dependencies

first you must generate a dataset, using eg. your osu!/Songs directory. This step only needs to be done once (unless you delete the generated dataset directory).

$ uv run python -m osu_dreamer.model generate-data [MAPS_DIR]

where [MAPS_DIR] is the path to eg. your osu!/Songs directory

after the dataset generation completes, you can start training. Training occurs in two stages:

$ uv run python -m osu_dreamer.model fit-latent

See osu_dreamer/latent_model/model.yml for all latent model training parameters.

At the end of every epoch, the model parameters will be checkpointed to runs/latent/version_{NUM}/checkpoints/epoch={EPOCH}-step={STEP}.ckpt. You can resume training from a saved checkpoint by adding --ckpt-path [PATH TO CHECKPOINT] to the fit-latent command.

run tensorboard --logdir=runs/latent in a new window to track training progress in Tensorboard

After training, copy/link the final checkpoint to the repo root:

ln runs/latent/version_{NUM}/checkpoints/epoch={EPOCH}-step={STEP}.ckpt latent.ckpt

Afterwards, proceed to the next training stage

$ uv run python -m osu_dreamer.model fit-denoiser

See osu_dreamer/diffusion_model/model.yml for all latent model training parameters.

At the end of every epoch, the model parameters will be checkpointed to runs/denoiser/version_{NUM}/checkpoints/epoch={EPOCH}-step={STEP}.ckpt. You can resume training from a saved checkpoint by adding --ckpt-path [PATH TO CHECKPOINT] to the fit-denoiser command.

run tensorboard --logdir=runs/denoiser in a new window to track training progress in Tensorboard

After training, copy/link the final checkpoint to the repo root:

ln runs/denoiser/version_{NUM}/checkpoints/epoch={EPOCH}-step={STEP}.ckpt denoiser.ckpt

After obtaining both training checkpoints, you must create an inference artifcat

# uv run python -m osu_dreamer export-inference

this will create an inference.pt file in the repo root.

$ uv run python -m osu_dreamer predict --help
Usage: python -m osu_dreamer predict [OPTIONS]

  generate osu!std maps from raw audio.

Options:
  --model-path FILE               inference artifact (.pt)  [required]
  --audio-file FILE               audio file to map  [required]
  --diff <FLOAT FLOAT FLOAT FLOAT FLOAT>...
                                  difficulty conditioning (sr, ar, od, cs, hp)
  --sample-steps INTEGER          number of diffusion steps to sample
  --title TEXT                    Song title - required if it cannot be
                                  determined from the audio metadata
  --artist TEXT                   Song artist - required if it cannot be
                                  determined from the audio metadata
  --help                          Show this message and exit.

you may specify --diff multiple times to generate multiple diffs at once.

The training process will generate one plot at the end of every epoch, using a sample from the validation set

• the first row is the spectrogram of the audio file
• the second row is the actual map associated with the audio file in its signal representation
• the third and fourth rows are signal representations of the maps produced by the model

⚠️ Support for training/evaluating the model locally on Windows is highly experimental and provided as-is

• 🐍 Python 3.9 (via Microsoft Store, or python.org)

Install the source code directly through github, or with the git clone command:

git clone https://github.com/jaswon/osu-dreamer

Setup from this point is pretty simple, navigate into the osu-dreamer directory and then into the windows_scripts folder, this is where all the batch scripts are stored.

First, you will need to run ! Install.bat, this will install osu-dreamer and all of its dependencies. Optionally you can install tensorboard and mathplotlib to view training statistics.

Now you're ready to begin training your own model! Here is a list of all the scripts and their functionality

• Install
  • Installs osu-dreamer and all of its dependencies.
• Run Training
  • Compiles the given songs directory and begins training a model
• Resume Training
  • Resumes training the given checkpoint
• Generate Beatmap
  • Generates a beatmap with the given information (requires a trained model and song)
• Tensorboard
  • Hosts tensorboard for tracking training statistics