feat: restore from hf

app/export.py

@@ -51,7 +51,12 @@ def download_expdir(exp_dir: str) -> str:
 
 
 def upload_to_huggingface(exp_dir: str, repo_id: str, token: str) -> str:
-    commit = upload_folder(repo_id=repo_id, folder_path=exp_dir, token=token)
+    commit = upload_folder(
+        repo_id=repo_id,
+        folder_path=exp_dir,
+        ignore_patterns=["_data", "*.zip", "tmp.wav"],
+        token=token if token.startswith("hf_") else None,
+    )
     return commit.commit_url
 
 

app/setup.py

@@ -5,6 +5,7 @@ import zipfile
 import tempfile
 from infer.modules.train.preprocess import PreProcess
 from typing import Tuple
+from huggingface_hub import snapshot_download
 
 
 def extract_audio_files(zip_file: str, target_dir: str) -> list[str]:
@@ -45,6 +46,14 @@ def restore_expdir(zip: str) -> str:
     return exp_dir
 
 
+def restore_from_huggingface(repo: str, token: str) -> str:
+    exp_dir = os.path.join(tempfile.mkdtemp(), repo.lower())
+    snapshot_download(
+        repo, local_dir=exp_dir, token=token if token.startswith("hf_") else None
+    )
+    return exp_dir
+
+
 def set_dir(dir_val: str) -> str:
     if not dir_val.startswith("/tmp/"):
         dir_val = os.path.join("/tmp", dir_val)
@@ -84,6 +93,18 @@ class SetupTab:
             )
             self.restore_btn = gr.Button(value="Restore Experiment", variant="primary")
 
+        gr.Markdown("You can also restore from a Hugging Face repo.")
+        with gr.Row():
+            self.hf_repo = gr.Textbox(
+                label="Restore from Hugging Face repo",
+                placeholder="username/repo",
+            )
+            self.hf_token = gr.Textbox(
+                label="Hugging Face token (optional)",
+                placeholder="hf_...",
+            )
+            self.restore_hf_btn = gr.Button(value="Restore from Hugging Face")
+
         with gr.Row():
             self.dir_val = gr.Textbox(
                 label="Manually set the experiment directory (don't touch it unless you know what you are doing)",
@@ -104,6 +125,12 @@ class SetupTab:
             outputs=[exp_dir],
         )
 
+        self.restore_hf_btn.click(
+            fn=restore_from_huggingface,
+            inputs=[self.hf_repo, self.hf_token],
+            outputs=[exp_dir],
+        )
+
         self.set_dir_btn.click(
             fn=set_dir,
             inputs=[self.dir_val],