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