Initial commit: AuroraMusic monorepo

packages/inference/aurora_inference.egg-info/PKG-INFO

@@ -0,0 +1,5 @@
+Metadata-Version: 2.4
+Name: aurora-inference
+Version: 0.1.0
+Requires-Python: >=3.10
+Requires-Dist: numpy

packages/inference/aurora_inference.egg-info/SOURCES.txt

@@ -0,0 +1,14 @@
+__init__.py
+base.py
+factory.py
+mock_provider.py
+pyproject.toml
+./__init__.py
+./base.py
+./factory.py
+./mock_provider.py
+aurora_inference.egg-info/PKG-INFO
+aurora_inference.egg-info/SOURCES.txt
+aurora_inference.egg-info/dependency_links.txt
+aurora_inference.egg-info/requires.txt
+aurora_inference.egg-info/top_level.txt

packages/inference/aurora_inference.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

packages/inference/aurora_inference.egg-info/requires.txt

@@ -0,0 +1 @@
+numpy

packages/inference/aurora_inference.egg-info/top_level.txt

@@ -0,0 +1 @@
+inference

packages/inference/base.py

@@ -0,0 +1,30 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass
+class GenerationRequest:
+    prompt: str
+    mode: str  # "instrumental" | "vocals"
+    lyrics: str | None
+    duration_seconds: int
+
+
+@dataclass
+class GenerationResult:
+    audio_path: Path
+    duration_seconds: float
+    sample_rate: int
+
+
+class MusicGenerationProvider(ABC):
+    """Pluggable backend for turning a prompt (+ optional lyrics) into a song.
+
+    Swap implementations (mock, local model, remote GPU endpoint) without
+    touching the worker or API — the worker only depends on this interface.
+    """
+
+    @abstractmethod
+    def generate(self, request: GenerationRequest, output_dir: Path) -> GenerationResult:
+        ...

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)

packages/inference/factory.py

@@ -0,0 +1,19 @@
+import os
+
+from .base import MusicGenerationProvider
+
+
+def get_provider() -> MusicGenerationProvider:
+    backend = os.environ.get("GENERATION_BACKEND", "mock")
+
+    if backend == "mock":
+        from .mock_provider import MockMusicGenerationProvider
+
+        return MockMusicGenerationProvider()
+
+    if backend == "diffrhythm":
+        from .diffrhythm_provider import DiffRhythmProvider
+
+        return DiffRhythmProvider()
+
+    raise ValueError(f"Unknown GENERATION_BACKEND: {backend}")

packages/inference/mock_provider.py

@@ -0,0 +1,95 @@
+import hashlib
+import uuid
+import wave
+from pathlib import Path
+
+import numpy as np
+
+from .base import GenerationRequest, GenerationResult, MusicGenerationProvider
+
+SAMPLE_RATE = 32000
+
+# A few simple chord progressions (root frequencies in Hz) so different
+# prompts audibly differ instead of producing the same tone every time.
+_PROGRESSIONS = [
+    [261.63, 329.63, 392.00, 261.63],  # C major-ish
+    [220.00, 261.63, 329.63, 220.00],  # A minor-ish
+    [196.00, 246.94, 293.66, 196.00],  # G-ish
+    [174.61, 220.00, 261.63, 174.61],  # F-ish
+]
+
+
+class MockMusicGenerationProvider(MusicGenerationProvider):
+    """Fast, dependency-light stand-in for a real model.
+
+    Synthesizes a short procedurally-generated audio clip (chord pad +
+    simple melody, with a gentle "vocal-ish" hum layer when mode is
+    "vocals") so the rest of the product — jobs, storage, playback,
+    library, feed — can be built and verified before wiring a real,
+    much slower/heavier local model.
+    """
+
+    def generate(self, request: GenerationRequest, output_dir: Path) -> GenerationResult:
+        seed = int(hashlib.sha256(request.prompt.encode()).hexdigest(), 16) % (2**32)
+        rng = np.random.default_rng(seed)
+        progression = _PROGRESSIONS[seed % len(_PROGRESSIONS)]
+
+        duration = max(5, min(request.duration_seconds, 30))
+        t = np.linspace(0, duration, int(SAMPLE_RATE * duration), endpoint=False)
+        audio = np.zeros_like(t)
+
+        bar_len = duration / len(progression)
+        for i, freq in enumerate(progression):
+            start = int(i * bar_len * SAMPLE_RATE)
+            end = int((i + 1) * bar_len * SAMPLE_RATE)
+            seg_t = t[start:end] - t[start]
+            pad = (
+                0.18 * np.sin(2 * np.pi * freq * seg_t)
+                + 0.12 * np.sin(2 * np.pi * freq * 1.5 * seg_t)
+                + 0.08 * np.sin(2 * np.pi * freq * 2 * seg_t)
+            )
+            envelope = np.clip(np.minimum(seg_t, bar_len - seg_t) * 4, 0, 1)
+            audio[start:end] += pad * envelope
+
+        # simple plucked melody riding on top, tempo/notes derived from the prompt
+        note_len = 60 / (90 + seed % 60)
+        n_notes = int(duration / note_len)
+        for i in range(n_notes):
+            start = int(i * note_len * SAMPLE_RATE)
+            end = min(int((i + 1) * note_len * SAMPLE_RATE), len(t))
+            if end <= start:
+                continue
+            freq = progression[i % len(progression)] * rng.choice([1, 1.25, 1.5, 2])
+            seg_t = t[start:end] - t[start]
+            note = 0.1 * np.sin(2 * np.pi * freq * seg_t) * np.exp(-3 * seg_t)
+            audio[start:end] += note
+
+        if request.mode == "vocals" and request.lyrics:
+            words = request.lyrics.split()
+            if words:
+                syl_len = duration / max(len(words), 1)
+                for i, _word in enumerate(words):
+                    start = int(i * syl_len * SAMPLE_RATE)
+                    end = min(int((i + 1) * syl_len * SAMPLE_RATE), len(t))
+                    if end <= start:
+                        continue
+                    base = progression[i % len(progression)] * 2
+                    vib = 1 + 0.01 * np.sin(2 * np.pi * 5 * (t[start:end] - t[start]))
+                    seg_t = t[start:end] - t[start]
+                    voice = 0.14 * np.sin(2 * np.pi * base * vib * seg_t)
+                    env = np.clip(np.minimum(seg_t, syl_len - seg_t) * 10, 0, 1)
+                    audio[start:end] += voice * env
+
+        audio = audio / max(np.abs(audio).max(), 1e-6) * 0.9
+        pcm = (audio * 32767).astype(np.int16)
+
+        output_dir.mkdir(parents=True, exist_ok=True)
+        filename = f"{uuid.uuid4().hex}.wav"
+        path = output_dir / filename
+        with wave.open(str(path), "wb") as f:
+            f.setnchannels(1)
+            f.setsampwidth(2)
+            f.setframerate(SAMPLE_RATE)
+            f.writeframes(pcm.tobytes())
+
+        return GenerationResult(audio_path=path, duration_seconds=duration, sample_rate=SAMPLE_RATE)

packages/inference/pyproject.toml

@@ -0,0 +1,16 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "aurora-inference"
+version = "0.1.0"
+requires-python = ">=3.10"
+dependencies = ["numpy"]
+
+[tool.setuptools]
+py-modules = []
+packages = ["inference"]
+
+[tool.setuptools.package-dir]
+inference = "."