auroramusic/packages/inference/diffrhythm_provider.py

"""Real generation backend, meant to run on a CUDA machine.

Wraps ASLP-lab/DiffRhythm (https://github.com/ASLP-lab/DiffRhythm, Apache-2.0),
an open-source latent-diffusion model that generates a full song — vocals sung
to your lyrics plus backing instrumentation — from a text style prompt in one
shot. Chosen over alternatives like YuE because it needs far less VRAM (8GB,
or ~6GB in fp16 with --chunked) and generates a full song in seconds, not
minutes, which matters a lot on a single consumer/prosumer GPU.

This provider shells out to DiffRhythm's own `infer/infer.py` CLI rather than
importing it as a library, because DiffRhythm is distributed as a standalone
script repo (not a pip package) with its own heavy, CUDA-specific dependency
set (torch, muq, etc). Keeping it out-of-process means Aurora's worker venv
stays lightweight and doesn't need to match DiffRhythm's exact dependency
pins. See infra/gpu-provider/README.md for setup instructions on the target
machine.
"""

import glob
import os
import shutil
import subprocess
import uuid
import wave
from pathlib import Path

from .base import GenerationRequest, GenerationResult, MusicGenerationProvider

# DiffRhythm-base tops out at 95s; the "full" checkpoint goes up to 285s (4:45).
# See DIFFRHYTHM_MAX_DURATION if you're running the -full checkpoint.
_DEFAULT_MAX_DURATION = 95
_MIN_DURATION = 10


def _seconds_to_lrc_timestamp(seconds: float) -> str:
    minutes = int(seconds // 60)
    secs = seconds - minutes * 60
    return f"[{minutes:02d}:{secs:05.2f}]"


def _write_lrc(lyrics: str | None, duration: float, path: Path) -> None:
    """Write a timestamped .lrc file, DiffRhythm's expected lyrics format.

    Lines are spread evenly across the requested duration. For instrumental
    mode (no lyrics), DiffRhythm accepts an empty/near-empty lrc file — verify
    this still holds for whatever DiffRhythm version you deploy; the repo's
    handling of instrumental-only generation has shifted across releases.
    """
    lines = [line.strip() for line in (lyrics or "").splitlines() if line.strip()]
    if not lines:
        path.write_text("")
        return

    step = duration / len(lines)
    with path.open("w") as f:
        for i, line in enumerate(lines):
            ts = _seconds_to_lrc_timestamp(i * step)
            f.write(f"{ts}{line}\n")


class DiffRhythmProvider(MusicGenerationProvider):
    def __init__(self) -> None:
        repo_dir = os.environ.get("DIFFRHYTHM_REPO_DIR")
        if not repo_dir:
            raise RuntimeError(
                "DIFFRHYTHM_REPO_DIR is not set. Clone "
                "https://github.com/ASLP-lab/DiffRhythm on the GPU machine and "
                "point this env var at it. See infra/gpu-provider/README.md."
            )
        self.repo_dir = Path(repo_dir)
        if not (self.repo_dir / "infer" / "infer.py").exists():
            raise RuntimeError(f"DIFFRHYTHM_REPO_DIR ({self.repo_dir}) doesn't look like a DiffRhythm checkout")

        self.python_bin = os.environ.get("DIFFRHYTHM_PYTHON", "python3")
        self.chunked = os.environ.get("DIFFRHYTHM_CHUNKED", "1") not in ("0", "false", "False")
        self.max_duration = int(os.environ.get("DIFFRHYTHM_MAX_DURATION", _DEFAULT_MAX_DURATION))
        self.default_style_prompt = os.environ.get(
            "DIFFRHYTHM_DEFAULT_STYLE_PROMPT", "melodic, well-produced, modern"
        )

    def generate(self, request: GenerationRequest, output_dir: Path) -> GenerationResult:
        duration = max(_MIN_DURATION, min(request.duration_seconds, self.max_duration))

        work_dir = Path("/tmp") / f"diffrhythm-{uuid.uuid4().hex}"
        work_dir.mkdir(parents=True, exist_ok=True)
        lrc_path = work_dir / "lyrics.lrc"
        infer_output_dir = work_dir / "output"

        lyrics = request.lyrics if request.mode == "vocals" else None
        _write_lrc(lyrics, duration, lrc_path)

        style_prompt = request.prompt.strip() or self.default_style_prompt

        cmd = [
            self.python_bin,
            "infer/infer.py",
            "--lrc-path",
            str(lrc_path),
            "--ref-prompt",
            style_prompt,
            "--audio-length",
            str(duration),
            "--output-dir",
            str(infer_output_dir),
            "--batch-infer-num",
            "1",
        ]
        if self.chunked:
            cmd.append("--chunked")

        try:
            proc = subprocess.run(
                cmd,
                cwd=str(self.repo_dir),
                capture_output=True,
                text=True,
                timeout=15 * 60,
                check=True,
            )
        except subprocess.CalledProcessError as exc:
            raise RuntimeError(f"DiffRhythm inference failed: {exc.stderr[-2000:]}") from exc
        except subprocess.TimeoutExpired as exc:
            raise RuntimeError("DiffRhythm inference timed out after 15 minutes") from exc

        wavs = sorted(glob.glob(str(infer_output_dir / "**" / "*.wav"), recursive=True))
        if not wavs:
            raise RuntimeError(f"DiffRhythm produced no output file. Last output: {proc.stdout[-1000:]}")

        source_wav = Path(wavs[-1])
        output_dir.mkdir(parents=True, exist_ok=True)
        final_path = output_dir / f"{uuid.uuid4().hex}.wav"
        shutil.copy(source_wav, final_path)

        actual_duration, rate = duration, 44100
        try:
            with wave.open(str(final_path), "rb") as f:
                rate = f.getframerate()
                actual_duration = f.getnframes() / float(rate)
        except wave.Error:
            # DiffRhythm may emit non-integer-PCM WAV (e.g. float32) that the
            # stdlib `wave` module can't parse; fall back to the requested
            # duration rather than failing the whole job over a metadata read.
            pass

        shutil.rmtree(work_dir, ignore_errors=True)

        return GenerationResult(audio_path=final_path, duration_seconds=actual_duration, sample_rate=rate)