diff --git a/modules/modelloader.py b/modules/modelloader.py index 1596e11b0..f5a2ff79c 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -10,6 +10,7 @@ import torch from modules import shared from modules.upscaler import Upscaler, UpscalerLanczos, UpscalerNearest, UpscalerNone +from modules.util import load_file_from_url # noqa, backwards compatibility if TYPE_CHECKING: import spandrel @@ -17,35 +18,6 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -def load_file_from_url( - url: str, - *, - model_dir: str, - progress: bool = True, - file_name: str | None = None, - hash_prefix: str | None = None, - re_download: bool = False, -) -> str: - """Download a file from `url` into `model_dir`, using the file present if possible. - Returns the path to the downloaded file. - - file_name: if specified, it will be used as the filename, otherwise the filename will be extracted from the url. - hash_prefix: is provided, the hash of the downloaded file will be checked against. - re_download: re-download the file even if it already exists. - - """ - os.makedirs(model_dir, exist_ok=True) - if not file_name: - parts = urlparse(url) - file_name = os.path.basename(parts.path) - cached_file = os.path.abspath(os.path.join(model_dir, file_name)) - if re_download or 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, hash_prefix=hash_prefix) - 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, hash_prefix=None) -> list: """ A one-and done loader to try finding the desired models in specified directories. diff --git a/modules/util.py b/modules/util.py index 7911b0db7..9e215d799 100644 --- a/modules/util.py +++ b/modules/util.py @@ -211,3 +211,78 @@ Requested path was: {path} subprocess.Popen(["explorer.exe", subprocess.check_output(["wslpath", "-w", path])]) else: subprocess.Popen(["xdg-open", path]) + + +def load_file_from_url( + url: str, + *, + model_dir: str, + progress: bool = True, + file_name: str | None = None, + hash_prefix: str | None = None, + re_download: bool = False, +) -> str: + """Download a file from `url` into `model_dir`, using the file present if possible. + Returns the path to the downloaded file. + + file_name: if specified, it will be used as the filename, otherwise the filename will be extracted from the url. + file is downloaded to {file_name}.tmp then moved to the final location after download is complete. + hash_prefix: sha256 hex string, if provided, the hash of the downloaded file will be checked against this prefix. + if the hash does not match, the temporary file is deleted and a ValueError is raised. + re_download: forcibly re-download the file even if it already exists. + """ + from urllib.parse import urlparse + import requests + try: + from tqdm import tqdm + except ImportError: + class tqdm: + def __init__(self, *args, **kwargs): + pass + + def update(self, n=1, *args, **kwargs): + pass + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + pass + + if not file_name: + parts = urlparse(url) + file_name = os.path.basename(parts.path) + + cached_file = os.path.abspath(os.path.join(model_dir, file_name)) + + if re_download or not os.path.exists(cached_file): + os.makedirs(model_dir, exist_ok=True) + temp_file = os.path.join(model_dir, f"{file_name}.tmp") + print(f'\nDownloading: "{url}" to {cached_file}') + response = requests.get(url, stream=True) + response.raise_for_status() + total_size = int(response.headers.get('content-length', 0)) + with tqdm(total=total_size, unit='B', unit_scale=True, desc=file_name, disable=not progress) as progress_bar: + with open(temp_file, 'wb') as file: + for chunk in response.iter_content(chunk_size=1024): + if chunk: + file.write(chunk) + progress_bar.update(len(chunk)) + + if hash_prefix and not compare_sha256(temp_file, hash_prefix): + print(f"Hash mismatch for {temp_file}. Deleting the temporary file.") + os.remove(temp_file) + raise ValueError(f"File hash does not match the expected hash prefix {hash_prefix}!") + + os.rename(temp_file, cached_file) + return cached_file + + +def compare_sha256(file_path: str, hash_prefix: str) -> bool: + """Check if the SHA256 hash of the file matches the given prefix.""" + import hashlib + sha256 = hashlib.sha256() + with open(file_path, 'rb') as f: + for chunk in iter(lambda: f.read(4096), b''): + sha256.update(chunk) + return sha256.hexdigest().startswith(hash_prefix)