Upload app.py with huggingface_hub

app.py

@@ -6,16 +6,72 @@ import os
 import time
 import sys
 import io
+import json
+import re
+from pathlib import Path
+
+from safetensors.torch import load_file
+from huggingface_hub import hf_hub_download
+
+from muscriptor.models.lm import LMModel, TorchAutocast
+from muscriptor.modules.conditioners import (
+    MelSpectrogramConditioner,
+    ClassConditioner,
+    ConditioningProvider,
+)
+from muscriptor.tokenizer.mt3 import MT3Tokenizer, MT3_FULL_PLUS_GROUP_NAMES
+from muscriptor.transcription_model import (
+    TranscriptionModel,
+    _resolve_source,
+    _resolve_config,
+    _remap_single_codebook_keys,
+    _build_model,
+)
+
+_SAMPLE_RATE = 16000
+
+
+def load_model_zerogpu():
+    """Load the MuScriptor model in a ZeroGPU-compatible way.
+
+    On ZeroGPU, safetensors.load_file(device="cuda") fails because there's no
+    real GPU at module scope. We load weights on CPU, build the model on CPU,
+    then call .to("cuda") — which ZeroGPU intercepts to pack weights to disk
+    and stream into VRAM on the first @spaces.GPU call.
+    """
+    device = torch.device("cuda")  # ZeroGPU intercepts this
+
+    source = _resolve_source("medium")
+    weights_path = Path(hf_hub_download(
+        repo_id="MuScriptor/muscriptor-medium",
+        filename="model.safetensors",
+    ))
+    cfg = _resolve_config(source, weights_path)
+
+    # Build the model on CPU first
+    cpu_device = torch.device("cpu")
+    model = _build_model(cpu_device, cfg)
+    model.eval()
+
+    # Load weights on CPU
+    state_dict = load_file(str(weights_path), device="cpu")
+    state_dict = _remap_single_codebook_keys(state_dict)
+    model.load_state_dict(state_dict)
+
+    # Now move to "cuda" — ZeroGPU intercepts this and packs to disk
+    model.to("cuda")
+
+    tokenizer = MT3Tokenizer(
+        instrument_vocabulary="MT3_FULL_PLUS",
+        max_shift_steps=1001,
+    )
+
+    return TranscriptionModel(model=model, tokenizer=tokenizer, device=device)
 
-from muscriptor.transcription_model import TranscriptionModel
-from muscriptor.events import NoteStartEvent, NoteEndEvent, ProgressEvent
-from muscriptor.tokenizer.mt3 import MT3_FULL_PLUS_GROUP_NAMES
 
-# Load the model at module scope — ZeroGPU intercepts .to("cuda") and packs
-# weights to disk, streaming them into VRAM on the first @spaces.GPU call.
 print("[muscriptor-space] Loading model...", file=sys.stderr, flush=True)
 t0 = time.perf_counter()
-model = TranscriptionModel.load_model("medium", device="cuda")
+model = load_model_zerogpu()
 print(f"[muscriptor-space] Model loaded in {time.perf_counter() - t0:.1f}s", file=sys.stderr, flush=True)
 
 # Build the instrument choices list (sorted by group ID for stable ordering)