diffusers/community-pipelines-mirror

1from typing import Optional, Tuple, Union2 3import torch4 5from diffusers import DDIMScheduler, DDPMScheduler, DiffusionPipeline, UNet2DConditionModel6from diffusers.pipeline_utils import ImagePipelineOutput7from diffusers.schedulers.scheduling_ddim import DDIMSchedulerOutput8from diffusers.schedulers.scheduling_ddpm import DDPMSchedulerOutput9from einops import rearrange, reduce10 11 12BITS = 813 14 15# convert to bit representations and back taken from https://github.com/lucidrains/bit-diffusion/blob/main/bit_diffusion/bit_diffusion.py16def decimal_to_bits(x, bits=BITS):17    """expects image tensor ranging from 0 to 1, outputs bit tensor ranging from -1 to 1"""18    device = x.device19 20    x = (x * 255).int().clamp(0, 255)21 22    mask = 2 ** torch.arange(bits - 1, -1, -1, device=device)23    mask = rearrange(mask, "d -> d 1 1")24    x = rearrange(x, "b c h w -> b c 1 h w")25 26    bits = ((x & mask) != 0).float()27    bits = rearrange(bits, "b c d h w -> b (c d) h w")28    bits = bits * 2 - 129    return bits30 31 32def bits_to_decimal(x, bits=BITS):33    """expects bits from -1 to 1, outputs image tensor from 0 to 1"""34    device = x.device35 36    x = (x > 0).int()37    mask = 2 ** torch.arange(bits - 1, -1, -1, device=device, dtype=torch.int32)38 39    mask = rearrange(mask, "d -> d 1 1")40    x = rearrange(x, "b (c d) h w -> b c d h w", d=8)41    dec = reduce(x * mask, "b c d h w -> b c h w", "sum")42    return (dec / 255).clamp(0.0, 1.0)43 44 45# modified scheduler step functions for clamping the predicted x_0 between -bit_scale and +bit_scale46def ddim_bit_scheduler_step(47    self,48    model_output: torch.FloatTensor,49    timestep: int,50    sample: torch.FloatTensor,51    eta: float = 0.0,52    use_clipped_model_output: bool = True,53    generator=None,54    return_dict: bool = True,55) -> Union[DDIMSchedulerOutput, Tuple]:56    """57    Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion58    process from the learned model outputs (most often the predicted noise).59    Args:60        model_output (`torch.FloatTensor`): direct output from learned diffusion model.61        timestep (`int`): current discrete timestep in the diffusion chain.62        sample (`torch.FloatTensor`):63            current instance of sample being created by diffusion process.64        eta (`float`): weight of noise for added noise in diffusion step.65        use_clipped_model_output (`bool`): TODO66        generator: random number generator.67        return_dict (`bool`): option for returning tuple rather than DDIMSchedulerOutput class68    Returns:69        [`~schedulers.scheduling_utils.DDIMSchedulerOutput`] or `tuple`:70        [`~schedulers.scheduling_utils.DDIMSchedulerOutput`] if `return_dict` is True, otherwise a `tuple`. When71        returning a tuple, the first element is the sample tensor.72    """73    if self.num_inference_steps is None:74        raise ValueError(75            "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"76        )77 78    # See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf79    # Ideally, read DDIM paper in-detail understanding80 81    # Notation (<variable name> -> <name in paper>82    # - pred_noise_t -> e_theta(x_t, t)83    # - pred_original_sample -> f_theta(x_t, t) or x_084    # - std_dev_t -> sigma_t85    # - eta -> η86    # - pred_sample_direction -> "direction pointing to x_t"87    # - pred_prev_sample -> "x_t-1"88 89    # 1. get previous step value (=t-1)90    prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps91 92    # 2. compute alphas, betas93    alpha_prod_t = self.alphas_cumprod[timestep]94    alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod95 96    beta_prod_t = 1 - alpha_prod_t97 98    # 3. compute predicted original sample from predicted noise also called99    # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf100    pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)101 102    # 4. Clip "predicted x_0"103    scale = self.bit_scale104    if self.config.clip_sample:105        pred_original_sample = torch.clamp(pred_original_sample, -scale, scale)106 107    # 5. compute variance: "sigma_t(η)" -> see formula (16)108    # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)109    variance = self._get_variance(timestep, prev_timestep)110    std_dev_t = eta * variance ** (0.5)111 112    if use_clipped_model_output:113        # the model_output is always re-derived from the clipped x_0 in Glide114        model_output = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)115 116    # 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf117    pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * model_output118 119    # 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf120    prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction121 122    if eta > 0:123        # randn_like does not support generator https://github.com/pytorch/pytorch/issues/27072124        device = model_output.device if torch.is_tensor(model_output) else "cpu"125        noise = torch.randn(model_output.shape, dtype=model_output.dtype, generator=generator).to(device)126        variance = self._get_variance(timestep, prev_timestep) ** (0.5) * eta * noise127 128        prev_sample = prev_sample + variance129 130    if not return_dict:131        return (prev_sample,)132 133    return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)134 135 136def ddpm_bit_scheduler_step(137    self,138    model_output: torch.FloatTensor,139    timestep: int,140    sample: torch.FloatTensor,141    prediction_type="epsilon",142    generator=None,143    return_dict: bool = True,144) -> Union[DDPMSchedulerOutput, Tuple]:145    """146    Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion147    process from the learned model outputs (most often the predicted noise).148    Args:149        model_output (`torch.FloatTensor`): direct output from learned diffusion model.150        timestep (`int`): current discrete timestep in the diffusion chain.151        sample (`torch.FloatTensor`):152            current instance of sample being created by diffusion process.153        prediction_type (`str`, default `epsilon`):154            indicates whether the model predicts the noise (epsilon), or the samples (`sample`).155        generator: random number generator.156        return_dict (`bool`): option for returning tuple rather than DDPMSchedulerOutput class157    Returns:158        [`~schedulers.scheduling_utils.DDPMSchedulerOutput`] or `tuple`:159        [`~schedulers.scheduling_utils.DDPMSchedulerOutput`] if `return_dict` is True, otherwise a `tuple`. When160        returning a tuple, the first element is the sample tensor.161    """162    t = timestep163 164    if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in ["learned", "learned_range"]:165        model_output, predicted_variance = torch.split(model_output, sample.shape[1], dim=1)166    else:167        predicted_variance = None168 169    # 1. compute alphas, betas170    alpha_prod_t = self.alphas_cumprod[t]171    alpha_prod_t_prev = self.alphas_cumprod[t - 1] if t > 0 else self.one172    beta_prod_t = 1 - alpha_prod_t173    beta_prod_t_prev = 1 - alpha_prod_t_prev174 175    # 2. compute predicted original sample from predicted noise also called176    # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf177    if prediction_type == "epsilon":178        pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)179    elif prediction_type == "sample":180        pred_original_sample = model_output181    else:182        raise ValueError(f"Unsupported prediction_type {prediction_type}.")183 184    # 3. Clip "predicted x_0"185    scale = self.bit_scale186    if self.config.clip_sample:187        pred_original_sample = torch.clamp(pred_original_sample, -scale, scale)188 189    # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t190    # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf191    pred_original_sample_coeff = (alpha_prod_t_prev ** (0.5) * self.betas[t]) / beta_prod_t192    current_sample_coeff = self.alphas[t] ** (0.5) * beta_prod_t_prev / beta_prod_t193 194    # 5. Compute predicted previous sample µ_t195    # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf196    pred_prev_sample = pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample197 198    # 6. Add noise199    variance = 0200    if t > 0:201        noise = torch.randn(202            model_output.size(), dtype=model_output.dtype, layout=model_output.layout, generator=generator203        ).to(model_output.device)204        variance = (self._get_variance(t, predicted_variance=predicted_variance) ** 0.5) * noise205 206    pred_prev_sample = pred_prev_sample + variance207 208    if not return_dict:209        return (pred_prev_sample,)210 211    return DDPMSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample)212 213 214class BitDiffusion(DiffusionPipeline):215    def __init__(216        self,217        unet: UNet2DConditionModel,218        scheduler: Union[DDIMScheduler, DDPMScheduler],219        bit_scale: Optional[float] = 1.0,220    ):221        super().__init__()222        self.bit_scale = bit_scale223        self.scheduler.step = (224            ddim_bit_scheduler_step if isinstance(scheduler, DDIMScheduler) else ddpm_bit_scheduler_step225        )226 227        self.register_modules(unet=unet, scheduler=scheduler)228 229    @torch.no_grad()230    def __call__(231        self,232        height: Optional[int] = 256,233        width: Optional[int] = 256,234        num_inference_steps: Optional[int] = 50,235        generator: Optional[torch.Generator] = None,236        batch_size: Optional[int] = 1,237        output_type: Optional[str] = "pil",238        return_dict: bool = True,239        **kwargs,240    ) -> Union[Tuple, ImagePipelineOutput]:241        latents = torch.randn(242            (batch_size, self.unet.in_channels, height, width),243            generator=generator,244        )245        latents = decimal_to_bits(latents) * self.bit_scale246        latents = latents.to(self.device)247 248        self.scheduler.set_timesteps(num_inference_steps)249 250        for t in self.progress_bar(self.scheduler.timesteps):251            # predict the noise residual252            noise_pred = self.unet(latents, t).sample253 254            # compute the previous noisy sample x_t -> x_t-1255            latents = self.scheduler.step(noise_pred, t, latents).prev_sample256 257        image = bits_to_decimal(latents)258 259        if output_type == "pil":260            image = self.numpy_to_pil(image)261 262        if not return_dict:263            return (image,)264 265        return ImagePipelineOutput(images=image)266