Pin the weights' revision, and fetch the fine-tuned models in the worker at startup rather than inside the GPU call.

app.py

@@ -45,6 +45,14 @@ init_model_choice = open("DEFAULT_MODEL").read().strip()
 
 # load the init model
 interface.load_finetuned(init_model_choice)
+
+# The other fine-tuned models the interface offers are fetched in the worker, which runs
+# process_fn, rather than inside the GPU call the first time a request picks one
+if __name__ != "__main__":
+    from vampnet import download_finetuned
+    for _name in interface.available_models():
+        if _name != "default":
+            download_finetuned(_name)
 print(f"Loaded initial model: {init_model_choice}")
 print(f"Coarse model path: {interface.coarse_path}")
 print(f"C2F model path: {interface.c2f_path}\n")

vampnet/__init__.py

@@ -15,6 +15,9 @@ from huggingface_hub import hf_hub_download, HfFileSystem
 DEFAULT_HF_MODEL_REPO_DIR = ROOT / "DEFAULT_HF_MODEL_REPO"
 DEFAULT_HF_MODEL_REPO = DEFAULT_HF_MODEL_REPO_DIR.read_text().strip()
 # DEFAULT_HF_MODEL_REPO = "hugggof/vampnet"
+# The weights' revision, which every download below names, so that the models cannot change
+# under the Space
+REVISION = "69e9c7dba72b54a45f93122ef91aa738d2e803ba"
 FS = HfFileSystem()
 
 def download_codec():
@@ -26,7 +29,8 @@ def download_codec():
         repo_id=repo_id,
         filename=filename,
         subfolder=None, 
-        local_dir=MODELS_DIR
+        local_dir=MODELS_DIR,
+        revision=REVISION,
     )
     return codec_path
     
@@ -39,7 +43,7 @@ def download_default():
         path = f"{MODELS_DIR}/{filename}"
         if not Path(path).exists():
             print(f"{path} does not exist, downloading")
-            FS.download(f"{repo_id}/{filename}", path)
+            FS.download(f"{repo_id}@{REVISION}/{filename}", path)
         paths.append(path)
     
     # load the models
@@ -53,14 +57,14 @@ def download_finetuned(name, repo_id=DEFAULT_HF_MODEL_REPO):
         path = f"{MODELS_DIR}/loras/{name}/{filename}"
         if not Path(path).exists():
             print(f"{path} does not exist, downloading")
-            FS.download(f"{repo_id}/loras/{name}/{filename}", path)
+            FS.download(f"{repo_id}@{REVISION}/loras/{name}/{filename}", path)
         paths.append(path)
     
     # load the models
     return paths[0], paths[1]
     
 def list_finetuned(repo_id=DEFAULT_HF_MODEL_REPO):
-    diritems = FS.listdir(f"{repo_id}/loras")
+    diritems = FS.listdir(f"{repo_id}@{REVISION}/loras")
     # iterate through all the names
     valid_diritems = []
     for item in diritems: