Pin muscriptor 0.3.0, whose 0.4.0 renamed transcribe_to_midi, fetch the weights at startup, run every size on the GPU, and check requests before the GPU is reserved.

app.py

@@ -17,11 +17,14 @@ except ImportError:
                 func = args[0]
                 return func
 
+import contextlib
 import tempfile
 from pathlib import Path
 
 import gradio as gr
+import soundfile as sf
 import torch
+from huggingface_hub import hf_hub_download, snapshot_download
 from pyharp import ModelCard, build_endpoint
 from pyharp.tags import Subcategory
 
@@ -43,18 +46,46 @@ model_card = ModelCard(
 
 _VALID_INSTRUMENTS = ", ".join(MT3_FULL_PLUS_GROUP_NAMES)
 
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Each size's weights, pinned, with the config.json that load_model reads beside them
+_WEIGHT_REVISIONS = {
+    "small": "8c127f603b807520fa465c838e9bfee8a91ada4e",
+    "medium": "f32236969308476e01fd3aae67357de5feb05a2d",
+    "large": "8809fdfbed2affa7ade94a7059e746e3880720e7",
+}
+_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",
+        ))
+
 
 def _get_model(variant: str) -> TranscriptionModel:
-    """Load and cache a TranscriptionModel for a given size, one per variant.
+    """Builds the TranscriptionModel of one size, on the GPU the call holds.
 
-    small/medium run on CPU; large runs on GPU. Built directly on its target
-    device inside this GPU-decorated call, per ZeroGPU rules.
+    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.
     """
     if variant not in _models:
-        device = "cuda" if variant == "large" else "cpu"
-        _models[variant] = TranscriptionModel.load_model(variant, device=device)
+        _models[variant] = TranscriptionModel.load_model(_weights[variant], device=DEVICE)
     return _models[variant]
 
 
@@ -88,8 +119,16 @@ def _resolve_instruments(text: str) -> tuple[list[str] | None, str]:
     return (resolved or None), note
 
 
-@spaces.GPU
-@torch.inference_mode()
+# The audio is transcribed in five-second segments, so the cost grows with its length, and with the
+# size of the model
+MAX_INPUT_SECONDS = 600
+SECONDS_PER_INPUT_SECOND = {"small": 0.5, "medium": 1.0, "large": 2.0}
+
+
+def gpu_seconds(input_audio_path, variant, instruments_text, use_sampling, temperature):
+    return 20 + SECONDS_PER_INPUT_SECOND[variant] * sf.info(input_audio_path).duration
+
+
 def process_fn(
     input_audio_path: str,
     variant: str,
@@ -98,6 +137,30 @@ def process_fn(
     temperature: float,
 ) -> tuple[str, str]:
     """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 not input_audio_path:
+        raise gr.Error("Provide the music to transcribe.")
+    try:
+        duration = sf.info(input_audio_path).duration
+    except Exception as exc:
+        raise gr.Error(f"Could not read the audio file: {exc}") from exc
+    if duration > MAX_INPUT_SECONDS:
+        raise gr.Error(
+            f"The music is {duration:.1f} s long. Trim it to {MAX_INPUT_SECONDS // 60} minutes "
+            "or less."
+        )
+    if gpu_seconds(input_audio_path, variant, instruments_text, use_sampling, temperature) > 200:
+        raise gr.Error(
+            f"The {variant} model would need more GPU time for {duration:.0f} s of music than a "
+            "request can reserve. Choose a smaller model or trim the music."
+        )
+    return transcribe(input_audio_path, variant, instruments_text, use_sampling, temperature)
+
+
+@spaces.GPU(duration=gpu_seconds)
+@torch.inference_mode()
+def transcribe(input_audio_path, variant, instruments_text, use_sampling, temperature):
     model = _get_model(variant)
     instruments, instrument_note = _resolve_instruments(instruments_text)
 

requirements.txt

@@ -1,3 +1,4 @@
 git+https://github.com/TEAMuP-dev/pyharp.git@v0.4.2
-muscriptor
+# 0.4.0 renamed transcribe_to_midi, which this app calls
+muscriptor==0.3.0
 spaces>=0.51.1