From 8f78025482e9f1fdd4b6678ee9f6b0cb50d70420 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Fri, 21 Jul 2023 03:52:09 +0000 Subject: [PATCH 1/8] add index_counter --- .../scheduling_dpmsolver_multistep.py | 29 ++++++++++++++----- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py index d7516fa601e1..7dbbe35f3d55 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py @@ -15,6 +15,7 @@ # DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver import math +from collections import defaultdict from typing import List, Optional, Tuple, Union import numpy as np @@ -274,11 +275,6 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc self.sigmas = torch.from_numpy(sigmas) - # when num_inference_steps == num_train_timesteps, we can end up with - # duplicates in timesteps. - _, unique_indices = np.unique(timesteps, return_index=True) - timesteps = timesteps[np.sort(unique_indices)] - self.timesteps = torch.from_numpy(timesteps).to(device) self.num_inference_steps = len(timesteps) @@ -288,6 +284,9 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc ] * self.config.solver_order self.lower_order_nums = 0 + # add an index counter for schedulers that allow duplicated timesteps + self._index_counter = defaultdict(int) + # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatTensor: """ @@ -660,11 +659,25 @@ def step( if isinstance(timestep, torch.Tensor): timestep = timestep.to(self.timesteps.device) - step_index = (self.timesteps == timestep).nonzero() - if len(step_index) == 0: + indices = (self.timesteps == timestep).nonzero() + timestep_int = timestep.cpu().item() if torch.is_tensor(timestep) else timestep + + if len(indices) == 0: step_index = len(self.timesteps) - 1 else: - step_index = step_index.item() + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + if len(self._index_counter) == 0: + pos = 1 if len(indices) > 1 else 0 + else: + pos = self._index_counter[timestep_int] + step_index = indices[pos].item() + + # advance index counter by 1 + self._index_counter[timestep_int] += 1 + prev_timestep = 0 if step_index == len(self.timesteps) - 1 else self.timesteps[step_index + 1] lower_order_final = ( (step_index == len(self.timesteps) - 1) and self.config.lower_order_final and len(self.timesteps) < 15 From 3b886af21bb070d48e984199778097221036aee5 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Fri, 21 Jul 2023 04:29:16 +0000 Subject: [PATCH 2/8] update test --- tests/schedulers/test_scheduler_dpm_multi.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/schedulers/test_scheduler_dpm_multi.py b/tests/schedulers/test_scheduler_dpm_multi.py index c9935780b983..86b24af24095 100644 --- a/tests/schedulers/test_scheduler_dpm_multi.py +++ b/tests/schedulers/test_scheduler_dpm_multi.py @@ -264,10 +264,10 @@ def test_fp16_support(self): assert sample.dtype == torch.float16 - def test_unique_timesteps(self, **config): + def test_duplicated_timesteps(self, **config): for scheduler_class in self.scheduler_classes: scheduler_config = self.get_scheduler_config(**config) scheduler = scheduler_class(**scheduler_config) scheduler.set_timesteps(scheduler.config.num_train_timesteps) - assert len(scheduler.timesteps.unique()) == scheduler.num_inference_steps + assert len(scheduler.timesteps) == scheduler.num_inference_steps From a05a13a9ab0a796d970c03e7b3e592948cf7b48a Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Mon, 24 Jul 2023 19:11:10 +0000 Subject: [PATCH 3/8] update add print lines add print lines and change --- src/diffusers/models/unet_2d_condition.py | 1 + .../pipeline_stable_diffusion.py | 9 +- .../pipeline_stable_diffusion_k_diffusion.py | 118 ++++++++++- .../scheduling_dpmsolver_multistep.py | 183 ++++++++++++------ .../schedulers/scheduling_euler_discrete.py | 8 + 5 files changed, 257 insertions(+), 62 deletions(-) diff --git a/src/diffusers/models/unet_2d_condition.py b/src/diffusers/models/unet_2d_condition.py index d7756ab5edb3..1bc493df638b 100644 --- a/src/diffusers/models/unet_2d_condition.py +++ b/src/diffusers/models/unet_2d_condition.py @@ -986,6 +986,7 @@ def forward( sample = self.conv_norm_out(sample) sample = self.conv_act(sample) sample = self.conv_out(sample) + print(f" - unet out (sample): {sample.shape},{sample[0,0,:3,:3]}") if not return_dict: return (sample,) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index 54927049571c..276ddae727b4 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -515,7 +515,10 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype latents = latents.to(device) # scale the initial noise by the standard deviation required by the scheduler + print(f" inside prepare_latents:") + print(f" - latents: {latents.shape},{latents[0,0,:3,:3]}") latents = latents * self.scheduler.init_noise_sigma + print(f" - latents * init_noise_sigma: {latents.shape},{latents[0,0,:3,:3]}") return latents @torch.no_grad() @@ -679,7 +682,9 @@ def __call__( for i, t in enumerate(timesteps): # expand the latents if we are doing classifier free guidance latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + print(f" - latent_model_input: {latent_model_input.shape},{latent_model_input[0,0,:3,:3]}") latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + print(f" - latent_model_input(scaled): {latent_model_input.shape},{latent_model_input[0,0,:3,:3]}") # predict the noise residual noise_pred = self.unet( @@ -689,12 +694,12 @@ def __call__( cross_attention_kwargs=cross_attention_kwargs, return_dict=False, )[0] - + print(f" - noise_pred: {noise_pred.shape},{noise_pred[0,0,:3,:3]}") # perform guidance if do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - + print(f" - noise_pred (cfg): {noise_pred.shape},{noise_pred[0,0,:3,:3]}") if do_classifier_free_guidance and guidance_rescale > 0.0: # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_k_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_k_diffusion.py index 29a57470a341..4291072b6e94 100755 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_k_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_k_diffusion.py @@ -18,7 +18,8 @@ from typing import Callable, List, Optional, Union import torch -from k_diffusion.external import CompVisDenoiser, CompVisVDenoiser +from k_diffusion.external import CompVisDenoiser, CompVisVDenoiser, DiscreteSchedule +from k_diffusion import utils from k_diffusion.sampling import BrownianTreeNoiseSampler, get_sigmas_karras from ...image_processor import VaeImageProcessor @@ -32,6 +33,97 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name +# yiyi testing +from tqdm.auto import trange +@torch.no_grad() +def sample_dpmpp_2m(model, x, sigmas, extra_args=None, callback=None, disable=None): + """DPM-Solver++(2M).""" + extra_args = {} if extra_args is None else extra_args + s_in = x.new_ones([x.shape[0]]) + sigma_fn = lambda t: t.neg().exp() + t_fn = lambda sigma: sigma.log().neg() + old_denoised = None + + for i in trange(len(sigmas) - 1, disable=disable): + print(f" - i :{i}, sigma: {sigmas[i]}") + denoised = model(x, sigmas[i] * s_in, **extra_args) + print(f" - denoised: {denoised.shape}, {denoised[0,0,:3,:3]}") + if callback is not None: + callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) + t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1]) + print(f" - sigma_t: {sigmas[i+1]}, sigma_s: {sigmas[i]}") + print(f" - t, t_next: {t},{t_next}") + h = t_next - t + print(f" - h: {h}") + if old_denoised is None or sigmas[i + 1] == 0: + print(f" first order") + print(f" - x/sample/latents: {x.shape},{x[0,0,:3,:3]}") + print(f" - sigma_fns(t_next): {sigma_fn(t_next)}, sigma_fn(t): {sigma_fn(t)}") + x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised + print(f" -> x: {x[0,0,:3,:3]}") + else: + print(" second order") + print(f" yiyi testing") + print(f" - sigmas: {sigmas[i]}, {sigmas[i+1]}") + print(f" - sigma_fns: {sigma_fn(t)}, {sigma_fn(t_next)}") + h_last = t - t_fn(sigmas[i - 1]) + r = h_last / h + denoised_d = (1 + 1 / (2 * r)) * denoised - (1 / (2 * r)) * old_denoised + x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_d + print(f" -> x: {x[0,0,:3,:3]}") + old_denoised = denoised + return x + + +class DiscreteEpsDDPMDenoiser(DiscreteSchedule): + """A wrapper for discrete schedule DDPM models that output eps (the predicted + noise).""" + + def __init__(self, model, alphas_cumprod, quantize): + super().__init__(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5, quantize) + self.inner_model = model + self.sigma_data = 1. + + def get_scalings(self, sigma): + c_out = -sigma + c_in = 1 / (sigma ** 2 + self.sigma_data ** 2) ** 0.5 + return c_out, c_in + + def get_eps(self, *args, **kwargs): + return self.inner_model(*args, **kwargs) + + def loss(self, input, noise, sigma, **kwargs): + c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] + noised_input = input + noise * utils.append_dims(sigma, input.ndim) + eps = self.get_eps(noised_input * c_in, self.sigma_to_t(sigma), **kwargs) + return (eps - noise).pow(2).flatten(1).mean(1) + + def forward(self, input, sigma, **kwargs): + c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] + print(f" arriving CompVisDenoiser.foward") + print(f" - input: {input.shape}, {input[0,0,:3,:3]}") + print(f" - c_in: {c_in.shape}, {c_in}") + print(f" - c_out:{c_out.shape}, {c_out}") + print(f" - sigma: {sigma}") + print(f" - t: {self.sigma_to_t(sigma)}") + print(f" - input * c_in : {(input * c_in).shape}, {(input * c_in)[0,0,:3,:3]}") + eps = self.get_eps(input * c_in, self.sigma_to_t(sigma), **kwargs) + print(f" - eps: {eps.shape}, {eps[0,0,:3,:3]}") + print(f" - eps * c_out : {(eps * c_out).shape}, {(eps * c_out)[0,0,:3,:3]}") + print(f" - input + eps * c_out: {(input + eps * c_out).shape}, {(input + eps * c_out)[0,0,:3,:3]}") + print(f" leaving CompVisDenoiser.foward") + return input + eps * c_out + + +class CompVisDenoiser(DiscreteEpsDDPMDenoiser): + """A wrapper for CompVis diffusion models.""" + + def __init__(self, model, quantize=False, device='cpu'): + super().__init__(model, model.alphas_cumprod, quantize=quantize) + + def get_eps(self, *args, **kwargs): + return self.inner_model.apply_model(*args, **kwargs) + class ModelWrapper: def __init__(self, model, alphas_cumprod): self.model = model @@ -123,9 +215,12 @@ def __init__( self.k_diffusion_model = CompVisDenoiser(model) def set_scheduler(self, scheduler_type: str): - library = importlib.import_module("k_diffusion") - sampling = getattr(library, "sampling") - self.sampler = getattr(sampling, scheduler_type) + #library = importlib.import_module("k_diffusion") + #sampling = getattr(library, "sampling") + #self.sampler = getattr(sampling, scheduler_type) + if scheduler_type == "sample_dpmpp_2m": + self.sampler = sample_dpmpp_2m + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_model_cpu_offload def enable_model_cpu_offload(self, gpu_id=0): @@ -530,9 +625,12 @@ def __call__( self.scheduler.set_timesteps(num_inference_steps, device=prompt_embeds.device) # 5. Prepare sigmas + if use_karras_sigmas: + print(f" - k_diffusion_model.sigmas :{self.k_diffusion_model.sigmas}") sigma_min: float = self.k_diffusion_model.sigmas[0].item() sigma_max: float = self.k_diffusion_model.sigmas[-1].item() + print(f" -sigma_max: {sigma_max}, sigma_min: {sigma_min}") sigmas = get_sigmas_karras(n=num_inference_steps, sigma_min=sigma_min, sigma_max=sigma_max) sigmas = sigmas.to(device) else: @@ -551,19 +649,28 @@ def __call__( generator, latents, ) + print(f" - prepare_latents -> {latents.shape},{latents[0,0,:3,:3]}") latents = latents * sigmas[0] + print(f" - latents * initial noise sigma: {latents.shape},{latents[0,0,:3,:3]}") self.k_diffusion_model.sigmas = self.k_diffusion_model.sigmas.to(latents.device) self.k_diffusion_model.log_sigmas = self.k_diffusion_model.log_sigmas.to(latents.device) # 7. Define model function def model_fn(x, t): + print(" ") + print(f" arriving model_fn") latent_model_input = torch.cat([x] * 2) + print(f" - latent_model_input: {latent_model_input.shape}, {latent_model_input[0,0,:3,:3]}") t = torch.cat([t] * 2) noise_pred = self.k_diffusion_model(latent_model_input, t, cond=prompt_embeds) + print(f" - noise_pred: {noise_pred.shape}, {noise_pred[0,0,:3,:3]}") noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + print(f" -> cfg -> {noise_pred.shape},{noise_pred[0,0,:3,:3]}") + print(" leaving model_fn") + print(" ") return noise_pred # 8. Run k-diffusion solver @@ -573,7 +680,8 @@ def model_fn(x, t): min_sigma, max_sigma = sigmas[sigmas > 0].min(), sigmas.max() noise_sampler = BrownianTreeNoiseSampler(latents, min_sigma, max_sigma, noise_sampler_seed) sampler_kwargs["noise_sampler"] = noise_sampler - + + print(f" sigmas: {sigmas}") latents = self.sampler(model_fn, latents, sigmas, **sampler_kwargs) if not output_type == "latent": diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py index 7dbbe35f3d55..d39417859da8 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py @@ -206,9 +206,6 @@ def __init__( self.sigma_t = torch.sqrt(1 - self.alphas_cumprod) self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t) - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - # settings for DPM-Solver if algorithm_type not in ["dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++"]: if algorithm_type == "deis": @@ -225,9 +222,27 @@ def __init__( # setable values self.num_inference_steps = None timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32)[::-1].copy() + self.timesteps = torch.from_numpy(timesteps) self.model_outputs = [None] * solver_order self.lower_order_nums = 0 + self._step_index = None + + @property + def init_noise_sigma(self): + # standard deviation of the initial noise distribution + if self.config.timestep_spacing in ["linspace", "trailing"]: + return self.sigmas.max() + + return (self.sigmas.max() ** 2 + 1) ** 0.5 + + @property + def step_index(self): + """ + TODO: Nice docstring + """ + return self._step_index + def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torch.device] = None): """ @@ -243,23 +258,25 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc # This is critical for cosine (squaredcos_cap_v2) noise schedule. clipped_idx = torch.searchsorted(torch.flip(self.lambda_t, [0]), self.config.lambda_min_clipped) last_timestep = ((self.config.num_train_timesteps - clipped_idx).numpy()).item() + print(f" - last_timestep: {last_timestep}") # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891 if self.config.timestep_spacing == "linspace": timesteps = ( - np.linspace(0, last_timestep - 1, num_inference_steps + 1).round()[::-1][:-1].copy().astype(np.int64) + np.linspace(0, last_timestep - 1, num_inference_steps, dtype=float).round()[::-1].copy() ) + print(f" - timesteps: {len(timesteps)},{timesteps[0]}:{timesteps[-1]}") elif self.config.timestep_spacing == "leading": - step_ratio = last_timestep // (num_inference_steps + 1) + step_ratio = last_timestep // self.num_inference_steps # creates integer timesteps by multiplying by ratio # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps + 1) * step_ratio).round()[::-1][:-1].copy().astype(np.int64) + timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(float) timesteps += self.config.steps_offset elif self.config.timestep_spacing == "trailing": step_ratio = self.config.num_train_timesteps / num_inference_steps # creates integer timesteps by multiplying by ratio # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = np.arange(last_timestep, 0, -step_ratio).round().copy().astype(np.int64) + timesteps = np.arange(last_timestep, 0, -step_ratio).round().copy().astype(float) timesteps -= 1 else: raise ValueError( @@ -267,13 +284,21 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc ) sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) + log_sigmas = np.log(sigmas) + sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) + print(f" sigmas: {sigmas[0]}: {sigmas[-1]}") + if self.config.use_karras_sigmas: - log_sigmas = np.log(sigmas) + print(f" sigmas (k): {sigmas[0]}, {sigmas[-1]}") sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round() - timesteps = np.flip(timesteps).copy().astype(np.int64) + timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) + + sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) + self.sigmas = torch.from_numpy(sigmas).to(device=device) - self.sigmas = torch.from_numpy(sigmas) + print(" set_timesteps") + print(f" - sigmas: {sigmas[0]}: {sigmas[-1]}") + print(sigmas) self.timesteps = torch.from_numpy(timesteps).to(device) @@ -283,9 +308,12 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc None, ] * self.config.solver_order self.lower_order_nums = 0 + print(f" - timesteps: {timesteps}") + print(f" - model_outputs: {self.model_outputs}") + print(f" - lower_order_nums: {self.lower_order_nums}") # add an index counter for schedulers that allow duplicated timesteps - self._index_counter = defaultdict(int) + self._step_index = None # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatTensor: @@ -384,13 +412,17 @@ def convert_model_output( """ # DPM-Solver++ needs to solve an integral of the data prediction model. + print(f" -algo_type: {self.config.algorithm_type}, pred_type: {self.config.prediction_type}, variance_type: {self.config.variance_type}") if self.config.algorithm_type in ["dpmsolver++", "sde-dpmsolver++"]: if self.config.prediction_type == "epsilon": # DPM-Solver and DPM-Solver++ only need the "mean" output. if self.config.variance_type in ["learned", "learned_range"]: model_output = model_output[:, :3] - alpha_t, sigma_t = self.alpha_t[timestep], self.sigma_t[timestep] - x0_pred = (sample - sigma_t * model_output) / alpha_t + # yiyi update/testing here: + sigma = self.sigmas[self.step_index] + x0_pred = sample - sigma * model_output + print(f" - sigma: {sigma}") + print(f" - x0_pred: {x0_pred.shape},{x0_pred[0,0,:3,:3]}") elif self.config.prediction_type == "sample": x0_pred = model_output elif self.config.prediction_type == "v_prediction": @@ -403,6 +435,7 @@ def convert_model_output( ) if self.config.thresholding: + print(f" threhold? ") x0_pred = self._threshold_sample(x0_pred) return x0_pred @@ -458,12 +491,18 @@ def dpm_solver_first_order_update( Returns: `torch.FloatTensor`: the sample tensor at the previous timestep. """ - lambda_t, lambda_s = self.lambda_t[prev_timestep], self.lambda_t[timestep] - alpha_t, alpha_s = self.alpha_t[prev_timestep], self.alpha_t[timestep] - sigma_t, sigma_s = self.sigma_t[prev_timestep], self.sigma_t[timestep] - h = lambda_t - lambda_s + def t_fn(_sigma): + return -torch.log(_sigma) + + print(f" - sample: {sample[0,0,:3,:3]}") + + sigma_t, sigma_s = self.sigmas[self.step_index +1], self.sigmas[self.step_index] + print(f" - sigma_t: {sigma_t}, sigma_s: {sigma_s}") + h = t_fn(sigma_t) - t_fn(sigma_s) + print(f" - h: {h}") if self.config.algorithm_type == "dpmsolver++": - x_t = (sigma_t / sigma_s) * sample - (alpha_t * (torch.exp(-h) - 1.0)) * model_output + x_t = (sigma_t / sigma_s) * sample - (torch.exp(-h) - 1.0) * model_output + print(f" -> prev_sample: {x_t[0,0,:3,:3]}") elif self.config.algorithm_type == "dpmsolver": x_t = (alpha_t / alpha_s) * sample - (sigma_t * (torch.exp(h) - 1.0)) * model_output elif self.config.algorithm_type == "sde-dpmsolver++": @@ -504,22 +543,32 @@ def multistep_dpm_solver_second_order_update( Returns: `torch.FloatTensor`: the sample tensor at the previous timestep. """ - t, s0, s1 = prev_timestep, timestep_list[-1], timestep_list[-2] + + def t_fn(_sigma): + return -torch.log(_sigma) + + sigma_t, sigma_s0, sigma_s1 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index], self.sigmas[self.step_index - 1] m0, m1 = model_output_list[-1], model_output_list[-2] - lambda_t, lambda_s0, lambda_s1 = self.lambda_t[t], self.lambda_t[s0], self.lambda_t[s1] - alpha_t, alpha_s0 = self.alpha_t[t], self.alpha_t[s0] - sigma_t, sigma_s0 = self.sigma_t[t], self.sigma_t[s0] - h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1 + + h, h_0 = t_fn(sigma_t) - t_fn(sigma_s0), t_fn(sigma_s0) - t_fn(sigma_s1) r0 = h_0 / h D0, D1 = m0, (1.0 / r0) * (m0 - m1) + print(f" sigma_t (current), sigma_s0 (prev), sigma_s1 (next): {sigma_t}, {sigma_s0}, {sigma_s1}") + # print(f" sigma_t, sigma_s0, {sigma_t},{sigma_s0}") + # print(f" sigma_s1: {self.sigma_t[s1]}") + # print(f" alpha_bar = {self.alphas_cumprod[s1]}") + print(f" sample: {sample[0,0,:3,:3]}") + # print(" -----") if self.config.algorithm_type == "dpmsolver++": # See https://arxiv.org/abs/2211.01095 for detailed derivations if self.config.solver_type == "midpoint": + print(f" dpmsolver++ /midpoint ") x_t = ( (sigma_t / sigma_s0) * sample - - (alpha_t * (torch.exp(-h) - 1.0)) * D0 - - 0.5 * (alpha_t * (torch.exp(-h) - 1.0)) * D1 + - (torch.exp(-h) - 1.0) * D0 + - 0.5 * (torch.exp(-h) - 1.0) * D1 ) + print(f" -> prev_sample: {x_t.shape}, {x_t[0,0,:3,:3]}") elif self.config.solver_type == "heun": x_t = ( (sigma_t / sigma_s0) * sample @@ -629,6 +678,23 @@ def multistep_dpm_solver_third_order_update( ) return x_t + def _init_step_index(self, timestep): + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + + index_candidates = (self.timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + if len(index_candidates) > 1: + step_index = index_candidates[1] + else: + step_index = index_candidates[0] + + self._step_index = step_index.item() + def step( self, model_output: torch.FloatTensor, @@ -657,37 +723,26 @@ def step( "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" ) - if isinstance(timestep, torch.Tensor): - timestep = timestep.to(self.timesteps.device) - indices = (self.timesteps == timestep).nonzero() - timestep_int = timestep.cpu().item() if torch.is_tensor(timestep) else timestep + if self.step_index is None: + self._init_step_index(timestep) + print(" ") + print(" ** inside step ***** ") - if len(indices) == 0: - step_index = len(self.timesteps) - 1 - else: - # The sigma index that is taken for the **very** first `step` - # is always the second index (or the last index if there is only 1) - # This way we can ensure we don't accidentally skip a sigma in - # case we start in the middle of the denoising schedule (e.g. for image-to-image) - if len(self._index_counter) == 0: - pos = 1 if len(indices) > 1 else 0 - else: - pos = self._index_counter[timestep_int] - step_index = indices[pos].item() - - # advance index counter by 1 - self._index_counter[timestep_int] += 1 + print(f" - timestep, step_index: {timestep}, {self.step_index}") + print(f" - sample/latents: {sample.shape},{sample[0,0,:3,:3]}") + print(f" - model_output (noise): {model_output.shape}, {model_output[0,0,:3,:3]}") - prev_timestep = 0 if step_index == len(self.timesteps) - 1 else self.timesteps[step_index + 1] - lower_order_final = ( - (step_index == len(self.timesteps) - 1) and self.config.lower_order_final and len(self.timesteps) < 15 - ) + prev_timestep = 0 if self.step_index == len(self.timesteps) - 1 else self.timesteps[self.step_index + 1] + lower_order_final = (self.step_index == len(self.timesteps) - 1) lower_order_second = ( - (step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15 + (self.step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15 ) + print(f" - prev_timestep: {prev_timestep}") model_output = self.convert_model_output(model_output, timestep, sample) + print(f" - model_output (x0): {model_output.shape}, {model_output[0,0,:3,:3]}") for i in range(self.config.solver_order - 1): + print(f" move outputs {i+1} -> {i}") self.model_outputs[i] = self.model_outputs[i + 1] self.model_outputs[-1] = model_output @@ -699,15 +754,18 @@ def step( noise = None if self.config.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final: + print(f" 1st order update: {self.config.solver_order}, {self.lower_order_nums}, {lower_order_final}") prev_sample = self.dpm_solver_first_order_update( model_output, timestep, prev_timestep, sample, noise=noise ) elif self.config.solver_order == 2 or self.lower_order_nums < 2 or lower_order_second: - timestep_list = [self.timesteps[step_index - 1], timestep] + print(f" 2nd order update") + timestep_list = [self.timesteps[self.step_index - 1], timestep] prev_sample = self.multistep_dpm_solver_second_order_update( self.model_outputs, timestep_list, prev_timestep, sample, noise=noise ) else: + print(f" 3rd order update") timestep_list = [self.timesteps[step_index - 2], self.timesteps[step_index - 1], timestep] prev_sample = self.multistep_dpm_solver_third_order_update( self.model_outputs, timestep_list, prev_timestep, sample @@ -715,23 +773,38 @@ def step( if self.lower_order_nums < self.config.solver_order: self.lower_order_nums += 1 + print(f" lower_order_nums +1 -> {self.lower_order_nums}") + + # upon completion increase step index by one + self._step_index += 1 if not return_dict: return (prev_sample,) return SchedulerOutput(prev_sample=prev_sample) - - def scale_model_input(self, sample: torch.FloatTensor, *args, **kwargs) -> torch.FloatTensor: + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.scale_model_input + def scale_model_input( + self, sample: torch.FloatTensor, timestep: Union[float, torch.FloatTensor] + ) -> torch.FloatTensor: """ - Ensures interchangeability with schedulers that need to scale the denoising model input depending on the - current timestep. + Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the Euler algorithm. Args: sample (`torch.FloatTensor`): input sample + timestep (`float` or `torch.FloatTensor`): the current timestep in the diffusion chain Returns: `torch.FloatTensor`: scaled input sample """ + if self.step_index is None: + self._init_step_index(timestep) + + sigma = self.sigmas[self.step_index] + + sample = sample / ((sigma**2 + 1) ** 0.5) + + self.is_scale_input_called = True return sample # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index cb126d4b953c..9dbe26e124f6 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -243,9 +243,13 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic ) sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) + print(f" -sigmas: {sigmas[0]}, {sigmas[-1]}") log_sigmas = np.log(sigmas) if self.config.interpolation_type == "linear": + print(f" - timesteps: {len(timesteps)}") + print(timesteps) + print(f" - sigmas: {len(sigmas)}") sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) elif self.config.interpolation_type == "log_linear": sigmas = torch.linspace(np.log(sigmas[-1]), np.log(sigmas[0]), num_inference_steps + 1).exp() @@ -256,8 +260,12 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic ) if self.use_karras_sigmas: + print(" use k_sigmas") + print(f" -sigmas: {sigmas[0]}, {sigmas[-1]}") sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=self.num_inference_steps) timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) + print(f" -sigmas: {sigmas}") + print(f" -timesteps: {timesteps}") sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) self.sigmas = torch.from_numpy(sigmas).to(device=device) From 670c782cb269ada256d0de4e5c5afb07858efb1b Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Mon, 21 Aug 2023 01:51:09 +0000 Subject: [PATCH 4/8] fix --- src/diffusers/models/unet_2d_condition.py | 1 - .../pipeline_stable_diffusion.py | 7 - .../pipeline_stable_diffusion_k_diffusion.py | 120 +----------------- .../scheduling_dpmsolver_multistep.py | 43 +------ .../schedulers/scheduling_euler_discrete.py | 8 -- 5 files changed, 7 insertions(+), 172 deletions(-) diff --git a/src/diffusers/models/unet_2d_condition.py b/src/diffusers/models/unet_2d_condition.py index 1bc493df638b..d7756ab5edb3 100644 --- a/src/diffusers/models/unet_2d_condition.py +++ b/src/diffusers/models/unet_2d_condition.py @@ -986,7 +986,6 @@ def forward( sample = self.conv_norm_out(sample) sample = self.conv_act(sample) sample = self.conv_out(sample) - print(f" - unet out (sample): {sample.shape},{sample[0,0,:3,:3]}") if not return_dict: return (sample,) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index 276ddae727b4..5ea7fd0d500a 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -515,10 +515,7 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype latents = latents.to(device) # scale the initial noise by the standard deviation required by the scheduler - print(f" inside prepare_latents:") - print(f" - latents: {latents.shape},{latents[0,0,:3,:3]}") latents = latents * self.scheduler.init_noise_sigma - print(f" - latents * init_noise_sigma: {latents.shape},{latents[0,0,:3,:3]}") return latents @torch.no_grad() @@ -682,9 +679,7 @@ def __call__( for i, t in enumerate(timesteps): # expand the latents if we are doing classifier free guidance latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - print(f" - latent_model_input: {latent_model_input.shape},{latent_model_input[0,0,:3,:3]}") latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - print(f" - latent_model_input(scaled): {latent_model_input.shape},{latent_model_input[0,0,:3,:3]}") # predict the noise residual noise_pred = self.unet( @@ -694,12 +689,10 @@ def __call__( cross_attention_kwargs=cross_attention_kwargs, return_dict=False, )[0] - print(f" - noise_pred: {noise_pred.shape},{noise_pred[0,0,:3,:3]}") # perform guidance if do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - print(f" - noise_pred (cfg): {noise_pred.shape},{noise_pred[0,0,:3,:3]}") if do_classifier_free_guidance and guidance_rescale > 0.0: # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_k_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_k_diffusion.py index 4291072b6e94..e272e9e3a505 100755 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_k_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_k_diffusion.py @@ -18,8 +18,7 @@ from typing import Callable, List, Optional, Union import torch -from k_diffusion.external import CompVisDenoiser, CompVisVDenoiser, DiscreteSchedule -from k_diffusion import utils +from k_diffusion.external import CompVisDenoiser, CompVisVDenoiser from k_diffusion.sampling import BrownianTreeNoiseSampler, get_sigmas_karras from ...image_processor import VaeImageProcessor @@ -33,97 +32,6 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name -# yiyi testing -from tqdm.auto import trange -@torch.no_grad() -def sample_dpmpp_2m(model, x, sigmas, extra_args=None, callback=None, disable=None): - """DPM-Solver++(2M).""" - extra_args = {} if extra_args is None else extra_args - s_in = x.new_ones([x.shape[0]]) - sigma_fn = lambda t: t.neg().exp() - t_fn = lambda sigma: sigma.log().neg() - old_denoised = None - - for i in trange(len(sigmas) - 1, disable=disable): - print(f" - i :{i}, sigma: {sigmas[i]}") - denoised = model(x, sigmas[i] * s_in, **extra_args) - print(f" - denoised: {denoised.shape}, {denoised[0,0,:3,:3]}") - if callback is not None: - callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) - t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1]) - print(f" - sigma_t: {sigmas[i+1]}, sigma_s: {sigmas[i]}") - print(f" - t, t_next: {t},{t_next}") - h = t_next - t - print(f" - h: {h}") - if old_denoised is None or sigmas[i + 1] == 0: - print(f" first order") - print(f" - x/sample/latents: {x.shape},{x[0,0,:3,:3]}") - print(f" - sigma_fns(t_next): {sigma_fn(t_next)}, sigma_fn(t): {sigma_fn(t)}") - x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised - print(f" -> x: {x[0,0,:3,:3]}") - else: - print(" second order") - print(f" yiyi testing") - print(f" - sigmas: {sigmas[i]}, {sigmas[i+1]}") - print(f" - sigma_fns: {sigma_fn(t)}, {sigma_fn(t_next)}") - h_last = t - t_fn(sigmas[i - 1]) - r = h_last / h - denoised_d = (1 + 1 / (2 * r)) * denoised - (1 / (2 * r)) * old_denoised - x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_d - print(f" -> x: {x[0,0,:3,:3]}") - old_denoised = denoised - return x - - -class DiscreteEpsDDPMDenoiser(DiscreteSchedule): - """A wrapper for discrete schedule DDPM models that output eps (the predicted - noise).""" - - def __init__(self, model, alphas_cumprod, quantize): - super().__init__(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5, quantize) - self.inner_model = model - self.sigma_data = 1. - - def get_scalings(self, sigma): - c_out = -sigma - c_in = 1 / (sigma ** 2 + self.sigma_data ** 2) ** 0.5 - return c_out, c_in - - def get_eps(self, *args, **kwargs): - return self.inner_model(*args, **kwargs) - - def loss(self, input, noise, sigma, **kwargs): - c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] - noised_input = input + noise * utils.append_dims(sigma, input.ndim) - eps = self.get_eps(noised_input * c_in, self.sigma_to_t(sigma), **kwargs) - return (eps - noise).pow(2).flatten(1).mean(1) - - def forward(self, input, sigma, **kwargs): - c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] - print(f" arriving CompVisDenoiser.foward") - print(f" - input: {input.shape}, {input[0,0,:3,:3]}") - print(f" - c_in: {c_in.shape}, {c_in}") - print(f" - c_out:{c_out.shape}, {c_out}") - print(f" - sigma: {sigma}") - print(f" - t: {self.sigma_to_t(sigma)}") - print(f" - input * c_in : {(input * c_in).shape}, {(input * c_in)[0,0,:3,:3]}") - eps = self.get_eps(input * c_in, self.sigma_to_t(sigma), **kwargs) - print(f" - eps: {eps.shape}, {eps[0,0,:3,:3]}") - print(f" - eps * c_out : {(eps * c_out).shape}, {(eps * c_out)[0,0,:3,:3]}") - print(f" - input + eps * c_out: {(input + eps * c_out).shape}, {(input + eps * c_out)[0,0,:3,:3]}") - print(f" leaving CompVisDenoiser.foward") - return input + eps * c_out - - -class CompVisDenoiser(DiscreteEpsDDPMDenoiser): - """A wrapper for CompVis diffusion models.""" - - def __init__(self, model, quantize=False, device='cpu'): - super().__init__(model, model.alphas_cumprod, quantize=quantize) - - def get_eps(self, *args, **kwargs): - return self.inner_model.apply_model(*args, **kwargs) - class ModelWrapper: def __init__(self, model, alphas_cumprod): self.model = model @@ -215,13 +123,10 @@ def __init__( self.k_diffusion_model = CompVisDenoiser(model) def set_scheduler(self, scheduler_type: str): - #library = importlib.import_module("k_diffusion") - #sampling = getattr(library, "sampling") - #self.sampler = getattr(sampling, scheduler_type) - if scheduler_type == "sample_dpmpp_2m": - self.sampler = sample_dpmpp_2m - - + library = importlib.import_module("k_diffusion") + sampling = getattr(library, "sampling") + self.sampler = getattr(sampling, scheduler_type) + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_model_cpu_offload def enable_model_cpu_offload(self, gpu_id=0): r""" @@ -625,12 +530,9 @@ def __call__( self.scheduler.set_timesteps(num_inference_steps, device=prompt_embeds.device) # 5. Prepare sigmas - if use_karras_sigmas: - print(f" - k_diffusion_model.sigmas :{self.k_diffusion_model.sigmas}") sigma_min: float = self.k_diffusion_model.sigmas[0].item() sigma_max: float = self.k_diffusion_model.sigmas[-1].item() - print(f" -sigma_max: {sigma_max}, sigma_min: {sigma_min}") sigmas = get_sigmas_karras(n=num_inference_steps, sigma_min=sigma_min, sigma_max=sigma_max) sigmas = sigmas.to(device) else: @@ -649,28 +551,19 @@ def __call__( generator, latents, ) - print(f" - prepare_latents -> {latents.shape},{latents[0,0,:3,:3]}") latents = latents * sigmas[0] - print(f" - latents * initial noise sigma: {latents.shape},{latents[0,0,:3,:3]}") self.k_diffusion_model.sigmas = self.k_diffusion_model.sigmas.to(latents.device) self.k_diffusion_model.log_sigmas = self.k_diffusion_model.log_sigmas.to(latents.device) # 7. Define model function def model_fn(x, t): - print(" ") - print(f" arriving model_fn") latent_model_input = torch.cat([x] * 2) - print(f" - latent_model_input: {latent_model_input.shape}, {latent_model_input[0,0,:3,:3]}") t = torch.cat([t] * 2) noise_pred = self.k_diffusion_model(latent_model_input, t, cond=prompt_embeds) - print(f" - noise_pred: {noise_pred.shape}, {noise_pred[0,0,:3,:3]}") noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - print(f" -> cfg -> {noise_pred.shape},{noise_pred[0,0,:3,:3]}") - print(" leaving model_fn") - print(" ") return noise_pred # 8. Run k-diffusion solver @@ -680,8 +573,7 @@ def model_fn(x, t): min_sigma, max_sigma = sigmas[sigmas > 0].min(), sigmas.max() noise_sampler = BrownianTreeNoiseSampler(latents, min_sigma, max_sigma, noise_sampler_seed) sampler_kwargs["noise_sampler"] = noise_sampler - - print(f" sigmas: {sigmas}") + latents = self.sampler(model_fn, latents, sigmas, **sampler_kwargs) if not output_type == "latent": diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py index d39417859da8..d1e66f8404af 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py @@ -258,14 +258,12 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc # This is critical for cosine (squaredcos_cap_v2) noise schedule. clipped_idx = torch.searchsorted(torch.flip(self.lambda_t, [0]), self.config.lambda_min_clipped) last_timestep = ((self.config.num_train_timesteps - clipped_idx).numpy()).item() - print(f" - last_timestep: {last_timestep}") # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891 if self.config.timestep_spacing == "linspace": timesteps = ( np.linspace(0, last_timestep - 1, num_inference_steps, dtype=float).round()[::-1].copy() ) - print(f" - timesteps: {len(timesteps)},{timesteps[0]}:{timesteps[-1]}") elif self.config.timestep_spacing == "leading": step_ratio = last_timestep // self.num_inference_steps # creates integer timesteps by multiplying by ratio @@ -286,20 +284,14 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) log_sigmas = np.log(sigmas) sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) - print(f" sigmas: {sigmas[0]}: {sigmas[-1]}") if self.config.use_karras_sigmas: - print(f" sigmas (k): {sigmas[0]}, {sigmas[-1]}") sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) self.sigmas = torch.from_numpy(sigmas).to(device=device) - print(" set_timesteps") - print(f" - sigmas: {sigmas[0]}: {sigmas[-1]}") - print(sigmas) - self.timesteps = torch.from_numpy(timesteps).to(device) self.num_inference_steps = len(timesteps) @@ -308,9 +300,6 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc None, ] * self.config.solver_order self.lower_order_nums = 0 - print(f" - timesteps: {timesteps}") - print(f" - model_outputs: {self.model_outputs}") - print(f" - lower_order_nums: {self.lower_order_nums}") # add an index counter for schedulers that allow duplicated timesteps self._step_index = None @@ -412,17 +401,13 @@ def convert_model_output( """ # DPM-Solver++ needs to solve an integral of the data prediction model. - print(f" -algo_type: {self.config.algorithm_type}, pred_type: {self.config.prediction_type}, variance_type: {self.config.variance_type}") if self.config.algorithm_type in ["dpmsolver++", "sde-dpmsolver++"]: if self.config.prediction_type == "epsilon": # DPM-Solver and DPM-Solver++ only need the "mean" output. if self.config.variance_type in ["learned", "learned_range"]: model_output = model_output[:, :3] - # yiyi update/testing here: sigma = self.sigmas[self.step_index] x0_pred = sample - sigma * model_output - print(f" - sigma: {sigma}") - print(f" - x0_pred: {x0_pred.shape},{x0_pred[0,0,:3,:3]}") elif self.config.prediction_type == "sample": x0_pred = model_output elif self.config.prediction_type == "v_prediction": @@ -435,7 +420,6 @@ def convert_model_output( ) if self.config.thresholding: - print(f" threhold? ") x0_pred = self._threshold_sample(x0_pred) return x0_pred @@ -494,15 +478,10 @@ def dpm_solver_first_order_update( def t_fn(_sigma): return -torch.log(_sigma) - print(f" - sample: {sample[0,0,:3,:3]}") - sigma_t, sigma_s = self.sigmas[self.step_index +1], self.sigmas[self.step_index] - print(f" - sigma_t: {sigma_t}, sigma_s: {sigma_s}") h = t_fn(sigma_t) - t_fn(sigma_s) - print(f" - h: {h}") if self.config.algorithm_type == "dpmsolver++": x_t = (sigma_t / sigma_s) * sample - (torch.exp(-h) - 1.0) * model_output - print(f" -> prev_sample: {x_t[0,0,:3,:3]}") elif self.config.algorithm_type == "dpmsolver": x_t = (alpha_t / alpha_s) * sample - (sigma_t * (torch.exp(h) - 1.0)) * model_output elif self.config.algorithm_type == "sde-dpmsolver++": @@ -553,22 +532,15 @@ def t_fn(_sigma): h, h_0 = t_fn(sigma_t) - t_fn(sigma_s0), t_fn(sigma_s0) - t_fn(sigma_s1) r0 = h_0 / h D0, D1 = m0, (1.0 / r0) * (m0 - m1) - print(f" sigma_t (current), sigma_s0 (prev), sigma_s1 (next): {sigma_t}, {sigma_s0}, {sigma_s1}") - # print(f" sigma_t, sigma_s0, {sigma_t},{sigma_s0}") - # print(f" sigma_s1: {self.sigma_t[s1]}") - # print(f" alpha_bar = {self.alphas_cumprod[s1]}") - print(f" sample: {sample[0,0,:3,:3]}") - # print(" -----") + if self.config.algorithm_type == "dpmsolver++": # See https://arxiv.org/abs/2211.01095 for detailed derivations if self.config.solver_type == "midpoint": - print(f" dpmsolver++ /midpoint ") x_t = ( (sigma_t / sigma_s0) * sample - (torch.exp(-h) - 1.0) * D0 - 0.5 * (torch.exp(-h) - 1.0) * D1 ) - print(f" -> prev_sample: {x_t.shape}, {x_t[0,0,:3,:3]}") elif self.config.solver_type == "heun": x_t = ( (sigma_t / sigma_s0) * sample @@ -725,24 +697,15 @@ def step( if self.step_index is None: self._init_step_index(timestep) - print(" ") - print(" ** inside step ***** ") - - print(f" - timestep, step_index: {timestep}, {self.step_index}") - print(f" - sample/latents: {sample.shape},{sample[0,0,:3,:3]}") - print(f" - model_output (noise): {model_output.shape}, {model_output[0,0,:3,:3]}") prev_timestep = 0 if self.step_index == len(self.timesteps) - 1 else self.timesteps[self.step_index + 1] lower_order_final = (self.step_index == len(self.timesteps) - 1) lower_order_second = ( (self.step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15 ) - print(f" - prev_timestep: {prev_timestep}") model_output = self.convert_model_output(model_output, timestep, sample) - print(f" - model_output (x0): {model_output.shape}, {model_output[0,0,:3,:3]}") for i in range(self.config.solver_order - 1): - print(f" move outputs {i+1} -> {i}") self.model_outputs[i] = self.model_outputs[i + 1] self.model_outputs[-1] = model_output @@ -754,18 +717,15 @@ def step( noise = None if self.config.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final: - print(f" 1st order update: {self.config.solver_order}, {self.lower_order_nums}, {lower_order_final}") prev_sample = self.dpm_solver_first_order_update( model_output, timestep, prev_timestep, sample, noise=noise ) elif self.config.solver_order == 2 or self.lower_order_nums < 2 or lower_order_second: - print(f" 2nd order update") timestep_list = [self.timesteps[self.step_index - 1], timestep] prev_sample = self.multistep_dpm_solver_second_order_update( self.model_outputs, timestep_list, prev_timestep, sample, noise=noise ) else: - print(f" 3rd order update") timestep_list = [self.timesteps[step_index - 2], self.timesteps[step_index - 1], timestep] prev_sample = self.multistep_dpm_solver_third_order_update( self.model_outputs, timestep_list, prev_timestep, sample @@ -773,7 +733,6 @@ def step( if self.lower_order_nums < self.config.solver_order: self.lower_order_nums += 1 - print(f" lower_order_nums +1 -> {self.lower_order_nums}") # upon completion increase step index by one self._step_index += 1 diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index 9dbe26e124f6..cb126d4b953c 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -243,13 +243,9 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic ) sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) - print(f" -sigmas: {sigmas[0]}, {sigmas[-1]}") log_sigmas = np.log(sigmas) if self.config.interpolation_type == "linear": - print(f" - timesteps: {len(timesteps)}") - print(timesteps) - print(f" - sigmas: {len(sigmas)}") sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) elif self.config.interpolation_type == "log_linear": sigmas = torch.linspace(np.log(sigmas[-1]), np.log(sigmas[0]), num_inference_steps + 1).exp() @@ -260,12 +256,8 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic ) if self.use_karras_sigmas: - print(" use k_sigmas") - print(f" -sigmas: {sigmas[0]}, {sigmas[-1]}") sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=self.num_inference_steps) timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - print(f" -sigmas: {sigmas}") - print(f" -timesteps: {timesteps}") sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) self.sigmas = torch.from_numpy(sigmas).to(device=device) From c95b545113b9b7db2f26a51f0f9ca8213f157b24 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Mon, 21 Aug 2023 03:01:34 +0000 Subject: [PATCH 5/8] style --- .../pipeline_stable_diffusion_k_diffusion.py | 2 +- .../scheduling_dpmsolver_multistep.py | 51 ++++++++++--------- 2 files changed, 29 insertions(+), 24 deletions(-) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_k_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_k_diffusion.py index e272e9e3a505..29a57470a341 100755 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_k_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_k_diffusion.py @@ -126,7 +126,7 @@ def set_scheduler(self, scheduler_type: str): library = importlib.import_module("k_diffusion") sampling = getattr(library, "sampling") self.sampler = getattr(sampling, scheduler_type) - + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_model_cpu_offload def enable_model_cpu_offload(self, gpu_id=0): r""" diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py index d1e66f8404af..82b77353223a 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py @@ -15,7 +15,6 @@ # DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver import math -from collections import defaultdict from typing import List, Optional, Tuple, Union import numpy as np @@ -222,12 +221,12 @@ def __init__( # setable values self.num_inference_steps = None timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32)[::-1].copy() - + self.timesteps = torch.from_numpy(timesteps) self.model_outputs = [None] * solver_order self.lower_order_nums = 0 self._step_index = None - + @property def init_noise_sigma(self): # standard deviation of the initial noise distribution @@ -235,7 +234,7 @@ def init_noise_sigma(self): return self.sigmas.max() return (self.sigmas.max() ** 2 + 1) ** 0.5 - + @property def step_index(self): """ @@ -243,7 +242,6 @@ def step_index(self): """ return self._step_index - def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torch.device] = None): """ Sets the timesteps used for the diffusion chain. Supporting function to be run before inference. @@ -261,20 +259,18 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891 if self.config.timestep_spacing == "linspace": - timesteps = ( - np.linspace(0, last_timestep - 1, num_inference_steps, dtype=float).round()[::-1].copy() - ) + timesteps = np.linspace(0, last_timestep - 1, num_inference_steps).round()[::-1].copy().astype(np.float32) elif self.config.timestep_spacing == "leading": step_ratio = last_timestep // self.num_inference_steps # creates integer timesteps by multiplying by ratio # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(float) + timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.float32) timesteps += self.config.steps_offset elif self.config.timestep_spacing == "trailing": step_ratio = self.config.num_train_timesteps / num_inference_steps # creates integer timesteps by multiplying by ratio # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = np.arange(last_timestep, 0, -step_ratio).round().copy().astype(float) + timesteps = np.arange(last_timestep, 0, -step_ratio).round().copy().astype(np.float32) timesteps -= 1 else: raise ValueError( @@ -288,10 +284,11 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc if self.config.use_karras_sigmas: sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - + sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) self.sigmas = torch.from_numpy(sigmas).to(device=device) + timesteps = timesteps.astype(np.int32) self.timesteps = torch.from_numpy(timesteps).to(device) self.num_inference_steps = len(timesteps) @@ -475,10 +472,14 @@ def dpm_solver_first_order_update( Returns: `torch.FloatTensor`: the sample tensor at the previous timestep. """ + def t_fn(_sigma): return -torch.log(_sigma) - sigma_t, sigma_s = self.sigmas[self.step_index +1], self.sigmas[self.step_index] + # YiYi notes: keep these for now so don't get an error, don't need once fully refactored + alpha_t, alpha_s = self.alpha_t[prev_timestep], self.alpha_t[timestep] + + sigma_t, sigma_s = self.sigmas[self.step_index + 1], self.sigmas[self.step_index] h = t_fn(sigma_t) - t_fn(sigma_s) if self.config.algorithm_type == "dpmsolver++": x_t = (sigma_t / sigma_s) * sample - (torch.exp(-h) - 1.0) * model_output @@ -525,8 +526,16 @@ def multistep_dpm_solver_second_order_update( def t_fn(_sigma): return -torch.log(_sigma) - - sigma_t, sigma_s0, sigma_s1 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index], self.sigmas[self.step_index - 1] + + # YiYi notes: keep these for now so don't get an error, not needed once fully refactored + t, s0 = prev_timestep, timestep_list[-1] + alpha_t, alpha_s0 = self.alpha_t[t], self.alpha_t[s0] + + sigma_t, sigma_s0, sigma_s1 = ( + self.sigmas[self.step_index + 1], + self.sigmas[self.step_index], + self.sigmas[self.step_index - 1], + ) m0, m1 = model_output_list[-1], model_output_list[-2] h, h_0 = t_fn(sigma_t) - t_fn(sigma_s0), t_fn(sigma_s0) - t_fn(sigma_s1) @@ -536,11 +545,7 @@ def t_fn(_sigma): if self.config.algorithm_type == "dpmsolver++": # See https://arxiv.org/abs/2211.01095 for detailed derivations if self.config.solver_type == "midpoint": - x_t = ( - (sigma_t / sigma_s0) * sample - - (torch.exp(-h) - 1.0) * D0 - - 0.5 * (torch.exp(-h) - 1.0) * D1 - ) + x_t = (sigma_t / sigma_s0) * sample - (torch.exp(-h) - 1.0) * D0 - 0.5 * (torch.exp(-h) - 1.0) * D1 elif self.config.solver_type == "heun": x_t = ( (sigma_t / sigma_s0) * sample @@ -699,7 +704,7 @@ def step( self._init_step_index(timestep) prev_timestep = 0 if self.step_index == len(self.timesteps) - 1 else self.timesteps[self.step_index + 1] - lower_order_final = (self.step_index == len(self.timesteps) - 1) + lower_order_final = self.step_index == len(self.timesteps) - 1 lower_order_second = ( (self.step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15 ) @@ -726,14 +731,14 @@ def step( self.model_outputs, timestep_list, prev_timestep, sample, noise=noise ) else: - timestep_list = [self.timesteps[step_index - 2], self.timesteps[step_index - 1], timestep] + timestep_list = [self.timesteps[self.step_index - 2], self.timesteps[self.step_index - 1], timestep] prev_sample = self.multistep_dpm_solver_third_order_update( self.model_outputs, timestep_list, prev_timestep, sample ) if self.lower_order_nums < self.config.solver_order: self.lower_order_nums += 1 - + # upon completion increase step index by one self._step_index += 1 @@ -741,7 +746,7 @@ def step( return (prev_sample,) return SchedulerOutput(prev_sample=prev_sample) - + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.scale_model_input def scale_model_input( self, sample: torch.FloatTensor, timestep: Union[float, torch.FloatTensor] From 515c1050409c5a383174a5c58ee691b36d9b20c5 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 22 Aug 2023 00:29:38 +0000 Subject: [PATCH 6/8] sde-dpmsolver+++ --- .../pipeline_stable_diffusion.py | 2 ++ .../scheduling_dpmsolver_multistep.py | 21 +++++++++---------- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index fa2a6715dea2..9bc2ad57fdcc 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -681,10 +681,12 @@ def __call__( cross_attention_kwargs=cross_attention_kwargs, return_dict=False, )[0] + # perform guidance if do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + if do_classifier_free_guidance and guidance_rescale > 0.0: # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py index 1c42dfd8c381..cd244956f68c 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py @@ -264,7 +264,6 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) self.sigmas = torch.from_numpy(sigmas).to(device=device) - timesteps = timesteps.astype(np.int32) self.timesteps = torch.from_numpy(timesteps).to(device) self.num_inference_steps = len(timesteps) @@ -460,7 +459,7 @@ def t_fn(_sigma): return -torch.log(_sigma) # YiYi notes: keep these for now so don't get an error, don't need once fully refactored - alpha_t, alpha_s = self.alpha_t[prev_timestep], self.alpha_t[timestep] + #alpha_t, alpha_s = self.alpha_t[prev_timestep], self.alpha_t[timestep] sigma_t, sigma_s = self.sigmas[self.step_index + 1], self.sigmas[self.step_index] h = t_fn(sigma_t) - t_fn(sigma_s) @@ -472,7 +471,7 @@ def t_fn(_sigma): assert noise is not None x_t = ( (sigma_t / sigma_s * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * model_output + + (1 - torch.exp(-2.0 * h)) * model_output + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise ) elif self.config.algorithm_type == "sde-dpmsolver": @@ -514,8 +513,8 @@ def t_fn(_sigma): return -torch.log(_sigma) # YiYi notes: keep these for now so don't get an error, not needed once fully refactored - t, s0 = prev_timestep, timestep_list[-1] - alpha_t, alpha_s0 = self.alpha_t[t], self.alpha_t[s0] + #t, s0 = prev_timestep, timestep_list[-1] + #alpha_t, alpha_s0 = self.alpha_t[t], self.alpha_t[s0] sigma_t, sigma_s0, sigma_s1 = ( self.sigmas[self.step_index + 1], @@ -535,8 +534,8 @@ def t_fn(_sigma): elif self.config.solver_type == "heun": x_t = ( (sigma_t / sigma_s0) * sample - - (alpha_t * (torch.exp(-h) - 1.0)) * D0 - + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1 + - (torch.exp(-h) - 1.0) * D0 + + ((torch.exp(-h) - 1.0) / h + 1.0) * D1 ) elif self.config.algorithm_type == "dpmsolver": # See https://arxiv.org/abs/2206.00927 for detailed derivations @@ -557,15 +556,15 @@ def t_fn(_sigma): if self.config.solver_type == "midpoint": x_t = ( (sigma_t / sigma_s0 * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 - + 0.5 * (alpha_t * (1 - torch.exp(-2.0 * h))) * D1 + + (1 - torch.exp(-2.0 * h)) * D0 + + 0.5 * (1 - torch.exp(-2.0 * h)) * D1 + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise ) elif self.config.solver_type == "heun": x_t = ( (sigma_t / sigma_s0 * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 - + (alpha_t * ((1.0 - torch.exp(-2.0 * h)) / (-2.0 * h) + 1.0)) * D1 + + (1 - torch.exp(-2.0 * h)) * D0 + + ((1.0 - torch.exp(-2.0 * h)) / (-2.0 * h) + 1.0) * D1 + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise ) elif self.config.algorithm_type == "sde-dpmsolver": From a85a18c0ce5c08e56ae6f1f7ab642ccab8c78676 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 22 Aug 2023 02:14:50 +0000 Subject: [PATCH 7/8] v_prediction --- src/diffusers/schedulers/scheduling_dpmsolver_multistep.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py index cd244956f68c..008331fab621 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py @@ -388,8 +388,8 @@ def convert_model_output( elif self.config.prediction_type == "sample": x0_pred = model_output elif self.config.prediction_type == "v_prediction": - alpha_t, sigma_t = self.alpha_t[timestep], self.sigma_t[timestep] - x0_pred = alpha_t * sample - sigma_t * model_output + sigma = self.sigmas[self.step_index] + x0_pred = model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1)) else: raise ValueError( f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" From f238e0d86280a0db5c4ae29ac6dfd49f5a81ef83 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Mon, 11 Sep 2023 01:36:54 +0000 Subject: [PATCH 8/8] remove round --- src/diffusers/schedulers/scheduling_dpmsolver_multistep.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py index 008331fab621..5e6869cb0f0c 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py @@ -235,7 +235,7 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891 if self.config.timestep_spacing == "linspace": - timesteps = np.linspace(0, last_timestep - 1, num_inference_steps).round()[::-1].copy().astype(np.float32) + timesteps = np.linspace(0, last_timestep - 1, num_inference_steps)[::-1].copy().astype(np.float32) elif self.config.timestep_spacing == "leading": step_ratio = last_timestep // self.num_inference_steps # creates integer timesteps by multiplying by ratio