redesign logging

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)

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

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)

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]))
 
 

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()
 

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:

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),

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 '<empty string>'}, using {upscaler.name} as a fallback")
+                log.warning(f"could not find upscaler named {upscaler_name or '<empty string>'}, 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')

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)

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

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"]

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()
 

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

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)
 

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)

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/<number>'
@@ -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: '<directory name>/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
-

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()
 

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}). '

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 <path-to-checkpoint> to force using existing checkpoint", file=sys.stderr)
+        shared.log.error("Cannot run without a checkpoint")
+        shared.log.error("Use --ckpt <path-to-checkpoint> 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
 
 

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

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):

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)

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:

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()

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)

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:

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:

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)

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

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)
 
     """