multimodalart/pix2pix-zero

1# Copyright 2022 Stanford University 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 code is strongly influenced by https://github.com/pesser/pytorch_diffusion16# and https://github.com/hojonathanho/diffusion17import os, sys, pdb18import math19from dataclasses import dataclass20from typing import List, Optional, Tuple, Union21 22import numpy as np23import torch24 25from diffusers.configuration_utils import ConfigMixin, register_to_config26from diffusers.utils import BaseOutput, randn_tensor27from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin28 29 30@dataclass31# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->DDIM32class DDIMSchedulerOutput(BaseOutput):33    """34    Output class for the scheduler's step function output.35 36    Args:37        prev_sample (`torch.FloatTensor` 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.FloatTensor` 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.FloatTensor46    pred_original_sample: Optional[torch.FloatTensor] = None47 48 49def betas_for_alpha_bar(num_diffusion_timesteps, max_beta=0.999) -> torch.Tensor:50    """51    Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of52    (1-beta) over time from t = [0,1].53 54    Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up55    to that part of the diffusion process.56 57 58    Args:59        num_diffusion_timesteps (`int`): the number of betas to produce.60        max_beta (`float`): the maximum beta to use; use values lower than 1 to61                     prevent singularities.62 63    Returns:64        betas (`np.ndarray`): the betas used by the scheduler to step the model outputs65    """66 67    def alpha_bar(time_step):68        return math.cos((time_step + 0.008) / 1.008 * math.pi / 2) ** 269 70    betas = []71    for i in range(num_diffusion_timesteps):72        t1 = i / num_diffusion_timesteps73        t2 = (i + 1) / num_diffusion_timesteps74        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))75    return torch.tensor(betas)76 77 78class DDIMInverseScheduler(SchedulerMixin, ConfigMixin):79    """80    Denoising diffusion implicit models is a scheduler that extends the denoising procedure introduced in denoising81    diffusion probabilistic models (DDPMs) with non-Markovian guidance.82 83    [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__`84    function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`.85    [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and86    [`~SchedulerMixin.from_pretrained`] functions.87 88    For more details, see the original paper: https://arxiv.org/abs/2010.0250289 90    Args:91        num_train_timesteps (`int`): number of diffusion steps used to train the model.92        beta_start (`float`): the starting `beta` value of inference.93        beta_end (`float`): the final `beta` value.94        beta_schedule (`str`):95            the beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from96            `linear`, `scaled_linear`, or `squaredcos_cap_v2`.97        trained_betas (`np.ndarray`, optional):98            option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc.99        clip_sample (`bool`, default `True`):100            option to clip predicted sample between -1 and 1 for numerical stability.101        set_alpha_to_one (`bool`, default `True`):102            each diffusion step uses the value of alphas product at that step and at the previous one. For the final103            step there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`,104            otherwise it uses the value of alpha at step 0.105        steps_offset (`int`, default `0`):106            an offset added to the inference steps. You can use a combination of `offset=1` and107            `set_alpha_to_one=False`, to make the last step use step 0 for the previous alpha product, as done in108            stable diffusion.109        prediction_type (`str`, default `epsilon`, optional):110            prediction type of the scheduler function, one of `epsilon` (predicting the noise of the diffusion111            process), `sample` (directly predicting the noisy sample`) or `v_prediction` (see section 2.4112            https://imagen.research.google/video/paper.pdf)113    """114 115    _compatibles = [e.name for e in KarrasDiffusionSchedulers]116    order = 1117 118    @register_to_config119    def __init__(120        self,121        num_train_timesteps: int = 1000,122        beta_start: float = 0.0001,123        beta_end: float = 0.02,124        beta_schedule: str = "linear",125        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,126        clip_sample: bool = True,127        set_alpha_to_one: bool = True,128        steps_offset: int = 0,129        prediction_type: str = "epsilon",130    ):131        if trained_betas is not None:132            self.betas = torch.tensor(trained_betas, dtype=torch.float32)133        elif beta_schedule == "linear":134            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)135        elif beta_schedule == "scaled_linear":136            # this schedule is very specific to the latent diffusion model.137            self.betas = (138                torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2139            )140        elif beta_schedule == "squaredcos_cap_v2":141            # Glide cosine schedule142            self.betas = betas_for_alpha_bar(num_train_timesteps)143        else:144            raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")145 146        self.alphas = 1.0 - self.betas147        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)148 149        # At every step in ddim, we are looking into the previous alphas_cumprod150        # For the final step, there is no previous alphas_cumprod because we are already at 0151        # `set_alpha_to_one` decides whether we set this parameter simply to one or152        # whether we use the final alpha of the "non-previous" one.153        self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0]154 155        # standard deviation of the initial noise distribution156        self.init_noise_sigma = 1.0157 158        # setable values159        self.num_inference_steps = None160        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64))161 162    def scale_model_input(self, sample: torch.FloatTensor, timestep: Optional[int] = None) -> torch.FloatTensor:163        """164        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the165        current timestep.166 167        Args:168            sample (`torch.FloatTensor`): input sample169            timestep (`int`, optional): current timestep170 171        Returns:172            `torch.FloatTensor`: scaled input sample173        """174        return sample175 176    def _get_variance(self, timestep, prev_timestep):177        alpha_prod_t = self.alphas_cumprod[timestep]178        alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod179        beta_prod_t = 1 - alpha_prod_t180        beta_prod_t_prev = 1 - alpha_prod_t_prev181 182        variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)183 184        return variance185 186    def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):187        """188        Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.189 190        Args:191            num_inference_steps (`int`):192                the number of diffusion steps used when generating samples with a pre-trained model.193        """194 195        if num_inference_steps > self.config.num_train_timesteps:196            raise ValueError(197                f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:"198                f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"199                f" maximal {self.config.num_train_timesteps} timesteps."200            )201 202        self.num_inference_steps = num_inference_steps203        step_ratio = self.config.num_train_timesteps // self.num_inference_steps204        # creates integer timesteps by multiplying by ratio205        # casting to int to avoid issues when num_inference_step is power of 3206        timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)207        self.timesteps = torch.from_numpy(timesteps).to(device)208        self.timesteps += self.config.steps_offset209 210    def step(211        self,212        model_output: torch.FloatTensor,213        timestep: int,214        sample: torch.FloatTensor,215        eta: float = 0.0,216        use_clipped_model_output: bool = False,217        generator=None,218        variance_noise: Optional[torch.FloatTensor] = None,219        return_dict: bool = True,220        reverse=False221    ) -> Union[DDIMSchedulerOutput, Tuple]:222        223 224        e_t = model_output225 226        x = sample227        prev_timestep = timestep + self.config.num_train_timesteps // self.num_inference_steps228        # print(timestep, prev_timestep)229        a_t = alpha_prod_t = self.alphas_cumprod[timestep-1]230        a_prev = alpha_t_prev = self.alphas_cumprod[prev_timestep-1] if prev_timestep >= 0 else self.final_alpha_cumprod231        beta_prod_t = 1 - alpha_prod_t232 233        pred_x0 = (x - (1-a_t)**0.5 * e_t) / a_t.sqrt()234        # direction pointing to x_t235        dir_xt = (1. - a_prev).sqrt() * e_t236        x = a_prev.sqrt()*pred_x0 + dir_xt237        if not return_dict:238            return (x,)239        return DDIMSchedulerOutput(prev_sample=x, pred_original_sample=pred_x0)240 241 242 243 244 245    def add_noise(246        self,247        original_samples: torch.FloatTensor,248        noise: torch.FloatTensor,249        timesteps: torch.IntTensor,250    ) -> torch.FloatTensor:251        # Make sure alphas_cumprod and timestep have same device and dtype as original_samples252        self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)253        timesteps = timesteps.to(original_samples.device)254 255        sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5256        sqrt_alpha_prod = sqrt_alpha_prod.flatten()257        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):258            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)259 260        sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5261        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()262        while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):263            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)264 265        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise266        return noisy_samples267 268    def get_velocity(269        self, sample: torch.FloatTensor, noise: torch.FloatTensor, timesteps: torch.IntTensor270    ) -> torch.FloatTensor:271        # Make sure alphas_cumprod and timestep have same device and dtype as sample272        self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device, dtype=sample.dtype)273        timesteps = timesteps.to(sample.device)274 275        sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5276        sqrt_alpha_prod = sqrt_alpha_prod.flatten()277        while len(sqrt_alpha_prod.shape) < len(sample.shape):278            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)279 280        sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5281        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()282        while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape):283            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)284 285        velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample286        return velocity287 288    def __len__(self):289        return self.config.num_train_timesteps290