diff --git a/CHANGELOG.md b/CHANGELOG.md index 06854d872..281fc7a01 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,8 @@ ## Update for 2025-01-17 +- **Upscale** + - code refactor to unify latent, resize and model based upscalers - **Gallery**: - add http fallback for slow/unreliable links - **Fixes**: diff --git a/modules/hypernetworks/hypernetwork.py b/modules/hypernetworks/hypernetwork.py index e15cf7724..fef95ce5d 100644 --- a/modules/hypernetworks/hypernetwork.py +++ b/modules/hypernetworks/hypernetwork.py @@ -9,6 +9,8 @@ from einops import rearrange, repeat from modules import devices, shared, hashes, errors, files_cache +loaded_hypernetworks = [] + class HypernetworkModule(torch.nn.Module): activation_dict = { "linear": torch.nn.Identity, @@ -280,10 +282,10 @@ def load_hypernetwork(name): def load_hypernetworks(names, multipliers=None): already_loaded = {} - for hypernetwork in shared.loaded_hypernetworks: - if hypernetwork.name in names: - already_loaded[hypernetwork.name] = hypernetwork - shared.loaded_hypernetworks.clear() + for hn in loaded_hypernetworks: + if hn.name in names: + already_loaded[hn.name] = hn + loaded_hypernetworks.clear() for i, name in enumerate(names): hypernetwork = already_loaded.get(name, None) if hypernetwork is None: @@ -291,7 +293,7 @@ def load_hypernetworks(names, multipliers=None): if hypernetwork is None: continue hypernetwork.set_multiplier(multipliers[i] if multipliers else 1.0) - shared.loaded_hypernetworks.append(hypernetwork) + loaded_hypernetworks.append(hypernetwork) def find_closest_hypernetwork_name(search: str): @@ -330,7 +332,7 @@ def attention_CrossAttention_forward(self, x, context=None, mask=None): h = self.heads q = self.to_q(x) context = default(context, x) - context_k, context_v = apply_hypernetworks(shared.loaded_hypernetworks, context, self) + context_k, context_v = apply_hypernetworks(loaded_hypernetworks, context, self) k = self.to_k(context_k) v = self.to_v(context_v) q, k, v = (rearrange(t, 'b n (h d) -> (b h) n d', h=h) for t in (q, k, v)) diff --git a/modules/images_resize.py b/modules/images_resize.py index 183e1d7f1..a549b5bf9 100644 --- a/modules/images_resize.py +++ b/modules/images_resize.py @@ -1,45 +1,50 @@ +from typing import Union import sys import time import numpy as np +import torch from PIL import Image -from modules import shared +from modules import shared, upscaler -def resize_image(resize_mode: int, im: Image.Image, width: int, height: int, upscaler_name: str=None, output_type: str='image', context: str=None): +def resize_image(resize_mode: int, im: Union[Image.Image, torch.Tensor], width: int, height: int, upscaler_name: str=None, output_type: str='image', context: str=None): upscaler_name = upscaler_name or shared.opts.upscaler_for_img2img - def latent(im, w, h, upscaler): - from modules.processing_vae import vae_encode, vae_decode - import torch - latents = vae_encode(im, shared.sd_model, full_quality=False) # TODO resize image: enable full VAE mode for resize-latent - latents = torch.nn.functional.interpolate(latents, size=(int(h // 8), int(w // 8)), mode=upscaler["mode"], antialias=upscaler["antialias"]) - im = vae_decode(latents, shared.sd_model, output_type='pil', full_quality=False)[0] - return im + def latent(im, scale: float, selected_upscaler: upscaler.UpscalerData): + if isinstance(im, torch.Tensor): + im = selected_upscaler.scaler.upscale(im, scale, selected_upscaler.name) + return im + else: + from modules.processing_vae import vae_encode, vae_decode + latents = vae_encode(im, shared.sd_model, full_quality=False) # TODO resize image: enable full VAE mode for resize-latent + latents = selected_upscaler.scaler.upscale(latents, scale, selected_upscaler.name) + im = vae_decode(latents, shared.sd_model, output_type='pil', full_quality=False)[0] + return im - def resize(im, w, h): - w = int(w) - h = int(h) - if upscaler_name is None or upscaler_name == "None" or im.mode == 'L': + def resize(im: Union[Image.Image, torch.Tensor], w, h): + w, h = int(w), int(h) + if upscaler_name is None or upscaler_name == "None" or (hasattr(im, 'mode') and im.mode == 'L'): return im.resize((w, h), resample=Image.Resampling.LANCZOS) # force for mask - scale = max(w / im.width, h / im.height) + if isinstance(im, torch.Tensor): + scale = max(w // 8 / im.shape[-1] , h // 8 / im.shape[-2]) + else: + scale = max(w / im.width, h / im.height) if scale > 1.0: upscalers = [x for x in shared.sd_upscalers if x.name.lower().replace('-', ' ') == upscaler_name.lower().replace('-', ' ')] if len(upscalers) > 0: - upscaler = upscalers[0] - im = upscaler.scaler.upscale(im, scale, upscaler.data_path) - else: - upscaler = shared.latent_upscale_modes.get(upscaler_name, None) - if upscaler is not None: - im = latent(im, w, h, upscaler) + selected_upscaler: upscaler.UpscalerData = upscalers[0] + if selected_upscaler.name.lower().startswith('latent'): + im = latent(im, scale, selected_upscaler) else: - upscaler = shared.sd_upscalers[0] - shared.log.warning(f"Resize upscaler: invalid={upscaler_name} fallback={upscaler.name}") - shared.log.debug(f"Resize upscaler: available={[u.name for u in shared.sd_upscalers]}") - if im.width != w or im.height != h: # probably downsample after upscaler created larger image + im = selected_upscaler.scaler.upscale(im, scale, selected_upscaler.name) + else: + shared.log.warning(f"Resize upscaler: invalid={upscaler_name} fallback={selected_upscaler.name}") + shared.log.debug(f"Resize upscaler: available={[u.name for u in shared.sd_upscalers]}") + if isinstance(im, Image.Image) and (im.width != w or im.height != h): # probably downsample after upscaler created larger image im = im.resize((w, h), resample=Image.Resampling.LANCZOS) return im - def crop(im): + def crop(im: Image.Image): ratio = width / height src_ratio = im.width / im.height src_w = width if ratio > src_ratio else im.width * height // im.height @@ -49,7 +54,7 @@ def resize_image(resize_mode: int, im: Image.Image, width: int, height: int, ups res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2)) return res - def fill(im, color=None): + def fill(im: Image.Image, color=None): color = color or shared.opts.image_background """ ratio = round(width / height, 1) @@ -77,7 +82,8 @@ def resize_image(resize_mode: int, im: Image.Image, width: int, height: int, ups res.paste(im, box=((width - im.width)//2, (height - im.height)//2)) return res - def context_aware(im, width, height, context): + def context_aware(im: Image.Image, width, height, context): + width, height = int(width), int(height) import seam_carving # https://github.com/li-plus/seam-carving if 'forward' in context.lower(): energy_mode = "forward" @@ -110,7 +116,10 @@ def resize_image(resize_mode: int, im: Image.Image, width: int, height: int, ups t0 = time.time() if resize_mode is None: resize_mode = 0 - if resize_mode == 0 or (im.width == width and im.height == height) or (width == 0 and height == 0): # none + if isinstance(im, torch.Tensor): # latent resize only supports fixed mode + res = resize(im, width, height) + return res + elif (resize_mode == 0) or (im.width == width and im.height == height) or (width == 0 and height == 0): # none res = im.copy() elif resize_mode == 1: # fixed res = resize(im, width, height) diff --git a/modules/modelloader.py b/modules/modelloader.py index b022b4fc6..0c5fa78d8 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -11,7 +11,7 @@ from PIL import Image import rich.progress as p import huggingface_hub as hf from modules import shared, errors, files_cache -from modules.upscaler import Upscaler, UpscalerLanczos, UpscalerNearest, UpscalerNone +from modules.upscaler import Upscaler from modules.paths import script_path, models_path @@ -575,7 +575,7 @@ def load_upscalers(): importlib.import_module(full_model) except Exception as e: shared.log.error(f'Error loading upscaler: {model_name} {e}') - datas = [] + upscalers = [] commandline_options = vars(shared.cmd_opts) # some of upscaler classes will not go away after reloading their modules, and we'll end up with two copies of those classes. The newest copy will always be the last in the list, so we go from end to beginning and ignore duplicates used_classes = {} @@ -583,7 +583,7 @@ def load_upscalers(): classname = str(cls) if classname not in used_classes: used_classes[classname] = cls - names = [] + upscaler_types = [] for cls in reversed(used_classes.values()): name = cls.__name__ cmd_name = f"{name.lower().replace('upscaler', '')}_models_path" @@ -591,9 +591,9 @@ def load_upscalers(): scaler = cls(commandline_model_path) scaler.user_path = commandline_model_path scaler.model_download_path = commandline_model_path or scaler.model_path - datas += scaler.scalers - names.append(name[8:]) - shared.sd_upscalers = sorted(datas, key=lambda x: x.name.lower() if not isinstance(x.scaler, (UpscalerNone, UpscalerLanczos, UpscalerNearest)) else "") # Special case for UpscalerNone keeps it at the beginning of the list. + upscalers += scaler.scalers + upscaler_types.append(name[8:]) + shared.sd_upscalers = upscalers t1 = time.time() - shared.log.info(f"Available Upscalers: items={len(shared.sd_upscalers)} downloaded={len([x for x in shared.sd_upscalers if x.data_path is not None and os.path.isfile(x.data_path)])} user={len([x for x in shared.sd_upscalers if x.custom])} time={t1-t0:.2f} types={names}") + shared.log.info(f"Available Upscalers: items={len(shared.sd_upscalers)} downloaded={len([x for x in shared.sd_upscalers if x.data_path is not None and os.path.isfile(x.data_path)])} user={len([x for x in shared.sd_upscalers if x.custom])} time={t1-t0:.2f} types={upscaler_types}") return [x.name for x in shared.sd_upscalers] diff --git a/modules/postprocess/sdupscaler_model.py b/modules/postprocess/sdupscaler_model.py index e5b1c8b45..5ec7168d3 100644 --- a/modules/postprocess/sdupscaler_model.py +++ b/modules/postprocess/sdupscaler_model.py @@ -4,7 +4,8 @@ from PIL import Image from modules import shared, devices from modules.upscaler import Upscaler, UpscalerData -class UpscalerSD(Upscaler): + +class UpscalerDiffusion(Upscaler): def __init__(self, dirname): # pylint: disable=super-init-not-called self.name = "SDUpscale" self.user_path = dirname @@ -12,8 +13,8 @@ class UpscalerSD(Upscaler): super().__init__() return self.scalers = [ - UpscalerData(name="SD Latent 2x", path="stabilityai/sd-x2-latent-upscaler", upscaler=self, model=None, scale=4), - UpscalerData(name="SD Latent 4x", path="stabilityai/stable-diffusion-x4-upscaler", upscaler=self, model=None, scale=4), + UpscalerData(name="Diffusion Latent Upscaler 2x", path="stabilityai/sd-x2-latent-upscaler", upscaler=self, model=None, scale=4), + UpscalerData(name="Diffusion Latent Upscaler 4x", path="stabilityai/stable-diffusion-x4-upscaler", upscaler=self, model=None, scale=4), ] self.pipelines = [ None, diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index c20ba85a8..240ac8670 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -178,9 +178,8 @@ def process_hires(p: processing.StableDiffusionProcessing, output): output.images = resize_hires(p, latents=output.images) if output is not None else [] sd_hijack_hypertile.hypertile_set(p, hr=True) - latent_upscale = shared.latent_upscale_modes.get(p.hr_upscaler, None) strength = p.hr_denoising_strength if p.hr_denoising_strength > 0 else p.denoising_strength - if (latent_upscale is not None or p.hr_force) and strength > 0: + if (p.hr_upscaler.lower().startswith('latent') or p.hr_force) and strength > 0: p.ops.append('hires') sd_models_compile.openvino_recompile_model(p, hires=True, refiner=False) if shared.sd_model.__class__.__name__ == "OnnxRawPipeline": diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py index 93896a02f..083d3489a 100644 --- a/modules/processing_helpers.py +++ b/modules/processing_helpers.py @@ -409,20 +409,19 @@ def resize_hires(p, latents): # input=latents output=pil if not latent_upscaler shared.log.warning('Hires: input is not tensor') first_pass_images = processing_vae.vae_decode(latents=latents, model=shared.sd_model, full_quality=p.full_quality, output_type='pil', width=p.width, height=p.height) return first_pass_images - latent_upscaler = shared.latent_upscale_modes.get(p.hr_upscaler, None) - # shared.log.info(f'Hires: upscaler={p.hr_upscaler} width={p.hr_upscale_to_x} height={p.hr_upscale_to_y} images={latents.shape[0]}') - if latent_upscaler is not None: - return torch.nn.functional.interpolate(latents, size=(p.hr_upscale_to_y // 8, p.hr_upscale_to_x // 8), mode=latent_upscaler["mode"], antialias=latent_upscaler["antialias"]) - first_pass_images = processing_vae.vae_decode(latents=latents, model=shared.sd_model, full_quality=p.full_quality, output_type='pil', width=p.width, height=p.height) - if p.hr_upscale_to_x == 0 or (p.hr_upscale_to_y == 0 and hasattr(p, 'init_hr')): + + if (p.hr_upscale_to_x == 0 or p.hr_upscale_to_y == 0) and hasattr(p, 'init_hr'): shared.log.error('Hires: missing upscaling dimensions') return first_pass_images + + if p.hr_upscaler.lower().startswith('latent'): + resized_image = images.resize_image(p.hr_resize_mode, latents, p.hr_upscale_to_x, p.hr_upscale_to_y, upscaler_name=p.hr_upscaler, context=p.hr_resize_context) + return resized_image + + first_pass_images = processing_vae.vae_decode(latents=latents, model=shared.sd_model, full_quality=p.full_quality, output_type='pil', width=p.width, height=p.height) resized_images = [] for img in first_pass_images: - if latent_upscaler is None: - resized_image = images.resize_image(p.hr_resize_mode, img, p.hr_upscale_to_x, p.hr_upscale_to_y, upscaler_name=p.hr_upscaler, context=p.hr_resize_context) - else: - resized_image = img + resized_image = images.resize_image(p.hr_resize_mode, img, p.hr_upscale_to_x, p.hr_upscale_to_y, upscaler_name=p.hr_upscaler, context=p.hr_resize_context) resized_images.append(resized_image) devices.torch_gc() return resized_images diff --git a/modules/processing_original.py b/modules/processing_original.py index 7a0af1b04..261fc8c13 100644 --- a/modules/processing_original.py +++ b/modules/processing_original.py @@ -72,14 +72,6 @@ def process_original(p: processing.StableDiffusionProcessing): def sample_txt2img(p: processing.StableDiffusionProcessingTxt2Img, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts): - latent_scale_mode = shared.latent_upscale_modes.get(p.hr_upscaler, None) if p.hr_upscaler is not None else shared.latent_upscale_modes.get(shared.latent_upscale_default_mode, "None") - if latent_scale_mode is not None: - p.hr_force = False # no need to force anything - if p.enable_hr and (latent_scale_mode is None or p.hr_force): - if len([x for x in shared.sd_upscalers if x.name == p.hr_upscaler]) == 0: - shared.log.warning(f"HiRes: upscaler={p.hr_upscaler} unknown") - p.enable_hr = False - p.ops.append('txt2img') hypertile_set(p) p.sampler = sd_samplers.create_sampler(p.sampler_name, p.sd_model) @@ -109,7 +101,16 @@ def sample_txt2img(p: processing.StableDiffusionProcessingTxt2Img, conditioning, info = processing.create_infotext(p, p.all_prompts, p.all_seeds, p.all_subseeds, [], iteration=p.iteration, position_in_batch=i) p.extra_generation_params, p.detailer_enabled = orig_extra_generation_params, orig_detailer images.save_image(image, p.outpath_samples, "", seeds[i], prompts[i], shared.opts.samples_format, info=info, suffix="-before-hires") - if latent_scale_mode is None or p.hr_force: # non-latent upscaling + + if p.hr_upscaler.lower().startswith('latent'): # non-latent upscaling + p.hr_force = True + shared.state.job = 'Upscale' + samples = images.resize_image(1, samples, target_width, target_height, upscaler_name=p.hr_upscaler) + if getattr(p, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) < 1.0: + image_conditioning = img2img_image_conditioning(p, decode_first_stage(p.sd_model, samples.to(dtype=devices.dtype_vae), p.full_quality), samples) + else: + image_conditioning = txt2img_image_conditioning(p, samples.to(dtype=devices.dtype_vae)) + else: shared.state.job = 'Upscale' if decoded_samples is None: decoded_samples = decode_first_stage(p.sd_model, samples.to(dtype=devices.dtype_vae), p.full_quality) @@ -130,15 +131,8 @@ def sample_txt2img(p: processing.StableDiffusionProcessingTxt2Img, conditioning, else: samples = p.sd_model.get_first_stage_encoding(p.sd_model.encode_first_stage(resized_samples)) image_conditioning = img2img_image_conditioning(p, resized_samples, samples) - else: - samples = torch.nn.functional.interpolate(samples, size=(target_height // 8, target_width // 8), mode=latent_scale_mode["mode"], antialias=latent_scale_mode["antialias"]) - if getattr(p, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) < 1.0: - image_conditioning = img2img_image_conditioning(p, decode_first_stage(p.sd_model, samples.to(dtype=devices.dtype_vae), p.full_quality), samples) - else: - image_conditioning = txt2img_image_conditioning(p, samples.to(dtype=devices.dtype_vae)) - if p.hr_sampler_name == "PLMS": - p.hr_sampler_name = 'UniPC' - if p.hr_force or latent_scale_mode is not None: + + if p.hr_force: shared.state.job = 'HiRes' if p.denoising_strength > 0: p.ops.append('hires') diff --git a/modules/processing_vae.py b/modules/processing_vae.py index 33319e174..46bf0ad26 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -162,6 +162,7 @@ def full_vae_decode(latents, model): def full_vae_encode(image, model): + t0 = time.time() if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'): log_debug('Moving to CPU: model=UNet') unet_device = model.unet.device @@ -170,9 +171,25 @@ def full_vae_encode(image, model): sd_models.move_model(model.vae, devices.device) vae_name = sd_vae.loaded_vae_file if sd_vae.loaded_vae_file is not None else "default" log_debug(f'Encode vae="{vae_name}" dtype={model.vae.dtype} upcast={model.vae.config.get("force_upcast", None)}') + + upcast = (model.vae.dtype == torch.float16) and (getattr(model.vae.config, 'force_upcast', False) or shared.opts.no_half_vae) + if upcast: + if hasattr(model, 'upcast_vae'): # this is done by diffusers automatically if output_type != 'latent' + model.upcast_vae() + else: # manual upcast and we restore it later + model.vae.orig_dtype = model.vae.dtype + model.vae = model.vae.to(dtype=torch.float32) + encoded = model.vae.encode(image.to(model.vae.device, model.vae.dtype)).latent_dist.sample() + + if hasattr(model.vae, "orig_dtype"): + model.vae = model.vae.to(dtype=model.vae.orig_dtype) + del model.vae.orig_dtype + if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'): sd_models.move_model(model.unet, unet_device) + t1 = time.time() + shared.log.debug(f'Encode: vae="{vae_name}" upcast={upcast} slicing={getattr(model.vae, "use_slicing", None)} tiling={getattr(model.vae, "use_tiling", None)} latents={encoded.shape}:{encoded.device}:{encoded.dtype} time={t1-t0:.3f}') return encoded diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py index 31089551e..33ff274ce 100644 --- a/modules/sd_hijack_optimizations.py +++ b/modules/sd_hijack_optimizations.py @@ -52,7 +52,7 @@ def split_cross_attention_forward_v1(self, x, context=None, mask=None): # pylint q_in = self.to_q(x) context = default(context, x) # pylint: disable=possibly-used-before-assignment - context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) + context_k, context_v = hypernetwork.apply_hypernetworks(hypernetwork.loaded_hypernetworks, context) k_in = self.to_k(context_k) v_in = self.to_v(context_v) del context, context_k, context_v, x @@ -90,7 +90,7 @@ def split_cross_attention_forward(self, x, context=None, mask=None): # pylint: d q_in = self.to_q(x) context = default(context, x) - context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) + context_k, context_v = hypernetwork.apply_hypernetworks(hypernetwork.loaded_hypernetworks, context) k_in = self.to_k(context_k) v_in = self.to_v(context_v) @@ -219,7 +219,7 @@ def split_cross_attention_forward_invokeAI(self, x, context=None, mask=None): # q = self.to_q(x) context = default(context, x) - context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) + context_k, context_v = hypernetwork.apply_hypernetworks(hypernetwork.loaded_hypernetworks, context) k = self.to_k(context_k) v = self.to_v(context_v) del context, context_k, context_v, x @@ -248,7 +248,7 @@ def sub_quad_attention_forward(self, x, context=None, mask=None): q = self.to_q(x) context = default(context, x) - context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) + context_k, context_v = hypernetwork.apply_hypernetworks(hypernetwork.loaded_hypernetworks, context) k = self.to_k(context_k) v = self.to_v(context_v) del context, context_k, context_v, x @@ -329,7 +329,7 @@ def xformers_attention_forward(self, x, context=None, mask=None): # pylint: disa q_in = self.to_q(x) context = default(context, x) - context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) + context_k, context_v = hypernetwork.apply_hypernetworks(hypernetwork.loaded_hypernetworks, context) k_in = self.to_k(context_k) v_in = self.to_v(context_v) @@ -360,7 +360,7 @@ def scaled_dot_product_attention_forward(self, x, context=None, mask=None): q_in = self.to_q(x) context = default(context, x) - context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) + context_k, context_v = hypernetwork.apply_hypernetworks(hypernetwork.loaded_hypernetworks, context) k_in = self.to_k(context_k) v_in = self.to_v(context_v) diff --git a/modules/shared.py b/modules/shared.py index 8c625bc24..720ca7c55 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -15,7 +15,6 @@ import gradio as gr import fasteners import orjson import diffusers -from rich.console import Console from modules import errors, devices, shared_items, shared_state, cmd_args, theme, history, files_cache from modules.paths import models_path, script_path, data_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, extensions_dir, extensions_builtin_dir # pylint: disable=W0611 from modules.dml import memory_providers, default_memory_provider, directml_do_hijack @@ -26,22 +25,19 @@ import modules.interrogate import modules.memmon import modules.styles import modules.paths as paths -from installer import print_dict -from installer import log as central_logger # pylint: disable=E0611 +from installer import log, print_dict, console # pylint: disable=unused-import errors.install([gr]) demo: gr.Blocks = None api = None -log = central_logger progress_print_out = sys.stdout parser = cmd_args.parser url = 'https://github.com/vladmandic/automatic' cmd_opts, _ = parser.parse_known_args() hide_dirs = {"visible": not cmd_opts.hide_ui_dir_config} xformers_available = False -locking_available = True -clip_model = None +locking_available = True # used by file read/write locking interrogator = modules.interrogate.InterrogateModels(os.path.join("models", "interrogate")) sd_upscalers = [] detailers = [] @@ -51,20 +47,7 @@ tab_names = [] extra_networks = [] options_templates = {} hypernetworks = {} -loaded_hypernetworks = [] settings_components = None -latent_upscale_default_mode = "None" -latent_upscale_modes = { - "Latent Nearest": {"mode": "nearest", "antialias": False}, - "Latent Nearest-exact": {"mode": "nearest-exact", "antialias": False}, - "Latent Area": {"mode": "area", "antialias": False}, - "Latent Bilinear": {"mode": "bilinear", "antialias": False}, - "Latent Bicubic": {"mode": "bicubic", "antialias": False}, - "Latent Bilinear antialias": {"mode": "bilinear", "antialias": True}, - "Latent Bicubic antialias": {"mode": "bicubic", "antialias": True}, - # "Latent Linear": {"mode": "linear", "antialias": False}, # not supported for latents with channels=4 - # "Latent Trilinear": {"mode": "trilinear", "antialias": False}, # not supported for latents with channels=4 -} restricted_opts = { "samples_filename_pattern", "directories_filename_pattern", @@ -80,7 +63,6 @@ restricted_opts = { } resize_modes = ["None", "Fixed", "Crop", "Fill", "Outpaint", "Context aware"] compatibility_opts = ['clip_skip', 'uni_pc_lower_order_final', 'uni_pc_order'] -console = Console(log_time=True, log_time_format='%H:%M:%S-%f') dir_timestamps = {} dir_cache = {} max_workers = 8 diff --git a/modules/ui_sections.py b/modules/ui_sections.py index def7e39b5..17436dd31 100644 --- a/modules/ui_sections.py +++ b/modules/ui_sections.py @@ -323,14 +323,6 @@ def create_hires_inputs(tab): with gr.Group(): with gr.Row(elem_id=f"{tab}_hires_row1"): enable_hr = gr.Checkbox(label='Enable refine pass', value=False, elem_id=f"{tab}_enable_hr") - """ - with gr.Row(elem_id=f"{tab}_hires_fix_row1", variant="compact"): - hr_upscaler = gr.Dropdown(label="Upscaler", elem_id=f"{tab}_hr_upscaler", choices=[*shared.latent_upscale_modes, *[x.name for x in shared.sd_upscalers]], value=shared.latent_upscale_default_mode) - hr_scale = gr.Slider(minimum=0.1, maximum=8.0, step=0.05, label="Rescale by", value=2.0, elem_id=f"{tab}_hr_scale") - with gr.Row(elem_id=f"{tab}_hires_fix_row3", variant="compact"): - hr_resize_x = gr.Slider(minimum=0, maximum=4096, step=8, label="Width resize", value=0, elem_id=f"{tab}_hr_resize_x") - hr_resize_y = gr.Slider(minimum=0, maximum=4096, step=8, label="Height resize", value=0, elem_id=f"{tab}_hr_resize_y") - """ hr_resize_mode, hr_upscaler, hr_resize_context, hr_resize_x, hr_resize_y, hr_scale, _selected_scale_tab = create_resize_inputs(tab, None, accordion=False, latent=True, non_zero=False) with gr.Row(elem_id=f"{tab}_hires_fix_row2", variant="compact"): hr_force = gr.Checkbox(label='Force HiRes', value=False, elem_id=f"{tab}_hr_force") @@ -355,8 +347,11 @@ def create_resize_inputs(tab, images, accordion=True, latent=False, non_zero=Tru prefix = f' {prefix}' with gr.Accordion(open=False, label="Resize", elem_classes=["small-accordion"], elem_id=f"{tab}_resize_group") if accordion else gr.Group(): with gr.Row(): + available_upscalers = [x.name for x in shared.sd_upscalers] + if not latent: + available_upscalers = [x for x in available_upscalers if not x.lower().startswith('latent')] resize_mode = gr.Dropdown(label=f"Mode{prefix}" if non_zero else "Resize mode", elem_id=f"{tab}_resize_mode", choices=shared.resize_modes, type="index", value='Fixed') - resize_name = gr.Dropdown(label=f"Method{prefix}", elem_id=f"{tab}_resize_name", choices=([] if not latent else list(shared.latent_upscale_modes)) + [x.name for x in shared.sd_upscalers], value=shared.latent_upscale_default_mode, visible=True) + resize_name = gr.Dropdown(label=f"Method{prefix}", elem_id=f"{tab}_resize_name", choices=available_upscalers, value=available_upscalers[0], visible=True) resize_context_choices = ["Add with forward", "Remove with forward", "Add with backward", "Remove with backward"] resize_context = gr.Dropdown(label=f"Context{prefix}", elem_id=f"{tab}_resize_context", choices=resize_context_choices, value=resize_context_choices[0], visible=False) ui_common.create_refresh_button(resize_name, modelloader.load_upscalers, lambda: {"choices": modelloader.load_upscalers()}, 'refresh_upscalers') diff --git a/modules/upscaler.py b/modules/upscaler.py index bda39f858..80c0ddaf9 100644 --- a/modules/upscaler.py +++ b/modules/upscaler.py @@ -8,10 +8,9 @@ from modules import devices, modelloader, shared from installer import setup_logging -LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.Resampling.LANCZOS) -NEAREST = (Image.Resampling.NEAREST if hasattr(Image, 'Resampling') else Image.Resampling.NEAREST) models = None + class Upscaler: name = None folder = None @@ -97,17 +96,24 @@ class Upscaler: 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) - for _ in range(3): - shape = (img.width, img.height) + if isinstance(img, Image.Image): + dest_w = int(img.width * scale) + dest_h = int(img.height * scale) + else: + dest_w = int(img.shape[-1] * scale) + dest_h = int(img.shape[-2] * scale) + if self.name.lower().startswith('latent'): img = self.do_upscale(img, selected_model) - if shape == (img.width, img.height): - break - if img.width >= dest_w and img.height >= dest_h: - break - if img.width != dest_w or img.height != dest_h: - img = img.resize((int(dest_w), int(dest_h)), resample=LANCZOS) + else: + for _ in range(3): + shape = (img.width, img.height) + img = self.do_upscale(img, selected_model) + if shape == (img.width, img.height): + break + if img.width >= dest_w and img.height >= dest_h: + break + if img.width != dest_w or img.height != dest_h: + img = img.resize((int(dest_w), int(dest_h)), resample=Image.Resampling.BICUBIC) shared.state.end() shared.state = orig_state return img @@ -125,7 +131,7 @@ class Upscaler: def find_model(self, path): info = None for scaler in self.scalers: - if scaler.data_path == path: + if (scaler.data_path == path) or (scaler.name == path): info = scaler break if info is None: @@ -157,50 +163,6 @@ class UpscalerData: self.model = model -class UpscalerNone(Upscaler): - name = "None" - scalers = [] - - def load_model(self, path): - pass - - def do_upscale(self, img, selected_model=None): - return img - - def __init__(self, dirname=None): # pylint: disable=unused-argument - super().__init__(False) - self.scalers = [UpscalerData("None", None, self)] - - -class UpscalerLanczos(Upscaler): - scalers = [] - - def do_upscale(self, img, selected_model=None): - return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=LANCZOS) - - def load_model(self, _): - pass - - def __init__(self, dirname=None): # pylint: disable=unused-argument - super().__init__(False) - self.name = "Lanczos" - self.scalers = [UpscalerData("Lanczos", None, self)] - - -class UpscalerNearest(Upscaler): - scalers = [] - - def do_upscale(self, img, selected_model=None): - return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=NEAREST) - - def load_model(self, _): - pass - - def __init__(self, dirname=None): # pylint: disable=unused-argument - super().__init__(False) - self.name = "Nearest" - self.scalers = [UpscalerData("Nearest", None, self)] - def compile_upscaler(model): try: if shared.opts.ipex_optimize and "Upscaler" in shared.opts.ipex_optimize: diff --git a/modules/upscaler_simple.py b/modules/upscaler_simple.py new file mode 100644 index 000000000..fcf4f2410 --- /dev/null +++ b/modules/upscaler_simple.py @@ -0,0 +1,91 @@ +from PIL import Image +from modules.upscaler import Upscaler, UpscalerData + + +class UpscalerNone(Upscaler): + def __init__(self, dirname=None): # pylint: disable=unused-argument + super().__init__(False) + self.name = "None" + self.scalers = [UpscalerData("None", None, self)] + + def load_model(self, path): + pass + + def do_upscale(self, img, selected_model=None): + return img + + +class UpscalerResize(Upscaler): + def __init__(self, dirname=None): # pylint: disable=unused-argument + super().__init__(False) + self.name = "Resize" + self.scalers = [ + UpscalerData("Resize Nearest", None, self), + UpscalerData("Resize Lanczos", None, self), + UpscalerData("Resize Bicubic", None, self), + UpscalerData("Resize Bilinear", None, self), + UpscalerData("Resize Hamming", None, self), + UpscalerData("Resize Box", None, self), + ] + + def do_upscale(self, img: Image, selected_model=None): + if selected_model is None: + return img + if selected_model == "Resize Nearest": + return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.NEAREST) + if selected_model == "Resize Lanczos": + return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.LANCZOS) + if selected_model == "Resize Bicubic": + return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.BICUBIC) + if selected_model == "Resize Bilinear": + return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.BILINEAR) + if selected_model == "Resize Hamming": + return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.HAMMING) + if selected_model == "Resize Box": + return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.BOX) + + + def load_model(self, _): + pass + + +class UpscalerLatent(Upscaler): + def __init__(self, dirname=None): # pylint: disable=unused-argument + super().__init__(False) + self.name = "Latent" + self.scalers = [ + UpscalerData("Latent Nearest", None, self), + UpscalerData("Latent Nearest exact", None, self), + UpscalerData("Latent Area", None, self), + UpscalerData("Latent Bilinear", None, self), + UpscalerData("Latent Bicubic", None, self), + UpscalerData("Latent Bilinear antialias", None, self), + UpscalerData("Latent Bicubic antialias", None, self), + ] + + def do_upscale(self, img: Image, selected_model=None): + import torch + import torch.nn.functional as F + if isinstance(img, torch.Tensor) and (len(img.shape) == 4): + _batch, _channel, h, w = img.shape + else: + raise ValueError(f"Latent upscale: image={img.shape if isinstance(img, torch.Tensor) else img} type={type(img)} if not supported") + h, w = int((8 * h * self.scale) // 8), int((8 * w * self.scale) // 8) + mode, antialias = '', '' + if selected_model == "Latent Nearest": + mode, antialias = 'nearest', False + elif selected_model == "Latent Nearest exact": + mode, antialias = 'nearest-exact', False + elif selected_model == "Latent Area": + mode, antialias = 'area', False + elif selected_model == "Latent Bilinear": + mode, antialias = 'bilinear', False + elif selected_model == "Latent Bicubic": + mode, antialias = 'bicubic', False + elif selected_model == "Latent Bilinear antialias": + mode, antialias = 'bilinear', True + elif selected_model == "Latent Bicubic antialias": + mode, antialias = 'bicubic', True + else: + raise ValueError(f"Latent upscale: model={selected_model} unknown") + return F.interpolate(img, size=(h, w), mode=mode, antialias=antialias) diff --git a/scripts/xyz_grid_classes.py b/scripts/xyz_grid_classes.py index 2b1fa2d59..cd4df56e8 100644 --- a/scripts/xyz_grid_classes.py +++ b/scripts/xyz_grid_classes.py @@ -128,7 +128,7 @@ axis_options = [ AxisOption("[Sampler] Shift", float, apply_setting("schedulers_shift")), AxisOption("[Sampler] eta delta", float, apply_setting("eta_noise_seed_delta")), AxisOption("[Sampler] eta multiplier", float, apply_setting("scheduler_eta")), - AxisOption("[Refine] Upscaler", str, apply_field("hr_upscaler"), cost=0.3, choices=lambda: [*shared.latent_upscale_modes, *[x.name for x in shared.sd_upscalers]]), + AxisOption("[Refine] Upscaler", str, apply_field("hr_upscaler"), cost=0.3, choices=lambda: [x.name for x in shared.sd_upscalers]), AxisOption("[Refine] Sampler", str, apply_hr_sampler_name, fmt=format_value_add_label, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]), AxisOption("[Refine] Denoising strength", float, apply_field("denoising_strength")), AxisOption("[Refine] Hires steps", int, apply_field("hr_second_pass_steps")), @@ -136,7 +136,7 @@ axis_options = [ AxisOption("[Refine] Guidance rescale", float, apply_field("diffusers_guidance_rescale")), AxisOption("[Refine] Refiner start", float, apply_field("refiner_start")), AxisOption("[Refine] Refiner steps", float, apply_field("refiner_steps")), - AxisOption("[Postprocess] Upscaler", str, apply_upscaler, cost=0.4, choices=lambda: [x.name for x in shared.sd_upscalers][1:]), + AxisOption("[Postprocess] Upscaler", str, apply_upscaler, cost=0.4, choices=lambda: [x.name for x in shared.sd_upscalers]), AxisOption("[Postprocess] Context", str, apply_context, choices=lambda: ["Add with forward", "Remove with forward", "Add with backward", "Remove with backward"]), AxisOption("[Postprocess] Detailer", str, apply_detailer, fmt=format_value_add_label), AxisOption("[Postprocess] Detailer strength", str, apply_field("detailer_strength")), diff --git a/webui.py b/webui.py index f104bbee8..23a9eba95 100644 --- a/webui.py +++ b/webui.py @@ -29,6 +29,7 @@ import modules.ui import modules.txt2img import modules.img2img import modules.upscaler +import modules.upscaler_simple import modules.extra_networks import modules.ui_extra_networks import modules.textual_inversion.textual_inversion diff --git a/wiki b/wiki index 6caea0521..3dcc0808d 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 6caea0521a48bd5b99a18566f21797bff1f2c244 +Subproject commit 3dcc0808db1e9c351a845184a3bf0fe7ca783190