import os import time import json import shutil import importlib from typing import Dict from urllib.parse import urlparse from PIL import Image import rich.progress as p from modules import shared, errors, files_cache from modules.upscaler import Upscaler, UpscalerLanczos, UpscalerNearest, UpscalerNone from modules.paths import script_path, models_path diffuser_repos = [] debug = shared.log.trace if os.environ.get('SD_DOWNLOAD_DEBUG', None) is not None else lambda *args, **kwargs: None def download_civit_meta(model_path: str, model_id): fn = os.path.splitext(model_path)[0] + '.json' url = f'https://civitai.com/api/v1/models/{model_id}' r = shared.req(url) if r.status_code == 200: try: shared.writefile(r.json(), filename=fn, mode='w', silent=True) msg = f'CivitAI download: id={model_id} url={url} file={fn}' shared.log.info(msg) return msg except Exception as e: msg = f'CivitAI download error: id={model_id} url={url} file={fn} {e}' errors.display(e, 'CivitAI download error') shared.log.error(msg) return msg return f'CivitAI download error: id={model_id} url={url} code={r.status_code}' def download_civit_preview(model_path: str, preview_url: str): ext = os.path.splitext(preview_url)[1] preview_file = os.path.splitext(model_path)[0] + ext if os.path.exists(preview_file): return '' res = f'CivitAI download: url={preview_url} file={preview_file}' r = shared.req(preview_url, stream=True) total_size = int(r.headers.get('content-length', 0)) block_size = 16384 # 16KB blocks written = 0 img = None shared.state.begin('civitai') try: with open(preview_file, 'wb') as f: with p.Progress(p.TextColumn('[cyan]{task.description}'), p.DownloadColumn(), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TransferSpeedColumn(), console=shared.console) as progress: task = progress.add_task(description="Download starting", total=total_size) for data in r.iter_content(block_size): written = written + len(data) f.write(data) progress.update(task, advance=block_size, description="Downloading") if written < 1024: # min threshold os.remove(preview_file) raise ValueError(f'removed invalid download: bytes={written}') img = Image.open(preview_file) except Exception as e: os.remove(preview_file) res += f' error={e}' shared.log.error(f'CivitAI download error: url={preview_url} file={preview_file} written={written} {e}') shared.state.end() if img is None: return res shared.log.info(f'{res} size={total_size} image={img.size}') img.close() return res download_pbar = None def download_civit_model_thread(model_name, model_url, model_path, model_type, token): import hashlib sha256 = hashlib.sha256() sha256.update(model_name.encode('utf-8')) temp_file = sha256.hexdigest()[:8] + '.tmp' if model_type == 'LoRA': model_file = os.path.join(shared.opts.lora_dir, model_path, model_name) temp_file = os.path.join(shared.opts.lora_dir, model_path, temp_file) elif model_type == 'Embedding': model_file = os.path.join(shared.opts.embeddings_dir, model_path, model_name) temp_file = os.path.join(shared.opts.embeddings_dir, model_path, temp_file) else: model_file = os.path.join(shared.opts.ckpt_dir, model_path, model_name) temp_file = os.path.join(shared.opts.ckpt_dir, model_path, temp_file) res = f'Model download: name="{model_name}" url="{model_url}" path="{model_path}" temp="{temp_file}"' if os.path.isfile(model_file): res += ' already exists' shared.log.warning(res) return res headers = {} starting_pos = 0 if os.path.isfile(temp_file): starting_pos = os.path.getsize(temp_file) res += f' resume={round(starting_pos/1024/1024)}Mb' headers['Range'] = f'bytes={starting_pos}-' if token is not None and len(token) > 0: headers['Authorization'] = f'Bearer {token}' r = shared.req(model_url, headers=headers, stream=True) total_size = int(r.headers.get('content-length', 0)) res += f' size={round((starting_pos + total_size)/1024/1024, 2)}Mb' shared.log.info(res) shared.state.begin('civitai') block_size = 16384 # 16KB blocks written = starting_pos global download_pbar # pylint: disable=global-statement if download_pbar is None: download_pbar = p.Progress(p.TextColumn('[cyan]{task.description}'), p.DownloadColumn(), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TransferSpeedColumn(), p.TextColumn('[cyan]{task.fields[name]}'), console=shared.console) with download_pbar: task = download_pbar.add_task(description="Download starting", total=starting_pos+total_size, name=model_name) try: with open(temp_file, 'ab') as f: for data in r.iter_content(block_size): if written == 0: try: # check if response is JSON message instead of bytes shared.log.error(f'Model download: response={json.loads(data.decode("utf-8"))}') raise ValueError('response: type=json expected=bytes') except Exception: # this is good pass written = written + len(data) f.write(data) download_pbar.update(task, description="Download", completed=written) if written < 1024: # min threshold os.remove(temp_file) raise ValueError(f'removed invalid download: bytes={written}') """ if preview is not None: preview_file = os.path.splitext(model_file)[0] + '.jpg' preview.save(preview_file) res += f' preview={preview_file}' """ except Exception as e: shared.log.error(f'{res} {e}') finally: download_pbar.stop_task(task) download_pbar.remove_task(task) if starting_pos+total_size != written: shared.log.warning(f'{res} written={round(written/1024/1024)}Mb incomplete download') elif os.path.exists(temp_file): shared.log.debug(f'Model download complete: temp="{temp_file}" path="{model_file}"') os.rename(temp_file, model_file) shared.state.end() return res def download_civit_model(model_url: str, model_name: str, model_path: str, model_type: str, token: str = None): import threading thread = threading.Thread(target=download_civit_model_thread, args=(model_name, model_url, model_path, model_type, token)) thread.start() return f'Model download: name={model_name} url={model_url} path={model_path}' def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config: Dict[str, str] = None, token = None, variant = None, revision = None, mirror = None, custom_pipeline = None): if hub_id is None or len(hub_id) == 0: return None from diffusers import DiffusionPipeline import huggingface_hub as hf shared.state.begin('huggingface') if download_config is None: download_config = { "force_download": False, "resume_download": True, "cache_dir": shared.opts.diffusers_dir, "load_connected_pipeline": True, } if cache_dir is not None: download_config["cache_dir"] = cache_dir if variant is not None and len(variant) > 0: download_config["variant"] = variant if revision is not None and len(revision) > 0: download_config["revision"] = revision if mirror is not None and len(mirror) > 0: download_config["mirror"] = mirror if custom_pipeline is not None and len(custom_pipeline) > 0: download_config["custom_pipeline"] = custom_pipeline shared.log.debug(f'Diffusers downloading: id="{hub_id}" args={download_config}') token = token or shared.opts.huggingface_token if token is not None and len(token) > 2: shared.log.debug(f"Diffusers authentication: {token}") hf.login(token) pipeline_dir = None ok = False err = None if not ok: try: pipeline_dir = DiffusionPipeline.download(hub_id, **download_config) ok = True except Exception as e: err = e ok = False debug(f'Diffusers download error: id="{hub_id}" {e}') if not ok and 'Repository Not Found' not in str(err): try: download_config.pop('load_connected_pipeline', None) download_config.pop('variant', None) pipeline_dir = hf.snapshot_download(hub_id, **download_config) except Exception as e: debug(f'Diffusers download error: id="{hub_id}" {e}') if 'gated' in str(e): shared.log.error(f'Diffusers download error: id="{hub_id}" model access requires login') return None if pipeline_dir is None: shared.log.error(f'Diffusers download error: id="{hub_id}" {err}') return None try: model_info_dict = hf.model_info(hub_id).cardData if pipeline_dir is not None else None except Exception: model_info_dict = None if model_info_dict is not None and "prior" in model_info_dict: # some checkpoints need to be downloaded as "hidden" as they just serve as pre- or post-pipelines of other pipelines download_dir = DiffusionPipeline.download(model_info_dict["prior"][0], **download_config) model_info_dict["prior"] = download_dir with open(os.path.join(download_dir, "hidden"), "w", encoding="utf-8") as f: # mark prior as hidden f.write("True") if pipeline_dir is not None: shared.writefile(model_info_dict, os.path.join(pipeline_dir, "model_info.json")) shared.state.end() return pipeline_dir def load_diffusers_models(clear=True): excluded_models = [ 'PhotoMaker', 'inswapper_128', 'IP-Adapter' ] t0 = time.time() place = shared.opts.diffusers_dir if place is None or len(place) == 0 or not os.path.isdir(place): place = os.path.join(models_path, 'Diffusers') if clear: diffuser_repos.clear() output = [] try: for folder in os.listdir(place): try: if any([x in folder for x in excluded_models]): # noqa:C419 continue if "--" not in folder: continue if folder.endswith("-prior"): continue _, name = folder.split("--", maxsplit=1) name = name.replace("--", "/") folder = os.path.join(place, folder) friendly = os.path.join(place, name) snapshots = os.listdir(os.path.join(folder, "snapshots")) if len(snapshots) == 0: shared.log.warning(f"Diffusers folder has no snapshots: location={place} folder={folder} name={name}") continue commit = os.path.join(folder, 'snapshots', snapshots[-1]) mtime = os.path.getmtime(commit) info = os.path.join(commit, "model_info.json") diffuser_repos.append({ 'name': name, 'filename': name, 'friendly': friendly, 'folder': folder, 'path': commit, 'hash': commit, 'mtime': mtime, 'model_info': info }) if os.path.exists(os.path.join(folder, 'hidden')): continue output.append(name) except Exception: # shared.log.error(f"Error analyzing diffusers model: {folder} {e}") pass except Exception as e: shared.log.error(f"Error listing diffusers: {place} {e}") shared.log.debug(f'Scanning diffusers cache: folder={place} items={len(output)} time={time.time()-t0:.2f}') return output def find_diffuser(name: str): repo = [r for r in diffuser_repos if name == r['name'] or name == r['friendly'] or name == r['path']] if len(repo) > 0: return repo['name'] import huggingface_hub as hf hf_api = hf.HfApi() hf_filter = hf.ModelFilter( model_name=name, # task='text-to-image', library=['diffusers'], ) models = list(hf_api.list_models(filter=hf_filter, full=True, limit=20, sort="downloads", direction=-1)) shared.log.debug(f'Searching diffusers models: {name} {len(models) > 0}') if len(models) > 0: return models[0].modelId return None def get_reference_opts(name: str): model_opts = {} for k, v in shared.reference_models.items(): model_name = os.path.splitext(v.get('path', '').split('@')[0])[0] if k == name or model_name == name: model_opts = v break if not model_opts: # shared.log.error(f'Reference: model="{name}" not found') return {} shared.log.debug(f'Reference: model="{name}" {model_opts.get("extras", None)}') return model_opts def load_reference(name: str, variant: str = None, revision: str = None, mirror: str = None, custom_pipeline: str = None): found = [r for r in diffuser_repos if name == r['name'] or name == r['friendly'] or name == r['path']] if len(found) > 0: # already downloaded model_opts = get_reference_opts(found[0]['name']) return True else: model_opts = get_reference_opts(name) if model_opts.get('skip', False): return True shared.log.debug(f'Reference: download="{name}"') model_dir = download_diffusers_model( hub_id=name, cache_dir=shared.opts.diffusers_dir, variant=variant or model_opts.get('variant', None), revision=revision or model_opts.get('revision', None), mirror=mirror or model_opts.get('mirror', None), custom_pipeline=custom_pipeline or model_opts.get('custom_pipeline', None) ) if model_dir is None: shared.log.error(f'Reference download: model="{name}"') return False else: shared.log.debug(f'Reference download complete: model="{name}"') model_opts = get_reference_opts(name) from modules import sd_models sd_models.list_models() return True def load_civitai(model: str, url: str): from modules import sd_models name, _ext = os.path.splitext(model) info = sd_models.get_closet_checkpoint_match(name) if info is not None: _model_opts = get_reference_opts(info.model_name) return name # already downloaded else: shared.log.debug(f'Reference download start: model="{name}"') download_civit_model_thread(model_name=model, model_url=url, model_path='', model_type='safetensors', token=None) shared.log.debug(f'Reference download complete: model="{name}"') sd_models.list_models() info = sd_models.get_closet_checkpoint_match(name) if info is not None: shared.log.debug(f'Reference: model="{name}"') return name # already downloaded else: shared.log.error(f'Reference model="{name}" not found') return None def download_url_to_file(url: str, dst: str): # based on torch.hub.download_url_to_file import uuid import tempfile from urllib.request import urlopen, Request from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn file_size = None req = Request(url, headers={"User-Agent": "sdnext"}) u = urlopen(req) # pylint: disable=R1732 meta = u.info() if hasattr(meta, 'getheaders'): content_length = meta.getheaders("Content-Length") else: content_length = meta.get_all("Content-Length") # pylint: disable=R1732 if content_length is not None and len(content_length) > 0: file_size = int(content_length[0]) dst = os.path.expanduser(dst) for _seq in range(tempfile.TMP_MAX): tmp_dst = dst + '.' + uuid.uuid4().hex + '.partial' try: f = open(tmp_dst, 'w+b') # pylint: disable=R1732 except FileExistsError: continue break else: shared.log.error('Error downloading: url={url} no usable temporary filename found') return try: with Progress(TextColumn('[cyan]{task.description}'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn(), console=shared.console) as progress: task = progress.add_task(description="Downloading", total=file_size) while True: buffer = u.read(8192) if len(buffer) == 0: break f.write(buffer) progress.update(task, advance=len(buffer)) f.close() shutil.move(f.name, dst) finally: f.close() if os.path.exists(f.name): os.remove(f.name) def load_file_from_url(url: str, *, model_dir: str, progress: bool = True, file_name = None): # pylint: disable=unused-argument """Download a file from url into model_dir, using the file present if possible. Returns the path to the downloaded file.""" if model_dir is None: shared.log.error('Download folder is none') 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 not os.path.exists(cached_file): shared.log.info(f'Downloading: url="{url}" file={cached_file}') download_url_to_file(url, cached_file) if os.path.exists(cached_file): return cached_file else: return None def load_models(model_path: str, model_url: str = None, command_path: str = None, ext_filter=None, download_name=None, ext_blacklist=None) -> list: """ A one-and done loader to try finding the desired models in specified directories. @param download_name: Specify to download from model_url immediately. @param model_url: If no other models are found, this will be downloaded on upscale. @param model_path: The location to store/find models in. @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 @return: A list of paths containing the desired model(s) """ places = [x for x in list(set([model_path, command_path])) if x is not None] # noqa:C405 output = [] try: output:list = [*files_cache.list_files(*places, ext_filter=ext_filter, ext_blacklist=ext_blacklist)] if model_url is not None and len(output) == 0: if download_name is not None: dl = load_file_from_url(model_url, model_dir=places[0], progress=True, file_name=download_name) if dl is not None: output.append(dl) else: output.append(model_url) except Exception as e: errors.display(e,f"Error listing models: {files_cache.unique_directories(places)}") return output def friendly_name(file: str): if "http" in file: file = urlparse(file).path file = os.path.basename(file) model_name, _extension = os.path.splitext(file) return model_name def friendly_fullname(file: str): if "http" in file: file = urlparse(file).path file = os.path.basename(file) return file def cleanup_models(): # This code could probably be more efficient if we used a tuple list or something to store the src/destinations # and then enumerate that, but this works for now. In the future, it'd be nice to just have every "model" scaler # somehow auto-register and just do these things... root_path = script_path src_path = models_path dest_path = os.path.join(models_path, "Stable-diffusion") # move_files(src_path, dest_path, ".ckpt") # move_files(src_path, dest_path, ".safetensors") src_path = os.path.join(root_path, "ESRGAN") dest_path = os.path.join(models_path, "ESRGAN") move_files(src_path, dest_path) src_path = os.path.join(models_path, "BSRGAN") dest_path = os.path.join(models_path, "ESRGAN") move_files(src_path, dest_path, ".pth") src_path = os.path.join(root_path, "gfpgan") dest_path = os.path.join(models_path, "GFPGAN") move_files(src_path, dest_path) src_path = os.path.join(root_path, "SwinIR") dest_path = os.path.join(models_path, "SwinIR") move_files(src_path, dest_path) src_path = os.path.join(root_path, "repositories/latent-diffusion/experiments/pretrained_models/") dest_path = os.path.join(models_path, "LDSR") move_files(src_path, dest_path) src_path = os.path.join(root_path, "SCUNet") dest_path = os.path.join(models_path, "SCUNet") move_files(src_path, dest_path) def move_files(src_path: str, dest_path: str, ext_filter: str = None): try: if not os.path.exists(dest_path): os.makedirs(dest_path) if os.path.exists(src_path): for file in os.listdir(src_path): fullpath = os.path.join(src_path, file) if os.path.isfile(fullpath): if ext_filter is not None: if ext_filter not in file: continue shared.log.warning(f"Moving {file} from {src_path} to {dest_path}.") try: shutil.move(fullpath, dest_path) except Exception: pass if len(os.listdir(src_path)) == 0: shared.log.info(f"Removing empty folder: {src_path}") shutil.rmtree(src_path, True) except Exception: pass def load_upscalers(): # We can only do this 'magic' method to dynamically load upscalers if they are referenced, so we'll try to import any _model.py files before looking in __subclasses__ t0 = time.time() modules_dir = os.path.join(shared.script_path, "modules", "postprocess") for file in os.listdir(modules_dir): if "_model.py" in file: model_name = file.replace("_model.py", "") full_model = f"modules.postprocess.{model_name}_model" try: importlib.import_module(full_model) except Exception as e: shared.log.error(f'Error loading upscaler: {model_name} {e}') datas = [] commandline_options = vars(shared.cmd_opts) # some of upscaler classes will not go away after reloading their modules, and we'll end up with two copies of those classes. The newest copy will always be the last in the list, so we go from end to beginning and ignore duplicates used_classes = {} for cls in reversed(Upscaler.__subclasses__()): classname = str(cls) if classname not in used_classes: used_classes[classname] = cls names = [] for cls in reversed(used_classes.values()): name = cls.__name__ cmd_name = f"{name.lower().replace('upscaler', '')}_models_path" commandline_model_path = commandline_options.get(cmd_name, None) scaler = cls(commandline_model_path) scaler.user_path = commandline_model_path scaler.model_download_path = commandline_model_path or scaler.model_path datas += scaler.scalers names.append(name[8:]) shared.sd_upscalers = sorted(datas, key=lambda x: x.name.lower() if not isinstance(x.scaler, (UpscalerNone, UpscalerLanczos, UpscalerNearest)) else "") # Special case for UpscalerNone keeps it at the beginning of the list. t1 = time.time() shared.log.debug(f"Load upscalers: total={len(shared.sd_upscalers)} downloaded={len([x for x in shared.sd_upscalers if x.data_path is not None and os.path.isfile(x.data_path)])} user={len([x for x in shared.sd_upscalers if x.custom])} time={t1-t0:.2f} {names}") return [x.name for x in shared.sd_upscalers]