added standalone schedulers

Author: divamgupta

backends/stable_diffusion/schedulers/__init__.py

@@ -0,0 +1,90 @@
+# source : https://github.com/huggingface/diffusers/
+
+from typing import Union
+import numpy as np
+
+
+class SchedulerMixin:
+ """
+ Mixin containing common functions for the schedulers.
+ """
+
+ ignore_for_config = ["tensor_format"]
+
+ def set_format(self, tensor_format="pt"):
+ self.tensor_format = tensor_format
+ if tensor_format == "pt":
+ for key, value in vars(self).items():
+ if isinstance(value, np.ndarray):
+ setattr(self, key, torch.from_numpy(value))
+
+ return self
+
+ def clip(self, tensor, min_value=None, max_value=None):
+ tensor_format = getattr(self, "tensor_format", "pt")
+
+ if tensor_format == "np":
+ return np.clip(tensor, min_value, max_value)
+ elif tensor_format == "pt":
+ return torch.clamp(tensor, min_value, max_value)
+
+ raise ValueError(f"`self.tensor_format`: {self.tensor_format} is not valid.")
+
+ def log(self, tensor):
+ tensor_format = getattr(self, "tensor_format", "pt")
+
+ if tensor_format == "np":
+ return np.log(tensor)
+ elif tensor_format == "pt":
+ return torch.log(tensor)
+
+ raise ValueError(f"`self.tensor_format`: {self.tensor_format} is not valid.")
+
+ def match_shape(self, values: Union[np.ndarray], broadcast_array: Union[np.ndarray]):
+ """
+ Turns a 1-D array into an array or tensor with len(broadcast_array.shape) dims.
+ Args:
+ values: an array or tensor of values to extract.
+ broadcast_array: an array with a larger shape of K dimensions with the batch
+ dimension equal to the length of timesteps.
+ Returns:
+ a tensor of shape [batch_size, 1, ...] where the shape has K dims.
+ """
+
+ tensor_format = getattr(self, "tensor_format", "pt")
+ values = values.flatten()
+
+ while len(values.shape) < len(broadcast_array.shape):
+ values = values[..., None]
+ if tensor_format == "pt":
+ values = values.to(broadcast_array.device)
+
+ return values
+
+ def norm(self, tensor):
+ tensor_format = getattr(self, "tensor_format", "pt")
+ if tensor_format == "np":
+ return np.linalg.norm(tensor)
+ elif tensor_format == "pt":
+ return torch.norm(tensor.reshape(tensor.shape[0], -1), dim=-1).mean()
+
+ raise ValueError(f"`self.tensor_format`: {self.tensor_format} is not valid.")
+
+ def randn_like(self, tensor, generator=None):
+ tensor_format = getattr(self, "tensor_format", "pt")
+ if tensor_format == "np":
+ return np.random.randn(*np.shape(tensor))
+ elif tensor_format == "pt":
+ # return torch.randn_like(tensor)
+ return torch.randn(tensor.shape, layout=tensor.layout, generator=generator).to(tensor.device)
+
+ raise ValueError(f"`self.tensor_format`: {self.tensor_format} is not valid.")
+
+ def zeros_like(self, tensor):
+ tensor_format = getattr(self, "tensor_format", "pt")
+ if tensor_format == "np":
+ return np.zeros_like(tensor)
+ elif tensor_format == "pt":
+ return torch.zeros_like(tensor)
+
+ raise ValueError(f"`self.tensor_format`: {self.tensor_format} is not valid.")

backends/stable_diffusion/schedulers/scheduler_mixin.py

@@ -0,0 +1,261 @@
+# source : https://github.com/huggingface/diffusers/
+
+
+# Copyright 2022 Stanford University Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion
+# and https://github.com/hojonathanho/diffusion
+
+import math
+from typing import Optional, Tuple, Union
+
+import numpy as np
+
+
+from .scheduler_mixin import SchedulerMixin
+
+
+def betas_for_alpha_bar(num_diffusion_timesteps, max_beta=0.999):
+ """
+ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
+ (1-beta) over time from t = [0,1].
+ Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up
+ to that part of the diffusion process.
+ Args:
+ num_diffusion_timesteps (`int`): the number of betas to produce.
+ max_beta (`float`): the maximum beta to use; use values lower than 1 to
+ prevent singularities.
+ Returns:
+ betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
+ """
+
+ def alpha_bar(time_step):
+ return math.cos((time_step + 0.008) / 1.008 * math.pi / 2) ** 2
+
+ betas = []
+ for i in range(num_diffusion_timesteps):
+ t1 = i / num_diffusion_timesteps
+ t2 = (i + 1) / num_diffusion_timesteps
+ betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+ return np.array(betas, dtype=np.float32)
+
+
+class DDIMScheduler(SchedulerMixin):
+ """
+ Denoising diffusion implicit models is a scheduler that extends the denoising procedure introduced in denoising
+ diffusion probabilistic models (DDPMs) with non-Markovian guidance.
+ [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__`
+ function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`.
+ [`~ConfigMixin`] also provides general loading and saving functionality via the [`~ConfigMixin.save_config`] and
+ [`~ConfigMixin.from_config`] functios.
+ For more details, see the original paper: https://arxiv.org/abs/2010.02502
+ Args:
+ num_train_timesteps (`int`): number of diffusion steps used to train the model.
+ beta_start (`float`): the starting `beta` value of inference.
+ beta_end (`float`): the final `beta` value.
+ beta_schedule (`str`):
+ the beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
+ `linear`, `scaled_linear`, or `squaredcos_cap_v2`.
+ trained_betas (`np.ndarray`, optional): TODO
+ timestep_values (`np.ndarray`, optional): TODO
+ clip_sample (`bool`, default `True`):
+ option to clip predicted sample between -1 and 1 for numerical stability.
+ set_alpha_to_one (`bool`, default `True`):
+ if alpha for final step is 1 or the final alpha of the "non-previous" one.
+ tensor_format (`str`): whether the scheduler expects pytorch or numpy arrays.
+ """
+
+ def __init__(
+ self,
+ num_train_timesteps: int = 1000,
+ beta_start: float = 0.0001,
+ beta_end: float = 0.02,
+ beta_schedule: str = "linear",
+ trained_betas: Optional[np.ndarray] = None,
+ timestep_values: Optional[np.ndarray] = None,
+ clip_sample: bool = True,
+ set_alpha_to_one: bool = True,
+ tensor_format: str = "pt",
+ ):
+ if trained_betas is not None:
+ self.betas = np.asarray(trained_betas)
+ if beta_schedule == "linear":
+ self.betas = np.linspace(beta_start, beta_end, num_train_timesteps, dtype=np.float32)
+ elif beta_schedule == "scaled_linear":
+ # this schedule is very specific to the latent diffusion model.
+ self.betas = np.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=np.float32) ** 2
+ elif beta_schedule == "squaredcos_cap_v2":
+ # Glide cosine schedule
+ self.betas = betas_for_alpha_bar(num_train_timesteps)
+ else:
+ raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")
+
+ self.alphas = 1.0 - self.betas
+ self.alphas_cumprod = np.cumprod(self.alphas, axis=0)
+
+ self.num_train_timesteps = num_train_timesteps
+
+ self.beta_start = beta_start
+ self.beta_end = beta_end
+ self.beta_schedule = beta_schedule
+ self.trained_betas = trained_betas
+ self.timestep_values = timestep_values
+ self.clip_sample = clip_sample
+ self.set_alpha_to_one = set_alpha_to_one
+ self.tensor_format = tensor_format
+
+ # 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 paratemer simply to one or
+ # whether we use the final alpha of the "non-previous" one.
+ self.final_alpha_cumprod = np.array(1.0) if set_alpha_to_one else self.alphas_cumprod[0]
+
+ # setable values
+ self.num_inference_steps = None
+ self.timesteps = np.arange(0, num_train_timesteps)[::-1].copy()
+
+ self.tensor_format = tensor_format
+ self.config = self
+ self.set_format(tensor_format=tensor_format)
+
+ def _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_cumprod
+ beta_prod_t = 1 - alpha_prod_t
+ beta_prod_t_prev = 1 - alpha_prod_t_prev
+
+ variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
+
+ return variance
+
+ def set_timesteps(self, num_inference_steps: int, offset: int = 0):
+ """
+ Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.
+ Args:
+ num_inference_steps (`int`):
+ the number of diffusion steps used when generating samples with a pre-trained model.
+ offset (`int`): TODO
+ """
+ self.num_inference_steps = num_inference_steps
+ self.timesteps = np.arange(
+ 0, self.config.num_train_timesteps, self.config.num_train_timesteps // self.num_inference_steps
+ )[::-1].copy()
+ self.timesteps += offset
+ self.set_format(tensor_format=self.tensor_format)
+
+ def step(
+ self,
+ model_output: Union[ np.ndarray],
+ timestep: int,
+ sample: Union[ np.ndarray],
+ eta: float = 0.0,
+ use_clipped_model_output: bool = False,
+ generator=None,
+ return_dict: bool = True,
+ ) -> Union[ Tuple]:
+ """
+ Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
+ process from the learned model outputs (most often the predicted noise).
+ Args:
+ model_output (`torch.FloatTensor` or `np.ndarray`): direct output from learned diffusion model.
+ timestep (`int`): current discrete timestep in the diffusion chain.
+ sample (`torch.FloatTensor` or `np.ndarray`):
+ current instance of sample being created by diffusion process.
+ eta (`float`): weight of noise for added noise in diffusion step.
+ use_clipped_model_output (`bool`): TODO
+ generator: random number generator.
+ return_dict (`bool`): option for returning tuple rather than SchedulerOutput class
+ Returns:
+ [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
+ [`~schedulers.scheduling_utils.SchedulerOutput`] if `return_dict` is True, otherwise a `tuple`. When
+ returning a tuple, the first element is the sample tensor.
+ """
+ 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"
+ )
+
+ # See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf
+ # Ideally, read DDIM paper in-detail understanding
+
+ # Notation (<variable name> -> <name in paper>
+ # - pred_noise_t -> e_theta(x_t, t)
+ # - pred_original_sample -> f_theta(x_t, t) or x_0
+ # - std_dev_t -> sigma_t
+ # - eta -> η
+ # - pred_sample_direction -> "direction pointingc to x_t"
+ # - pred_prev_sample -> "x_t-1"
+
+ # 1. get previous step value (=t-1)
+ prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps
+
+ # 2. compute alphas, betas
+ alpha_prod_t = self.alphas_cumprod[timestep]
+ alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
+ beta_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.pdf
+ pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
+
+ # 4. Clip "predicted x_0"
+ if self.config.clip_sample:
+ pred_original_sample = self.clip(pred_original_sample, -1, 1)
+
+ # 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 model_output is always re-derived from the clipped x_0 in Glide
+ model_output = (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.pdf
+ pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * model_output
+
+ # 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
+ prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction
+
+ if eta > 0:
+ noise = torch.randn(model_output.shape, generator=generator).to(device)
+ variance = self._get_variance(timestep, prev_timestep) ** (0.5) * eta * noise
+
+ if not torch.is_tensor(model_output):
+ variance = variance.numpy()
+
+ prev_sample = prev_sample + variance
+
+ if not return_dict:
+ return (prev_sample,)
+
+ return {"prev_sample":prev_sample} # SchedulerOutput(prev_sample=prev_sample)
+
+ def add_noise(
+ self,
+ original_samples: Union[ np.ndarray],
+ noise: Union[ np.ndarray],
+ timesteps: Union[np.ndarray],
+ ) -> Union[ np.ndarray]:
+ sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5
+ sqrt_alpha_prod = self.match_shape(sqrt_alpha_prod, original_samples)
+ sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5
+ sqrt_one_minus_alpha_prod = self.match_shape(sqrt_one_minus_alpha_prod, original_samples)
+
+ noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
+ return noisy_samples
+
+ def __len__(self):
+ return self.config.num_train_timesteps