Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 01:04:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

Cleanup upscalers

Browse Source
This commit is contained in:
Disty0
2024-01-27 02:54:39 +03:00
parent 1440e07f02
commit dbe4d2ff70
8 changed files with 85 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/intel/ipex/hijacks.py
+2
View File
@@ -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
modules/ldsr/ldsr_model_arch.py
+3 -3
View File
@@ -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
modules/postprocess/esrgan_model.py
+2 -2
View File
@@ -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]
modules/postprocess/realesrgan_model_arch.py
+3 -2
View File
@@ -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.
modules/postprocess/scunet_model.py
+6 -6
View File
@@ -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
modules/postprocess/swinir_model.py
+5 -5
View File
@@ -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))
modules/sd_models_compile.py
+6
View File
@@ -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:
modules/upscaler.py
+58 -38
View File
@@ -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