feat: integrate AFx-Rep model downloading using huggingface_hub

app.py

@@ -6,6 +6,8 @@ import torch
 from torch import Tensor
 import math
 import yaml
+import os
+from huggingface_hub import hf_hub_download
 import json
 import pyloudnorm as pyln
 from hydra.utils import instantiate
@@ -34,6 +36,21 @@ from st_ito.utils import (
     load_mir_feature_extractor,
 )
 
+# Download AFx-Rep model and config files
+os.makedirs("tmp", exist_ok=True)
+ckpt_path = hf_hub_download(
+    repo_id="csteinmetz1/afx-rep",
+    filename="afx-rep.ckpt",
+    local_dir="tmp",
+    local_files_only=False,
+)
+config_path = hf_hub_download(
+    repo_id="csteinmetz1/afx-rep",
+    filename="config.yaml",
+    local_dir="tmp",
+    local_files_only=False,
+)
+
 
 title_md = "# Vocal Effects Style Transfer Demo"
 description_md = """

st_ito/utils.py

@@ -5,6 +5,7 @@ import torchaudio
 import pyloudnorm as pyln
 from typing import Optional
 from importlib import import_module
+from huggingface_hub import hf_hub_download
 
 from modules.encoder import (
     StatisticReduction,
@@ -112,18 +113,28 @@ def get_param_embeds(
 def load_param_model(ckpt_path: Optional[str] = None):
 
     if ckpt_path is None:  # look in tmp direcory
-        ckpt_path = os.path.join(os.getcwd(), "tmp", "afx-rep.ckpt")
+        # ckpt_path = os.path.join(os.getcwd(), "tmp", "afx-rep.ckpt")
         os.makedirs("tmp", exist_ok=True)
-        if not os.path.isfile(ckpt_path):
-            # download from huggingfacehub
-            os.system(
-                "wget -O tmp/afx-rep.ckpt https://huggingface.co/csteinmetz1/afx-rep/resolve/main/afx-rep.ckpt"
-            )
-            os.system(
-                "wget -O tmp/config.yaml https://huggingface.co/csteinmetz1/afx-rep/resolve/main/config.yaml"
-            )
-
-    config_path = os.path.join(os.path.dirname(ckpt_path), "config.yaml")
+        # if not os.path.isfile(ckpt_path):
+        # download from huggingfacehub
+        # os.system(
+        #     "wget -O tmp/afx-rep.ckpt https://huggingface.co/csteinmetz1/afx-rep/resolve/main/afx-rep.ckpt"
+        # )
+        # os.system(
+        #     "wget -O tmp/config.yaml https://huggingface.co/csteinmetz1/afx-rep/resolve/main/config.yaml"
+        # )
+        ckpt_path = hf_hub_download(
+            repo_id="csteinmetz1/afx-rep",
+            filename="afx-rep.ckpt",
+            local_dir="tmp",
+        )
+        config_path = hf_hub_download(
+            repo_id="csteinmetz1/afx-rep",
+            filename="config.yaml",
+            local_dir="tmp",
+        )
+    else:
+        config_path = os.path.join(os.path.dirname(ckpt_path), "config.yaml")
 
     with open(config_path) as f:
         config = yaml.safe_load(f)