update

main.py

@@ -1,9 +1,11 @@
 #!/usr/bin/python3
 # -*- coding: utf-8 -*-
 import argparse
+from pathlib import Path
 import platform
 
 import gradio as gr
+from huggingface_hub import snapshot_download
 import numpy as np
 import torch
 
@@ -13,6 +15,22 @@ from toolbox.torchaudio.models.mpnet.inference_mpnet import InferenceMPNet
 
 def get_args():
     parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--examples_dir",
+        # default=(project_path / "data").as_posix(),
+        default=(project_path / "data/examples").as_posix(),
+        type=str
+    )
+    parser.add_argument(
+        "--models_repo_id",
+        default="qgyd2021/vm_sound_classification",
+        type=str
+    )
+    parser.add_argument(
+        "--trained_model_dir",
+        default=(project_path / "trained_models").as_posix(),
+        type=str
+    )
     parser.add_argument(
         "--hf_token",
         default=environment.get("hf_token"),
@@ -58,6 +76,18 @@ def when_click_denoise_button(noisy_audio_t, engine: str):
 def main():
     args = get_args()
 
+    examples_dir = Path(args.examples_dir)
+    trained_model_dir = Path(args.trained_model_dir)
+
+    # download models
+    if not trained_model_dir.exists():
+        trained_model_dir.mkdir(parents=True, exist_ok=True)
+        _ = snapshot_download(
+            repo_id=args.models_repo_id,
+            local_dir=trained_model_dir.as_posix(),
+            token=args.hf_token,
+        )
+
     # choices
     denoise_engine_choices = list(denoise_engines.keys())