| 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], | |
| ) | |