论文阅读笔记：Denoising Diffusion Implicit Models （4）

0、快速访问

论文阅读笔记：Denoising Diffusion Implicit Models （1）
论文阅读笔记：Denoising Diffusion Implicit Models （2）
论文阅读笔记：Denoising Diffusion Implicit Models （3）
论文阅读笔记：Denoising Diffusion Implicit Models （4）

4、接上文[论文阅读笔记：论文阅读笔记：Denoising Diffusion Implicit Models （3）

1. 已经知道跳$1$步时，$q_{\sigma }(x_{t-1}|x_t,x_0)$的分布满足公式（·）
   $$\begin{array}{c}\begin{array}{rl}{x}_{t-1}& =\sqrt{{\alpha }_{t-1}}\cdot \frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}+\sqrt{1-{\alpha }_{t-1}}\cdot z_t\end{array}\end{array}$$
2. 假设跳$n$步时，$q_{\sigma }(x_{t-n}|x_t,x_0)$的分布满足公式（2）
   $$\begin{array}{c}\begin{array}{rl}{x}_{t-n}& =\sqrt{{\alpha }_{t-n}}\cdot \frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}+\sqrt{1-{\alpha }_{t-n}}\cdot z_t\end{array}\end{array}$$
3. 证明：当跳$n+1$步时，分布$q_{\sigma }(x_{t-n-1}|x_t,x_0)$满足$x_{t-n-1}=\sqrt{{\alpha }_{t-n-1}}\cdot \frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}+\sqrt{1-{\alpha }_{t-n-1}}\cdot z_t$。
   由于$q_{\sigma }(x_{t-n-1}|x_t,x_0)$是$q_{\sigma }(x_{t-n-1},x_{t-n}|x_t,x_0)$的边缘分布，因此有
   $$\begin{array}{c}\begin{array}{rl}{q}_{\sigma }(x_{t-n-1}|x_t,x_0)& =\int q_{\sigma }(x_{t-n-1},x_{t-n}|x_t,x_0)\cdot d{x}_{t-n}\\ & =\int q_{\sigma }(x_{t-n-1},|x_{t-n},x_0)\cdot q_{\sigma }(x_{t-n}|x_t,x_0)\cdot d{x}_{t-n}\end{array}\end{array}$$
   $q_{\sigma }(x_{t-n-1},|x_{t-n},x_0)=N(x_{t-n-1}|\sqrt{\frac{1-{\alpha }_{t-n-1}}{1-{\alpha }_{t-n}}}\cdot x_{t-n}+[\sqrt{{\alpha }_{t-n-1}}-\frac{\sqrt{{\alpha }_{t-n}\cdot (1-{\alpha }_{t-n-1}})}{\sqrt{1-{\alpha }_{t-n}}}]\cdot x_0,0)$
   因此，分布$q_{\sigma }(x_{t-n-1}|x_t,x_0)$的均值${\mu }_{t-n-1}$如公式（4）所示。
   $$\begin{array}{c}\begin{array}{rl}{\mu }_{t-n-1}& =E(q_{\sigma }(x_{t-n-1}|x_t,x_0))\\ & =\int x_{t-n-1}\cdot q_{\sigma }(x_{t-n-1}|x_t,x_0)\cdot d{x}_{t-n-1}\\ & =\int x_{t-n-1}\cdot (\int q_{\sigma }(x_{t-n-1},|x_{t-n},x_0)\cdot q_{\sigma }(x_{t-n}|x_t,x_0)\cdot d{x}_{t-n})\cdot d{x}_{t-n-1}\\ & =\int \int x_{t-n-1}\cdot q_{\sigma }(x_{t-n-1},|x_{t-n},x_0)\cdot q_{\sigma }(x_{t-n}|x_t,x_0)\cdot d{x}_{t-n}\cdot d{x}_{t-n-1}\\ & =\int (\int x_{t-n-1}\cdot q_{\sigma }(x_{t-n-1},|x_{t-n},x_0)\cdot d{x}_{t-n-1})\cdot q_{\sigma }(x_{t-n}|x_t,x_0)\cdot d{x}_{t-n}\\ & =\int E(q_{\sigma }(x_{t-n-1},|x_{t-n},x_0))\cdot q_{\sigma }(x_{t-n}|x_t,x_0)\cdot d{x}_{t-n}\\ & =\int (\sqrt{\frac{1-{\alpha }_{t-n-1}}{1-{\alpha }_{t-n}}}\cdot x_{t-n}+[\sqrt{{\alpha }_{t-n-1}}-\frac{\sqrt{{\alpha }_{t-n}\cdot (1-{\alpha }_{t-n-1}})}{\sqrt{1-{\alpha }_{t-n}}}]\cdot x_0)\cdot q_{\sigma }(x_{t-n}|x_t,x_0)\cdot d{x}_{t-n}\\ & =\int \sqrt{\frac{1-{\alpha }_{t-n-1}}{1-{\alpha }_{t-n}}}\cdot x_{t-n}\cdot q_{\sigma }(x_{t-n}|x_t,x_0)\cdot d{x}_{t-n}+\int ([\sqrt{{\alpha }_{t-n-1}}-\frac{\sqrt{{\alpha }_{t-1}\cdot (1-{\alpha }_{t-n-1}})}{\sqrt{1-{\alpha }_{t-n}}}]\cdot x_0)\cdot q_{\sigma }(x_{t-n}|x_t,x_0)\cdot d{x}_{t-n}\\ & =\sqrt{\frac{1-{\alpha }_{t-n-1}}{1-{\alpha }_{t-n}}}\cdot \int q_{\sigma }(x_{t-n}|x_t,x_0)\cdot d{x}_{t-n}+[\sqrt{{\alpha }_{t-n-1}}-\frac{\sqrt{{\alpha }_{t-n}\cdot (1-{\alpha }_{t-n-1}})}{\sqrt{1-{\alpha }_{t-n}}}]\cdot x_0\\ & =\sqrt{\frac{1-{\alpha }_{t-n-1}}{1-{\alpha }_{t-n}}}\cdot E(q_{\sigma }(x_{t-n}|x_t,x_0))+[\sqrt{{\alpha }_{t-n-1}}-\frac{\sqrt{{\alpha }_{t-n}\cdot (1-{\alpha }_{t-n-1}})}{\sqrt{1-{\alpha }_{t-n}}}]\cdot \underset{x_0=\frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}}{\underbrace{x_0}}\\ & =\sqrt{\frac{1-{\alpha }_{t-n-1}}{1-{\alpha }_{t-n}}}\cdot (\sqrt{{\alpha }_{t-n}}\cdot \frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}+\sqrt{1-{\alpha }_{t-n}}\cdot z_t)+[\sqrt{{\alpha }_{t-n-1}}-\frac{\sqrt{{\alpha }_{t-n}\cdot (1-{\alpha }_{t-n-1}})}{\sqrt{1-{\alpha }_{t-n}}}]\cdot \frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}\\ & =\bcancel{\sqrt{\frac{1-{\alpha }_{t-n-1}}{1-{\alpha }_{t-n}}}\cdot \sqrt{{\alpha }_{t-n}}\cdot \frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}}+\sqrt{\frac{1-{\alpha }_{t-n-1}}{\bcancel{1-{\alpha }_{t-n}}}}\cdot \bcancel{\sqrt{1-{\alpha }_{t-n}}}\cdot z_t+\sqrt{{\alpha }_{t-n-1}}\cdot \frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}-\bcancel{\frac{\sqrt{{\alpha }_{t-n}\cdot (1-{\alpha }_{t-n-1}})}{\sqrt{1-{\alpha }_{t-n}}}\cdot \frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}}\\ & =\sqrt{{\alpha }_{t-n-1}}\cdot \frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}+\sqrt{1-{\alpha }_{t-n-1}}\cdot z_t\end{array}\end{array}$$
   证毕。综上所述，跳$n$步的公式为$$q_{\sigma }(x_{t-n}|x_t,x_0)=\sqrt{{\alpha }_{t-n}}\cdot \frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}+\sqrt{1-{\alpha }_{t-n}}\cdot z_t$$
   基于DDIM的多数论文，例如暗图像增强方法LightenDiffusion等，也都是令${\sigma }_t=0$。论文和代码中使用的跳$n$步的采样过程如公式（5）所示。
   $$\begin{array}{c}\begin{array}{rl}{x}_{t-n}& =\sqrt{{\alpha }_{t-n}}\cdot \underset{预测出z_t,进而计算出x_0}{\underbrace{\frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}}}+\sqrt{1-{\alpha }_{t-n}-{\sigma }_t^2}\cdot z_t+{\sigma }_t^2\underset{标准高斯分布}{\underbrace{{ϵ}_t}}\end{array}\end{array}$$

这里使用中的${\sigma }_t$是可以自己定义的量。有两种特殊的情况：
1、${\sigma }_t^2=0$：此时，
$x_{t-1}$满足公式（3）
$$\begin{array}{c}\begin{array}{rl}{x}_{t-1}& =\sqrt{{\alpha }_{t-1}}\cdot \frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}+\sqrt{1-{\alpha }_{t-1}-{\sigma }_t^2}\cdot z_t+{\sigma }_t^2{ϵ}_t\\ & =\sqrt{{\alpha }_{t-1}}\cdot x_0+\sqrt{1-{\alpha }_{t-1}}\cdot z_t\end{array}\end{array}$$
$x_{t-n}$满足
$$\begin{array}{c}\begin{array}{rl}{x}_{t-n}& =\sqrt{{\alpha }_{t-n}}\cdot \frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}+\sqrt{1-{\alpha }_{t-n}-{\sigma }_t^2}\cdot z_t+{\sigma }_t^2{ϵ}_t\\ & =\sqrt{{\alpha }_{t-n}}\cdot x_0+\sqrt{1-{\alpha }_{t-n}}\cdot z_t\end{array}\end{array}$$
可以看出，此时，$x_{t-1}$和$x_{t-n}$退化成上文论文阅读笔记：Denoising Diffusion Implicit Models （2）中的Lemma 1.

2、${\sigma }_t^2=\frac{1-{\alpha }_{t-1}}{1-{\alpha }_t}\cdot (1-\frac{{\alpha }_t}{{\alpha }_{t-1}})$：此时，$x_{t-1}$满足公式（4）
$$\begin{array}{c}\begin{array}{rl}{x}_{t-1}& =\sqrt{{\alpha }_{t-1}}\cdot \frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}+\sqrt{1-{\alpha }_{t-1}-{\sigma }_t^2}\cdot z_t+{\sigma }_t^2{ϵ}_t\\ & =\sqrt{{\alpha }_{t-1}}\cdot \frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}+\sqrt{1-{\alpha }_{t-1}-\frac{1-{\alpha }_{t-1}}{1-{\alpha }_t}\cdot (1-\frac{{\alpha }_t}{{\alpha }_{t-1}})}\cdot z_t+{\sigma }_t^2{ϵ}_t\\ & =\sqrt{{\alpha }_{t-1}}\cdot \frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}+\sqrt{(1-{\alpha }_{t-1})(1-\frac{1}{1-{\alpha }_t}\cdot \frac{{\alpha }_{t-1}-{\alpha }_t}{{\alpha }_{t-1}})}\cdot z_t+{\sigma }_t^2{ϵ}_t\\ & =\sqrt{{\alpha }_{t-1}}\cdot \frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}+\sqrt{(1-{\alpha }_{t-1})\frac{{\alpha }_{t-1}-{\alpha }_{t-1}\cdot {\alpha }_t-{\alpha }_{t-1}+{\alpha }_t}{{\alpha }_{t-1}\cdot (1-{\alpha }_t)}}\cdot z_t+{\sigma }_t^2{ϵ}_t\\ & =\sqrt{{\alpha }_{t-1}}\cdot \frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}+\sqrt{(1-{\alpha }_{t-1})\frac{-{\alpha }_{t-1}\cdot {\alpha }_t+{\alpha }_t}{{\alpha }_{t-1}\cdot (1-{\alpha }_t)}}\cdot z_t+{\sigma }_t^2{ϵ}_t\\ & =\sqrt{{\alpha }_{t-1}}\cdot \frac{x_t-\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}+(1-{\alpha }_{t-1})\sqrt{\frac{{\alpha }_t}{{\alpha }_{t-1}\cdot (1-{\alpha }_t)}}\cdot z_t+{\sigma }_t^2{ϵ}_t\\ & =\sqrt{{\alpha }_{t-1}}\cdot \frac{x_t}{\sqrt{{\alpha }_t}}-\frac{\sqrt{{\alpha }_{t-1}}\sqrt{1-{\alpha }_t}\cdot z_t}{\sqrt{{\alpha }_t}}+(1-{\alpha }_{t-1})\sqrt{\frac{{\alpha }_t}{{\alpha }_{t-1}\cdot (1-{\alpha }_t)}}\cdot z_t+{\sigma }_t^2{ϵ}_t\\ & =\sqrt{{\alpha }_{t-1}}\cdot \frac{x_t}{\sqrt{{\alpha }_t}}-\left(\frac{\sqrt{{\alpha }_{t-1}}\sqrt{1-{\alpha }_t}\cdot \sqrt{{\alpha }_{t-1}}\sqrt{1-{\alpha }_t}-(1-{\alpha }_{t-1})\cdot \sqrt{{\alpha }_t}\cdot \sqrt{{\alpha }_t}}{\sqrt{{\alpha }_t}\cdot \sqrt{{\alpha }_{t-1}\cdot (1-{\alpha }_t)}}\right)\cdot z_t+{\sigma }_t^2{ϵ}_t\\ & =\sqrt{{\alpha }_{t-1}}\cdot \frac{x_t}{\sqrt{{\alpha }_t}}-\left(\frac{\sqrt{{\alpha }_{t-1}}\sqrt{1-{\alpha }_t}\cdot \sqrt{{\alpha }_{t-1}}\sqrt{1-{\alpha }_t}-(1-{\alpha }_{t-1})\cdot \sqrt{{\alpha }_t}\cdot \sqrt{{\alpha }_t}}{\sqrt{{\alpha }_t}\cdot \sqrt{{\alpha }_{t-1}\cdot (1-{\alpha }_t)}}\right)\cdot z_t+{\sigma }_t^2{ϵ}_t\\ & =\sqrt{{\alpha }_{t-1}}\cdot \frac{x_t}{\sqrt{{\alpha }_t}}-\left(\frac{{\alpha }_{t-1}\cdot (1-{\alpha }_t)-(1-{\alpha }_{t-1})\cdot {\alpha }_t}{\sqrt{{\alpha }_t}\cdot \sqrt{{\alpha }_{t-1}\cdot (1-{\alpha }_t)}}\right)\cdot z_t+{\sigma }_t^2{ϵ}_t\\ & =\sqrt{{\alpha }_{t-1}}\cdot \frac{x_t}{\sqrt{{\alpha }_t}}-\left(\frac{{\alpha }_{t-1}-\bcancel{{\alpha }_t\cdot {\alpha }_{t-1}}-{\alpha }_t+\bcancel{{\alpha }_t\cdot {\alpha }_{t-1}}}{\sqrt{{\alpha }_t}\cdot \sqrt{{\alpha }_{t-1}\cdot (1-{\alpha }_t)}}\right)\cdot z_t\\ & =\sqrt{{\alpha }_{t-1}}\cdot \frac{x_t}{\sqrt{{\alpha }_t}}-\left(\frac{{\alpha }_{t-1}-{\alpha }_t}{\sqrt{{\alpha }_t}\cdot \sqrt{{\alpha }_{t-1}\cdot (1-{\alpha }_t)}}\right)\cdot z_t+{\sigma }_t^2{ϵ}_t\\ & =\sqrt{{\alpha }_{t-1}}\cdot \frac{x_t}{\sqrt{{\alpha }_t}}-\left(\frac{\sqrt{{\alpha }_{t-1}}\cdot ({\alpha }_{t-1}-{\alpha }_t)}{{\alpha }_{t-1}\cdot \sqrt{{\alpha }_t}\cdot \sqrt{(1-{\alpha }_t)}}\right)\cdot z_t+{\sigma }_t^2{ϵ}_t\\ & =\frac{\sqrt{{\alpha }_{t-1}}}{\sqrt{{\alpha }_t}}(x_t-\frac{{\alpha }_{t-1}-{\alpha }_t}{{\alpha }_{t-1}\cdot \sqrt{1-{\alpha }_t}})+{\sigma }_t^2{ϵ}_t\\ & =\frac{\sqrt{{\alpha }_{t-1}}}{\sqrt{{\alpha }_t}}(x_t-\frac{1}{\sqrt{1-{\alpha }_t}}\cdot (1-\frac{{\alpha }_t}{{\alpha }_{t-1}}))\cdot z_t+{\sigma }_t^2{ϵ}_t\\ & =\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}})\cdot z_t+{\sigma }_t^2{ϵ}_t（换成DDPM中的符号）\end{array}\end{array}$$
可以看出，此时，DDIM退化成了DDPM。
论文讨论了${\sigma }_t^2$选取$\eta \cdot \frac{1-{\alpha }_{t-1}}{1-{\alpha }_t}\cdot (1-\frac{{\alpha }_t}{{\alpha }_{t-1}}),\eta \in [0,1]$，即在0和DDPM之间变化时。不同$\eta$以及跳不同步时所对应的表现，如下图所示。

5、代码

class DDIMPipeline(DiffusionPipeline):model_cpu_offload_seq = "unet"def __init__(self, unet, scheduler):super().__init__()# make sure scheduler can always be converted to DDIMscheduler = DDIMScheduler.from_config(scheduler.config)self.register_modules(unet=unet, scheduler=scheduler)@torch.no_grad()def __call__(self,batch_size: int = 1,generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,eta: float = 0.0,num_inference_steps: int = 50,use_clipped_model_output: Optional[bool] = None,output_type: Optional[str] = "pil",return_dict: bool = True,) -> Union[ImagePipelineOutput, Tuple]:# Sample gaussian noise to begin loopif isinstance(self.unet.config.sample_size, int):image_shape = (batch_size,self.unet.config.in_channels,self.unet.config.sample_size,self.unet.config.sample_size,)else:image_shape = (batch_size, self.unet.config.in_channels, *self.unet.config.sample_size)if isinstance(generator, list) and len(generator) != batch_size:raise ValueError(f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"f" size of {batch_size}. Make sure the batch size matches the length of the generators.")# 随即生成噪音image = randn_tensor(image_shape, generator=generator, device=self._execution_device, dtype=self.unet.dtype)# 设置步数间隔。例如num_inference_steps = 50,然而总步长为1000,那么就是每次跳20步，例如在当前时刻， timestep=980, prev_timestep=960self.scheduler.set_timesteps(num_inference_steps)for t in self.progress_bar(self.scheduler.timesteps):# 1. 预测出timestep=980时刻对应噪音model_output = self.unet(image, t).sample# 2. 调用scheduler的方法step，执行公式（）得到prev_timestep=960时刻的图像image = self.scheduler.step(model_output, t, image, eta=eta, use_clipped_model_output=use_clipped_model_output, generator=generator).prev_sampleimage = (image / 2 + 0.5).clamp(0, 1)image = image.cpu().permute(0, 2, 3, 1).numpy()if output_type == "pil":image = self.numpy_to_pil(image)if not return_dict:return (image,)return ImagePipelineOutput(images=image)class DDIMScheduler(SchedulerMixin, ConfigMixin):_compatibles = [e.name for e in KarrasDiffusionSchedulers]order = 1@register_to_configdef __init__(self,num_train_timesteps: int = 1000,beta_start: float = 0.0001,beta_end: float = 0.02,beta_schedule: str = "linear",trained_betas: Optional[Union[np.ndarray, List[float]]] = None,clip_sample: bool = True,set_alpha_to_one: bool = True,steps_offset: int = 0,prediction_type: str = "epsilon",thresholding: bool = False,dynamic_thresholding_ratio: float = 0.995,clip_sample_range: float = 1.0,sample_max_value: float = 1.0,timestep_spacing: str = "leading",rescale_betas_zero_snr: bool = False,):if trained_betas is not None:self.betas = torch.tensor(trained_betas, dtype=torch.float32)elif beta_schedule == "linear":self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)elif beta_schedule == "scaled_linear":# this schedule is very specific to the latent diffusion model.self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2elif beta_schedule == "squaredcos_cap_v2":# Glide cosine scheduleself.betas = betas_for_alpha_bar(num_train_timesteps)else:raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")# Rescale for zero SNRif rescale_betas_zero_snr:self.betas = rescale_zero_terminal_snr(self.betas)self.alphas = 1.0 - self.betasself.alphas_cumprod = torch.cumprod(self.alphas, dim=0)# At every step in ddim, we are looking into the previous alphas_cumprod# For the final step, there is no previous alphas_cumprod because we are already at 0# `set_alpha_to_one` decides whether we set this parameter simply to one or# whether we use the final alpha of the "non-previous" one.self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0]# standard deviation of the initial noise distributionself.init_noise_sigma = 1.0# setable valuesself.num_inference_steps = Noneself.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64))def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:"""Ensures interchangeability with schedulers that need to scale the denoising model input depending on thecurrent timestep.Args:sample (`torch.Tensor`):The input sample.timestep (`int`, *optional*):The current timestep in the diffusion chain.Returns:`torch.Tensor`:A scaled input sample."""return sampledef _get_variance(self, timestep, prev_timestep):alpha_prod_t = self.alphas_cumprod[timestep]alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprodbeta_prod_t = 1 - alpha_prod_tbeta_prod_t_prev = 1 - alpha_prod_t_prevvariance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)return variance# Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sampledef _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:""""Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (theprediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide bys. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventingpixels from saturation at each step. We find that dynamic thresholding results in significantly betterphotorealism as well as better image-text alignment, especially when using very large guidance weights."https://arxiv.org/abs/2205.11487"""dtype = sample.dtypebatch_size, channels, *remaining_dims = sample.shapeif dtype not in (torch.float32, torch.float64):sample = sample.float()  # upcast for quantile calculation, and clamp not implemented for cpu half# Flatten sample for doing quantile calculation along each imagesample = sample.reshape(batch_size, channels * np.prod(remaining_dims))abs_sample = sample.abs()  # "a certain percentile absolute pixel value"s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)s = torch.clamp(s, min=1, max=self.config.sample_max_value)  # When clamped to min=1, equivalent to standard clipping to [-1, 1]s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0sample = torch.clamp(sample, -s, s) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"sample = sample.reshape(batch_size, channels, *remaining_dims)sample = sample.to(dtype)return sampledef set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):"""Sets the discrete timesteps used for the diffusion chain (to be run before inference).Args:num_inference_steps (`int`):The number of diffusion steps used when generating samples with a pre-trained model."""if num_inference_steps > self.config.num_train_timesteps:raise ValueError(f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:"f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"f" maximal {self.config.num_train_timesteps} timesteps.")self.num_inference_steps = num_inference_steps# "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891if self.config.timestep_spacing == "linspace":timesteps = (np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps).round()[::-1].copy().astype(np.int64))elif self.config.timestep_spacing == "leading":step_ratio = self.config.num_train_timesteps // self.num_inference_steps# creates integer timesteps by multiplying by ratio# casting to int to avoid issues when num_inference_step is power of 3timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)timesteps += self.config.steps_offsetelif self.config.timestep_spacing == "trailing":step_ratio = self.config.num_train_timesteps / self.num_inference_steps# creates integer timesteps by multiplying by ratio# casting to int to avoid issues when num_inference_step is power of 3timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64)timesteps -= 1else:raise ValueError(f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'leading' or 'trailing'.")self.timesteps = torch.from_numpy(timesteps).to(device)def step(self,model_output: torch.Tensor,timestep: int,sample: torch.Tensor,eta: float = 0.0,use_clipped_model_output: bool = False,generator=None,variance_noise: Optional[torch.Tensor] = None,return_dict: bool = True,) -> Union[DDIMSchedulerOutput, Tuple]:if self.num_inference_steps is None:raise ValueError("Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler")# 1. get previous step value (=t-1)；# timestep=980，self.config.num_train_timesteps=1000, self.num_inference_steps=50# prev_timestep = 960，步数的跳跃间隔为20prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps# 2. compute alphas, betasalpha_prod_t = self.alphas_cumprod[timestep]alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprodbeta_prod_t = 1 - alpha_prod_t# 3. compute predicted original sample from predicted noise also called# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdfif self.config.prediction_type == "epsilon":pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)pred_epsilon = model_outputelif self.config.prediction_type == "sample":pred_original_sample = model_outputpred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)elif self.config.prediction_type == "v_prediction":pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_outputpred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sampleelse:raise ValueError(f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"" `v_prediction`")# 4. Clip or threshold "predicted x_0"if self.config.thresholding:pred_original_sample = self._threshold_sample(pred_original_sample)elif self.config.clip_sample:pred_original_sample = pred_original_sample.clamp(-self.config.clip_sample_range, self.config.clip_sample_range)# 5. compute variance: "sigma_t(η)" -> see formula (16)# σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)variance = self._get_variance(timestep, prev_timestep)std_dev_t = eta * variance ** (0.5)if use_clipped_model_output:# the pred_epsilon is always re-derived from the clipped x_0 in Glidepred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)# 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdfpred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * pred_epsilon# 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdfprev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_directionif eta > 0:if variance_noise is not None and generator is not None:raise ValueError("Cannot pass both generator and variance_noise. Please make sure that either `generator` or"" `variance_noise` stays `None`.")if variance_noise is None:variance_noise = randn_tensor(model_output.shape, generator=generator, device=model_output.device, dtype=model_output.dtype)variance = std_dev_t * variance_noiseprev_sample = prev_sample + varianceif not return_dict:return (prev_sample,pred_original_sample,)return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)# Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noisedef add_noise(self,original_samples: torch.Tensor,noise: torch.Tensor,timesteps: torch.IntTensor,) -> torch.Tensor:# Make sure alphas_cumprod and timestep have same device and dtype as original_samples# Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement# for the subsequent add_noise callsself.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device)alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype)timesteps = timesteps.to(original_samples.device)sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5sqrt_alpha_prod = sqrt_alpha_prod.flatten()while len(sqrt_alpha_prod.shape) < len(original_samples.shape):sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noisereturn noisy_samples# Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocitydef get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor:# Make sure alphas_cumprod and timestep have same device and dtype as sampleself.alphas_cumprod = self.alphas_cumprod.to(device=sample.device)alphas_cumprod = self.alphas_cumprod.to(dtype=sample.dtype)timesteps = timesteps.to(sample.device)sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5sqrt_alpha_prod = sqrt_alpha_prod.flatten()while len(sqrt_alpha_prod.shape) < len(sample.shape):sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape):sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * samplereturn velocitydef __len__(self):return self.config.num_train_timesteps