OPENGIT/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2025-04-18 10:58:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

add option to check file hash after download

Browse Source 
if the sha256 hash does not match it will be automatically deleted
This commit is contained in:
w-e-w 2024-04-23 03:08:09 +09:00
parent ddb28b33a3
commit 246c269af8
1 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
modules/modelloader.py
View File

@ -23,6 +23,7 @@ def load_file_from_url(
model_dir: str, model_dir: str,
progress: bool = True, progress: bool = True,
file_name: str | None = None, file_name: str | None = None,
hash_prefix: str | None = None,
) -> str: ) -> str:
"""Download a file from `url` into `model_dir`, using the file present if possible. """Download a file from `url` into `model_dir`, using the file present if possible.
@ -36,11 +37,11 @@ def load_file_from_url(
if not os.path.exists(cached_file): if not os.path.exists(cached_file):
print(f'Downloading: "{url}" to {cached_file}\n') print(f'Downloading: "{url}" to {cached_file}\n')
from torch.hub import download_url_to_file from torch.hub import download_url_to_file
download_url_to_file(url, cached_file, progress=progress) download_url_to_file(url, cached_file, progress=progress, hash_prefix=hash_prefix)
return cached_file return cached_file
def load_models(model_path: str, model_url: str = None, command_path: str = None, ext_filter=None, download_name=None, ext_blacklist=None) -> list: def load_models(model_path: str, model_url: str = None, command_path: str = None, ext_filter=None, download_name=None, ext_blacklist=None, hash_prefix=None) -> list:
""" """
A one-and done loader to try finding the desired models in specified directories. A one-and done loader to try finding the desired models in specified directories.
@ -49,6 +50,7 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None
@param model_path: The location to store/find models in. @param model_path: The location to store/find models in.
@param command_path: A command-line argument to search for models in first. @param command_path: A command-line argument to search for models in first.
@param ext_filter: An optional list of filename extensions to filter by @param ext_filter: An optional list of filename extensions to filter by
@param hash_prefix: the expected sha256 of the model_url
@return: A list of paths containing the desired model(s) @return: A list of paths containing the desired model(s)
""" """
output = [] output = []
@ -78,7 +80,7 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None
if model_url is not None and len(output) == 0: if model_url is not None and len(output) == 0:
if download_name is not None: if download_name is not None:
output.append(load_file_from_url(model_url, model_dir=places[0], file_name=download_name)) output.append(load_file_from_url(model_url, model_dir=places[0], file_name=download_name, hash_prefix=hash_prefix))
else: else:
output.append(model_url) output.append(model_url)