From 0c70e6e595c2da92e7e567face3c44cc4ac27870 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 22 Oct 2023 00:29:37 +0300 Subject: [PATCH] Add compile support for upscalers --- modules/ldsr/ldsr_model_arch.py | 2 + modules/postprocess/esrgan_model.py | 4 +- modules/postprocess/realesrgan_model_arch.py | 6 ++- modules/postprocess/scunet_model.py | 3 +- modules/postprocess/swinir_model.py | 3 +- modules/shared.py | 6 ++- modules/upscaler.py | 44 ++++++++++++++++++++ 7 files changed, 61 insertions(+), 7 deletions(-) diff --git a/modules/ldsr/ldsr_model_arch.py b/modules/ldsr/ldsr_model_arch.py index 55a217589..9a4f1461f 100644 --- a/modules/ldsr/ldsr_model_arch.py +++ b/modules/ldsr/ldsr_model_arch.py @@ -10,6 +10,7 @@ import safetensors.torch from ldm.models.diffusion.ddim import DDIMSampler from ldm.util import instantiate_from_config, ismap from modules import devices, shared, sd_hijack +from modules.upscaler import compile_upscaler cached_ldsr_model: torch.nn.Module = None @@ -41,6 +42,7 @@ class LDSR: 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) cached_ldsr_model = model return {"model": model} diff --git a/modules/postprocess/esrgan_model.py b/modules/postprocess/esrgan_model.py index 8af099ea0..29a4cc8a2 100644 --- a/modules/postprocess/esrgan_model.py +++ b/modules/postprocess/esrgan_model.py @@ -4,7 +4,7 @@ from PIL import Image from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn import modules.postprocess.esrgan_model_arch as arch from modules import images, devices -from modules.upscaler import Upscaler, UpscalerData +from modules.upscaler import Upscaler, UpscalerData, compile_upscaler from modules.shared import opts, log, console @@ -154,6 +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) self.models[info.local_data_path] = model return self.models[info.local_data_path] @@ -168,6 +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) 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 fa290d5f9..06f8d9c1f 100644 --- a/modules/postprocess/realesrgan_model_arch.py +++ b/modules/postprocess/realesrgan_model_arch.py @@ -9,6 +9,7 @@ from torch import nn from torch.nn import functional as F from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn from modules.shared import log, console +from modules.upscaler import compile_upscaler ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) @@ -75,9 +76,10 @@ class RealESRGANer(): model.load_state_dict(loadnet[keyname], strict=True) model.eval() - self.model = model.to(self.device) if self.half: - self.model = self.model.half() + model = model.half() + model = compile_upscaler(model, name=self.name) + self.model = model.to(self.device) 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 83b492ef4..e48950afe 100644 --- a/modules/postprocess/scunet_model.py +++ b/modules/postprocess/scunet_model.py @@ -5,7 +5,7 @@ from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, T from modules import devices from modules.postprocess.scunet_model_arch import SCUNet as net from modules.shared import opts, log, console, device -from modules.upscaler import Upscaler +from modules.upscaler import Upscaler, compile_upscaler class UpscalerSCUNet(Upscaler): @@ -31,6 +31,7 @@ class UpscalerSCUNet(Upscaler): for _, v in model.named_parameters(): v.requires_grad = False model = model.to(device) + model = compile_upscaler(model, name=self.name) self.models[info.local_data_path] = model return model diff --git a/modules/postprocess/swinir_model.py b/modules/postprocess/swinir_model.py index fd0240ce4..3692cd836 100644 --- a/modules/postprocess/swinir_model.py +++ b/modules/postprocess/swinir_model.py @@ -5,7 +5,7 @@ from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, T from modules.postprocess.swinir_model_arch import SwinIR as net from modules.postprocess.swinir_model_arch_v2 import Swin2SR as net2 from modules import devices, script_callbacks, shared -from modules.upscaler import Upscaler +from modules.upscaler import Upscaler, compile_upscaler class UpscalerSwinIR(Upscaler): @@ -58,6 +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) self.models[info.local_data_path] = model return model except Exception as e: diff --git a/modules/shared.py b/modules/shared.py index 394c2c7c5..808eb7ae9 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -255,8 +255,8 @@ options_templates.update(options_section(('cuda', "Compute Settings"), { "math_sep": OptionInfo("

Execution precision

", "", gr.HTML), "precision": OptionInfo("Autocast", "Precision type", gr.Radio, {"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, {"choices": ["FP32", "FP16", "BF16"]}), - "no_half": OptionInfo(False, "Use full precision for model (--no-half)", None, None, None), - "no_half_vae": OptionInfo(False, "Use full precision for VAE (--no-half-vae)"), + "no_half": OptionInfo(True if cmd_opts.use_openvino else False, "Use full precision for model (--no-half)", None, None, None), + "no_half_vae": OptionInfo(True if cmd_opts.use_openvino else False, "Use full precision for VAE (--no-half-vae)"), "upcast_sampling": OptionInfo(True if sys.platform == "darwin" else False, "Enable upcast sampling"), "upcast_attn": OptionInfo(False, "Enable upcast cross attention layer"), "cuda_cast_unet": OptionInfo(False, "Use fixed UNet precision"), @@ -278,6 +278,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), { "cuda_compile_sep": OptionInfo("

Model Compile

", "", gr.HTML), "cuda_compile": OptionInfo(True if cmd_opts.use_openvino else False, "Enable model compile"), + "cuda_compile_upscaler": OptionInfo(False, "Enable upscaler compile"), "cuda_compile_backend": OptionInfo("openvino_fx" if cmd_opts.use_openvino else "none", "Model compile backend", gr.Radio, {"choices": ['none', 'inductor', 'cudagraphs', 'aot_ts_nvfuser', 'hidet', 'ipex', 'openvino_fx']}), "cuda_compile_mode": OptionInfo("default", "Model compile mode", gr.Radio, {"choices": ['default', 'reduce-overhead', 'max-autotune']}), "cuda_compile_fullgraph": OptionInfo(False, "Model compile fullgraph"), @@ -287,6 +288,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), { "ipex_sep": OptionInfo("

IPEX, DirectML and OpenVINO

", "", gr.HTML), "ipex_optimize": OptionInfo(True if devices.backend == "ipex" else False, "Enable IPEX Optimize for Intel GPUs"), + "ipex_optimize_upscaler": OptionInfo(True if devices.backend == "ipex" else False, "Enable IPEX Optimize for Intel GPUs with Upscalers"), "directml_memory_provider": OptionInfo(default_memory_provider, 'DirectML memory stats provider', gr.Radio, {"choices": memory_providers}), "openvino_disable_model_caching": OptionInfo(False, "OpenVINO disable model caching"), "openvino_hetero_gpu": OptionInfo(False, "OpenVINO use Hetero Device for single inference with multiple devices"), diff --git a/modules/upscaler.py b/modules/upscaler.py index b05862f78..8d82688f2 100644 --- a/modules/upscaler.py +++ b/modules/upscaler.py @@ -190,3 +190,47 @@ class UpscalerNearest(Upscaler): super().__init__(False) self.name = "Nearest" self.scalers = [UpscalerData("Nearest", None, self)] + +def compile_upscaler(model, name=""): + try: + if modules.shared.opts.ipex_optimize_upscaler: + import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import + from modules.devices import dtype as devices_dtype + 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 modules.shared.opts.cuda_compile_upscaler 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 + use_old_compiled_model_state = False + + if modules.shared.opts.cuda_compile_backend == "openvino_fx": + from modules.intel.openvino import openvino_fx, openvino_clear_caches # pylint: disable=unused-import + from modules.sd_models import CompiledModelState + + openvino_clear_caches() + torch._dynamo.eval_frame.check_if_dynamo_supported = lambda: True # pylint: disable=protected-access + + if modules.shared.compiled_model_state is not None: + use_old_compiled_model_state = True + old_compiled_model_state = modules.shared.compiled_model_state + modules.shared.compiled_model_state = CompiledModelState() + + log_level = logging.WARNING if modules.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 = 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 + + if use_old_compiled_model_state: + modules.shared.compiled_model_state = old_compiled_model_state + modules.shared.log.info("Upscaler: Complilation done.") + except Exception as err: + modules.shared.log.warning(f"Model compile not supported: {err}") + return model