Fetch the stem checkpoints from a model repository

app.py

@@ -1,3 +1,4 @@
+from huggingface_hub import hf_hub_download
 import math
 import librosa
 import tempfile
@@ -297,7 +298,33 @@ model_card = ModelCard(
     tags=[Subcategory.STEM_SEPARATION, 'drums', 'demucs', 'source-separation', 'pyharp', 'stems', 'multi-output'],
 )
 
-MODEL = LarsNet(wiener_filter=False, device="cpu", config="config.yaml")
+
+# The five stem checkpoints live in a model repository rather than in this
+# Space, which keeps the Space small enough to clone quickly. They are fetched
+# at module scope, so a failure shows up when the Space boots rather than
+# inside a job, and pinned to a revision so behavior cannot change under us.
+#
+# The vendored LarsNet class reads these paths out of config.yaml itself, so
+# the config is copied with the downloaded paths filled in rather than the
+# class being patched.
+_WEIGHTS_REPO = "teamup-tech/LarsNet-weights"
+_WEIGHTS_REVISION = "1da4e632b0c47b3002a7d2f2d23e54531ae6eee7"
+
+_config = yaml.safe_load(Path("config.yaml").read_text())
+_config["inference_models"] = {
+    stem: hf_hub_download(
+        repo_id=_WEIGHTS_REPO,
+        filename=f"pretrained_{stem}_unet.pth",
+        revision=_WEIGHTS_REVISION,
+    )
+    for stem in _config["inference_models"]
+}
+
+CONFIG_PATH = Path(tempfile.mkdtemp(prefix="larsnet-")) / "config.yaml"
+CONFIG_PATH.write_text(yaml.safe_dump(_config))
+
+MODEL = LarsNet(wiener_filter=False, device="cpu",
+                config=str(CONFIG_PATH))
 
 
 @torch.inference_mode()

pretrained_larsnet_models/cymbals/pretrained_cymbals_unet.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:889804c465e2c6fdbbe45febbd4daafef01d8d22f9f34fb968695ab9793858d0
-size 118037828

pretrained_larsnet_models/hihat/pretrained_hihat_unet.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ac004ac24a26e77f6d39671ed8ba45c3f476e956900d469ebb402386dad11dd7
-size 118037828

pretrained_larsnet_models/kick/pretrained_kick_unet.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ed821b6a69b1ef0413ac9ef7958ac9a37e5f4f056e66f0191496bf903ac2628d
-size 118037828

pretrained_larsnet_models/snare/pretrained_snare_unet.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:78bd75001ff6c52b23de6245cecd1606adbd45348e0216f6d8c116553013e4fd
-size 118037828

pretrained_larsnet_models/toms/pretrained_toms_unet.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:699112983948e14d805890a0723e5b402203a3a1db1cd0ccdd8a80058062ef77
-size 118037828