auroramusic/infra/gpu-provider

Real generation backend (DiffRhythm) — GPU machine setup

Aurora's worker (apps/worker) talks to generation backends through one interface (packages/inference/base.py). Locally you run it with GENERATION_BACKEND=mock, which synthesizes a placeholder clip instantly so you can build/test the product without a GPU. On your machine with the NVIDIA card, switch to GENERATION_BACKEND=diffrhythm, which shells out to DiffRhythm (Apache-2.0) — an open-source model that generates a full song (vocals + instrumentation) from a text style prompt and lyrics in one pass.

Why DiffRhythm over the other realistic open-source option (YuE): YuE needs ~20GB+ VRAM and multiple minutes per song; DiffRhythm runs on as little as 6–8GB VRAM and generates a full song in roughly the time it takes to read this sentence. With 128GB system RAM and an NVIDIA card, DiffRhythm will run comfortably; if your card has 24GB+ VRAM you can also try YuE for potentially higher-fidelity vocals — the same MusicGenerationProvider interface can host either.

1. Run Aurora's own services

Follow the root README to get Postgres, Redis, the API, and the worker running on this machine (same steps as the dev machine).

2. Install DiffRhythm

DiffRhythm is a script repo with its own heavy CUDA dependency set (torch, MuQ, etc.) — keep it in its own Python environment, separate from apps/worker/.venv, so version pins never collide.

git clone https://github.com/ASLP-lab/DiffRhythm.git
cd DiffRhythm

# espeak-ng is a hard dependency of the phonemizer used for lyrics
sudo apt-get install -y espeak-ng   # or your distro's equivalent

python3.10 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Confirm your CUDA build of torch actually sees the GPU before going further:

python -c "import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name(0))"

The first real inference run downloads model checkpoints from Hugging Face (ASLP-lab/DiffRhythm-1_2-full for the 4m45s-capable model, or the smaller -base checkpoint capped at 95s) — expect a multi-GB download the first time. Check infer/infer.py in the checked-out repo for how checkpoint selection works in the version you land on; this has moved around across DiffRhythm releases.

Sanity-check the install with the repo's own example before wiring it into Aurora:

bash scripts/infer_prompt_ref.sh
# should produce a .wav under infer/example/output

3. Point Aurora's worker at it

In apps/worker/.env on this machine:

GENERATION_BACKEND=diffrhythm
DIFFRHYTHM_REPO_DIR=/absolute/path/to/DiffRhythm
DIFFRHYTHM_PYTHON=/absolute/path/to/DiffRhythm/.venv/bin/python
DIFFRHYTHM_CHUNKED=1              # keeps VRAM usage low; drop if you have headroom
DIFFRHYTHM_MAX_DURATION=285       # 95 for the -base checkpoint, 285 for -full
DIFFRHYTHM_DEFAULT_STYLE_PROMPT=melodic, well-produced, modern

DIFFRHYTHM_PYTHON must point at the interpreter inside DiffRhythm's own venv, not Aurora's worker venv — the provider (packages/inference/diffrhythm_provider.py) invokes it as a subprocess, so the two dependency sets never need to reconcile.

Bump apps/worker/.env's DEFAULT_SONG_DURATION_SECONDS to whatever you want the default generated length to be (capped by DIFFRHYTHM_MAX_DURATION above).

Restart the worker; it'll route generation jobs through DiffRhythm instead of the mock provider. Generation will now take real wall-clock time (seconds to low minutes depending on length/GPU), so the frontend's library polling (already built to handle pending/processing states) is what makes this transition invisible to the product.

Known rough edges to expect

• Instrumental-only mode: the provider writes an empty .lrc file when no lyrics are given. Verify this actually produces an instrumental for whatever DiffRhythm version you land on — vocal-less generation has been handled differently across releases; check infer/infer.py if results look off.
• Vocal quality vs. Suno: DiffRhythm is a strong open-source result, not a drop-in Suno replacement. Expect rougher vocal fidelity, occasional lyric/timing drift, and more sensitivity to how the style prompt is worded. Iterate on DIFFRHYTHM_DEFAULT_STYLE_PROMPT and per-song prompts.
• Checkpoint choice: -full unlocks up to 4m45s songs; -base is faster/lighter but caps at 95s. Match DIFFRHYTHM_MAX_DURATION to whichever you configure DiffRhythm to load.