Create download/download.py

tabs/download/download.py

@@ -0,0 +1,80 @@
+import os, sys, shutil
+import tempfile
+import gradio as gr
+from core import run_download_script
+
+from assets.i18n.i18n import I18nAuto
+
+i18n = I18nAuto()
+
+now_dir = os.getcwd()
+sys.path.append(now_dir)
+
+gradio_temp_dir = os.path.join(tempfile.gettempdir(), "gradio")
+
+if os.path.exists(gradio_temp_dir):
+    shutil.rmtree(gradio_temp_dir)
+
+
+def save_drop_model(dropbox):
+    if "pth" not in dropbox and "index" not in dropbox:
+        raise gr.Error(
+            message="The file you dropped is not a valid model file. Please try again."
+        )
+    else:
+        file_name = os.path.basename(dropbox)
+        if ".pth" in dropbox:
+            model_name = file_name.split(".pth")[0]
+        else:
+            if "v2" not in dropbox:
+                model_name = file_name.split("_nprobe_1_")[1].split("_v1")[0]
+            else:
+                model_name = file_name.split("_nprobe_1_")[1].split("_v2")[0]
+        if " " in model_name:
+            model_name = model_name.replace(" ", "_")
+        model_path = os.path.join(now_dir, "logs", model_name)
+        if not os.path.exists(model_path):
+            os.makedirs(model_path)
+        if os.path.exists(os.path.join(model_path, file_name)):
+            os.remove(os.path.join(model_path, file_name))
+        os.rename(dropbox, os.path.join(model_path, file_name))
+        print(f"{file_name} saved in {model_path}")
+        gr.Info(f"{file_name} saved in {model_path}")
+    return None
+
+
+def download_tab():
+    with gr.Column():
+        gr.Markdown(value=i18n("## Download Model"))
+        model_link = gr.Textbox(
+            label=i18n("Model Link"),
+            placeholder=i18n("Introduce the model link"),
+            interactive=True,
+        )
+        model_download_output_info = gr.Textbox(
+            label=i18n("Output Information"),
+            value="",
+            max_lines=8,
+            interactive=False,
+        )
+        model_download_button = gr.Button(i18n("Download Model"))
+        model_download_button.click(
+            run_download_script,
+            [model_link],
+            model_download_output_info,
+            api_name="model_download",
+        )
+        gr.Markdown(value=i18n("## Drop files"))
+
+        dropbox = gr.File(
+            label=i18n(
+                "Drag your .pth file and .index file into this space. Drag one and then the other."
+            ),
+            type="filepath",
+        )
+
+        dropbox.upload(
+            fn=save_drop_model,
+            inputs=[dropbox],
+            outputs=[dropbox],
+        )