diff --git a/SUPIR/models/SUPIR_model.py b/SUPIR/models/SUPIR_model.py index e0669fb..0ed31a6 100644 --- a/SUPIR/models/SUPIR_model.py +++ b/SUPIR/models/SUPIR_model.py @@ -105,6 +105,7 @@ class SUPIRModel(DiffusionEngine): self.sampler_config.params.s_noise = s_noise self.sampler = instantiate_from_config(self.sampler_config) + print("Sampler: ", self.sampler_config.target) print("sampler_config: ", self.sampler_config.params) if seed == -1: diff --git a/nodes.py b/nodes.py index aced48c..e4ca677 100644 --- a/nodes.py +++ b/nodes.py @@ -143,6 +143,13 @@ class SUPIR_Upscale: "sampler_tile_stride": ("INT", {"default": 512, "min": 32, "max": 2048, "step": 32}), "fp8_unet": ("BOOLEAN", {"default": False}), "fp8_vae": ("BOOLEAN", {"default": False}), + "sampler": ( + [ + 'RestoreDPMPP2MSampler', + 'RestoreEDMSampler', + ], { + "default": 'RestoreEDMSampler' + }), } } @@ -156,7 +163,7 @@ class SUPIR_Upscale: encoder_tile_size_pixels, decoder_tile_size_latent, control_scale, cfg_scale_start, control_scale_start, restoration_scale, keep_model_loaded, a_prompt, n_prompt, sdxl_model, supir_model, use_tiled_vae, use_tiled_sampling=False, sampler_tile_size=128, sampler_tile_stride=64, captions="", diffusion_dtype="auto", - encoder_dtype="auto", batch_size=1, fp8_unet=False, fp8_vae=False): + encoder_dtype="auto", batch_size=1, fp8_unet=False, fp8_vae=False, sampler="RestoreEDMSampler"): device = mm.get_torch_device() mm.unload_all_models() @@ -176,7 +183,8 @@ class SUPIR_Upscale: 'supir_model': supir_model, 'use_tiled_sampling': use_tiled_sampling, 'fp8_unet': fp8_unet, - 'fp8_vae': fp8_vae + 'fp8_vae': fp8_vae, + 'sampler': sampler } if diffusion_dtype == 'auto': @@ -236,7 +244,7 @@ class SUPIR_Upscale: config.model.params.ae_dtype = vae_dtype config.model.params.diffusion_dtype = model_dtype - + config.model.params.sampler_config.target = f".sgm.modules.diffusionmodules.sampling.{sampler}" self.model = instantiate_from_config(config.model).cpu() try: diff --git a/sgm/modules/diffusionmodules/sampling.py b/sgm/modules/diffusionmodules/sampling.py index b34c769..212e756 100644 --- a/sgm/modules/diffusionmodules/sampling.py +++ b/sgm/modules/diffusionmodules/sampling.py @@ -8,7 +8,7 @@ from typing import Dict, Union import torch from omegaconf import ListConfig, OmegaConf from tqdm import tqdm - +from comfy.k_diffusion.sampling import BrownianTreeNoiseSampler, get_sigmas_karras from ...modules.diffusionmodules.sampling_utils import ( get_ancestral_step, linear_multistep_coeff, @@ -555,3 +555,169 @@ def _sliding_windows(h: int, w: int, tile_size: int, tile_stride: int): coords.append((hi, hi + tile_size, wi, wi + tile_size)) return coords +class RestoreDPMPP2MSampler(DPMPP2MSampler): + def __init__(self, s_churn=0.0, s_tmin=0.0, s_tmax=float("inf"), s_noise=1.0, restore_cfg=4.0, + restore_cfg_s_tmin=0.05, eta=1., *args, **kwargs): + self.s_noise = s_noise + self.eta = eta + super().__init__(*args, **kwargs) + + def denoise(self, x, denoiser, sigma, cond, uc, control_scale=1.0): + denoised = denoiser(*self.guider.prepare_inputs(x, sigma, cond, uc), control_scale) + denoised = self.guider(denoised, sigma) + return denoised + + def get_mult(self, h, r, t, t_next, previous_sigma): + eta_h = self.eta * h + mult1 = to_sigma(t_next) / to_sigma(t) * (-eta_h).exp() + mult2 = (-h -eta_h).expm1() + + if previous_sigma is not None: + mult3 = 1 + 1 / (2 * r) + mult4 = 1 / (2 * r) + return mult1, mult2, mult3, mult4 + else: + return mult1, mult2 + + + def sampler_step( + self, + old_denoised, + previous_sigma, + sigma, + next_sigma, + denoiser, + x, + cond, + uc=None, + eps_noise=None, + control_scale=1.0, + ): + denoised = self.denoise(x, denoiser, sigma, cond, uc, control_scale=control_scale) + + h, r, t, t_next = self.get_variables(sigma, next_sigma, previous_sigma) + eta_h = self.eta * h + mult = [ + append_dims(mult, x.ndim) + for mult in self.get_mult(h, r, t, t_next, previous_sigma) + ] + + x_standard = mult[0] * x - mult[1] * denoised + if old_denoised is None or torch.sum(next_sigma) < 1e-14: + # Save a network evaluation if all noise levels are 0 or on the first step + return x_standard, denoised + else: + denoised_d = mult[2] * denoised - mult[3] * old_denoised + x_advanced = mult[0] * x - mult[1] * denoised_d + + # apply correction if noise level is not 0 and not first step + x = torch.where( + append_dims(next_sigma, x.ndim) > 0.0, x_advanced, x_standard + ) + if self.eta: + x = x + eps_noise * next_sigma * (-2 * eta_h).expm1().neg().sqrt() * self.s_noise + + return x, denoised + + def __call__(self, denoiser, x, cond, uc=None, num_steps=None, control_scale=1.0, **kwargs): + x, s_in, sigmas, num_sigmas, cond, uc = self.prepare_sampling_loop( + x, cond, uc, num_steps + ) + sigmas_min, sigmas_max = sigmas[-2].cpu(), sigmas[0].cpu() + sigmas_new = get_sigmas_karras(self.num_steps, sigmas_min, sigmas_max, device=x.device) + sigmas = sigmas_new + + noise_sampler = BrownianTreeNoiseSampler(x, sigmas_min, sigmas_max) + + old_denoised = None + pbar_comfy = comfy.utils.ProgressBar(num_sigmas) + for i in self.get_sigma_gen(num_sigmas): + if i > 0 and torch.sum(s_in * sigmas[i + 1]) > 1e-14: + eps_noise = noise_sampler(s_in * sigmas[i], s_in * sigmas[i + 1]) + else: + eps_noise = None + x, old_denoised = self.sampler_step( + old_denoised, + None if i == 0 else s_in * sigmas[i - 1], + s_in * sigmas[i], + s_in * sigmas[i + 1], + denoiser, + x, + cond, + uc=uc, + eps_noise=eps_noise, + control_scale=control_scale, + ) + pbar_comfy.update(1) + + return x + +class TiledRestoreDPMPP2MSampler(RestoreDPMPP2MSampler): + def __init__(self, tile_size=128, tile_stride=64, *args, **kwargs): + + super().__init__(*args, **kwargs) + self.tile_size = tile_size + self.tile_stride = tile_stride + self.tile_weights = gaussian_weights(self.tile_size, self.tile_size, 1) + + def __call__(self, denoiser, x, cond, uc=None, num_steps=None, control_scale=1.0, **kwargs): + use_local_prompt = isinstance(cond, list) + b, _, h, w = x.shape + latent_tiles_iterator = _sliding_windows(h, w, self.tile_size, self.tile_stride) + tile_weights = self.tile_weights.repeat(b, 1, 1, 1) + if not use_local_prompt: + LQ_latent = cond['control'] + else: + assert len(cond) == len(latent_tiles_iterator), "Number of local prompts should be equal to number of tiles" + LQ_latent = cond[0]['control'] + x, s_in, sigmas, num_sigmas, cond, uc = self.prepare_sampling_loop( + x, cond, uc, num_steps + ) + sigmas_min, sigmas_max = sigmas[-2].cpu(), sigmas[0].cpu() + sigmas_new = get_sigmas_karras(self.num_steps, sigmas_min, sigmas_max, device=x.device) + sigmas = sigmas_new + + noise_sampler = BrownianTreeNoiseSampler(x, sigmas_min, sigmas_max) + + old_denoised = None + for _idx, i in enumerate(self.get_sigma_gen(num_sigmas)): + if i > 0 and torch.sum(s_in * sigmas[i + 1]) > 1e-14: + eps_noise = noise_sampler(s_in * sigmas[i], s_in * sigmas[i + 1]) + else: + eps_noise = torch.zeros_like(x) + x_next = torch.zeros_like(x) + old_denoised_next = torch.zeros_like(x) + count = torch.zeros_like(x) + for j, (hi, hi_end, wi, wi_end) in enumerate(latent_tiles_iterator): + x_tile = x[:, :, hi:hi_end, wi:wi_end] + _eps_noise = eps_noise[:, :, hi:hi_end, wi:wi_end] + if old_denoised is not None: + old_denoised_tile = old_denoised[:, :, hi:hi_end, wi:wi_end] + else: + old_denoised_tile = None + if use_local_prompt: + _cond = cond[j] + else: + _cond = cond + _cond['control'] = LQ_latent[:, :, hi:hi_end, wi:wi_end] + uc['control'] = LQ_latent[:, :, hi:hi_end, wi:wi_end] + _x, _old_denoised = self.sampler_step( + old_denoised_tile, + None if i == 0 else s_in * sigmas[i - 1], + s_in * sigmas[i], + s_in * sigmas[i + 1], + denoiser, + x_tile, + _cond, + uc=uc, + eps_noise=_eps_noise, + control_scale=control_scale, + ) + x_next[:, :, hi:hi_end, wi:wi_end] += _x * tile_weights + old_denoised_next[:, :, hi:hi_end, wi:wi_end] += _old_denoised * tile_weights + count[:, :, hi:hi_end, wi:wi_end] += tile_weights + old_denoised_next /= count + x_next /= count + x = x_next + old_denoised = old_denoised_next + return x \ No newline at end of file