Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

added the download model from url feature #2423

Open
wants to merge 9 commits into
base: main
Choose a base branch
from
3 changes: 3 additions & 0 deletions args_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down
15 changes: 14 additions & 1 deletion language/en.json
Original file line number Diff line number Diff line change
Expand Up @@ -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"

}
54 changes: 44 additions & 10 deletions modules/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please do it the other way round.
Instead of using an array and then assigning variables with keys it's more robust to assign variables first and then build an array from vars. This eliminates the need for array key string reading.

'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):
Expand Down Expand Up @@ -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(
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please use a more descriptive naming here, e.g. default_show_download_tab

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,
Expand Down
21 changes: 21 additions & 0 deletions modules/download_models.py
Original file line number Diff line number Diff line change
@@ -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."
codezeros marked this conversation as resolved.
Show resolved Hide resolved

return message
12 changes: 8 additions & 4 deletions modules/model_loader.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@codezeros downloading an invalid URL still reports as successful to the frontend despite raising an exception.
This has to be fixed (reverted).

Failed to download "invalid url" to /Volumes/data/private/development/ai/fooocus/models/checkpoints/invalid url, reason : unknown url type: 'invalid url'

Screenshot 2024-03-09 at 16 19 39

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@mashb1t thanks for notifying the problem i will look into it and also making some improvements based on previous feedback (in progress)


return cached_file
24 changes: 24 additions & 0 deletions webui.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
with gr.Tab(label='Download',visible=modules.config.default_download_tab_checkbox)as 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():
Expand Down Expand Up @@ -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')
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

set interactive to false when arg is set


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,
Expand Down Expand Up @@ -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) \
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do not process/output when arg is true

.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
Expand Down