Add ZeroGPU support for GPU-accelerated MusicGen on HF Spaces

Detect the ZeroGPU runtime via SPACES_ZERO_GPU and shim the spaces package to a no-op locally, so the same code runs on mps/cpu unchanged. On ZeroGPU the model loads on cuda at module level (the documented pattern) and generation runs inside an @spaces.GPU call with a dynamic duration budget; generate_song routes to the GPU path only on the Space, leaving local mps/cpu/stub paths and progress untouched. Bump the Space to python 3.12.12 (ZeroGPU requirement) and add spaces to requirements.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

README.md

@@ -5,7 +5,7 @@ colorFrom: yellow
 colorTo: green
 sdk: gradio
 sdk_version: 6.17.3
-python_version: '3.13'
+python_version: '3.12.12'
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -11,7 +11,7 @@ MusicGen ignores texture words in prompts, hence the separate bed.
 Env knobs:
   LOFINITY_ENGINE   musicgen (default) | stub
   LOFINITY_DURATION clip length in seconds (default 30, the single-shot max)
-  LOFINITY_DEVICE   mps | cpu (default: mps if available)
+  LOFINITY_DEVICE   cuda | mps | cpu (default: cuda on ZeroGPU, else mps if available)
   OLLAMA_URL        default http://localhost:11434
   OLLAMA_MODEL      default llama3.2:3b
 """
@@ -32,6 +32,24 @@ from gradio.server import Server
 ROOT = Path(__file__).parent
 FRONTEND = ROOT / "frontend"
 
+# ZeroGPU: on a Hugging Face ZeroGPU Space a GPU is attached only for the
+# duration of a function wrapped in @spaces.GPU, then released. The `spaces`
+# package exists only in that runtime (which sets SPACES_ZERO_GPU=true); locally
+# we shim @spaces.GPU to a no-op so the exact same code runs on mps/cpu untouched.
+IS_ZEROGPU = os.environ.get("SPACES_ZERO_GPU") == "true"
+try:
+    import spaces  # provided by the ZeroGPU Space runtime
+except ImportError:  # local dev / non-ZeroGPU — make the decorator harmless
+    class _SpacesShim:
+        @staticmethod
+        def GPU(*args, **kwargs):
+            # handle both bare @spaces.GPU and @spaces.GPU(duration=...)
+            if args and callable(args[0]):
+                return args[0]
+            return lambda fn: fn
+
+    spaces = _SpacesShim()
+
 ENGINE = os.getenv("LOFINITY_ENGINE", "musicgen")
 # 30s is musicgen-small's single-shot max (1500 tokens). Longer tapes are
 # stitched from 30s chunks: each one re-seeds the model with the last OVERLAP_S
@@ -139,7 +157,14 @@ def load_musicgen():
             from transformers import AutoProcessor, MusicgenForConditionalGeneration
 
             requested = os.getenv("LOFINITY_DEVICE")
-            device = requested or ("mps" if torch.backends.mps.is_available() else "cpu")
+            if requested:
+                device = requested
+            elif IS_ZEROGPU:
+                device = "cuda"
+            elif torch.backends.mps.is_available():
+                device = "mps"
+            else:
+                device = "cpu"
             print(f"[lofinity] loading musicgen-small on {device}…")
             processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
             model = MusicgenForConditionalGeneration.from_pretrained(
@@ -152,6 +177,14 @@ def load_musicgen():
     return _musicgen
 
 
+# ZeroGPU wants models resident on cuda at startup (module import time), not
+# lazily inside the @spaces.GPU call — outside the decorated function a CUDA
+# emulation layer lets this .to('cuda') succeed without a real GPU attached, and
+# placements done at startup are far more efficient than per-call transfers.
+if IS_ZEROGPU and ENGINE != "stub":
+    load_musicgen()
+
+
 def encode_wav(samples, rate: int) -> str:
     """Encode mono float samples as a base64 WAV data URI, entirely in memory.
 
@@ -259,6 +292,24 @@ def musicgen_engine(music_prompt: str, seconds: int = CHUNK_S, progress_cb=None)
     return samples, rate
 
 
+def _gpu_budget(music_prompt: str, seconds: int = CHUNK_S) -> int:
+    """GPU seconds to request from ZeroGPU for a brew of this length: per-chunk
+    render time plus headroom. Tighter budgets earn better queue priority, and
+    the signature must mirror gpu_musicgen so ZeroGPU can pass it the same args."""
+    chunks = max(1, round(int(seconds) / CHUNK_S))
+    return 20 + 25 * chunks  # 30s->45, 60s->70, 90s->95
+
+
+@spaces.GPU(duration=_gpu_budget)
+def gpu_musicgen(music_prompt: str, seconds: int = CHUNK_S) -> tuple:
+    """ZeroGPU entry point — runs MusicGen on the real GPU and returns the audio
+    to the web process. No progress_cb on purpose: this body executes in a
+    separate GPU worker, so _PROGRESS updates can't reach /api/progress yet (the
+    brewing garden is wired up in a later stage); on the Space the brew jumps to
+    done. Locally @spaces.GPU is a no-op, but this path is only taken on Spaces."""
+    return musicgen_engine(music_prompt, seconds)
+
+
 def stub_engine(_music_prompt: str, seconds: int = CHUNK_S, progress_cb=None) -> tuple:
     """A short audible tone — handy when developing without the heavy model.
     Honors `seconds` and fakes per-chunk timing so the length slider and the
@@ -295,10 +346,20 @@ def generate_song(prompt: str, seconds: int = DEFAULT_SECONDS) -> dict:
     _PROGRESS.update(done=0, total=chunks)
     music_prompt, title, bed = enrich_prompt(prompt)
     print(f"[lofinity] brewing {title!r} ({seconds}s) :: {music_prompt} [+ {bed}]")
-    engine = stub_engine if ENGINE == "stub" else musicgen_engine
-    samples, rate = engine(
-        music_prompt, seconds, progress_cb=lambda d, t: _PROGRESS.update(done=d, total=t)
-    )
+    if ENGINE == "stub":
+        samples, rate = stub_engine(
+            music_prompt, seconds,
+            progress_cb=lambda d, t: _PROGRESS.update(done=d, total=t),
+        )
+    elif IS_ZEROGPU:
+        # the GPU body runs in a separate worker process, so progress can't
+        # stream back here yet (Stage 3); the brewing bar jumps 0->100% on Space
+        samples, rate = gpu_musicgen(music_prompt, seconds)
+    else:
+        samples, rate = musicgen_engine(
+            music_prompt, seconds,
+            progress_cb=lambda d, t: _PROGRESS.update(done=d, total=t),
+        )
     _PROGRESS.update(done=chunks, total=chunks)
     try:
         samples = ambience.mix(samples, rate, bed)

requirements.txt

@@ -1,5 +1,8 @@
 # gradio itself is provided by the Space SDK (sdk_version in README.md).
-torch
+# `spaces` provides the @spaces.GPU decorator used on ZeroGPU. app.py shims it to
+# a no-op when it's absent, so it's only actually required on the Space.
+spaces
+torch  # ZeroGPU requires torch >=2.8; the Space runtime supplies the CUDA build
 transformers
 # Ollama must be running locally with the model below pulled:
 #   ollama pull llama3.2:3b