Merge 35ec483351 into d057f2fae9

2024-03-17 21:57:31 +07:00 · 2024-03-17 21:57:31 +07:00 · 066adb48dd
commit 066adb48dd
parent d057f2fae9 35ec483351
6 changed files with 113 additions and 14 deletions
--- a/args_manager.py
+++ b/args_manager.py
@ -4,6 +4,9 @@ import os
 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',
--- 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"
+
 }
--- 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

+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 = 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')
+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 @@ default_advanced_checkbox = get_config_item_or_set_default(
    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,
--- a/modules/download_models.py
+++ 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
--- 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
--- a/webui.py
+++ b/webui.py
@ -16,6 +16,7 @@ import modules.meta_parser
 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
@ -333,6 +334,25 @@ with shared.gradio_root:
                                                       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():
@ -434,6 +454,7 @@ with shared.gradio_root:
                        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,
@ -542,6 +563,9 @@ with shared.gradio_root:
        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