From dbe4d2ff70653a866ca308d7ef396a86ed0837b9 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sat, 27 Jan 2024 02:54:39 +0300 Subject: [PATCH] Cleanup upscalers --- modules/intel/ipex/hijacks.py | 2 + modules/ldsr/ldsr_model_arch.py | 6 +- modules/postprocess/esrgan_model.py | 4 +- modules/postprocess/realesrgan_model_arch.py | 5 +- modules/postprocess/scunet_model.py | 12 +-- modules/postprocess/swinir_model.py | 10 +- modules/sd_models_compile.py | 6 ++ modules/upscaler.py | 96 ++++++++++++-------- 8 files changed, 85 insertions(+), 56 deletions(-) diff --git a/modules/intel/ipex/hijacks.py b/modules/intel/ipex/hijacks.py index 2666c779a..cde138280 100644 --- a/modules/intel/ipex/hijacks.py +++ b/modules/intel/ipex/hijacks.py @@ -99,6 +99,8 @@ def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0. key = key.to(dtype=query.dtype) if query.dtype != value.dtype: value = value.to(dtype=query.dtype) + if attn_mask is not None and query.dtype != attn_mask.dtype: + attn_mask = attn_mask.to(dtype=query.dtype) return original_scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal) # A1111 FP16 diff --git a/modules/ldsr/ldsr_model_arch.py b/modules/ldsr/ldsr_model_arch.py index 9a4f1461f..7494fcfcc 100644 --- a/modules/ldsr/ldsr_model_arch.py +++ b/modules/ldsr/ldsr_model_arch.py @@ -35,14 +35,14 @@ class LDSR: config.model.target = "ldm.models.diffusion.ddpm.LatentDiffusionV1" model: torch.nn.Module = instantiate_from_config(config.model) model.load_state_dict(sd, strict=False) - model = model.to(shared.device) + model = model.to(devices.device) if half_attention: model = model.half() if shared.cmd_opts.opt_channelslast: model = model.to(memory_format=torch.channels_last) sd_hijack.model_hijack.hijack(model) # apply optimization model.eval() - model = compile_upscaler(model, name=self.modelPath) + model = compile_upscaler(model) cached_ldsr_model = model return {"model": model} @@ -151,7 +151,7 @@ def get_cond(selected_path): c = rearrange(c, '1 c h w -> 1 h w c') c = 2. * c - 1. - c = c.to(shared.device) + c = c.to(devices.device) example["LR_image"] = c example["image"] = c_up diff --git a/modules/postprocess/esrgan_model.py b/modules/postprocess/esrgan_model.py index c4d3e4ad4..74babe29e 100644 --- a/modules/postprocess/esrgan_model.py +++ b/modules/postprocess/esrgan_model.py @@ -154,7 +154,7 @@ class UpscalerESRGAN(Upscaler): model = arch.SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=num_conv, upscale=4, act_type='prelu') model.load_state_dict(state_dict) model.eval() - model = compile_upscaler(model, name=self.name) + model = compile_upscaler(model) self.models[info.local_data_path] = model return self.models[info.local_data_path] @@ -169,7 +169,7 @@ class UpscalerESRGAN(Upscaler): model = arch.RRDBNet(in_nc=in_nc, out_nc=out_nc, nf=nf, nb=nb, upscale=mscale, plus=plus) model.load_state_dict(state_dict) model.eval() - model = compile_upscaler(model, name=self.name) + model = compile_upscaler(model) self.models[info.local_data_path] = model return self.models[info.local_data_path] diff --git a/modules/postprocess/realesrgan_model_arch.py b/modules/postprocess/realesrgan_model_arch.py index 06f8d9c1f..30b8e65ac 100644 --- a/modules/postprocess/realesrgan_model_arch.py +++ b/modules/postprocess/realesrgan_model_arch.py @@ -8,6 +8,7 @@ import torch from torch import nn from torch.nn import functional as F from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn +from modules import devices from modules.shared import log, console from modules.upscaler import compile_upscaler @@ -54,7 +55,7 @@ class RealESRGANer(): self.device = torch.device( f'cuda:{gpu_id}' if torch.cuda.is_available() else 'cpu') if device is None else device else: - self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device is None else device + self.device = devices.device if device is None else device if isinstance(model_path, list): # dni @@ -78,8 +79,8 @@ class RealESRGANer(): model.eval() if self.half: model = model.half() - model = compile_upscaler(model, name=self.name) self.model = model.to(self.device) + self.model = compile_upscaler(self.model) def dni(self, net_a, net_b, dni_weight, key='params', loc='cpu'): """Deep network interpolation. diff --git a/modules/postprocess/scunet_model.py b/modules/postprocess/scunet_model.py index 566a8009c..d66dece39 100644 --- a/modules/postprocess/scunet_model.py +++ b/modules/postprocess/scunet_model.py @@ -4,7 +4,7 @@ import torch from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn from modules import devices from modules.postprocess.scunet_model_arch import SCUNet as net -from modules.shared import opts, log, console, device +from modules.shared import opts, log, console from modules.upscaler import Upscaler, compile_upscaler @@ -30,8 +30,8 @@ class UpscalerSCUNet(Upscaler): log.info(f"Upscaler loaded: type={self.name} model={info.local_data_path}") for _, v in model.named_parameters(): v.requires_grad = False - model = model.to(device) - model = compile_upscaler(model, name=self.name) + model = model.to(devices.device) + model = compile_upscaler(model) self.models[info.local_data_path] = model return model @@ -49,8 +49,8 @@ class UpscalerSCUNet(Upscaler): stride = tile - tile_overlap h_idx_list = list(range(0, h - tile, stride)) + [h - tile] w_idx_list = list(range(0, w - tile, stride)) + [w - tile] - E = torch.zeros(1, 3, h * sf, w * sf, dtype=img.dtype, device=device) - W = torch.zeros_like(E, dtype=devices.dtype, device=device) + E = torch.zeros(1, 3, h * sf, w * sf, dtype=img.dtype, device=devices.device) + W = torch.zeros_like(E, dtype=devices.dtype, device=devices.device) with Progress(TextColumn('[cyan]{task.description}'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn(), console=console) as progress: task = progress.add_task(description="Upscaling", total=len(h_idx_list) * len(w_idx_list)) for h_idx in h_idx_list: @@ -78,7 +78,7 @@ class UpscalerSCUNet(Upscaler): np_img = np.array(img) np_img = np_img[:, :, ::-1] # RGB to BGR np_img = np_img.transpose((2, 0, 1)) / 255 # HWC to CHW - torch_img = torch.from_numpy(np_img).float().unsqueeze(0).to(device) # type: ignore + torch_img = torch.from_numpy(np_img).float().unsqueeze(0).to(devices.device) # type: ignore if tile > h or tile > w: _img = torch.zeros(1, 3, max(h, tile), max(w, tile), dtype=torch_img.dtype, device=torch_img.device) _img[:, :, :h, :w] = torch_img # pad image diff --git a/modules/postprocess/swinir_model.py b/modules/postprocess/swinir_model.py index 3692cd836..4d4d1b476 100644 --- a/modules/postprocess/swinir_model.py +++ b/modules/postprocess/swinir_model.py @@ -58,7 +58,7 @@ class UpscalerSwinIR(Upscaler): else: model.load_state_dict(pretrained_model, strict=True) shared.log.info(f"Upscaler loaded: type={self.name} model={info.local_data_path} param={param}") - model = compile_upscaler(model, name=self.name) + model = compile_upscaler(model) self.models[info.local_data_path] = model return model except Exception as e: @@ -69,7 +69,7 @@ class UpscalerSwinIR(Upscaler): model = self.load_model(selected_model) if model is None: return img - model = model.to(shared.device, dtype=devices.dtype) + model = model.to(devices.device, dtype=devices.dtype) img = upscale(img, model) if shared.opts.upscaler_unload and selected_model in self.models: del self.models[selected_model] @@ -92,7 +92,7 @@ def upscale( img = img[:, :, ::-1] img = np.moveaxis(img, 2, 0) / 255 img = torch.from_numpy(img).float() - img = img.unsqueeze(0).to(shared.device, dtype=devices.dtype) + img = img.unsqueeze(0).to(devices.device, dtype=devices.dtype) with torch.no_grad(), devices.autocast(): _, _, h_old, w_old = img.size() h_pad = (h_old // window_size + 1) * window_size - h_old @@ -119,8 +119,8 @@ def inference(img, model, tile, tile_overlap, window_size, scale): stride = tile - tile_overlap h_idx_list = list(range(0, h - tile, stride)) + [h - tile] w_idx_list = list(range(0, w - tile, stride)) + [w - tile] - E = torch.zeros(b, c, h * sf, w * sf, dtype=devices.dtype, device=shared.device).type_as(img) - W = torch.zeros_like(E, dtype=devices.dtype, device=shared.device) + E = torch.zeros(b, c, h * sf, w * sf, dtype=devices.dtype, device=devices.device).type_as(img) + W = torch.zeros_like(E, dtype=devices.dtype, device=devices.device) with Progress(TextColumn('[cyan]{task.description}'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn(), console=shared.console) as progress: task = progress.add_task(description="Upscaling Initializing", total=len(h_idx_list) * len(w_idx_list)) diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py index 8fc9faa5c..f0dd85f8b 100644 --- a/modules/sd_models_compile.py +++ b/modules/sd_models_compile.py @@ -29,27 +29,33 @@ def ipex_optimize(sd_model): import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import if "Model" in shared.opts.ipex_optimize: if hasattr(sd_model, 'unet'): + sd_model.unet.eval() sd_model.unet.training = False sd_model.unet = ipex.optimize(sd_model.unet, dtype=devices.dtype_unet, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init elif hasattr(sd_model, 'transformer'): + sd_model.transformer.eval() sd_model.transformer.training = False sd_model.transformer = ipex.optimize(sd_model.transformer, dtype=devices.dtype_unet, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init else: shared.log.warning('IPEX Optimize enabled but model has no Unet or Transformer') if "VAE" in shared.opts.ipex_optimize: if hasattr(sd_model, 'vae'): + sd_model.vae.eval() sd_model.vae.training = False sd_model.vae = ipex.optimize(sd_model.vae, dtype=devices.dtype_vae, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init elif hasattr(sd_model, 'movq'): + sd_model.movq.eval() sd_model.movq.training = False sd_model.movq = ipex.optimize(sd_model.movq, dtype=devices.dtype_vae, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init else: shared.log.warning('Compress VAE Weights enabled but model has no VAE') if "Text Encoder" in shared.opts.ipex_optimize: if hasattr(sd_model, 'text_encoder'): + sd_model.text_encoder.eval() sd_model.text_encoder.training = False sd_model.text_encoder = ipex.optimize(sd_model.text_encoder, dtype=devices.dtype_unet, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init if hasattr(sd_model, 'text_encoder_2'): + sd_model.text_encoder_2.eval() sd_model.text_encoder_2.training = False sd_model.text_encoder_2 = ipex.optimize(sd_model.text_encoder_2, dtype=devices.dtype_unet, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init else: diff --git a/modules/upscaler.py b/modules/upscaler.py index 1a2b4ea9f..75b2485cc 100644 --- a/modules/upscaler.py +++ b/modules/upscaler.py @@ -1,9 +1,11 @@ import os import copy +import time +import logging from abc import abstractmethod from PIL import Image -import modules.shared -from modules import modelloader +from modules import devices, modelloader, shared +from installer import setup_logging LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.Resampling.LANCZOS) @@ -26,22 +28,22 @@ class Upscaler: def __init__(self, create_dirs=True): global models # pylint: disable=global-statement if models is None: - models = modules.shared.readfile('html/upscalers.json') + models = shared.readfile('html/upscalers.json') self.mod_pad_h = None - self.tile_size = modules.shared.opts.upscaler_tile_size - self.tile_pad = modules.shared.opts.upscaler_tile_overlap - self.device = modules.shared.device + self.tile_size = shared.opts.upscaler_tile_size + self.tile_pad = shared.opts.upscaler_tile_overlap + self.device = shared.device self.img = None self.output = None self.scale = 1 - self.half = not modules.shared.cmd_opts.no_half + self.half = not shared.cmd_opts.no_half self.pre_pad = 0 self.mod_scale = None self.model_download_path = None if self.user_path is not None and len(self.user_path) > 0 and not os.path.exists(self.user_path): - modules.shared.log.info(f'Upscaler create: folder="{self.user_path}"') + shared.log.info(f'Upscaler create: folder="{self.user_path}"') if self.model_path is None and self.name: - self.model_path = os.path.join(modules.shared.models_path, self.name) + self.model_path = os.path.join(shared.models_path, self.name) if self.model_path and create_dirs: os.makedirs(self.model_path, exist_ok=True) try: @@ -64,7 +66,7 @@ class Upscaler: scaler.custom = True scalers.append(scaler) loaded.append(file_name) - modules.shared.log.debug(f'Upscaler type={self.name} folder="{folder}" model="{model_name}" path="{file_name}"') + shared.log.debug(f'Upscaler type={self.name} folder="{folder}" model="{model_name}" path="{file_name}"') def find_scalers(self): scalers = [] @@ -78,7 +80,7 @@ class Upscaler: scaler = UpscalerData(name=f'{k} {model[0]}', path=model_path, upscaler=self) scalers.append(scaler) loaded.append(model_path) - # modules.shared.log.debug(f'Upscaler type={self.name} folder="{self.user_path}" model="{model[0]}" path="{model_path}"') + # shared.log.debug(f'Upscaler type={self.name} folder="{self.user_path}" model="{model[0]}" path="{model_path}"') if not os.path.exists(self.user_path): return scalers self.find_folder(self.user_path, scalers, loaded) @@ -89,8 +91,8 @@ class Upscaler: return img def upscale(self, img: Image, scale, selected_model: str = None): - orig_state = copy.deepcopy(modules.shared.state) - modules.shared.state.begin('upscale') + orig_state = copy.deepcopy(shared.state) + shared.state.begin('upscale') self.scale = scale dest_w = int(img.width * scale) dest_h = int(img.height * scale) @@ -103,8 +105,8 @@ class Upscaler: break if img.width != dest_w or img.height != dest_h: img = img.resize((int(dest_w), int(dest_h)), resample=LANCZOS) - modules.shared.state.end() - modules.shared.state = orig_state + shared.state.end() + shared.state = orig_state return img @abstractmethod @@ -115,7 +117,7 @@ class Upscaler: return modelloader.load_models(model_path=self.model_path, model_url=self.model_url, command_path=self.user_path) def update_status(self, prompt): - modules.shared.log.info(f'Upscaler: type={self.name} model="{prompt}"') + shared.log.info(f'Upscaler: type={self.name} model="{prompt}"') def find_model(self, path): info = None @@ -124,13 +126,13 @@ class Upscaler: info = scaler break if info is None: - modules.shared.log.error(f'Upscaler cannot match model: type={self.name} model="{path}"') + shared.log.error(f'Upscaler cannot match model: type={self.name} model="{path}"') return None if info.local_data_path.startswith("http"): from modules.modelloader import load_file_from_url info.local_data_path = load_file_from_url(url=info.data_path, model_dir=self.model_download_path, progress=True) if not os.path.isfile(info.local_data_path): - modules.shared.log.error(f'Upscaler cannot find model: type={self.name} model="{info.local_data_path}"') + shared.log.error(f'Upscaler cannot find model: type={self.name} model="{info.local_data_path}"') return None return info @@ -196,35 +198,53 @@ class UpscalerNearest(Upscaler): self.name = "Nearest" self.scalers = [UpscalerData("Nearest", None, self)] -def compile_upscaler(model, name=""): +def compile_upscaler(model): try: - if modules.shared.opts.ipex_optimize and "Upscaler" in modules.shared.opts.ipex_optimize: + if shared.opts.ipex_optimize and "Upscaler" in shared.opts.ipex_optimize: + t0 = time.time() import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import - from modules.devices import dtype as devices_dtype + model.eval() model.training = False - model = ipex.optimize(model, dtype=devices_dtype, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init - modules.shared.log.info("Applied Upscaler IPEX Optimize.") - except Exception as err: - modules.shared.log.warning(f"Upscaler IPEX Optimize not supported: {err}") - try: - if "Upscaler" in modules.shared.opts.cuda_compile and modules.shared.opts.cuda_compile_backend != 'none': - modules.shared.log.info(f"Upscaler Compiling: {name} mode={modules.shared.opts.cuda_compile_backend}") - import logging - import torch._dynamo # pylint: disable=unused-import,redefined-outer-name + model = ipex.optimize(model, dtype=devices.dtype, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init + t1 = time.time() + shared.log.info(f"Upscaler IPEX Optimize: time={t1-t0:.2f}") + except Exception as e: + shared.log.warning(f"Upscaler IPEX Optimize: error: {e}") - if modules.shared.opts.cuda_compile_backend == "openvino_fx": + try: + if "Upscaler" in shared.opts.cuda_compile and shared.opts.cuda_compile_backend != 'none': + import torch._dynamo # pylint: disable=unused-import,redefined-outer-name + torch._dynamo.reset() # pylint: disable=protected-access + shared.log.debug(f"Upscaler compile available backends: {torch._dynamo.list_backends()}") # pylint: disable=protected-access + + if shared.opts.cuda_compile_backend == "openvino_fx": from modules.intel.openvino import openvino_fx # pylint: disable=unused-import torch._dynamo.eval_frame.check_if_dynamo_supported = lambda: True # pylint: disable=protected-access - log_level = logging.WARNING if modules.shared.opts.cuda_compile_verbose else logging.CRITICAL # pylint: disable=protected-access + log_level = logging.WARNING if shared.opts.cuda_compile_verbose else logging.CRITICAL # pylint: disable=protected-access if hasattr(torch, '_logging'): torch._logging.set_logs(dynamo=log_level, aot=log_level, inductor=log_level) # pylint: disable=protected-access + torch._dynamo.config.verbose = shared.opts.cuda_compile_verbose # pylint: disable=protected-access + torch._dynamo.config.suppress_errors = shared.opts.cuda_compile_errors # pylint: disable=protected-access - torch._dynamo.config.verbose = modules.shared.opts.cuda_compile_verbose # pylint: disable=protected-access - torch._dynamo.config.suppress_errors = modules.shared.opts.cuda_compile_errors # pylint: disable=protected-access - model = torch.compile(model, mode=modules.shared.opts.cuda_compile_mode, backend=modules.shared.opts.cuda_compile_backend, fullgraph=modules.shared.opts.cuda_compile_fullgraph) # pylint: disable=attribute-defined-outside-init + try: + torch._inductor.config.conv_1x1_as_mm = True # pylint: disable=protected-access + torch._inductor.config.coordinate_descent_tuning = True # pylint: disable=protected-access + torch._inductor.config.epilogue_fusion = False # pylint: disable=protected-access + torch._inductor.config.coordinate_descent_check_all_directions = True # pylint: disable=protected-access + torch._inductor.config.use_mixed_mm = True # pylint: disable=protected-access + # torch._inductor.config.force_fuse_int_mm_with_mul = True # pylint: disable=protected-access + except Exception as e: + shared.log.error(f"Torch inductor config error: {e}") - modules.shared.log.info("Upscaler: Complilation done.") - except Exception as err: - modules.shared.log.warning(f"Model compile not supported: {err}") + t0 = time.time() + + model = torch.compile(model, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph) # pylint: disable=attribute-defined-outside-init + + setup_logging() # compile messes with logging so reset is needed + + t1 = time.time() + shared.log.info(f"Upscaler compile: time={t1-t0:.2f}") + except Exception as e: + shared.log.warning(f"Upscaler compile error: {e}") return model