mirror of
https://github.com/vladmandic/sdnext.git
synced 2026-01-27 15:02:48 +03:00
467 lines
21 KiB
Python
467 lines
21 KiB
Python
import io
|
|
import os
|
|
import time
|
|
import shutil
|
|
import importlib
|
|
import contextlib
|
|
from typing import Dict
|
|
from urllib.parse import urlparse
|
|
import huggingface_hub as hf
|
|
from installer import install, log
|
|
from modules import shared, errors, files_cache
|
|
from modules.upscaler import Upscaler
|
|
from modules.paths import script_path, models_path
|
|
|
|
|
|
loggedin = None
|
|
diffuser_repos = []
|
|
debug = log.trace if os.environ.get('SD_DOWNLOAD_DEBUG', None) is not None else lambda *args, **kwargs: None
|
|
pbar = None
|
|
|
|
|
|
def hf_login(token=None):
|
|
if shared.opts.offline_mode:
|
|
return False
|
|
global loggedin # pylint: disable=global-statement
|
|
token = token or shared.opts.huggingface_token
|
|
token = token.replace("\n", "").replace("\r", "").strip() if token is not None else None
|
|
install('hf_xet', quiet=True)
|
|
if token is None or len(token) <= 4:
|
|
log.debug('HF login: no token provided')
|
|
return False
|
|
if len(shared.opts.huggingface_mirror.strip()) > 0 and os.environ.get('HF_ENDPOINT', None) is None:
|
|
os.environ['HF_ENDPOINT'] = shared.opts.huggingface_mirror.strip()
|
|
if os.environ.get('HUGGING_FACE_HUB_TOKEN', None) is not None:
|
|
os.environ.pop('HUGGING_FACE_HUB_TOKEN', None)
|
|
os.unsetenv('HUGGING_FACE_HUB_TOKEN')
|
|
if os.environ.get('HF_TOKEN', None) is not None:
|
|
os.environ.pop('HF_TOKEN', None)
|
|
os.unsetenv('HF_TOKEN')
|
|
if loggedin != token:
|
|
stdout = io.StringIO()
|
|
try:
|
|
with contextlib.redirect_stdout(stdout):
|
|
hf.logout()
|
|
except Exception:
|
|
pass
|
|
with contextlib.redirect_stdout(stdout):
|
|
hf.login(token=token, add_to_git_credential=False)
|
|
os.environ['HF_TOKEN'] = token
|
|
text = stdout.getvalue() or ''
|
|
obfuscated_token = 'hf_...' + token[-4:]
|
|
line = [l for l in text.split('\n') if 'Token' in l]
|
|
log.info(f'HF login: token="{obfuscated_token}" fn="{hf.constants.HF_TOKEN_PATH}" {line[0] if len(line) > 0 else text}')
|
|
loggedin = token
|
|
return True
|
|
|
|
|
|
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
|
|
jobid = shared.state.begin('Download')
|
|
if hub_id.startswith('huggingface/'):
|
|
hub_id = hub_id.replace('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'HF download: id="{hub_id}" args={download_config}')
|
|
token = token or shared.opts.huggingface_token
|
|
if token is not None and len(token) > 2:
|
|
hf_login(token)
|
|
pipeline_dir = None
|
|
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'HF download error: id="{hub_id}" {e}')
|
|
if 'gated' in str(e):
|
|
shared.log.error(f'HF download error: id="{hub_id}" model access requires login')
|
|
shared.state.end(jobid)
|
|
return None
|
|
if pipeline_dir is None:
|
|
shared.log.error(f'HF download error: id="{hub_id}" no data')
|
|
shared.state.end(jobid)
|
|
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(jobid)
|
|
return pipeline_dir
|
|
|
|
|
|
def load_diffusers_models(clear=True):
|
|
# 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()
|
|
already_found = []
|
|
try:
|
|
for folder in os.listdir(place):
|
|
try:
|
|
name = folder[8:] if folder.startswith('models--') else folder
|
|
folder = os.path.join(place, folder)
|
|
if name.endswith("-prior"):
|
|
continue
|
|
if not os.path.isdir(folder):
|
|
continue
|
|
name = name.replace("--", "/")
|
|
friendly = os.path.join(place, name)
|
|
has_index = os.path.exists(os.path.join(folder, 'model_index.json'))
|
|
|
|
if has_index: # direct download of diffusers model
|
|
repo = { 'name': name, 'filename': name, 'friendly': friendly, 'folder': folder, 'path': folder, 'hash': None, 'mtime': os.path.getmtime(folder), 'model_info': os.path.join(folder, 'model_info.json'), 'model_index': os.path.join(folder, 'model_index.json') }
|
|
diffuser_repos.append(repo)
|
|
continue
|
|
|
|
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
|
|
for snapshot in snapshots: # download using from_pretrained which uses huggingface_hub or huggingface_hub directly and creates snapshot-like structure
|
|
commit = os.path.join(folder, 'snapshots', snapshot)
|
|
mtime = os.path.getmtime(commit)
|
|
info = os.path.join(commit, "model_info.json")
|
|
index = os.path.join(commit, "model_index.json")
|
|
config = os.path.join(commit, "config.json")
|
|
if (not os.path.exists(index)) and (not os.path.exists(info)) and (not os.path.exists(config)):
|
|
debug(f'Diffusers skip model no info: {name}')
|
|
continue
|
|
if name in already_found:
|
|
debug(f'Diffusers skip model already found: {name}')
|
|
continue
|
|
repo = { 'name': name, 'filename': name, 'friendly': friendly, 'folder': folder, 'path': commit, 'hash': snapshot, 'mtime': mtime, 'model_info': info, 'model_index': index, 'model_config': config }
|
|
already_found.append(name)
|
|
diffuser_repos.append(repo)
|
|
if os.path.exists(os.path.join(folder, 'hidden')):
|
|
continue
|
|
except Exception as e:
|
|
debug(f'Error analyzing diffusers model: "{folder}" {e}')
|
|
except Exception as e:
|
|
shared.log.error(f"Error listing diffusers: {place} {e}")
|
|
# shared.log.debug(f'Scanning diffusers cache: folder="{place}" items={len(list(diffuser_repos))} time={time.time()-t0:.2f}')
|
|
return diffuser_repos
|
|
|
|
|
|
def find_diffuser(name: str, full=False):
|
|
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[0]['name']]
|
|
hf_api = hf.HfApi()
|
|
suffix = ''
|
|
if len(name) > 3 and name.count('/') > 1:
|
|
parts = name.split('/')
|
|
name = '/'.join(parts[:2]) # only user/model
|
|
suffix = '/'.join(parts[2:]) # subfolder
|
|
if len(suffix) > 0:
|
|
suffix = '/' + suffix
|
|
models = list(hf_api.list_models(model_name=name, library=['diffusers'], full=True, limit=20, sort="downloads", direction=-1))
|
|
if len(models) == 0:
|
|
models = list(hf_api.list_models(model_name=name, full=True, limit=20, sort="downloads", direction=-1)) # widen search
|
|
models = [m for m in models if m.id.startswith(name)] # filter exact
|
|
shared.log.debug(f'Search model: repo="{name}" {len(models) > 0}')
|
|
if len(models) > 0:
|
|
if not full:
|
|
return models[0].id + suffix
|
|
else:
|
|
return [m.id + suffix for m in models]
|
|
return None
|
|
|
|
|
|
def get_reference_opts(name: str, quiet=False):
|
|
model_opts = {}
|
|
name = name.replace('Diffusers/', 'huggingface/')
|
|
for k, v in shared.reference_models.items():
|
|
model_name = v.get('path', '')
|
|
if k == name or model_name == name:
|
|
model_opts = v
|
|
break
|
|
model_name_split = os.path.splitext(model_name.split('@')[0])[0]
|
|
if k == name or model_name_split == name:
|
|
model_opts = v
|
|
break
|
|
model_name_replace = model_name.replace('huggingface/', '')
|
|
if k == name or model_name_replace == name:
|
|
model_opts = v
|
|
break
|
|
if not model_opts:
|
|
# shared.log.error(f'Reference: model="{name}" not found')
|
|
return {}
|
|
if not quiet:
|
|
desc = model_opts.copy()
|
|
desc.pop('desc', None)
|
|
shared.log.debug(f'Reference: model="{name}" {desc}')
|
|
return model_opts
|
|
|
|
|
|
def load_reference(name: str, variant: str = None, revision: str = None, mirror: str = None, custom_pipeline: str = None):
|
|
if '+' in name:
|
|
name = name.split('+')[0]
|
|
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_closest_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}"')
|
|
from modules.civitai.download_civitai import download_civit_model_thread
|
|
download_civit_model_thread(model_name=model, model_url=url, model_path='', model_type='safetensors', token=shared.opts.civitai_token)
|
|
shared.log.debug(f'Reference download complete: model="{name}"')
|
|
sd_models.list_models()
|
|
info = sd_models.get_closest_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 ssl
|
|
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"})
|
|
|
|
context = ssl._create_unverified_context() # pylint: disable=protected-access
|
|
u = urlopen(req, context=context) # 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(f'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}')
|
|
upscalers = []
|
|
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
|
|
upscaler_types = []
|
|
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
|
|
upscalers += scaler.scalers
|
|
upscaler_types.append(name[8:])
|
|
shared.sd_upscalers = upscalers
|
|
t1 = time.time()
|
|
shared.log.info(f"Available Upscalers: items={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} types={upscaler_types}")
|
|
return [x.name for x in shared.sd_upscalers]
|