Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-18 16:54:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add compile support for upscalers

Browse Source
This commit is contained in:
Disty0
2023-10-22 00:29:37 +03:00
parent 2f1f997421
commit 0c70e6e595
7 changed files with 61 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/ldsr/ldsr_model_arch.py
+2
View File
@@ -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}
modules/postprocess/esrgan_model.py
+3 -1
View File
@@ -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]
modules/postprocess/realesrgan_model_arch.py
+4 -2
View File
@@ -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.
modules/postprocess/scunet_model.py
+2 -1
View File
@@ -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
modules/postprocess/swinir_model.py
+2 -1
View File
@@ -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:
modules/shared.py
+4 -2
View File
@@ -255,8 +255,8 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
"math_sep": OptionInfo("<h2>Execution precision</h2>", "", 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("<h2>Model Compile</h2>", "", 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("<h2>IPEX, DirectML and OpenVINO</h2>", "", 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"),
modules/upscaler.py
+44
View File
@@ -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