diffusers/community-pipelines-mirror

1# Copyright 2024 UC Berkeley Team and The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15# DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim16 17import math18from dataclasses import dataclass19from typing import List, Optional, Tuple, Union20 21import numpy as np22import torch23 24from diffusers.configuration_utils import ConfigMixin, register_to_config25from diffusers.schedulers.scheduling_utils import SchedulerMixin26from diffusers.utils import BaseOutput27from diffusers.utils.torch_utils import randn_tensor28 29 30@dataclass31# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->UFOGen32class UFOGenSchedulerOutput(BaseOutput):33    """34    Output class for the scheduler's `step` function output.35 36    Args:37        prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images):38            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the39            denoising loop.40        pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images):41            The predicted denoised sample `(x_{0})` based on the model output from the current timestep.42            `pred_original_sample` can be used to preview progress or for guidance.43    """44 45    prev_sample: torch.Tensor46    pred_original_sample: Optional[torch.Tensor] = None47 48 49# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar50def betas_for_alpha_bar(51    num_diffusion_timesteps,52    max_beta=0.999,53    alpha_transform_type="cosine",54):55    """56    Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of57    (1-beta) over time from t = [0,1].58 59    Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up60    to that part of the diffusion process.61 62 63    Args:64        num_diffusion_timesteps (`int`): the number of betas to produce.65        max_beta (`float`): the maximum beta to use; use values lower than 1 to66                     prevent singularities.67        alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar.68                     Choose from `cosine` or `exp`69 70    Returns:71        betas (`np.ndarray`): the betas used by the scheduler to step the model outputs72    """73    if alpha_transform_type == "cosine":74 75        def alpha_bar_fn(t):76            return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 277 78    elif alpha_transform_type == "exp":79 80        def alpha_bar_fn(t):81            return math.exp(t * -12.0)82 83    else:84        raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")85 86    betas = []87    for i in range(num_diffusion_timesteps):88        t1 = i / num_diffusion_timesteps89        t2 = (i + 1) / num_diffusion_timesteps90        betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))91    return torch.tensor(betas, dtype=torch.float32)92 93 94# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr95def rescale_zero_terminal_snr(betas):96    """97    Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1)98 99 100    Args:101        betas (`torch.Tensor`):102            the betas that the scheduler is being initialized with.103 104    Returns:105        `torch.Tensor`: rescaled betas with zero terminal SNR106    """107    # Convert betas to alphas_bar_sqrt108    alphas = 1.0 - betas109    alphas_cumprod = torch.cumprod(alphas, dim=0)110    alphas_bar_sqrt = alphas_cumprod.sqrt()111 112    # Store old values.113    alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()114    alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()115 116    # Shift so the last timestep is zero.117    alphas_bar_sqrt -= alphas_bar_sqrt_T118 119    # Scale so the first timestep is back to the old value.120    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)121 122    # Convert alphas_bar_sqrt to betas123    alphas_bar = alphas_bar_sqrt**2  # Revert sqrt124    alphas = alphas_bar[1:] / alphas_bar[:-1]  # Revert cumprod125    alphas = torch.cat([alphas_bar[0:1], alphas])126    betas = 1 - alphas127 128    return betas129 130 131class UFOGenScheduler(SchedulerMixin, ConfigMixin):132    """133    `UFOGenScheduler` implements multistep and onestep sampling for a UFOGen model, introduced in134    [UFOGen: You Forward Once Large Scale Text-to-Image Generation via Diffusion GANs](https://arxiv.org/abs/2311.09257)135    by Yanwu Xu, Yang Zhao, Zhisheng Xiao, and Tingbo Hou. UFOGen is a varianet of the denoising diffusion GAN (DDGAN)136    model designed for one-step sampling.137 138    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic139    methods the library implements for all schedulers such as loading and saving.140 141    Args:142        num_train_timesteps (`int`, defaults to 1000):143            The number of diffusion steps to train the model.144        beta_start (`float`, defaults to 0.0001):145            The starting `beta` value of inference.146        beta_end (`float`, defaults to 0.02):147            The final `beta` value.148        beta_schedule (`str`, defaults to `"linear"`):149            The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from150            `linear`, `scaled_linear`, or `squaredcos_cap_v2`.151        clip_sample (`bool`, defaults to `True`):152            Clip the predicted sample for numerical stability.153        clip_sample_range (`float`, defaults to 1.0):154            The maximum magnitude for sample clipping. Valid only when `clip_sample=True`.155        set_alpha_to_one (`bool`, defaults to `True`):156            Each diffusion step uses the alphas product value at that step and at the previous one. For the final step157            there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`,158            otherwise it uses the alpha value at step 0.159        prediction_type (`str`, defaults to `epsilon`, *optional*):160            Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),161            `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen162            Video](https://imagen.research.google/video/paper.pdf) paper).163        thresholding (`bool`, defaults to `False`):164            Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such165            as Stable Diffusion.166        dynamic_thresholding_ratio (`float`, defaults to 0.995):167            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.168        sample_max_value (`float`, defaults to 1.0):169            The threshold value for dynamic thresholding. Valid only when `thresholding=True`.170        timestep_spacing (`str`, defaults to `"leading"`):171            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and172            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.173        steps_offset (`int`, defaults to 0):174            An offset added to the inference steps, as required by some model families.175        rescale_betas_zero_snr (`bool`, defaults to `False`):176            Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and177            dark samples instead of limiting it to samples with medium brightness. Loosely related to178            [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506).179        denoising_step_size (`int`, defaults to 250):180            The denoising step size parameter from the UFOGen paper. The number of steps used for training is roughly181            `math.ceil(num_train_timesteps / denoising_step_size)`.182    """183 184    order = 1185 186    @register_to_config187    def __init__(188        self,189        num_train_timesteps: int = 1000,190        beta_start: float = 0.0001,191        beta_end: float = 0.02,192        beta_schedule: str = "linear",193        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,194        clip_sample: bool = True,195        set_alpha_to_one: bool = True,196        prediction_type: str = "epsilon",197        thresholding: bool = False,198        dynamic_thresholding_ratio: float = 0.995,199        clip_sample_range: float = 1.0,200        sample_max_value: float = 1.0,201        timestep_spacing: str = "leading",202        steps_offset: int = 0,203        rescale_betas_zero_snr: bool = False,204        denoising_step_size: int = 250,205    ):206        if trained_betas is not None:207            self.betas = torch.tensor(trained_betas, dtype=torch.float32)208        elif beta_schedule == "linear":209            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)210        elif beta_schedule == "scaled_linear":211            # this schedule is very specific to the latent diffusion model.212            self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2213        elif beta_schedule == "squaredcos_cap_v2":214            # Glide cosine schedule215            self.betas = betas_for_alpha_bar(num_train_timesteps)216        elif beta_schedule == "sigmoid":217            # GeoDiff sigmoid schedule218            betas = torch.linspace(-6, 6, num_train_timesteps)219            self.betas = torch.sigmoid(betas) * (beta_end - beta_start) + beta_start220        else:221            raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")222 223        # Rescale for zero SNR224        if rescale_betas_zero_snr:225            self.betas = rescale_zero_terminal_snr(self.betas)226 227        self.alphas = 1.0 - self.betas228        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)229 230        # For the final step, there is no previous alphas_cumprod because we are already at 0231        # `set_alpha_to_one` decides whether we set this parameter simply to one or232        # whether we use the final alpha of the "non-previous" one.233        self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0]234 235        # standard deviation of the initial noise distribution236        self.init_noise_sigma = 1.0237 238        # setable values239        self.custom_timesteps = False240        self.num_inference_steps = None241        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy())242 243    def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:244        """245        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the246        current timestep.247 248        Args:249            sample (`torch.Tensor`):250                The input sample.251            timestep (`int`, *optional*):252                The current timestep in the diffusion chain.253 254        Returns:255            `torch.Tensor`:256                A scaled input sample.257        """258        return sample259 260    def set_timesteps(261        self,262        num_inference_steps: Optional[int] = None,263        device: Union[str, torch.device] = None,264        timesteps: Optional[List[int]] = None,265    ):266        """267        Sets the discrete timesteps used for the diffusion chain (to be run before inference).268 269        Args:270            num_inference_steps (`int`):271                The number of diffusion steps used when generating samples with a pre-trained model. If used,272                `timesteps` must be `None`.273            device (`str` or `torch.device`, *optional*):274                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.275            timesteps (`List[int]`, *optional*):276                Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default277                timestep spacing strategy of equal spacing between timesteps is used. If `timesteps` is passed,278                `num_inference_steps` must be `None`.279 280        """281        if num_inference_steps is not None and timesteps is not None:282            raise ValueError("Can only pass one of `num_inference_steps` or `custom_timesteps`.")283 284        if timesteps is not None:285            for i in range(1, len(timesteps)):286                if timesteps[i] >= timesteps[i - 1]:287                    raise ValueError("`custom_timesteps` must be in descending order.")288 289            if timesteps[0] >= self.config.num_train_timesteps:290                raise ValueError(291                    f"`timesteps` must start before `self.config.train_timesteps`:"292                    f" {self.config.num_train_timesteps}."293                )294 295            timesteps = np.array(timesteps, dtype=np.int64)296            self.custom_timesteps = True297        else:298            if num_inference_steps > self.config.num_train_timesteps:299                raise ValueError(300                    f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:"301                    f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"302                    f" maximal {self.config.num_train_timesteps} timesteps."303                )304 305            self.num_inference_steps = num_inference_steps306            self.custom_timesteps = False307 308            # TODO: For now, handle special case when num_inference_steps == 1 separately309            if num_inference_steps == 1:310                # Set the timestep schedule to num_train_timesteps - 1 rather than 0311                # (that is, the one-step timestep schedule is always trailing rather than leading or linspace)312                timesteps = np.array([self.config.num_train_timesteps - 1], dtype=np.int64)313            else:314                # TODO: For now, retain the DDPM timestep spacing logic315                # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891316                if self.config.timestep_spacing == "linspace":317                    timesteps = (318                        np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps)319                        .round()[::-1]320                        .copy()321                        .astype(np.int64)322                    )323                elif self.config.timestep_spacing == "leading":324                    step_ratio = self.config.num_train_timesteps // self.num_inference_steps325                    # creates integer timesteps by multiplying by ratio326                    # casting to int to avoid issues when num_inference_step is power of 3327                    timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)328                    timesteps += self.config.steps_offset329                elif self.config.timestep_spacing == "trailing":330                    step_ratio = self.config.num_train_timesteps / self.num_inference_steps331                    # creates integer timesteps by multiplying by ratio332                    # casting to int to avoid issues when num_inference_step is power of 3333                    timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64)334                    timesteps -= 1335                else:336                    raise ValueError(337                        f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'."338                    )339 340        self.timesteps = torch.from_numpy(timesteps).to(device)341 342    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample343    def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:344        """345        "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the346        prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by347        s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing348        pixels from saturation at each step. We find that dynamic thresholding results in significantly better349        photorealism as well as better image-text alignment, especially when using very large guidance weights."350 351        https://arxiv.org/abs/2205.11487352        """353        dtype = sample.dtype354        batch_size, channels, *remaining_dims = sample.shape355 356        if dtype not in (torch.float32, torch.float64):357            sample = sample.float()  # upcast for quantile calculation, and clamp not implemented for cpu half358 359        # Flatten sample for doing quantile calculation along each image360        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))361 362        abs_sample = sample.abs()  # "a certain percentile absolute pixel value"363 364        s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)365        s = torch.clamp(366            s, min=1, max=self.config.sample_max_value367        )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]368        s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0369        sample = torch.clamp(sample, -s, s) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"370 371        sample = sample.reshape(batch_size, channels, *remaining_dims)372        sample = sample.to(dtype)373 374        return sample375 376    def step(377        self,378        model_output: torch.Tensor,379        timestep: int,380        sample: torch.Tensor,381        generator: Optional[torch.Generator] = None,382        return_dict: bool = True,383    ) -> Union[UFOGenSchedulerOutput, Tuple]:384        """385        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion386        process from the learned model outputs (most often the predicted noise).387 388        Args:389            model_output (`torch.Tensor`):390                The direct output from learned diffusion model.391            timestep (`float`):392                The current discrete timestep in the diffusion chain.393            sample (`torch.Tensor`):394                A current instance of a sample created by the diffusion process.395            generator (`torch.Generator`, *optional*):396                A random number generator.397            return_dict (`bool`, *optional*, defaults to `True`):398                Whether or not to return a [`~schedulers.scheduling_ufogen.UFOGenSchedulerOutput`] or `tuple`.399 400        Returns:401            [`~schedulers.scheduling_ddpm.UFOGenSchedulerOutput`] or `tuple`:402                If return_dict is `True`, [`~schedulers.scheduling_ufogen.UFOGenSchedulerOutput`] is returned, otherwise a403                tuple is returned where the first element is the sample tensor.404 405        """406        # 0. Resolve timesteps407        t = timestep408        prev_t = self.previous_timestep(t)409 410        # 1. compute alphas, betas411        alpha_prod_t = self.alphas_cumprod[t]412        alpha_prod_t_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else self.final_alpha_cumprod413        beta_prod_t = 1 - alpha_prod_t414        # beta_prod_t_prev = 1 - alpha_prod_t_prev415        # current_alpha_t = alpha_prod_t / alpha_prod_t_prev416        # current_beta_t = 1 - current_alpha_t417 418        # 2. compute predicted original sample from predicted noise also called419        # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf420        if self.config.prediction_type == "epsilon":421            pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)422        elif self.config.prediction_type == "sample":423            pred_original_sample = model_output424        elif self.config.prediction_type == "v_prediction":425            pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output426        else:427            raise ValueError(428                f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or"429                " `v_prediction`  for UFOGenScheduler."430            )431 432        # 3. Clip or threshold "predicted x_0"433        if self.config.thresholding:434            pred_original_sample = self._threshold_sample(pred_original_sample)435        elif self.config.clip_sample:436            pred_original_sample = pred_original_sample.clamp(437                -self.config.clip_sample_range, self.config.clip_sample_range438            )439 440        # 4. Single-step or multi-step sampling441        # Noise is not used on the final timestep of the timestep schedule.442        # This also means that noise is not used for one-step sampling.443        if t != self.timesteps[-1]:444            # TODO: is this correct?445            # Sample prev sample x_{t - 1} ~ q(x_{t - 1} | x_0 =  G(x_t, t))446            device = model_output.device447            noise = randn_tensor(model_output.shape, generator=generator, device=device, dtype=model_output.dtype)448            sqrt_alpha_prod_t_prev = alpha_prod_t_prev**0.5449            sqrt_one_minus_alpha_prod_t_prev = (1 - alpha_prod_t_prev) ** 0.5450            pred_prev_sample = sqrt_alpha_prod_t_prev * pred_original_sample + sqrt_one_minus_alpha_prod_t_prev * noise451        else:452            # Simply return the pred_original_sample. If `prediction_type == "sample"`, this is equivalent to returning453            # the output of the GAN generator U-Net on the initial noisy latents x_T ~ N(0, I).454            pred_prev_sample = pred_original_sample455 456        if not return_dict:457            return (pred_prev_sample,)458 459        return UFOGenSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample)460 461    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise462    def add_noise(463        self,464        original_samples: torch.Tensor,465        noise: torch.Tensor,466        timesteps: torch.IntTensor,467    ) -> torch.Tensor:468        # Make sure alphas_cumprod and timestep have same device and dtype as original_samples469        alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)470        timesteps = timesteps.to(original_samples.device)471 472        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5473        sqrt_alpha_prod = sqrt_alpha_prod.flatten()474        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):475            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)476 477        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5478        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()479        while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):480            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)481 482        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise483        return noisy_samples484 485    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocity486    def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor:487        # Make sure alphas_cumprod and timestep have same device and dtype as sample488        alphas_cumprod = self.alphas_cumprod.to(device=sample.device, dtype=sample.dtype)489        timesteps = timesteps.to(sample.device)490 491        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5492        sqrt_alpha_prod = sqrt_alpha_prod.flatten()493        while len(sqrt_alpha_prod.shape) < len(sample.shape):494            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)495 496        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5497        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()498        while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape):499            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)500 501        velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample502        return velocity503 504    def __len__(self):505        return self.config.num_train_timesteps506 507    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.previous_timestep508    def previous_timestep(self, timestep):509        if self.custom_timesteps:510            index = (self.timesteps == timestep).nonzero(as_tuple=True)[0][0]511            if index == self.timesteps.shape[0] - 1:512                prev_t = torch.tensor(-1)513            else:514                prev_t = self.timesteps[index + 1]515        else:516            num_inference_steps = (517                self.num_inference_steps if self.num_inference_steps else self.config.num_train_timesteps518            )519            prev_t = timestep - self.config.num_train_timesteps // num_inference_steps520 521        return prev_t522