Spaces:
Runtime error
Runtime error
File size: 24,279 Bytes
c19ca42 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 |
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]
|