feat: upload expdir to hugging face

app/export.py

@@ -4,6 +4,7 @@ import shutil
 import gradio as gr
 from infer.lib.train.process_ckpt import extract_small_model
 from app.train import train_index
+from huggingface_hub import upload_folder
 
 
 def download_weight(exp_dir: str) -> str:
@@ -49,6 +50,11 @@ def download_expdir(exp_dir: str) -> str:
     return f"{exp_dir}.zip"
 
 
+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)
+    return commit.commit_url
+
+
 def remove_legacy_checkpoints(exp_dir: str):
     checkpoints = glob(f"{exp_dir}/G_*.pth")
     if not checkpoints:
@@ -106,6 +112,20 @@ class ExportTab:
             )
             self.download_expdir_output = gr.File(label="Archive experiment directory")
 
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("### Upload to Hugging Face")
+                gr.Markdown(
+                    "You can upload the entire experiment directory to Hugging Face."
+                )
+                self.commit_link = gr.Markdown("")
+            with gr.Column():
+                self.repo_id = gr.Textbox(label="Repository ID")
+                self.token = gr.Textbox(label="Personal access token")
+            self.upload_to_huggingface_btn = gr.Button(
+                value="Upload to Hugging Face", variant="primary"
+            )
+
         with gr.Row():
             self.remove_legacy_checkpoints_btn = gr.Button(
                 value="Remove legacy checkpoints"
@@ -135,6 +155,12 @@ class ExportTab:
             outputs=[self.download_expdir_output],
         )
 
+        self.upload_to_huggingface_btn.click(
+            fn=upload_to_huggingface,
+            inputs=[exp_dir, self.repo_id, self.token],
+            outputs=[self.commit_link],
+        )
+
         self.remove_legacy_checkpoints_btn.click(
             fn=remove_legacy_checkpoints,
             inputs=[exp_dir],