Initial commit: AuroraMusic monorepo

packages/inference/diffrhythm_provider.py

@@ -0,0 +1,147 @@
+"""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)