src/diffusion/noise_schedule.py

from __future__ import annotations

import math
from dataclasses import dataclass

import torch


@dataclass
class NoiseSchedule:
    """
    Precomputed DDPM noise schedule.

    Main variables:

        beta_t:
            amount of noise added at timestep t

        alpha_t:
            1 - beta_t

        alpha_bar_t:
            cumulative product of alphas up to t

        q(z_t | z_0):
            z_t = sqrt(alpha_bar_t) * z_0
                + sqrt(1 - alpha_bar_t) * eps
    """

    betas: torch.Tensor
    alphas: torch.Tensor
    alphas_cumprod: torch.Tensor
    alphas_cumprod_prev: torch.Tensor

    sqrt_alphas_cumprod: torch.Tensor
    sqrt_one_minus_alphas_cumprod: torch.Tensor

    log_one_minus_alphas_cumprod: torch.Tensor

    sqrt_recip_alphas_cumprod: torch.Tensor
    sqrt_recipm1_alphas_cumprod: torch.Tensor

    posterior_variance: torch.Tensor
    posterior_log_variance_clipped: torch.Tensor
    posterior_mean_coef1: torch.Tensor
    posterior_mean_coef2: torch.Tensor

    num_timesteps: int
    schedule_type: str

    def to(self, device: torch.device | str) -> "NoiseSchedule":
        device = torch.device(device)

        return NoiseSchedule(
            betas=self.betas.to(device),
            alphas=self.alphas.to(device),
            alphas_cumprod=self.alphas_cumprod.to(device),
            alphas_cumprod_prev=self.alphas_cumprod_prev.to(device),
            sqrt_alphas_cumprod=self.sqrt_alphas_cumprod.to(device),
            sqrt_one_minus_alphas_cumprod=self.sqrt_one_minus_alphas_cumprod.to(device),
            log_one_minus_alphas_cumprod=self.log_one_minus_alphas_cumprod.to(device),
            sqrt_recip_alphas_cumprod=self.sqrt_recip_alphas_cumprod.to(device),
            sqrt_recipm1_alphas_cumprod=self.sqrt_recipm1_alphas_cumprod.to(device),
            posterior_variance=self.posterior_variance.to(device),
            posterior_log_variance_clipped=self.posterior_log_variance_clipped.to(device),
            posterior_mean_coef1=self.posterior_mean_coef1.to(device),
            posterior_mean_coef2=self.posterior_mean_coef2.to(device),
            num_timesteps=self.num_timesteps,
            schedule_type=self.schedule_type,
        )


def make_beta_schedule(
    schedule_type: str = "cosine",
    num_timesteps: int = 1000,
    beta_start: float = 1e-4,
    beta_end: float = 2e-2,
    cosine_s: float = 0.008,
    max_beta: float = 0.999,
) -> torch.Tensor:
    """
    Create beta schedule.

    Supported:
        linear:
            Standard DDPM linear beta schedule.

        cosine:
            Improved DDPM cosine schedule.
            Usually better behaved and good default for v-prediction.

    Returns:
        betas: [num_timesteps], float32
    """
    schedule_type = schedule_type.lower()

    if schedule_type == "linear":
        betas = torch.linspace(
            beta_start,
            beta_end,
            num_timesteps,
            dtype=torch.float64,
        )

    elif schedule_type == "cosine":
        betas = cosine_beta_schedule(
            num_timesteps=num_timesteps,
            cosine_s=cosine_s,
            max_beta=max_beta,
        )

    else:
        raise ValueError(
            f"Unknown schedule_type={schedule_type}. "
            "Use 'linear' or 'cosine'."
        )

    return betas.float()


def cosine_beta_schedule(
    num_timesteps: int,
    cosine_s: float = 0.008,
    max_beta: float = 0.999,
) -> torch.Tensor:
    """
    Cosine beta schedule
    Instead of directly defining beta_t, we define alpha_bar(t)
    using a cosine curve, then derive beta_t.
    """
    steps = num_timesteps + 1

    x = torch.linspace(
        0,
        num_timesteps,
        steps,
        dtype=torch.float64,
    )

    alphas_cumprod = torch.cos(
        ((x / num_timesteps) + cosine_s)
        / (1.0 + cosine_s)
        * math.pi
        * 0.5
    ) ** 2

    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]

    betas = 1.0 - (
        alphas_cumprod[1:] / alphas_cumprod[:-1]
    )

    betas = torch.clamp(
        betas,
        min=1e-8,
        max=max_beta,
    )

    return betas


def create_noise_schedule(
    schedule_type: str = "cosine",
    num_timesteps: int = 1000,
    beta_start: float = 1e-4,
    beta_end: float = 2e-2,
    cosine_s: float = 0.008,
    max_beta: float = 0.999,
) -> NoiseSchedule:
    """
    all precomputed schedule tensors needed for DDPM training and sampling.
    """
    betas = make_beta_schedule(
        schedule_type=schedule_type,
        num_timesteps=num_timesteps,
        beta_start=beta_start,
        beta_end=beta_end,
        cosine_s=cosine_s,
        max_beta=max_beta,
    )

    alphas = 1.0 - betas

    alphas_cumprod = torch.cumprod(
        alphas,
        dim=0,
    )

    alphas_cumprod_prev = torch.cat(
        [
            torch.ones(1, dtype=alphas_cumprod.dtype),
            alphas_cumprod[:-1],
        ],
        dim=0,
    )

    sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)

    sqrt_one_minus_alphas_cumprod = torch.sqrt(
        1.0 - alphas_cumprod
    )

    log_one_minus_alphas_cumprod = torch.log(
        torch.clamp(
            1.0 - alphas_cumprod,
            min=1e-20,
        )
    )

    sqrt_recip_alphas_cumprod = torch.sqrt(
        1.0 / alphas_cumprod
    )

    sqrt_recipm1_alphas_cumprod = torch.sqrt(
        1.0 / alphas_cumprod - 1.0
    )

    # Posterior q(z_{t-1} | z_t, z_0)
    posterior_variance = (
        betas
        * (1.0 - alphas_cumprod_prev)
        / (1.0 - alphas_cumprod)
    )

    posterior_log_variance_clipped = torch.log(
        torch.clamp(
            posterior_variance,
            min=1e-20,
        )
    )

    posterior_mean_coef1 = (
        betas
        * torch.sqrt(alphas_cumprod_prev)
        / (1.0 - alphas_cumprod)
    )

    posterior_mean_coef2 = (
        (1.0 - alphas_cumprod_prev)
        * torch.sqrt(alphas)
        / (1.0 - alphas_cumprod)
    )

    return NoiseSchedule(
        betas=betas,
        alphas=alphas,
        alphas_cumprod=alphas_cumprod,
        alphas_cumprod_prev=alphas_cumprod_prev,
        sqrt_alphas_cumprod=sqrt_alphas_cumprod,
        sqrt_one_minus_alphas_cumprod=sqrt_one_minus_alphas_cumprod,
        log_one_minus_alphas_cumprod=log_one_minus_alphas_cumprod,
        sqrt_recip_alphas_cumprod=sqrt_recip_alphas_cumprod,
        sqrt_recipm1_alphas_cumprod=sqrt_recipm1_alphas_cumprod,
        posterior_variance=posterior_variance,
        posterior_log_variance_clipped=posterior_log_variance_clipped,
        posterior_mean_coef1=posterior_mean_coef1,
        posterior_mean_coef2=posterior_mean_coef2,
        num_timesteps=num_timesteps,
        schedule_type=schedule_type,
    )


def extract(
    values: torch.Tensor,
    timesteps: torch.Tensor,
    broadcast_shape: tuple[int, ...],
) -> torch.Tensor:
    """
    Extract values[t] and reshape for broadcasting.

    Args:
        values:
            Schedule tensor with shape [T].

        timesteps:
            Long tensor with shape [B].

        broadcast_shape:
            Shape of target tensor, e.g. z_t.shape = [B, C, H, W].

    Returns:
        Tensor with shape [B, 1, 1, 1], broadcastable to broadcast_shape.

    Example:
        sqrt_alpha_bar_t = extract(
            schedule.sqrt_alphas_cumprod,
            t,
            z_0.shape,
        )
    """
    if timesteps.dtype != torch.long:
        timesteps = timesteps.long()

    out = values.gather(
        dim=0,
        index=timesteps,
    )

    while len(out.shape) < len(broadcast_shape):
        out = out[..., None]

    return out