diff --git a/extensions-builtin/Lora/lora.py b/extensions-builtin/Lora/lora.py index 3cbd91646..40863787b 100644 --- a/extensions-builtin/Lora/lora.py +++ b/extensions-builtin/Lora/lora.py @@ -265,19 +265,32 @@ def lora_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.Mu return current_names = getattr(self, "lora_current_names", ()) + lora_prev_names = getattr(self, "lora_prev_names", ()) wanted_names = tuple((x.name, x.multiplier) for x in loaded_loras) weights_backup = getattr(self, "lora_weights_backup", None) - if weights_backup is None: + if weights_backup is None and len(loaded_loras): if isinstance(self, torch.nn.MultiheadAttention): weights_backup = (self.in_proj_weight.to(devices.cpu, copy=True), self.out_proj.weight.to(devices.cpu, copy=True)) else: weights_backup = self.weight.to(devices.cpu, copy=True) self.lora_weights_backup = weights_backup + elif lora_prev_names != current_names: + self.lora_weights_backup = None + weights_backup = None + elif len(loaded_loras) == 0: + self.lora_weights_backup = None - if current_names != wanted_names: - if weights_backup is not None: + if current_names != wanted_names or current_names != lora_prev_names: + if weights_backup is not None and current_names != lora_prev_names: + if isinstance(self, torch.nn.MultiheadAttention): + self.in_proj_weight.copy_(weights_backup[0]) + self.out_proj.weight.copy_(weights_backup[1]) + else: + self.weight.copy_(weights_backup) + elif weights_backup is not None and current_names == (): + # print('lora restore weight') if isinstance(self, torch.nn.MultiheadAttention): self.in_proj_weight.copy_(weights_backup[0]) self.out_proj.weight.copy_(weights_backup[1]) @@ -310,9 +323,11 @@ def lora_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.Mu print(f'failed to calculate lora weights for layer {lora_layer_name}') + setattr(self, "lora_prev_names", current_names) setattr(self, "lora_current_names", wanted_names) + def lora_reset_cached_weight(self: Union[torch.nn.Conv2d, torch.nn.Linear]): setattr(self, "lora_current_names", ()) setattr(self, "lora_weights_backup", None) diff --git a/extensions-builtin/sd-webui-controlnet b/extensions-builtin/sd-webui-controlnet index 4366c66e8..815b93021 160000 --- a/extensions-builtin/sd-webui-controlnet +++ b/extensions-builtin/sd-webui-controlnet @@ -1 +1 @@ -Subproject commit 4366c66e8393ca307dadf7431689416f61b91e1b +Subproject commit 815b930217873dc3bd72f7d9b518b017a6404ad8 diff --git a/launch.py b/launch.py index 76745b4c2..72e65eeb8 100644 --- a/launch.py +++ b/launch.py @@ -41,8 +41,7 @@ def commit_hash(): def run(command, desc=None, errdesc=None, custom_env=None, live=False): if desc is not None: - from rich import print # pylint: disable=redefined-builtin,wrong-import-order - print(desc) + installer.log(desc) if live: result = subprocess.run(command, check=False, shell=True, env=os.environ if custom_env is None else custom_env) if result.returncode != 0: @@ -97,7 +96,7 @@ if __name__ == "__main__": installer.extensions_preload(force=True) installer.log.info(f"Server arguments: {sys.argv[1:]}") installer.log.debug('Starting WebUI') - logging.disable(logging.INFO) + logging.disable(logging.NOTSET if args.debug else logging.DEBUG) if args.test: installer.log.info("Test only") import webui diff --git a/modules/api/api.py b/modules/api/api.py index fbe64a037..fa8110469 100644 --- a/modules/api/api.py +++ b/modules/api/api.py @@ -561,7 +561,7 @@ class Api: return TrainResponse(info=f"train embedding error: {error}") def shutdown(self): - print('shutdown request received') + shared.log.info('Shutdown request received') # from modules.shared import demo # demo.close() # time.sleep(0.5) diff --git a/modules/errors.py b/modules/errors.py index 181592d62..6b1efabc7 100644 --- a/modules/errors.py +++ b/modules/errors.py @@ -1,11 +1,10 @@ -import sys import logging import warnings -from rich import print # pylint: disable=redefined-builtin from rich.console import Console from rich.theme import Theme from rich.pretty import install as pretty_install from rich.traceback import install as traceback_install +from installer import log console = Console(log_time=True, log_time_format='%H:%M:%S-%f', theme=Theme({ "traceback.border": "black", @@ -23,18 +22,18 @@ def install(suppress=[]): pretty_install(console=console) traceback_install(console=console, extra_lines=1, width=console.width, word_wrap=False, indent_guides=False, suppress=suppress) logging.basicConfig(level=logging.INFO, format='%(asctime)s | %(levelname)s | %(pathname)s | %(message)s') - for handler in logging.getLogger().handlers: - handler.setLevel(logging.INFO) + # for handler in logging.getLogger().handlers: + # handler.setLevel(logging.INFO) def print_error_explanation(message): lines = message.strip().split("\n") for line in lines: - print(line, file=sys.stderr) + log.error(line) def display(e: Exception, task, suppress=[]): - print(f"{task or 'error'}: {type(e).__name__}", file=sys.stderr) + log.error(f"{task or 'error'}: {type(e).__name__}") console.print_exception(show_locals=False, max_frames=2, extra_lines=1, suppress=suppress, theme="ansi_dark", word_wrap=False, width=min([console.width, 200])) diff --git a/modules/extras.py b/modules/extras.py index 6c6704079..df0def4fd 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -53,9 +53,7 @@ def create_config(ckpt_result, config_source, a, b, c): filename, _ = os.path.splitext(ckpt_result) checkpoint_filename = filename + ".yaml" - print("Copying config:") - print(" from:", cfg) - print(" to:", checkpoint_filename) + shared.log.info("Copying config: {cfg} -> {checkpoint_filename}") shutil.copyfile(cfg, checkpoint_filename) @@ -134,14 +132,14 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_ if theta_func2: shared.state.textinfo = "Loading B" - print(f"Loading {secondary_model_info.filename}...") + shared.log.info(f"Loading {secondary_model_info.filename}...") theta_1 = sd_models.read_state_dict(secondary_model_info.filename) else: theta_1 = None if theta_func1: shared.state.textinfo = "Loading C" - print(f"Loading {tertiary_model_info.filename}...") + shared.log.info(f"Loading {tertiary_model_info.filename}...") theta_2 = sd_models.read_state_dict(tertiary_model_info.filename) shared.state.textinfo = 'Merging B and C' @@ -163,10 +161,10 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_ shared.state.nextjob() shared.state.textinfo = f"Loading {primary_model_info.filename}..." - print(f"Loading {primary_model_info.filename}...") + shared.log.info(f"Loading {primary_model_info.filename}...") theta_0 = sd_models.read_state_dict(primary_model_info.filename) - print("Merging...") + shared.log.info("Merging...") shared.state.textinfo = 'Merging A and B' shared.state.sampling_steps = len(theta_0.keys()) for key in tqdm.tqdm(theta_0.keys()): @@ -205,7 +203,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_ bake_in_vae_filename = sd_vae.vae_dict.get(bake_in_vae, None) if bake_in_vae_filename is not None: - print(f"Baking in VAE from {bake_in_vae_filename}") + shared.log.info(f"Baking in VAE from {bake_in_vae_filename}") shared.state.textinfo = 'Baking in VAE' vae_dict = sd_vae.load_vae_dict(bake_in_vae_filename) @@ -237,7 +235,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_ shared.state.nextjob() shared.state.textinfo = "Saving" - print(f"Saving to {output_modelname}...") + shared.log.info(f"Saving to {output_modelname}...") _, extension = os.path.splitext(output_modelname) if extension.lower() == ".safetensors": @@ -249,7 +247,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_ create_config(output_modelname, config_source, primary_model_info, secondary_model_info, tertiary_model_info) - print(f"Checkpoint saved to {output_modelname}.") + shared.log.info(f"Checkpoint saved to {output_modelname}.") shared.state.textinfo = "Checkpoint saved" shared.state.end() diff --git a/modules/generation_parameters_copypaste.py b/modules/generation_parameters_copypaste.py index b6609e224..3f8b5fbf0 100644 --- a/modules/generation_parameters_copypaste.py +++ b/modules/generation_parameters_copypaste.py @@ -56,7 +56,7 @@ def image_from_url_text(filedata): if is_in_right_dir: return Image.open(filename) else: - print(f'Attempted to open file outside of allowed directories: {filename}') + shared.log.warning(f'Attempted to open file outside of allowed directories: {filename}') if type(filedata) == list: if len(filedata) == 0: diff --git a/modules/hashes.py b/modules/hashes.py index 7a19e31c5..68c8422e3 100644 --- a/modules/hashes.py +++ b/modules/hashes.py @@ -1,13 +1,11 @@ import hashlib import json import os.path - import filelock - +from rich import progress from modules import shared from modules.paths import data_path - cache_filename = os.path.join(data_path, "cache.json") cache_data = None @@ -19,8 +17,7 @@ def dump_cache(): def cache(subsection): - global cache_data - + global cache_data # pylint: disable=global-statement if cache_data is None: with filelock.FileLock(cache_filename+".lock"): if not os.path.isfile(cache_filename): @@ -39,7 +36,7 @@ def calculate_sha256(filename): hash_sha256 = hashlib.sha256() blksize = 1024 * 1024 - with open(filename, "rb") as f: + with progress.open(filename, 'rb', description=f'Calculating model hash: [cyan]{filename}', auto_refresh=True) as f: for chunk in iter(lambda: f.read(blksize), b""): hash_sha256.update(chunk) @@ -72,9 +69,7 @@ def sha256(filename, title): if shared.cmd_opts.no_hashing: return None - print(f"Calculating sha256: {filename}", end='') sha256_value = calculate_sha256(filename) - print(f"{sha256_value}") hashes[title] = { "mtime": os.path.getmtime(filename), diff --git a/modules/images.py b/modules/images.py index ce384fd71..610fd7291 100644 --- a/modules/images.py +++ b/modules/images.py @@ -14,7 +14,7 @@ import piexif.helper from PIL import Image, ImageFont, ImageDraw, PngImagePlugin, ExifTags from modules import sd_samplers, shared, script_callbacks, errors -from modules.shared import opts, cmd_opts # pylint: disable=unused-import +from modules.shared import opts, log LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS) @@ -258,7 +258,7 @@ def resize_image(resize_mode, im, width, height, upscaler_name=None): upscalers = [x for x in shared.sd_upscalers if x.name == upscaler_name] if len(upscalers) == 0: upscaler = shared.sd_upscalers[0] - print(f"could not find upscaler named {upscaler_name or ''}, using {upscaler.name} as a fallback") + log.warning(f"could not find upscaler named {upscaler_name or ''}, using {upscaler.name} as a fallback") else: upscaler = upscalers[0] @@ -588,7 +588,7 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i return fullfn, txt_fullfn def safe_decode_string(s: bytes): - remove_prefix = lambda text, prefix: text[len(prefix):] if text.startswith(prefix) else text + remove_prefix = lambda text, prefix: text[len(prefix):] if text.startswith(prefix) else text # pylint: disable=unnecessary-lambda-assignment for encoding in ['utf-8', 'utf-16', 'ascii', 'latin_1', 'cp1252', 'cp437']: # try different encodings try: s = remove_prefix(s, b'UNICODE') diff --git a/modules/img2img.py b/modules/img2img.py index de889d4b7..612d68b79 100644 --- a/modules/img2img.py +++ b/modules/img2img.py @@ -4,10 +4,11 @@ from PIL import Image, ImageOps, ImageFilter, ImageEnhance, ImageChops, Unidenti import modules.scripts from modules import sd_samplers from modules.generation_parameters_copypaste import create_override_settings_dict -from modules.processing import Processed, StableDiffusionProcessingImg2Img, process_images, memory_stats -from modules.shared import opts, cmd_opts, log, state, listfiles, sd_model +from modules.processing import Processed, StableDiffusionProcessingImg2Img, process_images +from modules.shared import opts, debug, state, listfiles, sd_model, log from modules.ui import plaintext_to_html import modules.processing as processing +from modules.memstats import memory_stats def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args): @@ -18,8 +19,8 @@ def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args): inpaint_masks = listfiles(inpaint_mask_dir) is_inpaint_batch = len(inpaint_masks) > 0 if is_inpaint_batch: - print(f"\nInpaint batch is enabled. {len(inpaint_masks)} masks found.") - print(f"Will process {len(images)} images, creating {p.n_iter * p.batch_size} new images for each.") + log.info(f"\nInpaint batch is enabled. {len(inpaint_masks)} masks found.") + log.info(f"Will process {len(images)} images, creating {p.n_iter * p.batch_size} new images for each.") save_normally = output_dir == '' p.do_not_save_grid = True p.do_not_save_samples = not save_normally @@ -60,8 +61,7 @@ def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args): if processed_image.mode == 'RGBA': processed_image = processed_image.convert("RGB") processed_image.save(os.path.join(output_dir, filename)) - if cmd_opts.debug: - log.info(f'Processed: {len(images)} Memory: {memory_stats()} batch') + debug(f'Processed: {len(images)} Memory: {memory_stats()} batch') def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_styles, init_img, sketch, init_img_with_mask, inpaint_color_sketch, inpaint_color_sketch_orig, init_img_inpaint, init_mask_inpaint, steps: int, sampler_index: int, mask_blur: int, mask_alpha: float, inpainting_fill: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, image_cfg_scale: float, denoising_strength: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, resize_mode: int, inpaint_full_res: bool, inpaint_full_res_padding: int, inpainting_mask_invert: int, img2img_batch_input_dir: str, img2img_batch_output_dir: str, img2img_batch_inpaint_mask_dir: str, override_settings_texts, *args): # pylint: disable=unused-argument @@ -137,7 +137,6 @@ def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_s if mask: p.extra_generation_params["Mask blur"] = mask_blur if is_batch: - assert not cmd_opts.hide_ui_dir_config, "Launched with --hide-ui-dir-config, batch img2img disabled" process_batch(p, img2img_batch_input_dir, img2img_batch_output_dir, img2img_batch_inpaint_mask_dir, args) processed = Processed(p, [], p.seed, "") else: @@ -146,6 +145,5 @@ def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_s processed = process_images(p) p.close() generation_info_js = processed.js() - if cmd_opts.debug: - log.info(f'Processed: {len(processed.images)} Memory: {memory_stats()} img') + debug(f'Processed: {len(processed.images)} Memory: {memory_stats()} img') return processed.images, generation_info_js, plaintext_to_html(processed.info), plaintext_to_html(processed.comments) diff --git a/modules/import_hook.py b/modules/import_hook.py index 0cdac360f..c94c3b16b 100644 --- a/modules/import_hook.py +++ b/modules/import_hook.py @@ -1,15 +1,15 @@ import sys -from modules.shared import opts +from modules.shared import opts, log # this will break any attempt to import xformers which will prevent stability diffusion repo from trying to use it try: - import xformers # pylint: disable=unused-import - import xformers.ops # pylint: disable=unused-import + import xformers # pylint: disable=unused-import, import-error + import xformers.ops # pylint: disable=unused-import, import-error except: pass if opts.cross_attention_optimization != "xFormers": if sys.modules.get("xformers", None) is not None: - print('Unloading xFormers') + log.info('Unloading xFormers') sys.modules["xformers"] = None sys.modules["xformers.ops"] = None diff --git a/modules/interrogate.py b/modules/interrogate.py index 93bb08f20..d355d70eb 100644 --- a/modules/interrogate.py +++ b/modules/interrogate.py @@ -6,10 +6,10 @@ import re import torch try: - import intel_extension_for_pytorch as ipex + import intel_extension_for_pytorch as ipex # pylint: disable=unused-import,import-error except: pass -import torch.hub +import torch.hub # pylint: disable=ungrouped-imports from torchvision import transforms from torchvision.transforms.functional import InterpolationMode @@ -28,7 +28,7 @@ def category_types(): def download_default_clip_interrogate_categories(content_dir): - print("Downloading CLIP categories...") + shared.log.info("Downloading CLIP categories...") tmpdir = content_dir + "_tmp" cat_types = ["artists", "flavors", "mediums", "movements"] diff --git a/modules/memmon.py b/modules/memmon.py index 97199445d..c9d7131ef 100644 --- a/modules/memmon.py +++ b/modules/memmon.py @@ -64,27 +64,6 @@ class MemUsageMonitor(threading.Thread): self.data["min_free"] = min(self.data["min_free"], free) time.sleep(1 / self.opts.memmon_poll_rate) - def dump_debug(self): - try: - print(self, 'recorded data:') - for k, v in self.read().items(): - print(k, -(v // -(1024 ** 2))) - print(self, 'raw torch memory stats:') - if shared.cmd_opts.use_ipex: - tm = torch.xpu.memory_stats("xpu") - else: - tm = torch.cuda.memory_stats(self.device) - for k, v in tm.items(): - if 'bytes' not in k: - continue - print('\t' if 'peak' in k else '', k, -(v // -(1024 ** 2))) - if shared.cmd_opts.use_ipex: - print(torch.xpu.memory_summary()) - else: - print(torch.cuda.memory_summary()) - except: - self.disabled = True - def monitor(self): self.run_flag.set() diff --git a/modules/memstats.py b/modules/memstats.py new file mode 100644 index 000000000..5e59bb07d --- /dev/null +++ b/modules/memstats.py @@ -0,0 +1,40 @@ +import os +import psutil +import torch + +def memory_stats(): + def gb(val: float): + return round(val / 1024 / 1024 / 1024, 2) + mem = {} + try: + process = psutil.Process(os.getpid()) + res = process.memory_info() + ram_total = 100 * res.rss / process.memory_percent() + ram = { 'used': gb(res.rss), 'total': gb(ram_total) } + mem.update({ 'ram': ram }) + except Exception as e: + mem.update({ 'ram': e }) + try: + s = torch.cuda.mem_get_info() + gpu = { 'used': gb(s[1] - s[0]), 'total': gb(s[1]) } + s = dict(torch.cuda.memory_stats()) + mem.update({ + 'gpu': gpu, + 'retries': s['num_alloc_retries'], + 'oom': s['num_ooms'] + }) + return mem + except: + pass + try: + gpu = { 'used': gb(torch.xpu.memory_allocated()), 'total': gb(torch.xpu.get_device_properties("xpu").total_memory) } + s = dict(torch.xpu.memory_stats("xpu")) + mem.update({ + 'gpu': gpu, + 'retries': s['num_alloc_retries'], + 'oom': s['num_ooms'] + }) + return mem + except: + pass + return mem diff --git a/modules/middleware.py b/modules/middleware.py index 6520a0ed6..74c12303c 100644 --- a/modules/middleware.py +++ b/modules/middleware.py @@ -10,12 +10,13 @@ from starlette.responses import JSONResponse from fastapi import FastAPI, Request, Response from fastapi.exceptions import HTTPException from fastapi.encoders import jsonable_encoder +from installer import log import modules.errors as errors errors.install() def setup_middleware(app: FastAPI, cmd_opts): - print('Initializing middleware') + log.info('Initializing middleware') uvicorn_logger=logging.getLogger("uvicorn.error") uvicorn_logger.disabled = True from fastapi.middleware.cors import CORSMiddleware @@ -38,7 +39,7 @@ def setup_middleware(app: FastAPI, cmd_opts): res.headers["X-Process-Time"] = duration endpoint = req.scope.get('path', 'err') if cmd_opts.api_log and endpoint.startswith('/sdapi'): - print('API {t} {code} {prot}/{ver} {method} {endpoint} {cli} {duration}'.format( # pylint: disable=consider-using-f-string + log.info('API {t} {code} {prot}/{ver} {method} {endpoint} {cli} {duration}'.format( # pylint: disable=consider-using-f-string t = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"), code = res.status_code, ver = req.scope.get('http_version', '0.0'), @@ -57,7 +58,7 @@ def setup_middleware(app: FastAPI, cmd_opts): "body": vars(e).get('body', ''), "errors": str(e), } - print(f"API error: {req.method}: {req.url} {err}") + log.error(f"API error: {req.method}: {req.url} {err}") if not isinstance(e, HTTPException) and err['error'] != 'TypeError': # do not print backtrace on known httpexceptions errors.display(e, 'HTTP API', [anyio, fastapi, uvicorn, starlette]) return JSONResponse(status_code=vars(e).get('status_code', 500), content=jsonable_encoder(err)) @@ -67,7 +68,7 @@ def setup_middleware(app: FastAPI, cmd_opts): try: return await call_next(req) except CancelledError: - print('WebSocket closed (ignore asyncio.exceptions.CancelledError)') + log.warning('WebSocket closed (ignore asyncio.exceptions.CancelledError)') except BaseException as e: return handle_exception(req, e) diff --git a/modules/processing.py b/modules/processing.py index 41be44010..b8beaf314 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -1,12 +1,9 @@ import json import math import os -import sys import random import logging from typing import Any, Dict, List - -import psutil import torch try: import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import @@ -16,12 +13,10 @@ import numpy as np from PIL import Image, ImageFilter, ImageOps import cv2 from skimage import exposure - from ldm.data.util import AddMiDaS from ldm.models.diffusion.ddpm import LatentDepth2ImageDiffusion from einops import repeat, rearrange from blendmodes.blend import blendLayers, BlendType - import modules.sd_hijack from modules import devices, prompt_parser, masking, sd_samplers, lowvram, generation_parameters_copypaste, script_callbacks, extra_networks, sd_vae_approx, scripts # pylint: disable=unused-import from modules.sd_hijack import model_hijack @@ -39,41 +34,6 @@ opt_C = 4 opt_f = 8 -def memory_stats(): - def gb(val: float): - return round(val / 1024 / 1024 / 1024, 2) - mem = {} - try: - process = psutil.Process(os.getpid()) - res = process.memory_info() - ram_total = 100 * res.rss / process.memory_percent() - ram = { 'used': gb(res.rss), 'total': gb(ram_total) } - mem.update({ 'ram': ram }) - except Exception as e: - mem.update({ 'ram': e }) - try: - if cmd_opts.use_ipex: - gpu = { 'used': gb(torch.xpu.memory_allocated()), 'total': gb(torch.xpu.get_device_properties("xpu").total_memory) } - s = dict(torch.xpu.memory_stats("xpu")) - mem.update({ - 'gpu': gpu, - 'retries': s['num_alloc_retries'], - 'oom': s['num_ooms'] - }) - elif torch.cuda.is_available(): - s = torch.cuda.mem_get_info() - gpu = { 'used': gb(s[1] - s[0]), 'total': gb(s[1]) } - s = dict(torch.cuda.memory_stats(shared.device)) - mem.update({ - 'gpu': gpu, - 'retries': s['num_alloc_retries'], - 'oom': s['num_ooms'] - }) - except: - pass - return mem - - def setup_color_correction(image): logging.info("Calibrating color correction.") correction_target = cv2.cvtColor(np.asarray(image.copy()), cv2.COLOR_RGB2LAB) @@ -145,8 +105,6 @@ class StableDiffusionProcessing: The first set of paramaters: sd_models -> do_not_reload_embeddings represent the minimum required to create a StableDiffusionProcessing """ def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt: str = "", styles: List[str] = None, seed: int = -1, subseed: int = -1, subseed_strength: float = 0, seed_resize_from_h: int = -1, seed_resize_from_w: int = -1, seed_enable_extras: bool = True, sampler_name: str = None, batch_size: int = 1, n_iter: int = 1, steps: int = 20, cfg_scale: float = 6.0, width: int = 512, height: int = 512, restore_faces: bool = False, tiling: bool = False, do_not_save_samples: bool = False, do_not_save_grid: bool = False, extra_generation_params: Dict[Any, Any] = None, overlay_images: Any = None, negative_prompt: str = None, eta: float = None, do_not_reload_embeddings: bool = False, denoising_strength: float = 0, ddim_discretize: str = None, s_churn: float = 0.0, s_tmax: float = None, s_tmin: float = 0.0, s_noise: float = 1.0, override_settings: Dict[str, Any] = None, override_settings_restore_afterwards: bool = True, sampler_index: int = None, script_args: list = None): # pylint: disable=unused-argument - if sampler_index is not None: - print("sampler_index argument for StableDiffusionProcessing does not do anything; use sampler_name", file=sys.stderr) self.outpath_samples: str = outpath_samples self.outpath_grids: str = outpath_grids @@ -723,7 +681,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: devices.test_for_nans(x, "vae") except devices.NansException as e: if not shared.cmd_opts.no_half and not shared.cmd_opts.no_half_vae and shared.cmd_opts.rollback_vae: - print('\nA tensor with all NaNs was produced in VAE, try converting to bf16.') + log.warning('Tensor with all NaNs was produced in VAE') devices.dtype_vae = torch.bfloat16 vae_file, vae_source = sd_vae.resolve_vae(p.sd_model.sd_model_checkpoint) sd_vae.load_vae(p.sd_model, vae_file, vae_source) diff --git a/modules/safe.py b/modules/safe.py index dd463ccdd..3a4aec4d7 100644 --- a/modules/safe.py +++ b/modules/safe.py @@ -7,14 +7,14 @@ import re import torch try: - import intel_extension_for_pytorch as ipex + import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import except: pass import numpy import _codecs # PyTorch 1.13 and later have _TypedStorage renamed to TypedStorage -TypedStorage = torch.storage.TypedStorage if hasattr(torch.storage, 'TypedStorage') else torch.storage._TypedStorage +TypedStorage = torch.storage.TypedStorage if hasattr(torch.storage, 'TypedStorage') else torch.storage._TypedStorage # pylint: disable=protected-access def encode(*args): @@ -27,7 +27,10 @@ class RestrictedUnpickler(pickle.Unpickler): def persistent_load(self, saved_id): assert saved_id[0] == 'storage' - return TypedStorage() + try: + return TypedStorage(_internal=True) + except TypeError: + return TypedStorage() # PyTorch before 2.0 does not have the _internal argument def find_class(self, module, name): if self.extra_handler is not None: @@ -38,7 +41,7 @@ class RestrictedUnpickler(pickle.Unpickler): if module == 'collections' and name == 'OrderedDict': return getattr(collections, name) if module == 'torch._utils' and name in ['_rebuild_tensor_v2', '_rebuild_parameter', '_rebuild_device_tensor_from_numpy']: - return getattr(torch._utils, name) + return getattr(torch._utils, name) # pylint: disable=protected-access if module == 'torch' and name in ['FloatStorage', 'HalfStorage', 'IntStorage', 'LongStorage', 'DoubleStorage', 'ByteStorage', 'float32']: return getattr(torch, name) if module == 'torch.nn.modules.container' and name in ['ParameterDict']: @@ -59,7 +62,7 @@ class RestrictedUnpickler(pickle.Unpickler): return set # Forbid everything else. - raise Exception(f"global '{module}/{name}' is forbidden") + raise Exception(f"global '{module}/{name}' is forbidden") # pylint: disable=broad-exception-raised # Regular expression that accepts 'dirname/version', 'dirname/data.pkl', and 'dirname/data/' @@ -71,7 +74,7 @@ def check_zip_filenames(filename, names): if allowed_zip_names_re.match(name): continue - raise Exception(f"bad file inside {filename}: {name}") + raise Exception(f"bad file inside {filename}: {name}") # pylint: disable=broad-exception-raised def check_pt(filename, extra_handler): @@ -84,9 +87,9 @@ def check_pt(filename, extra_handler): # find filename of data.pkl in zip file: '/data.pkl' data_pkl_filenames = [f for f in z.namelist() if data_pkl_re.match(f)] if len(data_pkl_filenames) == 0: - raise Exception(f"data.pkl not found in {filename}") + raise Exception(f"data.pkl not found in {filename}") # pylint: disable=broad-exception-raised if len(data_pkl_filenames) > 1: - raise Exception(f"Multiple data.pkl found in {filename}") + raise Exception(f"Multiple data.pkl found in {filename}") # pylint: disable=broad-exception-raised with z.open(data_pkl_filenames[0]) as file: unpickler = RestrictedUnpickler(file) unpickler.extra_handler = extra_handler @@ -98,7 +101,7 @@ def check_pt(filename, extra_handler): with open(filename, "rb") as file: unpickler = RestrictedUnpickler(file) unpickler.extra_handler = extra_handler - for i in range(5): + for _i in range(5): unpickler.load() @@ -106,7 +109,7 @@ def load(filename, *args, **kwargs): return load_with_extra(filename, extra_handler=global_extra_handler, *args, **kwargs) -def load_with_extra(filename, extra_handler=None, *args, **kwargs): +def load_with_extra(filename, extra_handler=None, *args, **kwargs): # pylint: disable=keyword-arg-before-vararg """ this function is intended to be used by extensions that want to load models with some extra classes in them that the usual unpickler would find suspicious. @@ -164,13 +167,13 @@ with safe.Extra(handler): self.handler = handler def __enter__(self): - global global_extra_handler + global global_extra_handler # pylint: disable=global-statement assert global_extra_handler is None, 'already inside an Extra() block' global_extra_handler = self.handler def __exit__(self, exc_type, exc_val, exc_tb): - global global_extra_handler + global global_extra_handler # pylint: disable=global-statement global_extra_handler = None @@ -178,4 +181,3 @@ with safe.Extra(handler): unsafe_torch_load = torch.load torch.load = load global_extra_handler = None - diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index 49acdf3dd..e06692213 100644 --- a/modules/sd_hijack.py +++ b/modules/sd_hijack.py @@ -1,8 +1,7 @@ from types import MethodType -from rich import print # pylint: disable=redefined-builtin import torch try: - import intel_extension_for_pytorch as ipex + import intel_extension_for_pytorch as ipex # pylint: disable=unused-import except: pass from torch.nn.functional import silu @@ -41,49 +40,48 @@ def apply_optimizations(): can_use_sdp = hasattr(torch.nn.functional, "scaled_dot_product_attention") and callable(getattr(torch.nn.functional, "scaled_dot_product_attention")) if devices.device == torch.device("cpu"): if opts.cross_attention_optimization == "Scaled-Dot-Product": - print("Scaled dot product cross attention is not available on CPU") + shared.log.warning("Scaled dot product cross attention is not available on CPU") can_use_sdp = False if opts.cross_attention_optimization == "xFormers": - print("xFormers cross attention is not available on CPU") + shared.log.warning("xFormers cross attention is not available on CPU") shared.xformers_available = False if opts.cross_attention_optimization == "Disable cross-attention layer optimization": - print("Cross-attention optimization disabled") + shared.log.warning("Cross-attention optimization disabled") optimization_method = 'none' if can_use_sdp and opts.cross_attention_optimization == "Scaled-Dot-Product" and 'SDP disable memory attention' in opts.cross_attention_options: - print("Applying scaled dot product cross attention optimization (without memory efficient attention)") + shared.log.info("Applying scaled dot product cross attention optimization (without memory efficient attention)") ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.scaled_dot_product_no_mem_attention_forward ldm.modules.diffusionmodules.model.AttnBlock.forward = sd_hijack_optimizations.sdp_no_mem_attnblock_forward optimization_method = 'sdp-no-mem' elif can_use_sdp and opts.cross_attention_optimization == "Scaled-Dot-Product": - print("Applying scaled dot product cross attention optimization") + shared.log.info("Applying scaled dot product cross attention optimization") ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.scaled_dot_product_attention_forward ldm.modules.diffusionmodules.model.AttnBlock.forward = sd_hijack_optimizations.sdp_attnblock_forward optimization_method = 'sdp' if shared.xformers_available and opts.cross_attention_optimization == "xFormers": - print("Applying xformers cross attention optimization") + shared.log.info("Applying xformers cross attention optimization") ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.xformers_attention_forward ldm.modules.diffusionmodules.model.AttnBlock.forward = sd_hijack_optimizations.xformers_attnblock_forward optimization_method = 'xformers' if opts.cross_attention_optimization == "Sub-quadratic": - print("Applying sub-quadratic cross attention optimization") + shared.log.info("Applying sub-quadratic cross attention optimization") ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.sub_quad_attention_forward ldm.modules.diffusionmodules.model.AttnBlock.forward = sd_hijack_optimizations.sub_quad_attnblock_forward optimization_method = 'sub-quadratic' if opts.cross_attention_optimization == "Split attention": - print("Applying split attention optimization") + shared.log.info("Applying split attention optimization") ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward_v1 optimization_method = 'v1' if opts.cross_attention_optimization == "InvokeAI's": - print("Applying InvokeAI's cross attention optimization") + shared.log.info("Applying InvokeAI cross attention optimization") ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward_invokeAI optimization_method = 'invokeai' if opts.cross_attention_optimization == "Doggettx's": - print("Applying cross attention optimization (Doggettx).") + shared.log.info("Applying Doggettx cross attention optimization") ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward ldm.modules.diffusionmodules.model.AttnBlock.forward = sd_hijack_optimizations.cross_attention_attnblock_forward optimization_method = 'doggettx' - return optimization_method @@ -190,9 +188,9 @@ class StableDiffusionModelHijack: hidet.torch.dynamo_config.use_tensor_core(True) hidet.torch.dynamo_config.search_space(2) m.model = torch.compile(m.model, mode="default", backend=opts.cuda_compile_mode, fullgraph=False, dynamic=False) - print("Model compile enabled:", opts.cuda_compile_mode) + shared.log.info(f"Model compile enabled: {opts.cuda_compile_mode}") except Exception as err: - print(f"Model compile not supported: {err}") + shared.log.warning(f"Model compile not supported: {err}") self.optimization_method = apply_optimizations() diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py index 6c6824175..f4c89acc3 100644 --- a/modules/sd_hijack_optimizations.py +++ b/modules/sd_hijack_optimizations.py @@ -120,8 +120,6 @@ def split_cross_attention_forward(self, x, context=None, mask=None): steps = 1 if mem_required > mem_free_total: steps = 2 ** (math.ceil(math.log(mem_required / mem_free_total, 2))) - # print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB " - # f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}") if steps > 64: max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64 raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). ' diff --git a/modules/sd_models.py b/modules/sd_models.py index 332a19e65..ed570c909 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -6,10 +6,10 @@ import re import io from os import mkdir from urllib import request -from rich import print, progress # pylint: disable=redefined-builtin +from rich import progress # pylint: disable=redefined-builtin import torch try: - import intel_extension_for_pytorch as ipex + import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import except: pass import safetensors.torch @@ -20,11 +20,11 @@ from ldm.util import instantiate_from_config from modules import paths, shared, modelloader, devices, script_callbacks, sd_vae, sd_disable_initialization, errors, hashes, sd_models_config from modules.sd_hijack_inpainting import do_inpainting_hijack from modules.timer import Timer +from modules.memstats import memory_stats model_dir = "Stable-diffusion" model_path = os.path.abspath(os.path.join(paths.models_path, model_dir)) - checkpoints_list = {} checkpoint_aliases = {} checkpoints_loaded = collections.OrderedDict() @@ -34,27 +34,21 @@ class CheckpointInfo: def __init__(self, filename): self.filename = filename abspath = os.path.abspath(filename) - if shared.opts.ckpt_dir is not None and abspath.startswith(shared.opts.ckpt_dir): name = abspath.replace(shared.opts.ckpt_dir, '') elif abspath.startswith(model_path): name = abspath.replace(model_path, '') else: name = os.path.basename(filename) - if name.startswith("\\") or name.startswith("/"): name = name[1:] - self.name = name self.name_for_extra = os.path.splitext(os.path.basename(filename))[0] self.model_name = os.path.splitext(name.replace("/", "_").replace("\\", "_"))[0] self.hash = model_hash(filename) - self.sha256 = hashes.sha256_from_cache(self.filename, "checkpoint/" + name) self.shorthash = self.sha256[0:10] if self.sha256 else None - self.title = name if self.shorthash is None else f'{name} [{self.shorthash}]' - self.ids = [self.hash, self.model_name, self.title, name, f'{name} [{self.hash}]'] + ([self.shorthash, self.sha256, f'{self.name} [{self.shorthash}]'] if self.shorthash else []) def register(self): @@ -66,16 +60,12 @@ class CheckpointInfo: self.sha256 = hashes.sha256(self.filename, "checkpoint/" + self.name) if self.sha256 is None: return - self.shorthash = self.sha256[0:10] - if self.shorthash not in self.ids: self.ids += [self.shorthash, self.sha256, f'{self.name} [{self.shorthash}]'] - checkpoints_list.pop(self.title) self.title = f'{self.name} [{self.shorthash}]' self.register() - return self.shorthash @@ -90,7 +80,6 @@ except Exception: def setup_model(): if not os.path.exists(model_path): os.makedirs(model_path) - list_models() enable_midas_autodownload() @@ -98,10 +87,8 @@ def setup_model(): def checkpoint_tiles(): def convert(name): return int(name) if name.isdigit() else name.lower() - def alphanumeric_key(key): return [convert(c) for c in re.split('([0-9]+)', key)] - return sorted([x.title for x in checkpoints_list.values()], key=alphanumeric_key) @@ -116,11 +103,11 @@ def list_models(): checkpoint_info.register() shared.opts.data['sd_model_checkpoint'] = checkpoint_info.title elif shared.cmd_opts.ckpt != shared.default_sd_model_file: - print(f"Checkpoint not found: {shared.cmd_opts.ckpt}", file=sys.stderr) + shared.log.warning(f"Checkpoint not found: {shared.cmd_opts.ckpt}", file=sys.stderr) for filename in sorted(model_list, key=str.lower): checkpoint_info = CheckpointInfo(filename) checkpoint_info.register() - print(f'Available models: {shared.opts.ckpt_dir} {len(checkpoints_list)}') + shared.log.info(f'Available models: {shared.opts.ckpt_dir} {len(checkpoints_list)}') if len(checkpoints_list) == 0: if not shared.cmd_opts.no_download: key = input('Download the default model? (y/N) ') @@ -136,17 +123,14 @@ def get_closet_checkpoint_match(search_string): checkpoint_info = checkpoint_aliases.get(search_string, None) if checkpoint_info is not None: return checkpoint_info - found = sorted([info for info in checkpoints_list.values() if search_string in info.title], key=lambda x: len(x.title)) if found: return found[0] - return None def model_hash(filename): """old hash that only looks at a small part of the file and is prone to collisions""" - try: with open(filename, "rb") as file: import hashlib @@ -161,20 +145,16 @@ def model_hash(filename): def select_checkpoint(): model_checkpoint = shared.opts.sd_model_checkpoint - checkpoint_info = checkpoint_aliases.get(model_checkpoint, None) if checkpoint_info is not None: return checkpoint_info - if len(checkpoints_list) == 0: - print("Cannot run without a checkpoint", file=sys.stderr) - print("Use --ckpt to force using existing checkpoint", file=sys.stderr) + shared.log.error("Cannot run without a checkpoint") + shared.log.error("Use --ckpt to force using existing checkpoint") exit(1) - checkpoint_info = next(iter(checkpoints_list.values())) if model_checkpoint is not None: - print(f"Checkpoint {model_checkpoint} not found; loading fallback {checkpoint_info.title}", file=sys.stderr) - + shared.log.warning(f"Checkpoint {model_checkpoint} not found; loading fallback {checkpoint_info.title}") return checkpoint_info @@ -189,39 +169,31 @@ def transform_checkpoint_dict_key(k): for text, replacement in checkpoint_dict_replacements.items(): if k.startswith(text): k = replacement + k[len(text):] - return k def get_state_dict_from_checkpoint(pl_sd): pl_sd = pl_sd.pop("state_dict", pl_sd) pl_sd.pop("state_dict", None) - sd = {} for k, v in pl_sd.items(): new_key = transform_checkpoint_dict_key(k) - if new_key is not None: sd[new_key] = v - pl_sd.clear() pl_sd.update(sd) - return pl_sd def read_metadata_from_safetensors(filename): import json - with open(filename, mode="rb") as file: metadata_len = file.read(8) metadata_len = int.from_bytes(metadata_len, "little") json_start = file.read(2) - assert metadata_len > 2 and json_start in (b'{"', b"{'"), f"{filename} is not a safetensors file" json_data = json_start + file.read(metadata_len-2) json_obj = json.loads(json_data) - res = {} for k, v in json_obj.get("__metadata__", {}).items(): res[k] = v @@ -230,7 +202,6 @@ def read_metadata_from_safetensors(filename): res[k] = json.loads(v) except Exception: pass - return res @@ -239,7 +210,7 @@ def read_state_dict(checkpoint_file, map_location=None): # pylint: disable=unuse pl_sd = None with progress.open(checkpoint_file, 'rb', description=f'Loading weights: [cyan]{checkpoint_file}', auto_refresh=True) as f: _, extension = os.path.splitext(checkpoint_file) - if 'v1-5-pruned-emaonly.safetensors' or 'vae-ft-mse-840000-ema-pruned.ckpt' in checkpoint_file: + if 'v1-5-pruned-emaonly.safetensors' in checkpoint_file and not shared.opts.stream_load: if extension.lower() == ".safetensors": pl_sd = safetensors.torch.load_file(checkpoint_file, device='cpu') else: @@ -262,67 +233,52 @@ def read_state_dict(checkpoint_file, map_location=None): # pylint: disable=unuse def get_checkpoint_state_dict(checkpoint_info: CheckpointInfo, timer): if checkpoint_info in checkpoints_loaded: # use checkpoint cache - print("Loading weights from cache") + shared.log.info("Loading weights from cache") return checkpoints_loaded[checkpoint_info] - res = read_state_dict(checkpoint_info.filename) timer.record("load") - return res -def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer): +def load_model_weights(model: torch.nn.Module, checkpoint_info: CheckpointInfo, state_dict, timer): sd_model_hash = checkpoint_info.calculate_shorthash() timer.record("hash") - shared.opts.data["sd_model_checkpoint"] = checkpoint_info.title - if state_dict is None: state_dict = get_checkpoint_state_dict(checkpoint_info, timer) - model.load_state_dict(state_dict, strict=False) del state_dict timer.record("apply") - if shared.opts.sd_checkpoint_cache > 0: # cache newly loaded model checkpoints_loaded[checkpoint_info] = model.state_dict().copy() - if shared.opts.opt_channelslast: model.to(memory_format=torch.channels_last) timer.record("channels") - if not shared.cmd_opts.no_half: vae = model.first_stage_model depth_model = getattr(model, 'depth_model', None) - # with --no-half-vae, remove VAE from model when doing half() to prevent its weights from being converted to float16 if shared.cmd_opts.no_half_vae: model.first_stage_model = None # with --upcast-sampling, don't convert the depth model weights to float16 if shared.opts.upcast_sampling and depth_model: model.depth_model = None - model.half() model.first_stage_model = vae if depth_model: model.depth_model = depth_model - devices.set_cuda_params() devices.dtype_unet = model.model.diffusion_model.dtype model.first_stage_model.to(devices.dtype_vae) - # clean up cache if limit is reached while len(checkpoints_loaded) > shared.opts.sd_checkpoint_cache: checkpoints_loaded.popitem(last=False) - model.sd_model_hash = sd_model_hash model.sd_model_checkpoint = checkpoint_info.filename model.sd_checkpoint_info = checkpoint_info shared.opts.data["sd_checkpoint_hash"] = checkpoint_info.sha256 - model.logvar = model.logvar.to(devices.device) # fix for training - sd_vae.delete_base_vae() sd_vae.clear_loaded_vae() vae_file, vae_source = sd_vae.resolve_vae(checkpoint_info.filename) @@ -339,23 +295,16 @@ def enable_midas_autodownload(): This function applies a wrapper to download the model to the correct location automatically. """ - midas_path = os.path.join(paths.models_path, 'midas') - - # stable-diffusion-stability-ai hard-codes the midas model path to - # a location that differs from where other scripts using this model look. - # HACK: Overriding the path here. for k, v in midas.api.ISL_PATHS.items(): file_name = os.path.basename(v) midas.api.ISL_PATHS[k] = os.path.join(midas_path, file_name) - midas_urls = { "dpt_large": "https://github.com/intel-isl/DPT/releases/download/1_0/dpt_large-midas-2f21e586.pt", "dpt_hybrid": "https://github.com/intel-isl/DPT/releases/download/1_0/dpt_hybrid-midas-501f0c75.pt", "midas_v21": "https://github.com/AlexeyAB/MiDaS/releases/download/midas_dpt/midas_v21-f6b98070.pt", "midas_v21_small": "https://github.com/AlexeyAB/MiDaS/releases/download/midas_dpt/midas_v21_small-70d6b9c8.pt", } - midas.api.load_model_inner = midas.api.load_model def load_model_wrapper(model_type): @@ -363,29 +312,23 @@ def enable_midas_autodownload(): if not os.path.exists(path): if not os.path.exists(midas_path): mkdir(midas_path) - - print(f"Downloading midas model weights for {model_type} to {path}") + shared.log.info(f"Downloading midas model weights for {model_type} to {path}") request.urlretrieve(midas_urls[model_type], path) - print(f"{model_type} downloaded") - + shared.log.info(f"{model_type} downloaded") return midas.api.load_model_inner(model_type) midas.api.load_model = load_model_wrapper def repair_config(sd_config): - if not "use_ema" in sd_config.model.params: sd_config.model.params.use_ema = False - if shared.cmd_opts.no_half: sd_config.model.params.unet_config.params.use_fp16 = False elif shared.opts.upcast_sampling: sd_config.model.params.unet_config.params.use_fp16 = True - if getattr(sd_config.model.params.first_stage_config.params.ddconfig, "attn_type", None) == "vanilla-xformers" and not shared.xformers_available: sd_config.model.params.first_stage_config.params.ddconfig.attn_type = "vanilla" - # For UnCLIP-L, override the hardcoded karlo directory if "noise_aug_config" in sd_config.model.params and "clip_stats_path" in sd_config.model.params.noise_aug_config.params: karlo_path = os.path.join(paths.models_path, 'karlo') @@ -395,14 +338,13 @@ def repair_config(sd_config): sd1_clip_weight = 'cond_stage_model.transformer.text_model.embeddings.token_embedding.weight' sd2_clip_weight = 'cond_stage_model.model.transformer.resblocks.0.attn.in_proj_weight' -def load_model(checkpoint_info=None, already_loaded_state_dict=None): + +def load_model(checkpoint_info=None, already_loaded_state_dict=None, timer=None): from modules import lowvram, sd_hijack checkpoint_info = checkpoint_info or select_checkpoint() - do_inpainting_hijack() - - timer = Timer() - + if timer is None: + timer = Timer() current_checkpoint_info = None if shared.sd_model: current_checkpoint_info = shared.sd_model.sd_checkpoint_info @@ -410,122 +352,100 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None): shared.sd_model = None gc.collect() devices.torch_gc() - + shared.debug(f'Model unloaded: {memory_stats()}') if already_loaded_state_dict is not None: state_dict = already_loaded_state_dict else: state_dict = get_checkpoint_state_dict(checkpoint_info, timer) - checkpoint_config = sd_models_config.find_checkpoint_config(state_dict, checkpoint_info) if state_dict is None or checkpoint_config is None: - print(f"Failed to load checkpooint: {checkpoint_info.filename}") + shared.log.error(f"Failed to load checkpooint: {checkpoint_info.filename}") if current_checkpoint_info is not None: - print(f"Restoring previous checkpoint: {current_checkpoint_info.filename}") + shared.log.info(f"Restoring previous checkpoint: {current_checkpoint_info.filename}") load_model(current_checkpoint_info, None) return - + shared.debug(f'Model dict loaded: {memory_stats()}') clip_is_included_into_sd = sd1_clip_weight in state_dict or sd2_clip_weight in state_dict - sd_config = OmegaConf.load(checkpoint_config) repair_config(sd_config) - timer.record("config") - - print(f"Creating model from config: {checkpoint_config}") - + shared.debug(f'Model config loaded: {memory_stats()}') + shared.log.info(f"Creating model from config: {checkpoint_config}") sd_model = None try: with sd_disable_initialization.DisableInitialization(disable_clip=clip_is_included_into_sd): sd_model = instantiate_from_config(sd_config.model) except Exception: sd_model = instantiate_from_config(sd_config.model) - sd_model.used_config = checkpoint_config - timer.record("create") - load_model_weights(sd_model, checkpoint_info, state_dict, timer) - + timer.record("load") + shared.debug(f'Model weights loaded: {memory_stats()}') if shared.cmd_opts.lowvram or shared.cmd_opts.medvram: lowvram.setup_for_low_vram(sd_model, shared.cmd_opts.medvram) else: sd_model.to(shared.device) - timer.record("move") - + shared.debug(f'Model weights moved: {memory_stats()}') sd_hijack.model_hijack.hijack(sd_model) - timer.record("hijack") - sd_model.eval() shared.sd_model = sd_model - sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings(force_reload=True) # Reload embeddings after model load as they may or may not fit the model - timer.record("embeddings") - script_callbacks.model_loaded_callback(sd_model) - timer.record("callbacks") - - print(f"Model loaded in {timer.summary()}") - + shared.log.info(f"Model loaded in {timer.summary()}") + shared.debug(f'Model load finished: {memory_stats()}') return sd_model def reload_model_weights(sd_model=None, info=None): from modules import lowvram, sd_hijack checkpoint_info = info or select_checkpoint() - if not sd_model: sd_model = shared.sd_model - + if not shared.opts.model_reuse_dict and sd_model is not None: + sd_model = None + else: + shared.log.info('Reusing previous model dictionary') if sd_model is None: # previous model load failed current_checkpoint_info = None else: current_checkpoint_info = sd_model.sd_checkpoint_info if sd_model.sd_model_checkpoint == checkpoint_info.filename: return - if shared.cmd_opts.lowvram or shared.cmd_opts.medvram: lowvram.send_everything_to_cpu() else: sd_model.to(devices.cpu) - sd_hijack.model_hijack.undo_hijack(sd_model) - timer = Timer() - state_dict = get_checkpoint_state_dict(checkpoint_info, timer) - checkpoint_config = sd_models_config.find_checkpoint_config(state_dict, checkpoint_info) - - timer.record("find config") - + timer.record("config") if sd_model is None or checkpoint_config != sd_model.used_config: del sd_model checkpoints_loaded.clear() - load_model(checkpoint_info, already_loaded_state_dict=state_dict) + load_model(checkpoint_info, already_loaded_state_dict=state_dict, timer=timer) return shared.sd_model - try: load_model_weights(sd_model, checkpoint_info, state_dict, timer) except Exception: - print("Failed to load checkpoint, restoring previous") + shared.log.error("Failed to load checkpoint, restoring previous") load_model_weights(sd_model, current_checkpoint_info, None, timer) raise finally: sd_hijack.model_hijack.hijack(sd_model) timer.record("hijack") - script_callbacks.model_loaded_callback(sd_model) timer.record("callbacks") - if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram: sd_model.to(devices.device) timer.record("device") + shared.log.info(f"Weights loaded in {timer.summary()}") - print(f"Weights loaded in {timer.summary()}") def unload_model_weights(sd_model=None, _info=None): from modules import sd_hijack @@ -539,7 +459,7 @@ def unload_model_weights(sd_model=None, _info=None): sd_model = None gc.collect() devices.torch_gc() - print(f"Unloaded weights {timer.summary()}") + shared.log.info(f"Unloaded weights {timer.summary()}") return sd_model diff --git a/modules/sd_vae.py b/modules/sd_vae.py index a13d73be7..bc5820b70 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -2,16 +2,13 @@ import os import collections import glob from copy import deepcopy -from rich import print # pylint: disable=redefined-builtin -from modules import shared import torch +from modules import shared, paths, devices, script_callbacks, sd_models try: - import intel_extension_for_pytorch as ipex + import intel_extension_for_pytorch as ipex # pylint: disable=unused-import except: if shared.cmd_opts.use_ipex: - print("Failed to import IPEX") -from modules import paths, devices, script_callbacks, sd_models - + shared.log.error("Failed to import IPEX") vae_ignore_keys = {"model_ema.decay", "model_ema.num_updates"} vae_dict = {} @@ -44,7 +41,7 @@ def delete_base_vae(): def restore_base_vae(model): global loaded_vae_file # pylint: disable=global-statement if base_vae is not None and checkpoint_info == model.sd_checkpoint_info: - print("Restoring base VAE") + shared.log.info("Restoring base VAE") _load_vae_dict(model, base_vae) loaded_vae_file = None delete_base_vae() @@ -120,7 +117,7 @@ def resolve_vae(checkpoint_file): return vae_from_options, 'specified in settings' if not is_automatic: - print(f"VAE not found: {shared.opts.sd_vae}") + shared.log.warning(f"VAE not found: {shared.opts.sd_vae}") return None, None @@ -140,7 +137,7 @@ def load_vae(model, vae_file=None, vae_source="from unknown source"): if vae_file: if cache_enabled and vae_file in checkpoints_loaded: # use vae checkpoint cache - print(f"Loading VAE weights {vae_source}: cached {get_filename(vae_file)}") + shared.log.info(f"Loading VAE weights {vae_source}: cached {get_filename(vae_file)}") store_base_vae(model) _load_vae_dict(model, checkpoints_loaded[vae_file]) else: @@ -220,5 +217,5 @@ def reload_vae_weights(sd_model=None, vae_file=unspecified): if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram: sd_model.to(devices.device) - print("VAE weights loaded.") + shared.log.info("VAE weights loaded.") return sd_model diff --git a/modules/shared.py b/modules/shared.py index 77d5b3e14..32e8b0ab7 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -162,6 +162,12 @@ state.server_start = time.time() interrogator = modules.interrogate.InterrogateModels("interrogate") face_restorers = [] + +def debug(message): + if cmd_opts.debug: + log.info(message) + + class OptionInfo: def __init__(self, default=None, label="", component=None, component_args=None, onchange=None, section=None, refresh=None): self.default = default @@ -215,9 +221,9 @@ def refresh_themes(): with open(os.path.join('javascript', 'themes.json'), mode='w', encoding='utf=8') as f: f.write(json.dumps(res)) else: - print('Error refreshing UI themes') + log.error('Error refreshing UI themes') except: - print('Exception refreshing UI themes') + log.error('Exception refreshing UI themes') hide_dirs = {"visible": not cmd_opts.hide_ui_dir_config} tab_names = [] @@ -229,6 +235,8 @@ options_templates.update(options_section(('sd', "Stable Diffusion"), { "sd_checkpoint_cache": OptionInfo(0, "Model checkpoints to cache in RAM", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}), "sd_vae_checkpoint_cache": OptionInfo(0, "VAE checkpoints to cache in RAM", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}), "sd_vae": OptionInfo("Automatic", "Select VAE", gr.Dropdown, lambda: {"choices": shared_items.sd_vae_items()}, refresh=shared_items.refresh_vae_list), + "stream_load": OptionInfo(False, "When loading models attempt stream loading optimized for slow or network storage"), + "model_reuse_dict": OptionInfo(False, "When loading models attempt to reuse previous model dictionary"), "inpainting_mask_weight": OptionInfo(1.0, "Inpainting conditioning mask strength", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}), "initial_noise_multiplier": OptionInfo(1.0, "Noise multiplier for img2img", gr.Slider, {"minimum": 0.5, "maximum": 1.5, "step": 0.01}), "img2img_color_correction": OptionInfo(False, "Apply color correction to img2img results to match original colors."), @@ -314,7 +322,7 @@ options_templates.update(options_section(('cuda', "CUDA Settings"), { "memmon_poll_rate": OptionInfo(2, "VRAM usage polls per second during generation. Set to 0 to disable.", gr.Slider, {"minimum": 0, "maximum": 40, "step": 1}), "precision": OptionInfo("Autocast", "Precision type", gr.Radio, lambda: {"choices": ["Autocast", "Full"]}), "cuda_dtype": OptionInfo("FP32" if sys.platform == "darwin" else "FP16", "Device precision type", gr.Radio, lambda: {"choices": ["FP32", "FP16", "BF16"]}), - "no_half": OptionInfo(True if is_device_dml else False, "Use full precision for model (--no-half)", None, None, lambda: print("Warning: Most of DirectML devices do not fully support half mode. Recommend to use full precision to model.") if is_device_dml else None), + "no_half": OptionInfo(True if is_device_dml else False, "Use full precision for model (--no-half)", None, None, None), "no_half_vae": OptionInfo(True if is_device_dml else False, "Use full precision for VAE (--no-half-vae)"), "upcast_sampling": OptionInfo(True if sys.platform == "darwin" or cmd_opts.use_ipex else False, "Enable upcast sampling. Usually produces similar results to --no-half with better performance while using less memory"), "disable_nan_check": OptionInfo(True, "Do not check if produced images/latent spaces have NaN values"), @@ -468,10 +476,10 @@ class Options: if self.data is not None: if key in self.data or key in self.data_labels: if cmd_opts.freeze: - print(f'Settings are frozen: {key}') + log.warning(f'Settings are frozen: {key}') return if cmd_opts.hide_ui_dir_config and key in restricted_opts: - print(f'Settings key is restricted: {key}') + log.warning(f'Settings key is restricted: {key}') return else: self.data[key] = value @@ -531,11 +539,11 @@ class Options: for k, v in self.data.items(): info = self.data_labels.get(k, None) if info is not None and not self.same_type(info.default, v): - log.error(f"Warning: bad setting value: {k}: {v} ({type(v).__name__}; expected {type(info.default).__name__})", file=sys.stderr) + log.error(f"Warning: bad setting value: {k}: {v} ({type(v).__name__}; expected {type(info.default).__name__})") bad_settings += 1 if bad_settings > 0: - log.error(f"The program is likely to not work with bad settings.\nSettings file: {filename}\nEither fix the file, or delete it and restart.", file=sys.stderr) + log.error(f"The program is likely to not work with bad settings.\nSettings file: {filename}\nEither fix the file, or delete it and restart.") def onchange(self, key, func, call=True): item = self.data_labels.get(key) @@ -630,9 +638,9 @@ def reload_gradio_theme(theme_name=None): try: gradio_theme = gr.themes.ThemeClass.from_hub(theme_name) except: - print("Theme download error accessing HuggingFace") + log.error("Theme download error accessing HuggingFace") gradio_theme = gr.themes.Default() - print(f'Loading theme: {theme_name}') + log.info(f'Loading theme: {theme_name}') class TotalTQDM: @@ -678,14 +686,14 @@ def restart_server(): return try: import logging - logging.disable(logging.CRITICAL) + log.setLevel(logging.DEBUG if cmd_opts.debug else logging.CRITICAL) demo.server.should_exit = True demo.server.force_exit = True demo.close(verbose=False) demo.server.close() except: pass - print('Server shutdown') + log.info('Server shutdown') def listfiles(dirname): diff --git a/modules/shared_items.py b/modules/shared_items.py index 4c2924518..c329e73e9 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -1,5 +1,3 @@ - - def realesrgan_models_names(): import modules.realesrgan_model return [x.name for x in modules.realesrgan_model.get_realesrgan_models(None)] @@ -32,6 +30,3 @@ def list_crossattention(): "Sub-quadratic", "Split attention" ] -# parser.add_argument("--sub-quad-q-chunk-size", type=int, help="query chunk size for the sub-quadratic cross-attention layer optimization to use", default=1024) -# parser.add_argument("--sub-quad-kv-chunk-size", type=int, help="kv chunk size for the sub-quadratic cross-attention layer optimization to use", default=None) -# parser.add_argument("--sub-quad-chunk-threshold", type=int, help="the percentage of VRAM threshold for the sub-quadratic cross-attention layer optimization to use chunking", default=None) diff --git a/modules/styles.py b/modules/styles.py index e7c33025d..fb151da96 100644 --- a/modules/styles.py +++ b/modules/styles.py @@ -1,11 +1,9 @@ # We need this so Python doesn't complain about the unknown StableDiffusionProcessing-typehint at runtime from __future__ import annotations - import csv import os import os.path import typing -import collections.abc as abc import tempfile import shutil @@ -49,7 +47,6 @@ class StyleDatabase: self.styles.clear() if not os.path.exists(self.path): - print(f'Creating styles database: {self.path}') self.save_styles(self.path) with open(self.path, "r", encoding="utf-8-sig", newline='') as file: diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py index 343e55539..d4e210e62 100644 --- a/modules/textual_inversion/textual_inversion.py +++ b/modules/textual_inversion/textual_inversion.py @@ -9,7 +9,6 @@ except: pass import tqdm import safetensors.torch -from rich import print # pylint: disable=redefined-builtin import numpy as np from PIL import Image, PngImagePlugin from torch.utils.tensorboard import SummaryWriter @@ -131,7 +130,7 @@ class EmbeddingDatabase: def get_expected_shape(self): if shared.sd_model is None: - print('Model not loaded') + shared.log.error('Model not loaded') return 0 vec = shared.sd_model.cond_stage_model.encode_embedding_init_text(",", 1) return vec.shape[1] @@ -234,9 +233,9 @@ class EmbeddingDatabase: displayed_embeddings = (tuple(self.word_embeddings.keys()), tuple(self.skipped_embeddings.keys())) if self.previously_displayed_embeddings != displayed_embeddings: self.previously_displayed_embeddings = displayed_embeddings - print(f"Embeddings loaded: {', '.join(self.word_embeddings.keys())} ({len(self.word_embeddings)})") + shared.log.info(f"Embeddings loaded: {', '.join(self.word_embeddings.keys())} ({len(self.word_embeddings)})") if len(self.skipped_embeddings) > 0: - print(f"Textual inversion embeddings skipped({len(self.skipped_embeddings)}): {', '.join(self.skipped_embeddings.keys())}") + shared.log.info(f"Textual inversion embeddings skipped({len(self.skipped_embeddings)}): {', '.join(self.skipped_embeddings.keys())}") def find_embedding_at_position(self, tokens, offset): token = tokens[offset] @@ -271,12 +270,12 @@ def create_embedding(name, num_vectors_per_token, overwrite_old, init_text='*'): name = "".join( x for x in name if (x.isalnum() or x in "._- ")) fn = os.path.join(shared.opts.embeddings_dir, f"{name}.pt") if not overwrite_old and os.path.exists(fn): - print(f"Embedding already exists: {fn}") + shared.log.warning(f"Embedding already exists: {fn}") else: embedding = Embedding(vec, name) embedding.step = 0 embedding.save(fn) - print(f'Created embedding: {fn} vectors {num_vectors_per_token} init {init_text}') + shared.log.info(f'Created embedding: {fn} vectors {num_vectors_per_token} init {init_text}') return fn @@ -424,9 +423,9 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st optimizer_state_dict = optimizer_saved_dict.get('optimizer_state_dict', None) if optimizer_state_dict is not None: optimizer.load_state_dict(optimizer_state_dict) - print("Loaded existing optimizer from checkpoint") + shared.log.info("Loaded existing optimizer from checkpoint") else: - print("No saved optimizer exists in checkpoint") + shared.log.info("No saved optimizer exists in checkpoint") if shared.cmd_opts.use_ipex: scaler = torch.xpu.amp.GradScaler() diff --git a/modules/txt2img.py b/modules/txt2img.py index 17e5ce909..36bf6dfb6 100644 --- a/modules/txt2img.py +++ b/modules/txt2img.py @@ -1,9 +1,10 @@ import modules.scripts from modules import sd_samplers from modules.generation_parameters_copypaste import create_override_settings_dict -from modules.processing import StableDiffusionProcessingTxt2Img, process_images, memory_stats -from modules.shared import opts, sd_model, cmd_opts, log +from modules.processing import StableDiffusionProcessingTxt2Img, process_images +from modules.shared import opts, sd_model, debug from modules.ui import plaintext_to_html +from modules.memstats import memory_stats def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, enable_hr: bool, denoising_strength: float, hr_scale: float, hr_upscaler: str, hr_second_pass_steps: int, hr_resize_x: int, hr_resize_y: int, override_settings_texts, *args): # pylint: disable=unused-argument @@ -46,6 +47,5 @@ def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, step processed = process_images(p) p.close() generation_info_js = processed.js() - if cmd_opts.debug: - log.info(f'Processed: {len(processed.images)} Memory: {memory_stats()} txt') + debug(f'Processed: {len(processed.images)} Memory: {memory_stats()} txt') return processed.images, generation_info_js, plaintext_to_html(processed.info), plaintext_to_html(processed.comments) diff --git a/modules/ui.py b/modules/ui.py index e9d6c6efc..d2397bb26 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -109,9 +109,7 @@ def process_interrogate(interrogation_function, mode, ii_input_dir, ii_output_di elif mode == 2: return [interrogation_function(ii_singles[mode]["image"]), None] elif mode == 5: - assert not shared.cmd_opts.hide_ui_dir_config, "Launched with --hide-ui-dir-config, batch img2img disabled" images = shared.listfiles(ii_input_dir) - print(f"Will process {len(images)} images.") if ii_output_dir != "": os.makedirs(ii_output_dir, exist_ok=True) else: @@ -183,8 +181,7 @@ def connect_reuse_seed(seed: gr.Number, reuse_seed: gr.Button, generation_info: res = all_seeds[index if 0 <= index < len(all_seeds) else 0] except json.decoder.JSONDecodeError: if gen_info_string != '': - print("Error parsing JSON generation info:", file=sys.stderr) - print(gen_info_string, file=sys.stderr) + shared.log.error(f"Error parsing JSON generation info: {gen_info_string}") return [res, gr_show(False)] reuse_seed.click(fn=copy_seed, _js="(x, y) => [x, selected_gallery_index()]", show_progress=False, inputs=[generation_info, dummy_component], outputs=[seed, dummy_component]) @@ -1496,9 +1493,6 @@ def create_ui(): ui_settings[key] = getattr(obj, field) elif condition and not condition(saved_value): pass - - # this warning is generally not useful; - # print(f'Warning: Bad ui setting value: {key}: {saved_value}; Default value "{getattr(obj, field)}" will be used instead.') else: setattr(obj, field, saved_value) if init_field is not None: diff --git a/modules/ui_common.py b/modules/ui_common.py index cb9870196..d823e7d45 100644 --- a/modules/ui_common.py +++ b/modules/ui_common.py @@ -2,7 +2,6 @@ import json import html import os import platform -import sys import subprocess as sp import gradio as gr @@ -107,15 +106,10 @@ def create_output_panel(tabname, outdir): def open_folder(f): if not os.path.exists(f): - print(f'Folder "{f}" does not exist. After you create an image, the folder will be created.') + shared.log.warning(f'Folder "{f}" does not exist. After you create an image, the folder will be created.') return elif not os.path.isdir(f): - print(f""" -WARNING -An open_folder request was made with an argument that is not a folder. -This could be an error or a malicious attempt to run code on your computer. -Requested path was: {f} -""", file=sys.stderr) + shared.log.warning(f"An open_folder request was made with an argument that is not a folder: {f}") return if not shared.cmd_opts.hide_ui_dir_config: diff --git a/modules/ui_extensions.py b/modules/ui_extensions.py index 66754f58c..fb7624bff 100644 --- a/modules/ui_extensions.py +++ b/modules/ui_extensions.py @@ -4,11 +4,8 @@ import time import shutil import errno import html - import git import gradio as gr - -from rich import print # pylint: disable=redefined-builtin from modules import extensions, shared, paths, errors from modules.call_queue import wrap_gradio_gpu_call @@ -46,7 +43,7 @@ def apply_and_restart(disable_list, update_list, disable_all): # shared.state.interrupt() # shared.state.need_restart = True # shared.restart_server() - print('Extension list updated - please restart the server') + shared.log.warning('Extension list updated - please restart the server') def check_updates(_id_task, disable_list): @@ -135,12 +132,12 @@ def install_extension_from_url(dirname, url): assert url, 'No URL specified' if dirname is None or dirname == "": - *parts, last_part = url.split('/') + *parts, last_part = url.split('/') # pylint: disable=unused-variable last_part = normalize_git_url(last_part) dirname = last_part target_dir = os.path.join(extensions.extensions_dir, dirname) - print(f'Installing extension: {url} into {target_dir}') + shared.log.info(f'Installing extension: {url} into {target_dir}') assert not os.path.exists(target_dir), f'Extension directory already exists: {target_dir}' normalized_url = normalize_git_url(url) diff --git a/modules/upscaler.py b/modules/upscaler.py index 9ca6e4596..560ca6167 100644 --- a/modules/upscaler.py +++ b/modules/upscaler.py @@ -40,7 +40,7 @@ class Upscaler: os.makedirs(self.model_path, exist_ok=True) try: - import cv2 + import cv2 # pylint: disable=unused-import self.can_tile = True except: pass diff --git a/webui.py b/webui.py index d3afac97f..60b75915e 100644 --- a/webui.py +++ b/webui.py @@ -6,7 +6,6 @@ import signal import asyncio import logging import warnings -from rich import print # pylint: disable=W0622 from modules import timer, errors startup_timer = timer.Timer() @@ -18,6 +17,9 @@ except: pass import torchvision # pylint: disable=W0611,C0411 import pytorch_lightning # pytorch_lightning should be imported after torch, but it re-enables warnings on import so import once to disable them # pylint: disable=W0611,C0411 +if ".dev" in torch.__version__ or "+git" in torch.__version__: + torch.__long_version__ = torch.__version__ + torch.__version__ = re.search(r'[\d.]+[\d]', torch.__version__).group(0) logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage()) logging.getLogger("pytorch_lightning").disabled = True warnings.filterwarnings(action="ignore", category=DeprecationWarning, module="pytorch_lightning") @@ -34,12 +36,6 @@ from modules import extra_networks, ui_extra_networks_checkpoints # pylint: disa from modules import extra_networks_hypernet, ui_extra_networks_hypernets, ui_extra_networks_textual_inversion from modules.call_queue import wrap_queued_call, queue_lock, wrap_gradio_gpu_call # pylint: disable=W0611,C0411 from modules.paths import create_paths - -# Truncate version number of nightly/local build of PyTorch to not cause exceptions with CodeFormer or Safetensors -if ".dev" in torch.__version__ or "+git" in torch.__version__: - torch.__long_version__ = torch.__version__ - torch.__version__ = re.search(r'[\d.]+[\d]', torch.__version__).group(0) - from modules import shared, extensions, ui_tempdir, ui_extra_networks import modules.devices import modules.sd_samplers @@ -59,11 +55,14 @@ import modules.textual_inversion.textual_inversion import modules.progress import modules.ui from modules import modelloader -from modules.shared import cmd_opts, opts +from modules.shared import cmd_opts, opts, log import modules.hypernetworks.hypernetwork from modules.middleware import setup_middleware startup_timer.record("libraries") + +log.setLevel(logging.DEBUG if cmd_opts.debug else logging.INFO) +logging.disable(logging.NOTSET if cmd_opts.debug else logging.DEBUG) if cmd_opts.server_name: server_name = cmd_opts.server_name else: @@ -73,17 +72,18 @@ else: def check_rollback_vae(): if shared.cmd_opts.rollback_vae: if not torch.cuda.is_available(): - print("Rollback VAE functionality requires compatible GPU") + log.error("Rollback VAE functionality requires compatible GPU") shared.cmd_opts.rollback_vae = False elif not torch.__version__.startswith('2.1'): - print("Rollback VAE functionality requires Torch 2.1 or higher") + log.error("Rollback VAE functionality requires Torch 2.1 or higher") shared.cmd_opts.rollback_vae = False elif 0 < torch.cuda.get_device_capability()[0] < 8: - print('Rollback VAE functionality device capabilities not met') + log.error('Rollback VAE functionality device capabilities not met') shared.cmd_opts.rollback_vae = False def initialize(): + log.debug('Entering Initialize') check_rollback_vae() extensions.list_extensions() @@ -127,19 +127,19 @@ def initialize(): if cmd_opts.tls_keyfile is not None and cmd_opts.tls_keyfile is not None: try: if not os.path.exists(cmd_opts.tls_keyfile): - print("Invalid path to TLS keyfile given") + log.error("Invalid path to TLS keyfile given") if not os.path.exists(cmd_opts.tls_certfile): - print(f"Invalid path to TLS certfile: '{cmd_opts.tls_certfile}'") + log.error(f"Invalid path to TLS certfile: '{cmd_opts.tls_certfile}'") except TypeError: cmd_opts.tls_keyfile = cmd_opts.tls_certfile = None - print("TLS setup invalid, running webui without TLS") + log.error("TLS setup invalid, running webui without TLS") else: - print("Running with TLS") + log.info("Running with TLS") startup_timer.record("TLS") # make the program just exit at ctrl+c without waiting for anything def sigint_handler(_sig, _frame): - print('Exiting') + log.info('Exiting') os._exit(0) signal.signal(signal.SIGINT, sigint_handler) @@ -152,10 +152,10 @@ def load_model(): modules.sd_models.load_model() except Exception as e: errors.display(e, "loading stable diffusion model") - print("Stable diffusion model failed to load") + log.error("Stable diffusion model failed to load") exit(1) if shared.sd_model is None: - print("No stable diffusion model loaded") + log.error("No stable diffusion model loaded") exit(1) shared.opts.data["sd_model_checkpoint"] = shared.sd_model.sd_checkpoint_info.title shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights())) @@ -185,7 +185,8 @@ def async_policy(): def start_ui(): - logging.disable(logging.INFO) + log.debug('Entering StartUI') + logging.disable(logging.NOTSET if cmd_opts.debug else logging.DEBUG) create_paths(opts) async_policy() initialize() @@ -197,7 +198,7 @@ def start_ui(): shared.demo = modules.ui.create_ui() startup_timer.record("ui") if cmd_opts.disable_queue: - print('Server queues disabled') + log.info('Server queues disabled') shared.demo.progress_tracking = False else: shared.demo.queue(concurrency_count=16) @@ -238,10 +239,10 @@ def start_ui(): def webui(): + log.debug('Entering WebUI') start_ui() load_model() - print(f"Startup time: {startup_timer.summary()}") - logging.disable(logging.DEBUG) + log.info(f"Startup time: {startup_timer.summary()}") while True: try: @@ -249,10 +250,10 @@ def webui(): except: alive = False if not alive: - print('Server restart') + log.warning('Server restart') startup_timer.reset() start_ui() - print(f"Startup time: {startup_timer.summary()}") + log.info(f"Startup time: {startup_timer.summary()}") time.sleep(1) """