diff --git a/args_manager.py b/args_manager.py index c7c1b7ab1..2815c8e81 100644 --- a/args_manager.py +++ b/args_manager.py @@ -4,6 +4,9 @@ from tempfile import gettempdir args_parser.parser.add_argument("--share", action='store_true', help="Set whether to share on Gradio.") +args_parser.parser.add_argument("--disable-download-tab", action='store_true', + help="Disables Download Tabs in the UI.") + args_parser.parser.add_argument("--preset", type=str, default=None, help="Apply specified UI preset.") args_parser.parser.add_argument("--language", type=str, default='default', diff --git a/language/en.json b/language/en.json index f61255c96..f007fa690 100644 --- a/language/en.json +++ b/language/en.json @@ -381,5 +381,18 @@ "Metadata Scheme": "Metadata Scheme", "Image Prompt parameters are not included. Use png and a1111 for compatibility with Civitai.": "Image Prompt parameters are not included. Use png and a1111 for compatibility with Civitai.", "fooocus (json)": "fooocus (json)", - "a1111 (plain text)": "a1111 (plain text)" + "a1111 (plain text)": "a1111 (plain text)", + "Download": "Download", + "URL": "URL", + "File Name with its Extension (Optional)": "File Name with its Extension (Optional)", + "Select Path": "Select Path", + "Output": "Output", + "download": "download", + "Download successful! Model saved to": "Download successful! Model saved to", + "Download failed! Please check the URL and try again.": "Download failed! Please check the URL and try again.", + "Downloading: ": "Downloading: ", + "to": "to", + "Failed to download": "Failed to download", + "reason": "reason" + } \ No newline at end of file diff --git a/modules/config.py b/modules/config.py index a68bd2187..9fbd598f8 100644 --- a/modules/config.py +++ b/modules/config.py @@ -166,16 +166,45 @@ def get_dir_or_set_default(key, default_value, as_array=False, make_directory=Fa config_dict[key] = dp return dp - -paths_checkpoints = get_dir_or_set_default('path_checkpoints', ['../models/checkpoints/'], True) -paths_loras = get_dir_or_set_default('path_loras', ['../models/loras/'], True) -path_embeddings = get_dir_or_set_default('path_embeddings', '../models/embeddings/') -path_vae_approx = get_dir_or_set_default('path_vae_approx', '../models/vae_approx/') -path_upscale_models = get_dir_or_set_default('path_upscale_models', '../models/upscale_models/') -path_inpaint = get_dir_or_set_default('path_inpaint', '../models/inpaint/') -path_controlnet = get_dir_or_set_default('path_controlnet', '../models/controlnet/') -path_clip_vision = get_dir_or_set_default('path_clip_vision', '../models/clip_vision/') -path_fooocus_expansion = get_dir_or_set_default('path_fooocus_expansion', '../models/prompt_expansion/fooocus_expansion') +config_paths = { + 'checkpoints': get_dir_or_set_default('path_checkpoints', ['../models/checkpoints/'], True), + 'clip': get_dir_or_set_default('path_clip', ['../models/clip/']), + 'config': get_dir_or_set_default('path_config', '../models/configs/'), + 'diffusers': get_dir_or_set_default('path_diffusers', ['../models/diffusers/']), + 'gligen': get_dir_or_set_default('path_gligen', ['../models/gligen/']), + 'hypernetworks': get_dir_or_set_default('path_hypernetworks', ['../models/hypernetworks/']), + 'prompt_expansion': get_dir_or_set_default('path_prompt_expansion', '../models/prompt_expansion/'), + 'style_models': get_dir_or_set_default('path_style_models', '../models/style_models/'), + 'unet': get_dir_or_set_default('path_unet', '../models/unet/'), + 'vae': get_dir_or_set_default('path_vae', '../models/vae/'), + 'loras': get_dir_or_set_default('path_loras', ['../models/loras/'], True), + 'embeddings': get_dir_or_set_default('path_embeddings', '../models/embeddings/'), + 'vae_approx': get_dir_or_set_default('path_vae_approx', '../models/vae_approx/'), + 'upscale_models': get_dir_or_set_default('path_upscale_models', '../models/upscale_models/'), + 'inpaint': get_dir_or_set_default('path_inpaint', '../models/inpaint/'), + 'controlnet': get_dir_or_set_default('path_controlnet', '../models/controlnet/'), + 'clip_vision': get_dir_or_set_default('path_clip_vision', '../models/clip_vision/'), + 'fooocus_expansion': get_dir_or_set_default('path_fooocus_expansion', '../models/prompt_expansion/fooocus_expansion') +} + +paths_checkpoints = config_paths['checkpoints'] +paths_clip = config_paths['clip'] +paths_config = config_paths['config'] +paths_diffusers = config_paths['diffusers'] +paths_gligen = config_paths['gligen'] +paths_hypernetworks = config_paths['hypernetworks'] +paths_prompt_expansion = config_paths['prompt_expansion'] +paths_style_models = config_paths['style_models'] +paths_unet = config_paths['unet'] +paths_vae = config_paths['vae'] +paths_loras = config_paths['loras'] +path_embeddings = config_paths['embeddings'] +path_vae_approx = config_paths['vae_approx'] +path_upscale_models = config_paths['upscale_models'] +path_inpaint = config_paths['inpaint'] +path_controlnet = config_paths['controlnet'] +path_clip_vision = config_paths['clip_vision'] +path_fooocus_expansion = config_paths['fooocus_expansion'] path_outputs = get_path_output() def get_config_item_or_set_default(key, default_value, validator, disable_empty_as_none=False): @@ -318,6 +347,11 @@ def get_config_item_or_set_default(key, default_value, validator, disable_empty_ default_value=False, validator=lambda x: isinstance(x, bool) ) +default_download_tab_checkbox = get_config_item_or_set_default( + key='default_download_tab_checkbox', + default_value=True, + validator=lambda x: isinstance(x, bool) +) default_max_image_number = get_config_item_or_set_default( key='default_max_image_number', default_value=32, diff --git a/modules/download_models.py b/modules/download_models.py new file mode 100644 index 000000000..97b26ad32 --- /dev/null +++ b/modules/download_models.py @@ -0,0 +1,21 @@ +# download_models.py +import os + +from modules import config +from modules.model_loader import load_file_from_url + + +def download_models(url, selected, file_name=None): + model_paths = config.config_paths + paths = model_paths[selected] + if isinstance(paths, list): + path = os.path.join(*paths) + else: + path = paths + try: + load_file_from_url(url, model_dir=path, progress=True, file_name=file_name) + message = f"Download successful! Model saved to {path}." + except Exception as e: + message = f"Download failed! Please check the URL and try again." + + return message diff --git a/modules/model_loader.py b/modules/model_loader.py index 8ba336a91..f8343ad1e 100644 --- a/modules/model_loader.py +++ b/modules/model_loader.py @@ -20,7 +20,11 @@ def load_file_from_url( file_name = os.path.basename(parts.path) cached_file = os.path.abspath(os.path.join(model_dir, file_name)) if not os.path.exists(cached_file): - print(f'Downloading: "{url}" to {cached_file}\n') - from torch.hub import download_url_to_file - download_url_to_file(url, cached_file, progress=progress) - return cached_file + try: + print(f'Downloading: "{url}" to {cached_file}\n') + from torch.hub import download_url_to_file + download_url_to_file(url, cached_file, progress=progress) + except Exception as e: + print(f"Failed to download \"{url}\" to {cached_file}, reason : {e}") + + return cached_file \ No newline at end of file diff --git a/webui.py b/webui.py index 42dd890f7..588ba3145 100644 --- a/webui.py +++ b/webui.py @@ -16,6 +16,7 @@ import args_manager import copy +from modules.download_models import download_models from modules.sdxl_styles import legal_style_names from modules.private_logger import get_current_html_path from modules.ui_gradio_extensions import reload_javascript @@ -332,6 +333,25 @@ def update_history_link(): show_progress=False).then( lambda: None, _js='()=>{refresh_style_localization();}') + if not args_manager.args.disable_download_tab: + with gr.Tab(label='Download',visible=modules.config.default_download_tab_checkbox)as download_tab: + with gr.Group(): + choices = list(modules.config.config_paths.keys()) + with gr.Row(): + url_input = gr.Textbox(label="URL") + with gr.Row(): + file_name = gr.Textbox(label="File Name with its Extension (Optional)") + with gr.Row(): + selected_path = gr.Dropdown(label='Select Path', choices=choices, + show_label=True) + with gr.Row(): + output = gr.Textbox(value="", label="Output") + with gr.Row(): + start_download = gr.Button("download", label="Download") + + start_download.click(download_models, inputs=[url_input, selected_path, file_name], + outputs=[output], queue=True, show_progress=True) + with gr.Tab(label='Model'): with gr.Group(): with gr.Row(): @@ -433,6 +453,7 @@ def update_history_link(): disable_seed_increment = gr.Checkbox(label='Disable seed increment', info='Disable automatic seed increment when image number is > 1.', value=False) + disable_download_checkbox = gr.Checkbox(label='Disable Download Tab',info='Disable the download tab when clicked',value=modules.config.default_download_tab_checkbox,elem_classes='min_check') if not args_manager.args.disable_metadata: save_metadata_to_images = gr.Checkbox(label='Save Metadata to Images', value=modules.config.default_save_metadata_to_images, @@ -541,6 +562,9 @@ def model_refresh_clicked(): advanced_checkbox.change(lambda x: gr.update(visible=x), advanced_checkbox, advanced_column, queue=False, show_progress=False) \ .then(fn=lambda: None, _js='refresh_grid_delayed', queue=False, show_progress=False) + disable_download_checkbox.change(lambda x: gr.update(visible=x), disable_download_checkbox, download_tab) \ + .then(fn=lambda: None, _js='refresh_grid_delayed', queue=False, show_progress=False) + def inpaint_mode_change(mode): assert mode in modules.flags.inpaint_options