Build the models at startup, check inputs before reserving the GPU, bound the retries, and reserve 170 s rather than 300.

app.py

@@ -42,7 +42,9 @@ MAX_CFG = 5.0          # classifier-free guidance ceiling (default: 5.0, per pap
 MIN_CFG = 1.0          # classifier-free guidance floor, annealed down to this (default: 1.0, per paper)
 RESCORE_WEIGHT = 0.5   # MAGNeT/MusicGen rescoring interpolation weight (default: 0.5, per paper)
 HINT_DURATION = 30.0   # seconds of MusicGen hint audio generated for beat detection (default: 30.0, per repo)
-MAX_RETRIES = 5        # beat-alignment retries before falling back to a fixed length (default: 5, per repo)
+# Beat-alignment retries before falling back to a fixed length. The repo's default is 5, but each
+# hint costs about 46 s of the caller's GPU allowance, and two keep a request inside its duration
+MAX_RETRIES = 2
 
 beat_checkpoint_path = None
 beat_checkpoint_ready = threading.Event()
@@ -51,7 +53,6 @@ beat_checkpoint_error = None
 musicgen = None
 loopgen = None
 file2beats = None
-model_lock = threading.Lock()
 
 
 def download_beat_checkpoint():
@@ -84,6 +85,14 @@ _loader.start()
 if __name__ != "__main__":
     _loader.join()
 
+    # The worker, which runs process_fn, builds the models on CUDA here, at module level on its main
+    # thread, where ZeroGPU's emulation records the placement, rather than inside every call,
+    # which ZeroGPU runs in a fresh fork. The server never runs process_fn, so it builds none.
+    if beat_checkpoint_error is None:
+        musicgen = MusicGen.get_pretrained(MUSICGEN_MODEL, device=DEVICE)
+        loopgen = LoopGen.get_pretrained(MAGNET_MODEL, device=DEVICE)
+        file2beats = File2Beats(checkpoint_path=beat_checkpoint_path, device=DEVICE)
+
 
 model_card = ModelCard(
     name="LoopGen",
@@ -167,33 +176,25 @@ def _beat_aligned_length(path, min_length, max_length, unit_size=0.02, length_re
     return length
 
 
-@spaces.GPU(duration=300)
-@torch.inference_mode()
 def process_fn(description: str, min_length: float, max_length: float, seed: int) -> str:
     """Generates a seamlessly loopable music clip: a MusicGen hint is used to
     find a beat-aligned loop length (retrying up to MAX_RETRIES times), then
     MAGNeT generates the final loop as a continuation of that hint. Adapted
-    from pipelines.py's loopgen_loop.
-
-    Model construction happens here, inside the GPU call, rather than in the
-    background-download thread above: LoopGen/MusicGen.get_pretrained()
-    bakes device="cuda" in immediately, and ZeroGPU only allows CUDA calls
-    made inside an @spaces.GPU-decorated call, not from a background thread."""
-    global musicgen, loopgen, file2beats
-
-    if not beat_checkpoint_ready.wait(timeout=90):
-        raise gr.Error("Beat-tracking checkpoint is still downloading, please try again shortly.")
+    from pipelines.py's loopgen_loop."""
+    # Checked before the GPU is reserved, so a request that cannot run spends none of the
+    # caller's ZeroGPU allowance
     if beat_checkpoint_error is not None:
         raise gr.Error(f"Beat-tracking checkpoint failed to download: {beat_checkpoint_error}")
+    if musicgen is None or loopgen is None or file2beats is None:
+        raise gr.Error("The models are still loading, please try again shortly.")
+
+    return generate_loop(description, min_length, max_length, seed)
 
-    with model_lock:
-        if musicgen is None:
-            musicgen = MusicGen.get_pretrained(MUSICGEN_MODEL, device=DEVICE)
-        if loopgen is None:
-            loopgen = LoopGen.get_pretrained(MAGNET_MODEL, device=DEVICE)
-        if file2beats is None:
-            file2beats = File2Beats(checkpoint_path=beat_checkpoint_path, device=DEVICE)
 
+# About 46 s per hint, at most MAX_RETRIES of them, then about 61 s for MAGNeT's loop
+@spaces.GPU(duration=170)
+@torch.inference_mode()
+def generate_loop(description: str, min_length: float, max_length: float, seed: int) -> str:
     torch.manual_seed(seed)
     np.random.seed(seed)
 
@@ -268,7 +269,7 @@ if __name__ == "__main__":
                       info="Random seed for generation (default: 42, per repo)."),
         ]
         output_components = [
-            gr.Audio(type="filepath", label="Output Audio").set_info("Generated loopable music clip."),
+            gr.Audio(type="filepath", label="Loop").set_info("Generated loopable music clip."),
         ]
 
         build_endpoint(