Spaces:
Runtime error
Runtime error
Create download/download.py
Browse files- tabs/download/download.py +80 -0
tabs/download/download.py
ADDED
@@ -0,0 +1,80 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os, sys, shutil
|
2 |
+
import tempfile
|
3 |
+
import gradio as gr
|
4 |
+
from core import run_download_script
|
5 |
+
|
6 |
+
from assets.i18n.i18n import I18nAuto
|
7 |
+
|
8 |
+
i18n = I18nAuto()
|
9 |
+
|
10 |
+
now_dir = os.getcwd()
|
11 |
+
sys.path.append(now_dir)
|
12 |
+
|
13 |
+
gradio_temp_dir = os.path.join(tempfile.gettempdir(), "gradio")
|
14 |
+
|
15 |
+
if os.path.exists(gradio_temp_dir):
|
16 |
+
shutil.rmtree(gradio_temp_dir)
|
17 |
+
|
18 |
+
|
19 |
+
def save_drop_model(dropbox):
|
20 |
+
if "pth" not in dropbox and "index" not in dropbox:
|
21 |
+
raise gr.Error(
|
22 |
+
message="The file you dropped is not a valid model file. Please try again."
|
23 |
+
)
|
24 |
+
else:
|
25 |
+
file_name = os.path.basename(dropbox)
|
26 |
+
if ".pth" in dropbox:
|
27 |
+
model_name = file_name.split(".pth")[0]
|
28 |
+
else:
|
29 |
+
if "v2" not in dropbox:
|
30 |
+
model_name = file_name.split("_nprobe_1_")[1].split("_v1")[0]
|
31 |
+
else:
|
32 |
+
model_name = file_name.split("_nprobe_1_")[1].split("_v2")[0]
|
33 |
+
if " " in model_name:
|
34 |
+
model_name = model_name.replace(" ", "_")
|
35 |
+
model_path = os.path.join(now_dir, "logs", model_name)
|
36 |
+
if not os.path.exists(model_path):
|
37 |
+
os.makedirs(model_path)
|
38 |
+
if os.path.exists(os.path.join(model_path, file_name)):
|
39 |
+
os.remove(os.path.join(model_path, file_name))
|
40 |
+
os.rename(dropbox, os.path.join(model_path, file_name))
|
41 |
+
print(f"{file_name} saved in {model_path}")
|
42 |
+
gr.Info(f"{file_name} saved in {model_path}")
|
43 |
+
return None
|
44 |
+
|
45 |
+
|
46 |
+
def download_tab():
|
47 |
+
with gr.Column():
|
48 |
+
gr.Markdown(value=i18n("## Download Model"))
|
49 |
+
model_link = gr.Textbox(
|
50 |
+
label=i18n("Model Link"),
|
51 |
+
placeholder=i18n("Introduce the model link"),
|
52 |
+
interactive=True,
|
53 |
+
)
|
54 |
+
model_download_output_info = gr.Textbox(
|
55 |
+
label=i18n("Output Information"),
|
56 |
+
value="",
|
57 |
+
max_lines=8,
|
58 |
+
interactive=False,
|
59 |
+
)
|
60 |
+
model_download_button = gr.Button(i18n("Download Model"))
|
61 |
+
model_download_button.click(
|
62 |
+
run_download_script,
|
63 |
+
[model_link],
|
64 |
+
model_download_output_info,
|
65 |
+
api_name="model_download",
|
66 |
+
)
|
67 |
+
gr.Markdown(value=i18n("## Drop files"))
|
68 |
+
|
69 |
+
dropbox = gr.File(
|
70 |
+
label=i18n(
|
71 |
+
"Drag your .pth file and .index file into this space. Drag one and then the other."
|
72 |
+
),
|
73 |
+
type="filepath",
|
74 |
+
)
|
75 |
+
|
76 |
+
dropbox.upload(
|
77 |
+
fn=save_drop_model,
|
78 |
+
inputs=[dropbox],
|
79 |
+
outputs=[dropbox],
|
80 |
+
)
|