From 9a09b2eef30264490a0b6f155f8044a47cfe2a04 Mon Sep 17 00:00:00 2001 From: cool-bigdogs-tshirt <131823432+cool-bigdogs-tshirt@users.noreply.github.com> Date: Thu, 27 Apr 2023 11:24:26 -0500 Subject: [PATCH 1/4] attempt at unipc latent upscaling i should have taken linear algebra before i dropped out... --- modules/models/diffusion/uni_pc/sampler.py | 74 ++++++++++++++++++++++ modules/processing.py | 3 +- 2 files changed, 76 insertions(+), 1 deletion(-) diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py index a241c8a7c..41b8c9a5b 100644 --- a/modules/models/diffusion/uni_pc/sampler.py +++ b/modules/models/diffusion/uni_pc/sampler.py @@ -4,6 +4,7 @@ import torch from .uni_pc import NoiseScheduleVP, model_wrapper, UniPC from modules import shared, devices +from ldm.modules.diffusionmodules.util import extract_into_tensor class UniPCSampler(object): @@ -15,6 +16,79 @@ class UniPCSampler(object): self.after_sample = None self.register_buffer('alphas_cumprod', to_torch(model.alphas_cumprod)) + def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True): + # persist steps so we can eventually find denoising strength + self.inflated_steps = ddim_num_steps + + @torch.no_grad() + def stochastic_encode(self, x0, t, use_original_steps=False, noise=None): + if noise is None: + noise = torch.randn_like(x0) + + # first time we have all the info to get the real parameters from the ui + hires_steps = t[0] + 1 + inflated_steps = self.inflated_steps + self.denoising_strength = hires_steps/inflated_steps + + adjusted_steps = int(hires_steps * self.denoising_strength) + self.steps = max(adjusted_steps, shared.opts.uni_pc_order+1) + + t = torch.full(t.shape, self.steps).to(t.device) + + timesteps = torch.asarray(list(range( + t, + self.model.num_timesteps, + self.model.num_timesteps // hires_steps, + ))) + 1 + alphas = self.model.alphas_cumprod[timesteps] + sqrt_one_minus_alphas = torch.sqrt(1. - alphas) + a = extract_into_tensor(torch.sqrt(alphas), t, x0.shape) * x0 + b = extract_into_tensor(sqrt_one_minus_alphas, t, x0.shape) * noise + + return (a+b) + + def decode(self, x_latent, conditioning, t_start, unconditional_guidance_scale=1.0, unconditional_conditioning=None, + use_original_steps=False, callback=None): + #print(f'steps {self.steps} denoising {self.denoising_strength}') + + noise_schedule = NoiseScheduleVP("discrete", alphas_cumprod=self.alphas_cumprod) + + # same as in .sample(), i guess + model_type = "v" if self.model.parameterization == "v" else "noise" + + model_fn = model_wrapper( + lambda x, t, c: self.model.apply_model(x, t, c), + noise_schedule, + model_type=model_type, + guidance_type="classifier-free", + #condition=conditioning, + #unconditional_condition=unconditional_conditioning, + guidance_scale=unconditional_guidance_scale, + ) + + self.uni_pc = UniPC( + model_fn, + noise_schedule, + predict_x0=True, + thresholding=False, + variant=shared.opts.uni_pc_variant, + condition=conditioning, + unconditional_condition=unconditional_conditioning, + before_sample=self.before_sample, + after_sample=self.after_sample, + after_update=self.after_update, + ) + + return self.uni_pc.sample( + x_latent, + steps=self.steps, + skip_type=shared.opts.uni_pc_skip_type, + method="multistep", + order=shared.opts.uni_pc_order, + lower_order_final=shared.opts.uni_pc_lower_order_final, + t_start=self.denoising_strength, + ) + def register_buffer(self, name, attr): if type(attr) == torch.Tensor: if attr.device != devices.device: diff --git a/modules/processing.py b/modules/processing.py index 04379fabe..c3ea4b2b8 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -970,7 +970,8 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing): shared.state.nextjob() img2img_sampler_name = self.sampler_name - if self.sampler_name in ['PLMS', 'UniPC']: # PLMS/UniPC do not support img2img so we just silently switch to DDIM + if self.sampler_name in ['PLMS']: + # PLMS does not support img2img, use fallback instead img2img_sampler_name = shared.opts.fallback_sampler self.sampler = sd_samplers.create_sampler(img2img_sampler_name, self.sd_model) From 42e30bfc3cf6b10664fbd66cabcd3c65f69e9cd4 Mon Sep 17 00:00:00 2001 From: cool-bigdogs-tshirt <131823432+cool-bigdogs-tshirt@users.noreply.github.com> Date: Fri, 28 Apr 2023 22:31:53 -0500 Subject: [PATCH 2/4] unipc img2img - add a bunch of code to get a single value that maybe performs slightly better? --- modules/models/diffusion/uni_pc/sampler.py | 51 ++++++++++++++++------ 1 file changed, 37 insertions(+), 14 deletions(-) diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py index 41b8c9a5b..3b468cf3f 100644 --- a/modules/models/diffusion/uni_pc/sampler.py +++ b/modules/models/diffusion/uni_pc/sampler.py @@ -1,5 +1,6 @@ """SAMPLING ONLY.""" +import numpy as np import torch from .uni_pc import NoiseScheduleVP, model_wrapper, UniPC @@ -26,26 +27,48 @@ class UniPCSampler(object): noise = torch.randn_like(x0) # first time we have all the info to get the real parameters from the ui - hires_steps = t[0] + 1 + # value from the hires steps slider: + num_inference_steps = t[0] + 1 + # (num_inference_steps // denoising_strength): inflated_steps = self.inflated_steps - self.denoising_strength = hires_steps/inflated_steps + # not exact: + self.denoising_strength = num_inference_steps/inflated_steps - adjusted_steps = int(hires_steps * self.denoising_strength) - self.steps = max(adjusted_steps, shared.opts.uni_pc_order+1) + # values used for timesteps that generate noise in diffusers repo + init_timestep = min( + int(num_inference_steps * self.denoising_strength), + num_inference_steps, + ) + t_start = max(num_inference_steps - init_timestep, 0) + + # actual number of steps we'll run + self.steps = max( + num_inference_steps - init_timestep, + shared.opts.uni_pc_order+1, + ) t = torch.full(t.shape, self.steps).to(t.device) - timesteps = torch.asarray(list(range( - t, - self.model.num_timesteps, - self.model.num_timesteps // hires_steps, - ))) + 1 - alphas = self.model.alphas_cumprod[timesteps] - sqrt_one_minus_alphas = torch.sqrt(1. - alphas) - a = extract_into_tensor(torch.sqrt(alphas), t, x0.shape) * x0 - b = extract_into_tensor(sqrt_one_minus_alphas, t, x0.shape) * noise + scheduler_timesteps = np.linspace( + 0, + self.model.num_timesteps-1, + num_inference_steps + 1, + ).round()[::-1][:-1].copy().astype(np.int64) + _, unique_indices = np.unique(scheduler_timesteps, return_index=True) + scheduler_timesteps = scheduler_timesteps[np.sort(unique_indices)] + scheduler_timesteps = torch.from_numpy(scheduler_timesteps).to(t.device) - return (a+b) + sample_timesteps = scheduler_timesteps[t_start:] + latent_timestep = sample_timesteps[:1].repeat(x0.shape[0]) + + alphas_cumprod = self.alphas_cumprod + sqrt_alphas_prod = alphas_cumprod[latent_timestep] ** 0.5 + sqrt_alphas_prod = sqrt_alphas_prod.flatten() + + sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[latent_timestep]) ** 0.5 + sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() + + return (sqrt_alphas_prod * x0 + sqrt_one_minus_alpha_prod * noise) def decode(self, x_latent, conditioning, t_start, unconditional_guidance_scale=1.0, unconditional_conditioning=None, use_original_steps=False, callback=None): From a78ce0a3ca4521a5661042c47e3dd56eee7a1eeb Mon Sep 17 00:00:00 2001 From: cool-bigdogs-tshirt <131823432+cool-bigdogs-tshirt@users.noreply.github.com> Date: Thu, 27 Apr 2023 12:04:14 -0500 Subject: [PATCH 3/4] xyz override for latent upscaler fallback --- modules/processing.py | 5 +++-- scripts/xyz_grid.py | 9 +++++++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/modules/processing.py b/modules/processing.py index c3ea4b2b8..a83d9fa12 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -970,9 +970,10 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing): shared.state.nextjob() img2img_sampler_name = self.sampler_name - if self.sampler_name in ['PLMS']: + force_latent_upscaler = shared.opts.data.get('xyz_fallback_sampler') + if self.sampler_name in ['PLMS'] or force_latent_upscaler is not None: # PLMS does not support img2img, use fallback instead - img2img_sampler_name = shared.opts.fallback_sampler + img2img_sampler_name = force_latent_upscaler or shared.opts.fallback_sampler self.sampler = sd_samplers.create_sampler(img2img_sampler_name, self.sd_model) samples = samples[:, :, self.truncate_y//2:samples.shape[2]-(self.truncate_y+1)//2, self.truncate_x//2:samples.shape[3]-(self.truncate_x+1)//2] diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py index 52ae1c6e1..9a5a67241 100644 --- a/scripts/xyz_grid.py +++ b/scripts/xyz_grid.py @@ -128,6 +128,14 @@ def apply_styles(p: StableDiffusionProcessingTxt2Img, x: str, _): p.styles.extend(x.split(',')) +def apply_fallback(p, x, xs): + sampler_name = sd_samplers.samplers_map.get(x.lower(), None) + if sampler_name is None: + raise RuntimeError(f"Unknown sampler: {x}") + + opts.data["xyz_fallback_sampler"] = sampler_name + + def apply_uni_pc_order(p, x, xs): opts.data["uni_pc_order"] = min(x, p.steps - 1) @@ -220,6 +228,7 @@ axis_options = [ AxisOption("Clip skip", int, apply_clip_skip), AxisOption("Denoising", float, apply_field("denoising_strength")), AxisOptionTxt2Img("Hires upscaler", str, apply_field("hr_upscaler"), choices=lambda: [*shared.latent_upscale_modes, *[x.name for x in shared.sd_upscalers]]), + AxisOptionTxt2Img("Fallback latent upscaler sampler", str, apply_fallback, format_value=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]), AxisOptionImg2Img("Cond. Image Mask Weight", float, apply_field("inpainting_mask_weight")), AxisOption("VAE", str, apply_vae, cost=0.7, choices=lambda: list(sd_vae.vae_dict)), AxisOption("Styles", str, apply_styles, choices=lambda: list(shared.prompt_styles.styles)), From 5148c5b0ad6a5a9522e57952b428f89d29f64375 Mon Sep 17 00:00:00 2001 From: cool-bigdogs-tshirt <131823432+cool-bigdogs-tshirt@users.noreply.github.com> Date: Fri, 28 Apr 2023 22:50:56 -0500 Subject: [PATCH 4/4] fix batching issue --- modules/models/diffusion/uni_pc/sampler.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py index 3b468cf3f..3100522ab 100644 --- a/modules/models/diffusion/uni_pc/sampler.py +++ b/modules/models/diffusion/uni_pc/sampler.py @@ -47,8 +47,6 @@ class UniPCSampler(object): shared.opts.uni_pc_order+1, ) - t = torch.full(t.shape, self.steps).to(t.device) - scheduler_timesteps = np.linspace( 0, self.model.num_timesteps-1, @@ -62,13 +60,17 @@ class UniPCSampler(object): latent_timestep = sample_timesteps[:1].repeat(x0.shape[0]) alphas_cumprod = self.alphas_cumprod - sqrt_alphas_prod = alphas_cumprod[latent_timestep] ** 0.5 - sqrt_alphas_prod = sqrt_alphas_prod.flatten() + sqrt_alpha_prod = alphas_cumprod[latent_timestep] ** 0.5 + sqrt_alpha_prod = sqrt_alpha_prod.flatten() + while len(sqrt_alpha_prod.shape) < len(x0.shape): + sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[latent_timestep]) ** 0.5 sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() + while len(sqrt_one_minus_alpha_prod.shape) < len(x0.shape): + sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - return (sqrt_alphas_prod * x0 + sqrt_one_minus_alpha_prod * noise) + return (sqrt_alpha_prod * x0 + sqrt_one_minus_alpha_prod * noise) def decode(self, x_latent, conditioning, t_start, unconditional_guidance_scale=1.0, unconditional_conditioning=None, use_original_steps=False, callback=None):