add light theme, enhance profilng and logging

modules/call_queue.py

@@ -47,10 +47,8 @@ def wrap_gradio_gpu_call(func, extra_outputs=None):
 
 def wrap_gradio_call(func, extra_outputs=None, add_stats=False):
     def f(*args, extra_outputs_array=extra_outputs, **kwargs):
-        run_memmon = shared.opts.memmon_poll_rate > 0 and not shared.mem_mon.disabled and add_stats
-        if run_memmon:
-            shared.mem_mon.monitor()
         t = time.perf_counter()
+        shared.mem_mon.reset()
         try:
             if shared.cmd_opts.profile:
                 pr = cProfile.Profile()
@@ -83,18 +81,11 @@ def wrap_gradio_call(func, extra_outputs=None, add_stats=False):
         elapsed = time.perf_counter() - t
         elapsed_m = int(elapsed // 60)
         elapsed_s = elapsed % 60
-        elapsed_text = f"{elapsed_s:.2f}s"
-        if elapsed_m > 0:
-            elapsed_text = f"{elapsed_m}m "+elapsed_text
-        if run_memmon:
-            mem_stats = {k: -(v//-(1024*1024)) for k, v in shared.mem_mon.stop().items()}
-            active_peak = mem_stats['active_peak']
-            reserved_peak = mem_stats['reserved_peak']
-            sys_peak = mem_stats['system_peak']
-            sys_total = mem_stats['total']
-            vram_html = f" | <p class='vram'>GPU active {active_peak} MB reserved {reserved_peak} MB | System peak {sys_peak} MB total {sys_total} MB</p>"
-        else:
-            vram_html = ''
-        res[-1] += f"<div class='performance'><p class='time'>Time taken: {elapsed_text}</p>{vram_html}</div>"
+        elapsed_text = f"{elapsed_m}m {elapsed_s:.2f}s" if elapsed_m > 0 else f"{elapsed_s:.2f}s"
+        vram_html = ''
+        if not shared.mem_mon.disabled:
+            vram = {k: -(v//-(1024*1024)) for k, v in shared.mem_mon.read().items()}
+            vram_html += f" | <p class='vram'>GPU active {max(vram['active_peak'], vram['reserved_peak'])} MB reserved {vram['reserved']} | used {vram['used']} MB free {vram['free']} MB total {vram['total']} MB | retries {vram['retries']} oom {vram['oom']}</p>"
+        res[-1] += f"<div class='performance'><p class='time'>Time: {elapsed_text}</p>{vram_html}</div>"
         return tuple(res)
     return f

modules/images.py

@@ -451,7 +451,7 @@ def atomically_save_image():
             image_format = 'JPEG'
         if shared.opts.image_watermark_enabled:
             image = set_watermark(image, shared.opts.image_watermark)
-        shared.log.debug(f'Saving image: type={image_format} size={image.size} {fn}')
+        shared.log.debug(f'Saving: image={fn} type={image_format} size={image.width}x{image.height}')
         # actual save
         exifinfo = (exifinfo or "") if shared.opts.image_metadata else ""
         if image_format == 'PNG':

modules/img2img.py

@@ -74,7 +74,7 @@ def process_batch(p, input_files, input_dir, output_dir, inpaint_mask_dir, args)
             for k, v in items.items():
                 image.info[k] = v
             images.save_image(image, path=output_dir, basename=basename, seed=None, prompt=None, extension=ext, info=geninfo, short_filename=True, no_prompt=True, grid=False, pnginfo_section_name="extras", existing_info=image.info, forced_filename=None)
-        shared.log.debug(f'Processed: {len(image_files)} Memory: {memory_stats()} batch')
+        shared.log.debug(f'Processed: images={len(image_files)} memory={memory_stats()} op=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, latent_index: int, mask_blur: int, mask_alpha: float, inpainting_fill: int, full_quality: bool, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, image_cfg_scale: float, diffusers_guidance_rescale: float, refiner_steps: int, refiner_start: float, clip_skip: int, denoising_strength: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, selected_scale_tab: int, height: int, width: int, scale_by: float, resize_mode: int, inpaint_full_res: bool, inpaint_full_res_padding: int, inpainting_mask_invert: int, img2img_batch_files: list, img2img_batch_input_dir: str, img2img_batch_output_dir: str, img2img_batch_inpaint_mask_dir: str, override_settings_texts, *args): # pylint: disable=unused-argument
@@ -195,5 +195,4 @@ def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_s
             processed = processing.process_images(p)
     p.close()
     generation_info_js = processed.js()
-    shared.log.debug(f'Processed: {len(processed.images)} Memory: {memory_stats()} img')
     return processed.images, generation_info_js, processed.info, plaintext_to_html(processed.comments)

modules/memmon.py

@@ -1,73 +1,52 @@
-import threading
-import time
 from collections import defaultdict
 import torch
-from modules import devices
 
 
-class MemUsageMonitor(threading.Thread):
-    run_flag = None
+class MemUsageMonitor():
     device = None
     disabled = False
     opts = None
     data = None
 
-    def __init__(self, name, device, opts):
-        threading.Thread.__init__(self)
+    def __init__(self, name, device):
         self.name = name
         self.device = device
-        self.opts = opts
-        self.daemon = True
-        self.run_flag = threading.Event()
         self.data = defaultdict(int)
         if not torch.cuda.is_available():
             self.disabled = True
         else:
             try:
-                self.cuda_mem_get_info()
+                torch.cuda.mem_get_info(self.device.index if self.device.index is not None else torch.cuda.current_device())
                 torch.cuda.memory_stats(self.device)
             except Exception:
                 self.disabled = True
 
-    def cuda_mem_get_info(self):
-        index = self.device.index if self.device.index is not None else torch.cuda.current_device()
-        return torch.cuda.mem_get_info(index)
-
-    def run(self):
+    def cuda_mem_get_info(self): # legacy for extensions only
         if self.disabled:
-            return
-        while True:
-            self.run_flag.wait()
-            torch.cuda.reset_peak_memory_stats()
-            self.data.clear()
-            if self.opts.memmon_poll_rate <= 0:
-                self.run_flag.clear()
-                continue
-            self.data["min_free"] = self.cuda_mem_get_info()[0]
-            while self.run_flag.is_set():
-                free, _total = self.cuda_mem_get_info()
-                self.data["min_free"] = min(self.data["min_free"], free)
-                time.sleep(1 / self.opts.memmon_poll_rate)
+            return 0, 0
+        return torch.cuda.mem_get_info(self.device.index if self.device.index is not None else torch.cuda.current_device())
 
-    def monitor(self):
-        self.run_flag.set()
+    def reset(self):
+        if not self.disabled:
+            torch.cuda.reset_peak_memory_stats(self.device)
+            self.data['retries'] = 0
+            self.data['oom'] = 0
+            # torch.cuda.reset_accumulated_memory_stats(self.device)
+            # torch.cuda.reset_max_memory_allocated(self.device)
+            # torch.cuda.reset_max_memory_cached(self.device)
 
     def read(self):
         if not self.disabled:
-            free, total = self.cuda_mem_get_info()
-            self.data["free"] = free
-            self.data["total"] = total
             try:
+                self.data["free"], self.data["total"] = torch.cuda.mem_get_info(self.device.index if self.device.index is not None else torch.cuda.current_device())
                 torch_stats = torch.cuda.memory_stats(self.device)
                 self.data["active"] = torch_stats["active.all.current"]
                 self.data["active_peak"] = torch_stats["active_bytes.all.peak"]
                 self.data["reserved"] = torch_stats["reserved_bytes.all.current"]
                 self.data["reserved_peak"] = torch_stats["reserved_bytes.all.peak"]
-                self.data["system_peak"] = total - self.data["min_free"]
+                self.data['retries'] = torch_stats["num_alloc_retries"]
+                self.data['oom'] = torch_stats["num_ooms"]
+                self.data["used"] = self.data["total"] - self.data["free"]
             except Exception:
                 self.disabled = True
         return self.data
-
-    def stop(self):
-        self.run_flag.clear()
-        return self.read()

modules/paths.py

@@ -42,10 +42,6 @@ for d, must_exist, what, _options in path_dirs:
         print(f"Warning: {what} not found at path {must_exist_path}", file=sys.stderr)
     else:
         d = os.path.abspath(d)
-        # if "atstart" in options:
-        #    sys.path.insert(0, d)
-        # else:
-        #    sys.path.append(d)
         sys.path.append(d)
         paths[what] = d
 

modules/processing.py

@@ -1,29 +1,39 @@
+import os
 import json
 import math
-import os
+import time
 import hashlib
 import random
 from contextlib import nullcontext
 from typing import Any, Dict, List
 import torch
 import numpy as np
-from PIL import Image, ImageFilter, ImageOps
 import cv2
+from PIL import Image, ImageFilter, ImageOps
 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
 from installer import git_commit
-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, sd_samplers_common # pylint: disable=unused-import
-import modules.shared as shared
-import modules.paths as paths
+from modules import shared, devices
+import modules.memstats
+import modules.lowvram
+import modules.masking
+import modules.paths
+import modules.scripts
+import modules.prompt_parser
+import modules.extra_networks
 import modules.face_restoration
 import modules.images as images
 import modules.styles
-import modules.sd_models as sd_models
-import modules.sd_vae as sd_vae
+import modules.sd_hijack
+import modules.sd_samplers
+import modules.sd_samplers_common
+import modules.sd_models
+import modules.sd_vae
+import modules.sd_vae_approx
+import modules.generation_parameters_copypaste
 
 
 opt_C = 4
@@ -451,7 +461,7 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts, all_seeds, all_su
         index = position_in_batch + iteration * p.batch_size
     if all_negative_prompts is None:
         all_negative_prompts = p.all_negative_prompts
-    vae = (None if not shared.opts.add_model_name_to_info or sd_vae.loaded_vae_file is None else os.path.splitext(os.path.basename(sd_vae.loaded_vae_file))[0]) if p.full_quality else 'TAESD'
+    vae = (None if not shared.opts.add_model_name_to_info or modules.sd_vae.loaded_vae_file is None else os.path.splitext(os.path.basename(modules.sd_vae.loaded_vae_file))[0]) if p.full_quality else 'TAESD'
     comment = ', '.join(comments) if comments is not None and type(comments) is list else None
 
     args = {
@@ -473,7 +483,7 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts, all_seeds, all_su
         "Prompt2": p.refiner_prompt if len(p.refiner_prompt) > 0 else None,
         "Negative2": p.refiner_negative if len(p.refiner_negative) > 0 else None,
         # other
-        "ENSD": shared.opts.eta_noise_seed_delta if shared.opts.eta_noise_seed_delta != 0 and sd_samplers_common.is_sampler_using_eta_noise_seed_delta(p) else None,
+        "ENSD": shared.opts.eta_noise_seed_delta if shared.opts.eta_noise_seed_delta != 0 and modules.sd_samplers_common.is_sampler_using_eta_noise_seed_delta(p) else None,
         "Tiling": p.tiling if p.tiling else None,
         # sdnext
         "Backend": 'Diffusers' if shared.backend == shared.Backend.DIFFUSERS else 'Original',
@@ -524,7 +534,7 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts, all_seeds, all_su
     args['Token merging ratio hr'] = token_merging_ratio_hr if token_merging_ratio_hr != 0 else None
 
     args.update(p.extra_generation_params)
-    params_text = ", ".join([k if k == v else f'{k}: {generation_parameters_copypaste.quote(v)}' for k, v in args.items() if v is not None])
+    params_text = ", ".join([k if k == v else f'{k}: {modules.generation_parameters_copypaste.quote(v)}' for k, v in args.items() if v is not None])
     negative_prompt_text = f"\nNegative prompt: {all_negative_prompts[index]}" if all_negative_prompts[index] else ""
     infotext = f"{all_prompts[index]}{negative_prompt_text}\n{params_text}".strip()
     return infotext
@@ -568,18 +578,18 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
         stored_opts[k] = shared.opts.data.get(k, None) or shared.opts.data_labels[k].default
     try:
         # if no checkpoint override or the override checkpoint can't be found, remove override entry and load opts checkpoint
-        if p.override_settings.get('sd_model_checkpoint', None) is not None and sd_models.checkpoint_aliases.get(p.override_settings.get('sd_model_checkpoint')) is None:
+        if p.override_settings.get('sd_model_checkpoint', None) is not None and modules.sd_models.checkpoint_aliases.get(p.override_settings.get('sd_model_checkpoint')) is None:
             p.override_settings.pop('sd_model_checkpoint', None)
-            sd_models.reload_model_weights()
+            modules.sd_models.reload_model_weights()
         for k, v in p.override_settings.items():
             setattr(shared.opts, k, v)
             if k == 'sd_model_checkpoint':
-                sd_models.reload_model_weights()
+                modules.sd_models.reload_model_weights()
             if k == 'sd_vae':
-                sd_vae.reload_vae_weights()
+                modules.sd_vae.reload_vae_weights()
 
         if not shared.opts.cuda_compile:
-            sd_models.apply_token_merging(p.sd_model, p.get_token_merging_ratio())
+            modules.sd_models.apply_token_merging(p.sd_model, p.get_token_merging_ratio())
 
         if shared.cmd_opts.profile:
             """
@@ -598,16 +608,16 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
             res = process_images_inner(p)
     finally:
         if not shared.opts.cuda_compile:
-            sd_models.apply_token_merging(p.sd_model, 0)
+            modules.sd_models.apply_token_merging(p.sd_model, 0)
         if p.override_settings_restore_afterwards: # restore opts to original state
             for k, v in stored_opts.items():
                 setattr(shared.opts, k, v)
                 if k == 'sd_model_checkpoint':
-                    sd_models.reload_model_weights()
+                    modules.sd_models.reload_model_weights()
                 if k == 'sd_model_refiner':
-                    sd_models.reload_model_weights()
+                    modules.sd_models.reload_model_weights()
                 if k == 'sd_vae':
-                    sd_vae.reload_vae_weights()
+                    modules.sd_vae.reload_vae_weights()
     return res
 
 
@@ -672,15 +682,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
     cached_c = [None, None]
 
     def get_conds_with_caching(function, required_prompts, steps, cache):
-        """
-        Returns the result of calling function(shared.sd_model, required_prompts, steps)
-        using a cache to store the result if the same arguments have been used before.
-
-        cache is an array containing two elements. The first element is a tuple
-        representing the previously used arguments, or None if no arguments
-        have been used before. The second element is where the previously
-        computed result is stored.
-        """
         if cache[0] is not None and (required_prompts, steps) == cache[0]:
             return cache[1]
         with devices.autocast():
@@ -692,11 +693,12 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
         return ''
 
     ema_scope_context = p.sd_model.ema_scope if shared.backend == shared.Backend.ORIGINAL else nullcontext
-    with torch.no_grad(), ema_scope_context():
+    with torch.inference_mode(), ema_scope_context():
+        t0 = time.time()
         with devices.autocast():
             p.init(p.all_prompts, p.all_seeds, p.all_subseeds)
             if shared.opts.live_previews_enable and shared.opts.show_progress_type == "Approximate NN" and shared.backend == shared.Backend.ORIGINAL:
-                sd_vae_approx.model()
+                modules.sd_vae_approx.model()
         if shared.state.job_count == -1:
             shared.state.job_count = p.n_iter
         extra_network_data = None
@@ -717,25 +719,25 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                 p.scripts.before_process_batch(p, batch_number=n, prompts=p.prompts, seeds=p.seeds, subseeds=p.subseeds)
             if len(p.prompts) == 0:
                 break
-            p.prompts, extra_network_data = extra_networks.parse_prompts(p.prompts)
+            p.prompts, extra_network_data = modules.extra_networks.parse_prompts(p.prompts)
             if not p.disable_extra_networks:
                 with devices.autocast():
-                    extra_networks.activate(p, extra_network_data)
+                    modules.extra_networks.activate(p, extra_network_data)
             if p.scripts is not None:
                 p.scripts.process_batch(p, batch_number=n, prompts=p.prompts, seeds=p.seeds, subseeds=p.subseeds)
             if n == 0:
-                with open(os.path.join(paths.data_path, "params.txt"), "w", encoding="utf8") as file:
+                with open(os.path.join(modules.paths.data_path, "params.txt"), "w", encoding="utf8") as file:
                     processed = Processed(p, [], p.seed, "")
                     file.write(processed.infotext(p, 0))
             step_multiplier = 1
-            sampler_config = sd_samplers.find_sampler_config(p.sampler_name)
+            sampler_config = modules.sd_samplers.find_sampler_config(p.sampler_name)
             step_multiplier = 2 if sampler_config and sampler_config.options.get("second_order", False) else 1
             if p.n_iter > 1:
                 shared.state.job = f"Batch {n+1} out of {p.n_iter}"
 
             if shared.backend == shared.Backend.ORIGINAL:
-                uc = get_conds_with_caching(prompt_parser.get_learned_conditioning, p.negative_prompts, p.steps * step_multiplier, cached_uc)
-                c = get_conds_with_caching(prompt_parser.get_multicond_learned_conditioning, p.prompts, p.steps * step_multiplier, cached_c)
+                uc = get_conds_with_caching(modules.prompt_parser.get_learned_conditioning, p.negative_prompts, p.steps * step_multiplier, cached_uc)
+                c = get_conds_with_caching(modules.prompt_parser.get_multicond_learned_conditioning, p.prompts, p.steps * step_multiplier, cached_c)
                 if len(modules.sd_hijack.model_hijack.comments) > 0:
                     for comment in modules.sd_hijack.model_hijack.comments:
                         comments[comment] = 1
@@ -749,8 +751,8 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                     if not shared.opts.no_half and not shared.opts.no_half_vae and shared.cmd_opts.rollback_vae:
                         shared.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)
+                        vae_file, vae_source = modules.sd_vae.resolve_vae(p.sd_model.sd_model_checkpoint)
+                        modules.sd_vae.load_vae(p.sd_model, vae_file, vae_source)
                         x_samples_ddim = [decode_first_stage(p.sd_model, samples_ddim[i:i+1].to(dtype=devices.dtype_vae))[0].cpu() for i in range(samples_ddim.size(0))]
                         for x in x_samples_ddim:
                             devices.test_for_nans(x, "vae")
@@ -768,14 +770,14 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                 raise ValueError(f"Unknown backend {shared.backend}")
 
             if shared.cmd_opts.lowvram or shared.cmd_opts.medvram and shared.backend == shared.Backend.ORIGINAL:
-                lowvram.send_everything_to_cpu()
+                modules.lowvram.send_everything_to_cpu()
                 devices.torch_gc()
             if p.scripts is not None:
                 p.scripts.postprocess_batch(p, x_samples_ddim, batch_number=n)
             if p.scripts is not None:
                 p.prompts = p.all_prompts[n * p.batch_size:(n + 1) * p.batch_size]
                 p.negative_prompts = p.all_negative_prompts[n * p.batch_size:(n + 1) * p.batch_size]
-                batch_params = scripts.PostprocessBatchListArgs(list(x_samples_ddim))
+                batch_params = modules.scripts.PostprocessBatchListArgs(list(x_samples_ddim))
                 p.scripts.postprocess_batch_list(p, batch_params, batch_number=n)
                 x_samples_ddim = batch_params.images
 
@@ -797,7 +799,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                     x_sample = modules.face_restoration.restore_faces(x_sample)
                 image = Image.fromarray(x_sample)
                 if p.scripts is not None:
-                    pp = scripts.PostprocessImageArgs(image)
+                    pp = modules.scripts.PostprocessImageArgs(image)
                     p.scripts.postprocess_image(p, pp)
                     image = pp.image
                 if p.color_corrections is not None and i < len(p.color_corrections):
@@ -832,6 +834,9 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             devices.torch_gc()
             shared.state.nextjob()
 
+        t1 = time.time()
+        shared.log.info(f'Processed: images={len(output_images)} time={t1 - t0:.2f}s its={(p.steps * len(output_images)) / (t1 - t0):.2f} memory={modules.memstats.memory_stats()}')
+
         p.color_corrections = None
         index_of_first_image = 0
         unwanted_grid_because_of_img_count = len(output_images) < 2 and shared.opts.grid_only_if_multiple
@@ -848,7 +853,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                     images.save_image(grid, p.outpath_grids, "grid", p.all_seeds[0], p.all_prompts[0], shared.opts.grid_format, info=infotext(), short_filename=not shared.opts.grid_extended_filename, p=p, grid=True)
 
     if not p.disable_extra_networks and extra_network_data:
-        extra_networks.deactivate(p, extra_network_data)
+        modules.extra_networks.deactivate(p, extra_network_data)
 
     res = Processed(
         p,
@@ -905,7 +910,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 
     def init(self, all_prompts, all_seeds, all_subseeds):
         if shared.backend == shared.Backend.DIFFUSERS:
-            sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE)
+            modules.sd_models.set_diffuser_pipe(self.sd_model, modules.sd_models.DiffusersTaskType.TEXT_2_IMAGE)
 
         self.width = self.width or 512
         self.height = self.height or 512
@@ -953,7 +958,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
             if not shared.opts.save or self.do_not_save_samples or not shared.opts.save_images_before_highres_fix:
                 return
             if not isinstance(image, Image.Image):
-                image = sd_samplers.sample_to_image(image, index, approximation=0)
+                image = modules.sd_samplers.sample_to_image(image, index, approximation=0)
             orig1 = self.extra_generation_params
             orig2 = self.restore_faces
             self.extra_generation_params = {}
@@ -964,10 +969,10 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
             images.save_image(image, self.outpath_samples, "", seeds[index], prompts[index], shared.opts.samples_format, info=info, suffix="-before-highres-fix")
 
         if shared.backend == shared.Backend.DIFFUSERS:
-            sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE)
+            modules.sd_models.set_diffuser_pipe(self.sd_model, modules.sd_models.DiffusersTaskType.TEXT_2_IMAGE)
 
         self.ops.append('txt2img')
-        self.sampler = sd_samplers.create_sampler(self.sampler_name, self.sd_model)
+        self.sampler = modules.sd_samplers.create_sampler(self.sampler_name, self.sd_model)
         latent_scale_mode = shared.latent_upscale_modes.get(self.hr_upscaler, None) if self.hr_upscaler is not None else shared.latent_upscale_modes.get(shared.latent_upscale_default_mode, "None")
         if self.enable_hr and latent_scale_mode is None:
             if len([x for x in shared.sd_upscalers if x.name == self.hr_upscaler]) == 0:
@@ -1020,14 +1025,14 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
             shared.state.nextjob()
             if self.latent_sampler == "PLMS":
                 self.latent_sampler = 'UniPC'
-            self.sampler = sd_samplers.create_sampler(self.latent_sampler or self.sampler_name, self.sd_model)
+            self.sampler = modules.sd_samplers.create_sampler(self.latent_sampler or self.sampler_name, self.sd_model)
             samples = samples[:, :, self.truncate_y//2:samples.shape[2]-(self.truncate_y+1)//2, self.truncate_x//2:samples.shape[3]-(self.truncate_x+1)//2]
             noise = create_random_tensors(samples.shape[1:], seeds=seeds, subseeds=subseeds, subseed_strength=subseed_strength, p=self)
             x = None
             devices.torch_gc() # GC now before running the next img2img to prevent running out of memory
-            sd_models.apply_token_merging(self.sd_model, self.get_token_merging_ratio(for_hr=True))
+            modules.sd_models.apply_token_merging(self.sd_model, self.get_token_merging_ratio(for_hr=True))
             samples = self.sampler.sample_img2img(self, samples, noise, conditioning, unconditional_conditioning, steps=self.hr_second_pass_steps or self.steps, image_conditioning=image_conditioning)
-            sd_models.apply_token_merging(self.sd_model, self.get_token_merging_ratio())
+            modules.sd_models.apply_token_merging(self.sd_model, self.get_token_merging_ratio())
             self.is_hr_pass = False
 
         return samples
@@ -1065,14 +1070,14 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
 
     def init(self, all_prompts, all_seeds, all_subseeds):
         if shared.backend == shared.Backend.DIFFUSERS and self.image_mask is None:
-            sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE)
+            modules.sd_models.set_diffuser_pipe(self.sd_model, modules.sd_models.DiffusersTaskType.IMAGE_2_IMAGE)
         elif shared.backend == shared.Backend.DIFFUSERS and self.image_mask is not None:
-            sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.INPAINTING)
+            modules.sd_models.set_diffuser_pipe(self.sd_model, modules.sd_models.DiffusersTaskType.INPAINTING)
             self.sd_model.dtype = self.sd_model.unet.dtype
 
         if self.sampler_name == "PLMS":
             self.sampler_name = 'UniPC'
-        self.sampler = sd_samplers.create_sampler(self.sampler_name, self.sd_model)
+        self.sampler = modules.sd_samplers.create_sampler(self.sampler_name, self.sd_model)
 
         if self.image_mask is not None:
             self.ops.append('inpaint')
@@ -1089,8 +1094,8 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
             if self.inpaint_full_res:
                 self.mask_for_overlay = image_mask
                 mask = image_mask.convert('L')
-                crop_region = masking.get_crop_region(np.array(mask), self.inpaint_full_res_padding)
-                crop_region = masking.expand_crop_region(crop_region, self.width, self.height, mask.width, mask.height)
+                crop_region = modules.masking.get_crop_region(np.array(mask), self.inpaint_full_res_padding)
+                crop_region = modules.masking.expand_crop_region(crop_region, self.width, self.height, mask.width, mask.height)
                 x1, y1, x2, y2 = crop_region
                 mask = mask.crop(crop_region)
                 image_mask = images.resize_image(3, mask, self.width, self.height)
@@ -1137,7 +1142,7 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
             self.init_images = [image] # assign early for diffusers
             if image_mask is not None:
                 if self.inpainting_fill != 1:
-                    image = masking.fill(image, latent_mask)
+                    image = modules.masking.fill(image, latent_mask)
             if add_color_corrections:
                 self.color_corrections.append(setup_color_correction(image))
             image = np.array(image).astype(np.float32) / 255.0
@@ -1183,9 +1188,9 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
     def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts):
         if shared.backend == shared.Backend.DIFFUSERS:
             if self.init_mask is None: # pylint: disable=no-member
-                sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE)
+                modules.sd_models.set_diffuser_pipe(self.sd_model, modules.sd_models.DiffusersTaskType.IMAGE_2_IMAGE)
             else:
-                sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.INPAINTING)
+                modules.sd_models.set_diffuser_pipe(self.sd_model, modules.sd_models.DiffusersTaskType.INPAINTING)
                 self.sd_model.dtype = self.sd_model.unet.dtype
 
         x = create_random_tensors([4, self.height // 8, self.width // 8], seeds=seeds, subseeds=subseeds, subseed_strength=self.subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self)

modules/scripts.py

@@ -276,6 +276,25 @@ def wrap_call(func, filename, funcname, *args, default=None, **kwargs):
     return default
 
 
+class ScriptSummary:
+    def __init__(self, op):
+        self.start = time.time()
+        self.update = time.time()
+        self.op = op
+        self.time = {}
+
+    def record(self, script):
+        self.update = time.time()
+        self.time[script] = round(time.time() - self.update, 2)
+
+    def report(self):
+        total = sum(self.time.values())
+        if total == 0:
+            return
+        scripts = [f'{k}:{v}s' for k, v in self.time.items() if v > 0]
+        log.debug(f'Script: op={self.op} total={total}s scripts={scripts}')
+
+
 class ScriptRunner:
     def __init__(self):
         self.scripts = []
@@ -405,6 +424,7 @@ class ScriptRunner:
         return inputs
 
     def run(self, p, *args):
+        s = ScriptSummary('run')
         script_index = args[0]
         if script_index == 0:
             return None
@@ -412,117 +432,111 @@ class ScriptRunner:
         if script is None:
             return None
         parsed = p.per_script_args.get(script.title(), args[script.args_from:script.args_to])
-        t0 = time.time()
         processed = script.run(p, *parsed)
-        log.debug(f'Script run: {script.title()}:{round(time.time()-t0, 2)}s')
+        s.record(script.title())
+        s.report()
         return processed
 
     def process(self, p, **kwargs):
-        s = []
+        s = ScriptSummary('process')
         for script in self.alwayson_scripts:
             try:
-                t0 = time.time()
                 args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
                 script.process(p, *args, **kwargs)
-                s.append(f'{script.title()}:{round(time.time()-t0, 2)}s')
             except Exception as e:
                 errors.display(e, f'Running script process: {script.filename}')
-        log.debug(f'Script process: {s}')
+            s.record(script.title())
+        s.report()
 
     def before_process_batch(self, p, **kwargs):
-        s = []
+        s = ScriptSummary('before-process-batch')
         for script in self.alwayson_scripts:
             try:
-                t0 = time.time()
                 args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
                 script.before_process_batch(p, *args, **kwargs)
-                s.append(f'{script.title()}:{round(time.time()-t0, 2)}s')
             except Exception as e:
                 errors.display(e, f'Running script before process batch: {script.filename}')
-        log.debug(f'Script before-process-batch: {s}')
+            s.record(script.title())
+        s.report()
 
     def process_batch(self, p, **kwargs):
-        s = []
+        s = ScriptSummary('process-batch')
         for script in self.alwayson_scripts:
             try:
-                t0 = time.time()
                 args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
                 script.process_batch(p, *args, **kwargs)
-                s.append(f'{script.title()}:{round(time.time()-t0, 2)}s')
             except Exception as e:
                 errors.display(e, f'Running script process batch: {script.filename}')
-        log.debug(f'Script process-batch: {s}')
+            s.record(script.title())
+        s.report()
 
     def postprocess(self, p, processed):
-        s = []
+        s = ScriptSummary('postprocess')
         for script in self.alwayson_scripts:
             try:
-                t0 = time.time()
                 args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
                 script.postprocess(p, processed, *args)
-                s.append(f'{script.title()}:{round(time.time()-t0, 2)}s')
             except Exception as e:
                 errors.display(e, f'Running script postprocess: {script.filename}')
-        log.debug(f'Script postprocess: {s}')
+            s.record(script.title())
+        s.report()
 
     def postprocess_batch(self, p, images, **kwargs):
-        s = []
+        s = ScriptSummary('postprocess-batch')
         for script in self.alwayson_scripts:
             try:
-                t0 = time.time()
                 args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
                 script.postprocess_batch(p, *args, images=images, **kwargs)
-                s.append(f'{script.title()}:{round(time.time()-t0, 2)}s')
             except Exception as e:
                 errors.display(e, f'Running script before postprocess batch: {script.filename}')
-        log.debug(f'Script postprocess-batch: {s}')
+            s.record(script.title())
+        s.report()
 
     def postprocess_batch_list(self, p, pp: PostprocessBatchListArgs, **kwargs):
-        s = []
+        s = ScriptSummary('postprocess-batch-list')
         for script in self.alwayson_scripts:
             try:
-                t0 = time.time()
                 args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
                 script.postprocess_batch_list(p, pp, *args, **kwargs)
-                s.append(f'{script.title()}:{round(time.time()-t0, 2)}s')
             except Exception as e:
                 errors.display(e, f'Running script before postprocess batch list: {script.filename}')
-        log.debug(f'Script postprocess-batch-list: {s}')
+            s.record(script.title())
+        s.report()
 
     def postprocess_image(self, p, pp: PostprocessImageArgs):
-        s = []
+        s = ScriptSummary('postprocess-image')
         for script in self.alwayson_scripts:
             try:
-                t0 = time.time()
                 args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
                 script.postprocess_image(p, pp, *args)
-                s.append(f'{script.title()}:{round(time.time()-t0, 2)}s')
             except Exception as e:
                 errors.display(e, f'Running script postprocess image: {script.filename}')
-        log.debug(f'Script postprocess-image: {s}')
+            s.record(script.title())
+        s.report()
 
     def before_component(self, component, **kwargs):
+        s = ScriptSummary('before-component')
         for script in self.scripts:
             try:
-                t0 = time.time()
                 script.before_component(component, **kwargs)
-                time_component[script.title()] = time_component.get(script.title(), 0) + (time.time()-t0)
             except Exception as e:
                 errors.display(e, f'Running script before component: {script.filename}')
+            s.record(script.title())
+        s.report()
 
     def after_component(self, component, **kwargs):
+        s = ScriptSummary('after-component')
         for script in self.scripts:
             try:
-                t0 = time.time()
                 script.after_component(component, **kwargs)
-                time_component[script.title()] = time_component.get(script.title(), 0) + (time.time()-t0)
             except Exception as e:
                 errors.display(e, f'Running script after component: {script.filename}')
+            s.record(script.title())
+        s.report()
 
     def reload_sources(self, cache):
-        s = []
+        s = ScriptSummary('reload-sources')
         for si, script in list(enumerate(self.scripts)):
-            t0 = time.time()
             args_from = script.args_from
             args_to = script.args_to
             filename = script.filename
@@ -536,8 +550,8 @@ class ScriptRunner:
                     self.scripts[si].filename = filename
                     self.scripts[si].args_from = args_from
                     self.scripts[si].args_to = args_to
-            s.append(f'{script.title()}:{round(time.time()-t0, 2)}s')
-        log.debug(f'Script reload-sources: {s}')
+            s.record(script.title())
+        s.report()
 
 
 scripts_txt2img: ScriptRunner = None

modules/sd_models.py

@@ -355,7 +355,7 @@ def read_state_dict(checkpoint_file, map_location=None): # pylint: disable=unuse
         return None
     try:
         pl_sd = None
-        with progress.open(checkpoint_file, 'rb', description=f'Loading weights: [cyan]{checkpoint_file}', auto_refresh=True) as f:
+        with progress.open(checkpoint_file, 'rb', description=f'[cyan]Loading weights: [yellow]{checkpoint_file}', auto_refresh=True) as f:
             _, extension = os.path.splitext(checkpoint_file)
             if extension.lower() == ".ckpt" and shared.opts.sd_disable_ckpt:
                 shared.log.warning(f"Checkpoint loading disabled: {checkpoint_file}")

modules/sd_samplers.py

@@ -47,14 +47,14 @@ def create_sampler(name, model):
         sampler = config.constructor(model)
         sampler.config = config
         sampler.name = name
-        shared.log.debug(f'Sampler: {sampler.name} {sampler.config.options}')
+        shared.log.debug(f'Sampler: sampler={sampler.name} config={sampler.config.options}')
         return sampler
     elif shared.backend == shared.Backend.DIFFUSERS:
         sampler = config.constructor(model)
         if not hasattr(model, 'scheduler_config'):
             model.scheduler_config = sampler.sampler.config.copy()
         model.scheduler = sampler.sampler
-        shared.log.debug(f'Sampler: {sampler.name} {sampler.config}')
+        shared.log.debug(f'Sampler: sampler={sampler.name} config={sampler.config}')
         return sampler.sampler
     else:
         return None

modules/shared.py

@@ -378,7 +378,7 @@ options_templates.update(options_section(('optimizations', "Optimizations"), {
 }))
 
 options_templates.update(options_section(('cuda', "Compute Settings"), {
-    "memmon_poll_rate": OptionInfo(2, "VRAM usage polls per second during generation", gr.Slider, {"minimum": 0, "maximum": 40, "step": 1}),
+    # "memmon_poll_rate": OptionInfo(2, "VRAM usage polls per second during generation", 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" or cmd_opts.use_openvino else "BF16" if devices.backend == "ipex" else "FP16", "Device precision type", gr.Radio, lambda: {"choices": ["FP32", "FP16", "BF16"]}),
     "no_half": OptionInfo(False, "Use full precision for model (--no-half)", None, None, None),
@@ -825,8 +825,7 @@ devices.device, devices.device_interrogate, devices.device_gfpgan, devices.devic
 device = devices.device
 batch_cond_uncond = opts.always_batch_cond_uncond or not (cmd_opts.lowvram or cmd_opts.medvram)
 parallel_processing_allowed = not cmd_opts.lowvram
-mem_mon = modules.memmon.MemUsageMonitor("MemMon", device, opts)
-mem_mon.start()
+mem_mon = modules.memmon.MemUsageMonitor("MemMon", devices.device)
 if devices.backend == "directml":
     directml_do_hijack()
 
@@ -916,7 +915,7 @@ def restart_server(restart=True):
         demo.server.close()
         demo.fns = []
         # os._exit(0)
-    except Exception as e:
+    except (Exception, BaseException) as e:
         log.error(f'Server shutdown error: {e}')
     if restart:
         log.info('Server will restart')

modules/txt2img.py

@@ -2,7 +2,6 @@ import modules.scripts
 from modules import sd_samplers, shared, processing
 from modules.generation_parameters_copypaste import create_override_settings_dict
 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, latent_index: int, full_quality: bool, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, image_cfg_scale: float, diffusers_guidance_rescale: float, clip_skip: int, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, 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, refiner_steps: int, refiner_start: int, refiner_prompt: str, refiner_negative: str, override_settings_texts, *args): # pylint: disable=unused-argument
@@ -68,5 +67,4 @@ def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, step
     if processed is None:
         return [], '', '', 'Error: processing failed'
     generation_info_js = processed.js()
-    shared.log.debug(f'Processed: {len(processed.images)} Memory: {memory_stats()} txt')
     return processed.images, generation_info_js, processed.info, plaintext_to_html(processed.comments)

modules/ui_extra_networks.py

@@ -17,7 +17,7 @@ import modules.ui_symbols as symbols
 
 
 extra_pages = []
-allowed_dirs = set()
+allowed_dirs = []
 dir_cache = {} # key=path, value=(mtime, listdir(path))
 refresh_time = None
 
@@ -39,16 +39,21 @@ def register_page(page):
     # registers extra networks page for the UI; recommend doing it in on_before_ui() callback for extensions
     extra_pages.append(page)
     allowed_dirs.clear()
-    allowed_dirs.update(set(sum([x.allowed_directories_for_previews() for x in extra_pages], [])))
+    for page in extra_pages:
+        for folder in page.allowed_directories_for_previews():
+            if folder not in allowed_dirs:
+                allowed_dirs.append(os.path.abspath(folder))
 
 
 def fetch_file(filename: str = ""):
-    if filename.startswith('html/'):
+    if not os.path.exists(filename):
+        return JSONResponse({ "error": f"file {filename}: not found" }, status_code=404)
+    if filename.startswith('html/') or filename.startswith('models/'):
         return FileResponse(filename, headers={"Accept-Ranges": "bytes"})
-    if not any(Path(x).absolute() in Path(filename).absolute().parents for x in allowed_dirs):
-        return JSONResponse({"error": f"File cannot be fetched: {filename}. Must be in one of directories registered by extra pages."})
+    if not any(Path(folder).absolute() in Path(filename).absolute().parents for folder in allowed_dirs):
+        return JSONResponse({ "error": f"file {filename}: must be in one of allowed directories" }, status_code=403)
     if os.path.splitext(filename)[1].lower() not in (".png", ".jpg", ".jpeg", ".webp"):
-        return JSONResponse({"error": f"File cannot be fetched: {filename}. Only png and jpg and webp."})
+        return JSONResponse({"error": f"file {filename}: not an image file"}, status_code=403)
     return FileResponse(filename, headers={"Accept-Ranges": "bytes"})
 
 
@@ -362,7 +367,7 @@ def create_ui(container, button, tabname, skip_indexing = False):
     def refresh(title):
         res = []
         for page in extra_pages:
-            if title == '' or title == page.title:
+            if title == '' or title == page.title or len(page.html) == 0:
                 shared.log.debug(f"Refreshing Extra networks: page={page.title} tab={ui.tabname}")
                 page.refresh()
                 page.create_page(ui.tabname)