diff --git a/launch.py b/launch.py index e00da58c7..5c8a6051a 100755 --- a/launch.py +++ b/launch.py @@ -192,6 +192,9 @@ def main(): global args # pylint: disable=global-statement installer.ensure_base_requirements() init_args() # setup argparser and default folders + if args.malloc: + import tracemalloc + tracemalloc.start() installer.args = args installer.setup_logging() installer.log.info('Starting SD.Next') diff --git a/modules/cmd_args.py b/modules/cmd_args.py index 752ad02c0..cb4e5fc16 100644 --- a/modules/cmd_args.py +++ b/modules/cmd_args.py @@ -26,6 +26,7 @@ def main_args(): group_diag.add_argument("--no-hashing", default=os.environ.get("SD_NOHASHING", False), action='store_true', help="Disable hashing of checkpoints, default: %(default)s") group_diag.add_argument("--no-metadata", default=os.environ.get("SD_NOMETADATA", False), action='store_true', help="Disable reading of metadata from models, default: %(default)s") group_diag.add_argument("--profile", default=os.environ.get("SD_PROFILE", False), action='store_true', help="Run profiler, default: %(default)s") + group_diag.add_argument("--malloc", default=os.environ.get("SD_PROFILE", False), action='store_true', help="Trace memory ops, default: %(default)s") group_diag.add_argument("--disable-queue", default=os.environ.get("SD_DISABLEQUEUE", False), action='store_true', help="Disable queues, default: %(default)s") group_diag.add_argument('--debug', default=os.environ.get("SD_DEBUG", False), action='store_true', help = "Run installer with debug logging, default: %(default)s") diff --git a/modules/lora/networks.py b/modules/lora/networks.py index f211149bd..23c45ff2a 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -22,7 +22,7 @@ from modules import shared, devices, sd_models, sd_models_compile, errors, files debug = os.environ.get('SD_LORA_DEBUG', None) is not None -pbar = p.Progress(p.TextColumn('[cyan]LoRA apply'), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TextColumn('[cyan]{task.description}'), console=shared.console) +pbar = p.Progress(p.TextColumn('[cyan]{task.description}'), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), console=shared.console) extra_network_lora = None available_networks = {} available_network_aliases = {} @@ -50,6 +50,13 @@ def total_time(): return sum(timer.values()) +def get_timers(): + t = { 'total': round(sum(timer.values()), 2) } + for k, v in timer.items(): + t[k] = round(v, 2) + return t + + def assign_network_names_to_compvis_modules(sd_model): if sd_model is None: return @@ -362,7 +369,8 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn network_layer_name = getattr(self, 'network_layer_name', None) current_names = getattr(self, "network_current_names", ()) wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in loaded_networks) - maybe_backup_weights(self, wanted_names) + if network_layer_name is not None and any([net.modules.get(network_layer_name, None) for net in loaded_networks]): # noqa: C419 + maybe_backup_weights(self, wanted_names) if current_names != wanted_names: batch_updown = None batch_ex_bias = None @@ -414,30 +422,23 @@ def network_load(): # called from processing if shared.opts.diffusers_offload_mode != "none": sd_models.disable_offload(sd_model) sd_models.move_model(sd_model, device=devices.cpu) + modules = [] + for component_name in ['text_encoder','text_encoder_2', 'unet', 'transformer']: + component = getattr(sd_model, component_name, None) + if component is not None and hasattr(component, 'named_modules'): + modules += list(component.named_modules()) with pbar: - for component_name in ['text_encoder','text_encoder_2', 'unet', 'transformer']: - component = getattr(sd_model, component_name, None) - if component is not None: - applied = 0 - modules = list(component.named_modules()) - task_start = time.time() - task = pbar.add_task(description=component_name , total=len(modules), visible=False) - for _, module in modules: - layer_name = getattr(module, 'network_layer_name', None) - if layer_name is None: - continue - present = any([net.modules.get(layer_name, None) for net in loaded_networks]) # noqa: C419 - if present: - network_apply_weights(module) - applied += 1 - pbar.update(task, advance=1, visible=(time.time() - task_start) > 1) # progress bar becomes visible if operation takes more than 1sec - pbar.remove_task(task) - if debug: - shared.log.debug(f'Load network: type=LoRA component={component_name} modules={len(modules)} applied={applied}') + task = pbar.add_task(description='Apply network: type=LoRA' , total=len(modules), visible=len(loaded_networks) > 0) + for _, module in modules: + network_apply_weights(module) + # pbar.update(task, advance=1) # progress bar becomes visible if operation takes more than 1sec + pbar.remove_task(task) + if debug: + shared.log.debug(f'Load network: type=LoRA modules={len(modules)}') if shared.opts.diffusers_offload_mode != "none": sd_models.set_diffuser_offload(sd_model, op="model") if debug: - shared.log.debug(f'Load network: type=LoRA total={total_time():.2f} timers={timer}') + shared.log.debug(f'Load network: type=LoRA timers{get_timers()}') def list_available_networks(): diff --git a/modules/processing.py b/modules/processing.py index 16e7a9213..92faaee8d 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -473,4 +473,13 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: p.scripts.postprocess(p, processed) timer.process.record('post') shared.log.info(f'Processed: images={len(output_images)} its={(p.steps * len(output_images)) / (t1 - t0):.2f} time={t1-t0:.2f} timers={timer.process.dct(min_time=0.02)} memory={memstats.memory_stats()}') + + if shared.cmd_opts.malloc: + import tracemalloc + snapshot = tracemalloc.take_snapshot() + stats = snapshot.statistics('lineno') + shared.log.debug('Profile malloc:') + for stat in stats[:20]: + frame = stat.traceback[0] + shared.log.debug(f' file="{frame.filename}":{frame.lineno} size={stat.size}') return processed diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py index 22acf296c..ab08d4cc8 100644 --- a/modules/processing_helpers.py +++ b/modules/processing_helpers.py @@ -1,4 +1,5 @@ import os +import time import math import random import warnings @@ -9,7 +10,7 @@ import cv2 from PIL import Image from skimage import exposure from blendmodes.blend import blendLayers, BlendType -from modules import shared, devices, images, sd_models, sd_samplers, sd_hijack_hypertile, processing_vae +from modules import shared, devices, images, sd_models, sd_samplers, sd_hijack_hypertile, processing_vae, timer debug = shared.log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -352,6 +353,7 @@ def img2img_image_conditioning(p, source_image, latent_image, image_mask=None): def validate_sample(tensor): + t0 = time.time() if not isinstance(tensor, np.ndarray) and not isinstance(tensor, torch.Tensor): return tensor dtype = tensor.dtype @@ -377,6 +379,8 @@ def validate_sample(tensor): if upcast is not None and not upcast: setattr(shared.sd_model.vae.config, 'force_upcast', True) # noqa: B010 shared.log.warning('Decode: upcast=True set, retry operation') + t1 = time.time() + timer.process.add('validate', t1 - t0) return cast