# 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](https://github.com/ASLP-lab/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](../../README.md) 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.

```bash
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:

```bash
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
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:

```bash
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.