Build the default size in the worker, and check the variant

Builds the medium model, the default, in the worker rather than only downloading every size there, builds another size when a request picks it, moves the beat checkpoint link into the worker, and refuses an unknown variant. It also imports spaces plainly, since it is always installed.

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

app.py

@@ -1,22 +1,9 @@
+import spaces  # imported first, as ZeroGPU requires
+
 import sys
 
 sys.stdout.reconfigure(line_buffering=True)
 
-try:
-    import spaces
-except ImportError:
-    # keep @spaces.GPU usable as a no-op; ZeroGPU requires this exact name.
-    class spaces:
-        class GPU:
-            def __init__(self, func=None, duration=60):
-                self.func = func
-
-            def __call__(self, *args, **kwargs):
-                if self.func is not None:
-                    return self.func(*args, **kwargs)
-                func = args[0]
-                return func
-
 import contextlib
 import tempfile
 from pathlib import Path
@@ -57,38 +44,46 @@ _WEIGHT_REVISIONS = {
 _weights: dict[str, Path] = {}
 _models: dict[str, TranscriptionModel] = {}
 
-# Fetched in the worker, which runs process_fn and starts as the server comes up, so that no
-# request downloads inside its GPU reservation. The server never runs process_fn.
-if __name__ != "__main__":
-    for _size, _revision in _WEIGHT_REVISIONS.items():
-        _weights[_size] = Path(snapshot_download(
-            f"MuScriptor/muscriptor-{_size}", revision=_revision
-        )) / "model.safetensors"
-
-# Tempo detection runs beat_this, which downloads its checkpoint from the JKU cloud on first use,
-# and that has failed both inside a GPU call (no name resolution) and at startup (503). The same
-# file is mirrored on the Hub and linked where torch.hub looks before downloading.
-_beat_checkpoint = Path(torch.hub.get_dir()) / "checkpoints" / "beat_this-final0.ckpt"
-if not _beat_checkpoint.exists():
-    _beat_checkpoint.parent.mkdir(parents=True, exist_ok=True)
-    with contextlib.suppress(FileExistsError):
-        _beat_checkpoint.symlink_to(hf_hub_download(
-            "teamup-tech/Beat-This-weights", "final0.ckpt",
-            revision="4afbe376535aaab71b76ae2ca2d7c16b08877321",
-        ))
+DEFAULT_SIZE = "medium"
 
 
 def _get_model(variant: str) -> TranscriptionModel:
-    """Builds the TranscriptionModel of one size, on the GPU the call holds.
+    """Builds the TranscriptionModel of one size, once.
 
-    Built inside the call, so ZeroGPU builds it again on every call, which takes seconds from local
-    weights. Holding all three sizes in the worker instead would keep about 7 GB there.
+    The default size is built in the worker at startup. Holding all three would keep about 7 GB
+    there, so another size is built when a request picks it, inside the call, where ZeroGPU builds
+    it again on every call, which takes seconds from local weights.
     """
     if variant not in _models:
         _models[variant] = TranscriptionModel.load_model(_weights[variant], device=DEVICE)
     return _models[variant]
 
 
+# Fetched in the worker, which runs process_fn and starts as the server comes up, so that no
+# request downloads inside its GPU reservation, and the default size is built and placed there, on
+# its main thread, where ZeroGPU's emulation records the placement. The server never runs
+# process_fn.
+if __name__ != "__main__":
+    for _size, _revision in _WEIGHT_REVISIONS.items():
+        _weights[_size] = Path(snapshot_download(
+            f"MuScriptor/muscriptor-{_size}", revision=_revision
+        )) / "model.safetensors"
+
+    # Tempo detection runs beat_this, which downloads its checkpoint from the JKU cloud on first
+    # use, and that has failed both inside a GPU call (no name resolution) and at startup (503).
+    # The same file is mirrored on the Hub and linked where torch.hub looks before downloading.
+    _beat_checkpoint = Path(torch.hub.get_dir()) / "checkpoints" / "beat_this-final0.ckpt"
+    if not _beat_checkpoint.exists():
+        _beat_checkpoint.parent.mkdir(parents=True, exist_ok=True)
+        with contextlib.suppress(FileExistsError):
+            _beat_checkpoint.symlink_to(hf_hub_download(
+                "teamup-tech/Beat-This-weights", "final0.ckpt",
+                revision="4afbe376535aaab71b76ae2ca2d7c16b08877321",
+            ))
+
+    _get_model(DEFAULT_SIZE)
+
+
 def _resolve_instruments(text: str) -> tuple[list[str] | None, str]:
     """Resolve the comma-separated instruments box into exact group names.
 
@@ -143,6 +138,8 @@ def process_fn(
     """Transcribe the input audio to MIDI, plus a note on the Instruments field."""
     # Checked before the GPU is reserved, so a request the model cannot take spends none of the
     # caller's ZeroGPU allowance
+    if variant not in _WEIGHT_REVISIONS:
+        raise gr.Error(f"Unknown model size: {variant}")
     if not input_audio_path:
         raise gr.Error("Provide the music to transcribe.")
     try:
@@ -192,7 +189,7 @@ if __name__ == "__main__":
             )),
             gr.Dropdown(
                 choices=["small", "medium", "large"],
-                value="medium",
+                value=DEFAULT_SIZE,
                 label="Model Size",
                 info="small = fastest, least accurate. medium = balanced (default, per repo). large = most accurate, slower.",
             ),