RePaintScheduler

RePaintScheduler is a DDPM-based inpainting scheduler for unsupervised inpainting with extreme masks. It is designed to be used with the [RePaintPipeline], and it is based on the paper RePaint: Inpainting using Denoising Diffusion Probabilistic Models by Andreas Lugmayr et al.

The abstract from the paper is:

Free-form inpainting is the task of adding new content to an image in the regions specified by an arbitrary binary mask. Most existing approaches train for a certain distribution of masks, which limits their generalization capabilities to unseen mask types. Furthermore, training with pixel-wise and perceptual losses often leads to simple textural extensions towards the missing areas instead of semantically meaningful generation. In this work, we propose RePaint: A Denoising Diffusion Probabilistic Model (DDPM) based inpainting approach that is applicable to even extreme masks. We employ a pretrained unconditional DDPM as the generative prior. To condition the generation process, we only alter the reverse diffusion iterations by sampling the unmasked regions using the given image information. Since this technique does not modify or condition the original DDPM network itself, the model produces high-quality and diverse output images for any inpainting form. We validate our method for both faces and general-purpose image inpainting using standard and extreme masks. RePaint outperforms state-of-the-art Autoregressive, and GAN approaches for at least five out of six mask distributions. GitHub Repository: this http URL.

The original implementation can be found at andreas128/RePaint.

mindone.diffusers.RePaintScheduler

Bases: SchedulerMixin, ConfigMixin

RePaintScheduler is a scheduler for DDPM inpainting inside a given mask.

This model inherits from [SchedulerMixin] and [ConfigMixin]. Check the superclass documentation for the generic methods the library implements for all schedulers such as loading and saving.

PARAMETER	DESCRIPTION
num_train_timesteps	The number of diffusion steps to train the model. TYPE: `int`, defaults to 1000 DEFAULT: 1000
beta_start	The starting beta value of inference. TYPE: `float`, defaults to 0.0001 DEFAULT: 0.0001
beta_end	The final beta value. TYPE: `float`, defaults to 0.02 DEFAULT: 0.02
beta_schedule	The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from linear, scaled_linear, squaredcos_cap_v2, or sigmoid. TYPE: `str`, defaults to `"linear"` DEFAULT: 'linear'
eta	The weight of noise for added noise in diffusion step. If its value is between 0.0 and 1.0 it corresponds to the DDIM scheduler, and if its value is between -0.0 and 1.0 it corresponds to the DDPM scheduler. TYPE: `float` DEFAULT: 0.0
trained_betas	Pass an array of betas directly to the constructor to bypass beta_start and beta_end. TYPE: `np.ndarray`, *optional* DEFAULT: None
clip_sample	Clip the predicted sample between -1 and 1 for numerical stability. TYPE: `bool`, defaults to `True` DEFAULT: True

Source code in mindone/diffusers/schedulers/scheduling_repaint.py

class RePaintScheduler(SchedulerMixin, ConfigMixin):
    """
    `RePaintScheduler` is a scheduler for DDPM inpainting inside a given mask.

    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
    methods the library implements for all schedulers such as loading and saving.

    Args:
        num_train_timesteps (`int`, defaults to 1000):
            The number of diffusion steps to train the model.
        beta_start (`float`, defaults to 0.0001):
            The starting `beta` value of inference.
        beta_end (`float`, defaults to 0.02):
            The final `beta` value.
        beta_schedule (`str`, defaults to `"linear"`):
            The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
            `linear`, `scaled_linear`, `squaredcos_cap_v2`, or `sigmoid`.
        eta (`float`):
            The weight of noise for added noise in diffusion step. If its value is between 0.0 and 1.0 it corresponds
            to the DDIM scheduler, and if its value is between -0.0 and 1.0 it corresponds to the DDPM scheduler.
        trained_betas (`np.ndarray`, *optional*):
            Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
        clip_sample (`bool`, defaults to `True`):
            Clip the predicted sample between -1 and 1 for numerical stability.

    """

    order = 1

    @register_to_config
    def __init__(
        self,
        num_train_timesteps: int = 1000,
        beta_start: float = 0.0001,
        beta_end: float = 0.02,
        beta_schedule: str = "linear",
        eta: float = 0.0,
        trained_betas: Optional[np.ndarray] = None,
        clip_sample: bool = True,
    ):
        if trained_betas is not None:
            self.betas = ms.tensor(trained_betas)
        elif beta_schedule == "linear":
            self.betas = ms.tensor(np.linspace(beta_start, beta_end, num_train_timesteps), dtype=ms.float32)
        elif beta_schedule == "scaled_linear":
            # this schedule is very specific to the latent diffusion model.
            self.betas = (
                ms.tensor(np.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps), dtype=ms.float32) ** 2
            )
        elif beta_schedule == "squaredcos_cap_v2":
            # Glide cosine schedule
            self.betas = betas_for_alpha_bar(num_train_timesteps)
        elif beta_schedule == "sigmoid":
            # GeoDiff sigmoid schedule
            betas = ms.tensor(np.linspace(-6, 6, num_train_timesteps))
            self.betas = ops.sigmoid(betas) * (beta_end - beta_start) + beta_start
        else:
            raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")

        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = ops.cumprod(self.alphas, dim=0)
        self.one = ms.tensor(1.0)

        self.final_alpha_cumprod = ms.tensor(1.0)

        # standard deviation of the initial noise distribution
        self.init_noise_sigma = 1.0

        # setable values
        self.num_inference_steps = None
        self.timesteps = ms.tensor(np.arange(0, num_train_timesteps)[::-1].copy())

        self.eta = eta

    def scale_model_input(self, sample: ms.Tensor, timestep: Optional[int] = None) -> ms.Tensor:
        """
        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
        current timestep.

        Args:
            sample (`ms.Tensor`):
                The input sample.
            timestep (`int`, *optional*):
                The current timestep in the diffusion chain.

        Returns:
            `ms.Tensor`:
                A scaled input sample.
        """
        return sample

    def set_timesteps(
        self,
        num_inference_steps: int,
        jump_length: int = 10,
        jump_n_sample: int = 10,
    ):
        """
        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 used,
                `timesteps` must be `None`.
            jump_length (`int`, defaults to 10):
                The number of steps taken forward in time before going backward in time for a single jump (“j” in
                RePaint paper). Take a look at Figure 9 and 10 in the paper.
            jump_n_sample (`int`, defaults to 10):
                The number of times to make a forward time jump for a given chosen time sample. Take a look at Figure 9
                and 10 in the paper.

        """
        num_inference_steps = min(self.config.num_train_timesteps, num_inference_steps)
        self.num_inference_steps = num_inference_steps

        timesteps = []

        jumps = {}
        for j in range(0, num_inference_steps - jump_length, jump_length):
            jumps[j] = jump_n_sample - 1

        t = num_inference_steps
        while t >= 1:
            t = t - 1
            timesteps.append(t)

            if jumps.get(t, 0) > 0:
                jumps[t] = jumps[t] - 1
                for _ in range(jump_length):
                    t = t + 1
                    timesteps.append(t)

        timesteps = np.array(timesteps) * (self.config.num_train_timesteps // self.num_inference_steps)
        self.timesteps = ms.tensor(timesteps)

    def _get_variance(self, t):
        prev_timestep = t - self.config.num_train_timesteps // self.num_inference_steps

        alpha_prod_t = self.alphas_cumprod[t]
        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

        # For t > 0, compute predicted variance βt (see formula (6) and (7) from
        # https://arxiv.org/pdf/2006.11239.pdf) and sample from it to get
        # previous sample x_{t-1} ~ N(pred_prev_sample, variance) == add
        # variance to pred_sample
        # Is equivalent to formula (16) in https://arxiv.org/pdf/2010.02502.pdf
        # without eta.
        # variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * self.betas[t]
        variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)

        return variance

    def step(
        self,
        model_output: ms.Tensor,
        timestep: int,
        sample: ms.Tensor,
        original_image: ms.Tensor,
        mask: ms.Tensor,
        generator: Optional[np.random.Generator] = None,
        return_dict: bool = False,
    ) -> Union[RePaintSchedulerOutput, Tuple]:
        """
        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
        process from the learned model outputs (most often the predicted noise).

        Args:
            model_output (`ms.Tensor`):
                The direct output from learned diffusion model.
            timestep (`int`):
                The current discrete timestep in the diffusion chain.
            sample (`ms.Tensor`):
                A current instance of a sample created by the diffusion process.
            original_image (`ms.Tensor`):
                The original image to inpaint on.
            mask (`ms.Tensor`):
                The mask where a value of 0.0 indicates which part of the original image to inpaint.
            generator (`np.random.Generator`, *optional*):
                A random number generator.
            return_dict (`bool`, *optional*, defaults to `False`):
                Whether or not to return a [`~schedulers.scheduling_repaint.RePaintSchedulerOutput`] or `tuple`.

        Returns:
            [`~schedulers.scheduling_repaint.RePaintSchedulerOutput`] or `tuple`:
                If return_dict is `True`, [`~schedulers.scheduling_repaint.RePaintSchedulerOutput`] is returned,
                otherwise a tuple is returned where the first element is the sample tensor.

        """
        dtype = sample.dtype
        t = timestep
        prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps

        # 1. compute alphas, betas
        alpha_prod_t = self.alphas_cumprod[t]
        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

        # 2. compute predicted original sample from predicted noise also called
        # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
        pred_original_sample = ((sample - (beta_prod_t**0.5).to(dtype) * model_output) / alpha_prod_t**0.5).to(
            dtype
        )

        # 3. Clip "predicted x_0"
        if self.config.clip_sample:
            pred_original_sample = ops.clamp(pred_original_sample, -1, 1)

        # We choose to follow RePaint Algorithm 1 to get x_{t-1}, however we
        # substitute formula (7) in the algorithm coming from DDPM paper
        # (formula (4) Algorithm 2 - Sampling) with formula (12) from DDIM paper.
        # DDIM schedule gives the same results as DDPM with eta = 1.0
        # Noise is being reused in 7. and 8., but no impact on quality has
        # been observed.

        # 5. Add noise
        noise = randn_tensor(model_output.shape, generator=generator, dtype=model_output.dtype)
        std_dev_t = self.eta * self._get_variance(timestep) ** 0.5

        variance = 0
        if t > 0 and self.eta > 0:
            variance = std_dev_t * noise

        # 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).to(dtype) * model_output

        # 7. compute x_{t-1} of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
        prev_unknown_part = (
            (alpha_prod_t_prev**0.5).to(dtype) * pred_original_sample + pred_sample_direction + variance
        )

        # 8. Algorithm 1 Line 5 https://arxiv.org/pdf/2201.09865.pdf
        prev_known_part = (alpha_prod_t_prev**0.5).to(dtype) * original_image + ((1 - alpha_prod_t_prev) ** 0.5).to(
            dtype
        ) * noise

        # 9. Algorithm 1 Line 8 https://arxiv.org/pdf/2201.09865.pdf
        pred_prev_sample = (mask * prev_known_part + (1.0 - mask) * prev_unknown_part).to(dtype)

        if not return_dict:
            return (
                pred_prev_sample,
                pred_original_sample,
            )

        return RePaintSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample)

    def undo_step(self, sample, timestep, generator=None):
        n = self.config.num_train_timesteps // self.num_inference_steps

        for i in range(n):
            beta = self.betas[timestep + i]
            noise = randn_tensor(sample.shape, generator=generator, dtype=sample.dtype)

            # 10. Algorithm 1 Line 10 https://arxiv.org/pdf/2201.09865.pdf
            sample = ((1 - beta) ** 0.5).to(sample.dtype) * sample + (beta**0.5).to(noise.dtype) * noise

        return sample

    def add_noise(
        self,
        original_samples: ms.Tensor,
        noise: ms.Tensor,
        timesteps: ms.Tensor,
    ) -> ms.Tensor:
        raise NotImplementedError("Use `DDPMScheduler.add_noise()` to train for sampling with RePaint.")

    def __len__(self):
        return self.config.num_train_timesteps

mindone.diffusers.RePaintScheduler.scale_model_input(sample, timestep=None)

Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep.

PARAMETER	DESCRIPTION
sample	The input sample. TYPE: `ms.Tensor`
timestep	The current timestep in the diffusion chain. TYPE: `int`, *optional* DEFAULT: None

RETURNS	DESCRIPTION
Tensor	ms.Tensor: A scaled input sample.

Source code in mindone/diffusers/schedulers/scheduling_repaint.py

def scale_model_input(self, sample: ms.Tensor, timestep: Optional[int] = None) -> ms.Tensor:
    """
    Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
    current timestep.

    Args:
        sample (`ms.Tensor`):
            The input sample.
        timestep (`int`, *optional*):
            The current timestep in the diffusion chain.

    Returns:
        `ms.Tensor`:
            A scaled input sample.
    """
    return sample

mindone.diffusers.RePaintScheduler.set_timesteps(num_inference_steps, jump_length=10, jump_n_sample=10)

Sets the discrete timesteps used for the diffusion chain (to be run before inference).

PARAMETER	DESCRIPTION
num_inference_steps	The number of diffusion steps used when generating samples with a pre-trained model. If used, timesteps must be None. TYPE: `int`
jump_length	The number of steps taken forward in time before going backward in time for a single jump (“j” in RePaint paper). Take a look at Figure 9 and 10 in the paper. TYPE: `int`, defaults to 10 DEFAULT: 10
jump_n_sample	The number of times to make a forward time jump for a given chosen time sample. Take a look at Figure 9 and 10 in the paper. TYPE: `int`, defaults to 10 DEFAULT: 10

Source code in mindone/diffusers/schedulers/scheduling_repaint.py

def set_timesteps(
    self,
    num_inference_steps: int,
    jump_length: int = 10,
    jump_n_sample: int = 10,
):
    """
    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 used,
            `timesteps` must be `None`.
        jump_length (`int`, defaults to 10):
            The number of steps taken forward in time before going backward in time for a single jump (“j” in
            RePaint paper). Take a look at Figure 9 and 10 in the paper.
        jump_n_sample (`int`, defaults to 10):
            The number of times to make a forward time jump for a given chosen time sample. Take a look at Figure 9
            and 10 in the paper.

    """
    num_inference_steps = min(self.config.num_train_timesteps, num_inference_steps)
    self.num_inference_steps = num_inference_steps

    timesteps = []

    jumps = {}
    for j in range(0, num_inference_steps - jump_length, jump_length):
        jumps[j] = jump_n_sample - 1

    t = num_inference_steps
    while t >= 1:
        t = t - 1
        timesteps.append(t)

        if jumps.get(t, 0) > 0:
            jumps[t] = jumps[t] - 1
            for _ in range(jump_length):
                t = t + 1
                timesteps.append(t)

    timesteps = np.array(timesteps) * (self.config.num_train_timesteps // self.num_inference_steps)
    self.timesteps = ms.tensor(timesteps)

mindone.diffusers.RePaintScheduler.step(model_output, timestep, sample, original_image, mask, generator=None, return_dict=False)

Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise).

PARAMETER	DESCRIPTION
model_output	The direct output from learned diffusion model. TYPE: `ms.Tensor`
timestep	The current discrete timestep in the diffusion chain. TYPE: `int`
sample	A current instance of a sample created by the diffusion process. TYPE: `ms.Tensor`
original_image	The original image to inpaint on. TYPE: `ms.Tensor`
mask	The mask where a value of 0.0 indicates which part of the original image to inpaint. TYPE: `ms.Tensor`
generator	A random number generator. TYPE: `np.random.Generator`, *optional* DEFAULT: None
return_dict	Whether or not to return a [~schedulers.scheduling_repaint.RePaintSchedulerOutput] or tuple. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False

RETURNS	DESCRIPTION
Union[RePaintSchedulerOutput, Tuple]	[~schedulers.scheduling_repaint.RePaintSchedulerOutput] or tuple: If return_dict is True, [~schedulers.scheduling_repaint.RePaintSchedulerOutput] is returned, otherwise a tuple is returned where the first element is the sample tensor.

Source code in mindone/diffusers/schedulers/scheduling_repaint.py

def step(
    self,
    model_output: ms.Tensor,
    timestep: int,
    sample: ms.Tensor,
    original_image: ms.Tensor,
    mask: ms.Tensor,
    generator: Optional[np.random.Generator] = None,
    return_dict: bool = False,
) -> Union[RePaintSchedulerOutput, Tuple]:
    """
    Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
    process from the learned model outputs (most often the predicted noise).

    Args:
        model_output (`ms.Tensor`):
            The direct output from learned diffusion model.
        timestep (`int`):
            The current discrete timestep in the diffusion chain.
        sample (`ms.Tensor`):
            A current instance of a sample created by the diffusion process.
        original_image (`ms.Tensor`):
            The original image to inpaint on.
        mask (`ms.Tensor`):
            The mask where a value of 0.0 indicates which part of the original image to inpaint.
        generator (`np.random.Generator`, *optional*):
            A random number generator.
        return_dict (`bool`, *optional*, defaults to `False`):
            Whether or not to return a [`~schedulers.scheduling_repaint.RePaintSchedulerOutput`] or `tuple`.

    Returns:
        [`~schedulers.scheduling_repaint.RePaintSchedulerOutput`] or `tuple`:
            If return_dict is `True`, [`~schedulers.scheduling_repaint.RePaintSchedulerOutput`] is returned,
            otherwise a tuple is returned where the first element is the sample tensor.

    """
    dtype = sample.dtype
    t = timestep
    prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps

    # 1. compute alphas, betas
    alpha_prod_t = self.alphas_cumprod[t]
    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

    # 2. compute predicted original sample from predicted noise also called
    # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
    pred_original_sample = ((sample - (beta_prod_t**0.5).to(dtype) * model_output) / alpha_prod_t**0.5).to(
        dtype
    )

    # 3. Clip "predicted x_0"
    if self.config.clip_sample:
        pred_original_sample = ops.clamp(pred_original_sample, -1, 1)

    # We choose to follow RePaint Algorithm 1 to get x_{t-1}, however we
    # substitute formula (7) in the algorithm coming from DDPM paper
    # (formula (4) Algorithm 2 - Sampling) with formula (12) from DDIM paper.
    # DDIM schedule gives the same results as DDPM with eta = 1.0
    # Noise is being reused in 7. and 8., but no impact on quality has
    # been observed.

    # 5. Add noise
    noise = randn_tensor(model_output.shape, generator=generator, dtype=model_output.dtype)
    std_dev_t = self.eta * self._get_variance(timestep) ** 0.5

    variance = 0
    if t > 0 and self.eta > 0:
        variance = std_dev_t * noise

    # 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).to(dtype) * model_output

    # 7. compute x_{t-1} of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
    prev_unknown_part = (
        (alpha_prod_t_prev**0.5).to(dtype) * pred_original_sample + pred_sample_direction + variance
    )

    # 8. Algorithm 1 Line 5 https://arxiv.org/pdf/2201.09865.pdf
    prev_known_part = (alpha_prod_t_prev**0.5).to(dtype) * original_image + ((1 - alpha_prod_t_prev) ** 0.5).to(
        dtype
    ) * noise

    # 9. Algorithm 1 Line 8 https://arxiv.org/pdf/2201.09865.pdf
    pred_prev_sample = (mask * prev_known_part + (1.0 - mask) * prev_unknown_part).to(dtype)

    if not return_dict:
        return (
            pred_prev_sample,
            pred_original_sample,
        )

    return RePaintSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample)

mindone.diffusers.schedulers.scheduling_repaint.RePaintSchedulerOutput dataclass

Bases: BaseOutput

Output class for the scheduler's step function output.

PARAMETER	DESCRIPTION
prev_sample	Computed sample (x_{t-1}) of previous timestep. prev_sample should be used as next model input in the denoising loop. TYPE: `ms.Tensor` of shape `(batch_size, num_channels, height, width)` for images
pred_original_sample	The predicted denoised sample (x_{0}) based on the model output from the current timestep. pred_original_sample can be used to preview progress or for guidance. TYPE: `ms.Tensor` of shape `(batch_size, num_channels, height, width)` for images

Source code in mindone/diffusers/schedulers/scheduling_repaint.py

@dataclass
class RePaintSchedulerOutput(BaseOutput):
    """
    Output class for the scheduler's step function output.

    Args:
        prev_sample (`ms.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
            Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the
            denoising loop.
        pred_original_sample (`ms.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
            The predicted denoised sample (x_{0}) based on the model output from
             the current timestep. `pred_original_sample` can be used to preview progress or for guidance.
    """

    prev_sample: ms.Tensor
    pred_original_sample: ms.Tensor