TorchDiff 2.0.0__py3-none-any.whl

ldm/forward_idm.py

@@ -0,0 +1,100 @@
+import torch
+import torch.nn as nn
+
+
+
+
+class ForwardSDE(nn.Module):
+    def __init__(self, hyper_params, method):
+        super().__init__()
+        self.hyper_params = hyper_params
+        self.method = method
+
+    def forward(self, x0, noise, time_steps):
+
+        dt = self.hyper_params.dt
+        if self.method == "ve":
+            sigma_t = self.hyper_params.sigmas[time_steps]
+            sigma_t_prev = self.hyper_params.sigmas[time_steps - 1] if time_steps.min() > 0 else torch.zeros_like(sigma_t)
+            sigma_diff = torch.sqrt(torch.clamp(sigma_t ** 2 - sigma_t_prev ** 2, min=0))
+            x0 = x0 + noise * sigma_diff.view(-1, 1, 1, 1)
+
+        elif self.method == "vp":
+            betas = self.hyper_params.betas[time_steps].view(-1, 1, 1, 1)
+            drift = -0.5 * betas * x0 * dt
+            diffusion = torch.sqrt(betas * dt) * noise
+            x0 = x0 + drift + diffusion
+
+        elif self.method == "sub-vp":
+            betas = self.hyper_params.betas[time_steps].view(-1, 1, 1, 1)
+            cum_betas = self.hyper_params.cum_betas[time_steps].view(-1, 1, 1, 1)
+            drift = -0.5 * betas * x0 * dt
+            diffusion = torch.sqrt(betas * (1 - torch.exp(-2 * cum_betas)) * dt) * noise
+            x0 = x0 + drift + diffusion
+
+        elif self.method == "ode":
+            if self.method == "ve":
+                x0 = x0
+            else:
+                betas = self.hyper_params.betas[time_steps].view(-1, 1, 1, 1)
+                drift = -0.5 * betas * x0 * dt
+                x0 = x0 + drift
+        else:
+            raise ValueError(f"Unknown method: {self.method}")
+        return x0
+
+
+
+class ForwardDDPM(nn.Module):
+    """forward diffusion process of DDPM."""
+    def __init__(self, hyper_params):
+        super().__init__()
+        self.hyper_params = hyper_params
+
+    def forward(self, x0, noise, time_steps):
+        if not torch.all((time_steps >= 0) & (time_steps < self.hyper_params.num_steps)):
+            raise ValueError(f"time_steps must be between 0 and {self.hyper_params.num_steps - 1}")
+
+        if self.hyper_params.trainable_beta:
+            _, _, _, sqrt_alpha_bar_t, sqrt_one_minus_alpha_bar_t = self.hyper_params.compute_schedule(
+                self.hyper_params.betas
+            )
+            sqrt_alpha_bar_t = sqrt_alpha_bar_t[time_steps].to(x0.device)
+            sqrt_one_minus_alpha_bar_t = sqrt_one_minus_alpha_bar_t[time_steps].to(x0.device)
+        else:
+            sqrt_alpha_bar_t = self.hyper_params.sqrt_alpha_bars[time_steps].to(x0.device)
+            sqrt_one_minus_alpha_bar_t = self.hyper_params.sqrt_one_minus_alpha_bars[time_steps].to(x0.device)
+
+        sqrt_alpha_bar_t = sqrt_alpha_bar_t.view(-1, 1, 1, 1)
+        sqrt_one_minus_alpha_bar_t = sqrt_one_minus_alpha_bar_t.view(-1, 1, 1, 1)
+
+        xt = sqrt_alpha_bar_t * x0 + sqrt_one_minus_alpha_bar_t * noise
+        return xt
+
+
+
+class ForwardDDIM(nn.Module):
+    def __init__(self, hyper_params):
+        """forward diffusion process of DDIM"""
+        super().__init__()
+        self.hyper_params = hyper_params
+
+    def forward(self, x0, noise, time_steps):
+        if not torch.all((time_steps >= 0) & (time_steps < self.hyper_params.num_steps)):
+            raise ValueError(f"time_steps must be between 0 and {self.hyper_params.num_steps - 1}")
+
+        if self.hyper_params.trainable_beta:
+            _, _, _, sqrt_alpha_cumprod_t, sqrt_one_minus_alpha_cumprod_t = self.hyper_params.compute_schedule(
+                self.hyper_params.betas
+            )
+            sqrt_alpha_cumprod_t = sqrt_alpha_cumprod_t[time_steps].to(x0.device)
+            sqrt_one_minus_alpha_cumprod_t = sqrt_one_minus_alpha_cumprod_t[time_steps].to(x0.device)
+        else:
+            sqrt_alpha_cumprod_t = self.hyper_params.sqrt_alpha_cumprod[time_steps].to(x0.device)
+            sqrt_one_minus_alpha_cumprod_t = self.hyper_params.sqrt_one_minus_alpha_cumprod[time_steps].to(x0.device)
+
+        sqrt_alpha_cumprod_t = sqrt_alpha_cumprod_t.view(-1, 1, 1, 1)
+        sqrt_one_minus_alpha_cumprod_t = sqrt_one_minus_alpha_cumprod_t.view(-1, 1, 1, 1)
+
+        xt = sqrt_alpha_cumprod_t * x0 + sqrt_one_minus_alpha_cumprod_t * noise
+        return xt

ldm/hyper_param.py

@@ -0,0 +1,239 @@
+import torch
+import torch.nn as nn
+
+
+
+class HyperParamsDDPM(nn.Module):
+    """hyperparameters for DDPM noise schedule with flexible beta computation."""
+    def __init__(self, num_steps=1000, beta_start=1e-4, beta_end=0.02, trainable_beta=False, beta_method="linear"):
+        super().__init__()
+        self.num_steps = num_steps
+        self.beta_start = beta_start
+        self.beta_end = beta_end
+        self.trainable_beta = trainable_beta
+        self.beta_method = beta_method
+
+        # validate inputs
+        if not (0 < beta_start < beta_end < 1):
+            raise ValueError(f"beta_start ({beta_start}) and beta_end ({beta_end}) must satisfy 0 < start < end < 1")
+        if num_steps <= 0:
+            raise ValueError(f"num_steps ({num_steps}) must be positive")
+
+        # compute initial beta schedule
+        beta_range = (beta_start, beta_end)
+        betas_init = self.compute_beta_schedule(beta_range, num_steps, beta_method)
+
+        # initialize betas
+        if trainable_beta:
+            self.betas = nn.Parameter(betas_init)  # Trainable parameter
+        else:
+            self.register_buffer('betas', betas_init)  # Fixed buffer
+            self.register_buffer('alphas', 1 - self.betas)
+            self.register_buffer('alpha_bars', torch.cumprod(self.alphas, dim=0))
+            self.register_buffer('sqrt_alpha_bars', torch.sqrt(self.alpha_bars))
+            self.register_buffer('sqrt_one_minus_alpha_bars', torch.sqrt(1 - self.alpha_bars))
+
+    def compute_beta_schedule(self, beta_range, num_steps, method):
+        """
+        Computes the beta schedule based on the selected method.
+        Args:
+            beta_range: Tuple of (min_beta, max_beta) values
+            num_steps: Number of diffusion steps
+            method: Method for beta schedule ("linear", "sigmoid", "quadratic", "constant", "inverse_time")
+        Returns:
+            Tensor of beta values, shape [num_steps]
+        """
+        beta_min, beta_max = beta_range
+        if method == "sigmoid":
+            x = torch.linspace(-6, 6, num_steps)
+            beta = torch.sigmoid(x) * (beta_max - beta_min) + beta_min
+        elif method == "quadratic":
+            x = torch.linspace(beta_min**0.5, beta_max**0.5, num_steps)
+            beta = x**2
+        elif method == "constant":
+            beta = torch.full((num_steps,), beta_max)
+        elif method == "inverse_time":
+            beta = 1.0 / torch.linspace(num_steps, 1, num_steps)
+            # scale to beta_range
+            beta = beta_min + (beta_max - beta_min) * (beta - beta.min()) / (beta.max() - beta.min())
+        elif method == "linear":
+            beta = torch.linspace(beta_min, beta_max, num_steps)
+        else:
+            raise ValueError(f"Unknown beta_method: {method}. Supported: linear, sigmoid, quadratic, constant, inverse_time")
+
+        beta = torch.clamp(beta, min=beta_min, max=beta_max)
+        return beta
+
+    @staticmethod
+    def compute_schedule(betas):
+        """computes noise schedule parameters dynamically from betas."""
+        alphas = 1 - betas
+        alpha_bars = torch.cumprod(alphas, dim=0)
+        return betas, alphas, alpha_bars, torch.sqrt(alpha_bars), torch.sqrt(1 - alpha_bars)
+
+    def constrain_betas(self):
+        """constrains betas to a valid range during training."""
+        if self.trainable_beta:
+            with torch.no_grad():
+                self.betas.clamp_(min=self.beta_start, max=self.beta_end)
+
+
+
+
+
+class HyperParamsDDIM(nn.Module):
+    """Hyperparameters for DDIM noise schedule with flexible beta computation."""
+    def __init__(self, eta=None, num_steps=1000, tau_num_steps=100, beta_start=1e-4, beta_end=0.02,
+                 trainable_beta=False, beta_method="linear"):
+        super().__init__()
+        self.eta = eta or 0
+        self.num_steps = num_steps
+        self.tau_num_steps = tau_num_steps
+        self.beta_start = beta_start
+        self.beta_end = beta_end
+        self.trainable_beta = trainable_beta
+        self.beta_method = beta_method
+
+        if not (0 < beta_start < beta_end < 1):
+            raise ValueError(f"beta_start ({beta_start}) and beta_end ({beta_end}) must satisfy 0 < start < end < 1")
+        if num_steps <= 0:
+            raise ValueError(f"num_steps ({num_steps}) must be positive")
+
+        beta_range = (beta_start, beta_end)
+        betas_init = self.compute_beta_schedule(beta_range, num_steps, beta_method)
+
+        if trainable_beta:
+            self.betas = nn.Parameter(betas_init)
+        else:
+            self.register_buffer('betas', betas_init)
+            self.register_buffer('alphas', 1 - self.betas)
+            self.register_buffer('alpha_cumprod', torch.cumprod(self.alphas, dim=0))
+            self.register_buffer('sqrt_alpha_cumprod', torch.sqrt(self.alpha_cumprod))
+            self.register_buffer('sqrt_one_minus_alpha_cumprod', torch.sqrt(1 - self.alpha_cumprod))
+
+        self.register_buffer('tau_indices', torch.linspace(0, num_steps - 1, tau_num_steps, dtype=torch.long))
+
+    def compute_beta_schedule(self, beta_range, num_steps, method):
+
+        beta_min, beta_max = beta_range
+        if method == "sigmoid":
+            x = torch.linspace(-6, 6, num_steps)
+            beta = torch.sigmoid(x) * (beta_max - beta_min) + beta_min
+        elif method == "quadratic":
+            x = torch.linspace(beta_min ** 0.5, beta_max ** 0.5, num_steps)
+            beta = x ** 2
+        elif method == "constant":
+            beta = torch.full((num_steps,), beta_max)
+        elif method == "inverse_time":
+            beta = 1.0 / torch.linspace(num_steps, 1, num_steps)
+            # scale to beta_range
+            beta = beta_min + (beta_max - beta_min) * (beta - beta.min()) / (beta.max() - beta.min())
+        elif method == "linear":
+            beta = torch.linspace(beta_min, beta_max, num_steps)
+        else:
+            raise ValueError(
+                f"Unknown beta_method: {method}. Supported: linear, sigmoid, quadratic, constant, inverse_time")
+
+        beta = torch.clamp(beta, min=beta_min, max=beta_max)
+        return beta
+
+    def get_tau_schedule(self):
+        if self.trainable_beta:
+            betas, alphas, alpha_cumprod, sqrt_alpha_cumprod, sqrt_one_minus_alpha_cumprod = self.compute_schedule(self.betas)
+        else:
+            betas = self.betas
+            alphas = self.alphas
+            alpha_cumprod = self.alpha_cumprod
+            sqrt_alpha_cumprod = self.sqrt_alpha_cumprod
+            sqrt_one_minus_alpha_cumprod = self.sqrt_one_minus_alpha_cumprod
+
+        tau_betas = betas[self.tau_indices]
+        tau_alphas = alphas[self.tau_indices]
+        tau_alpha_cumprod = alpha_cumprod[self.tau_indices]
+        tau_sqrt_alpha_cumprod = sqrt_alpha_cumprod[self.tau_indices]
+        tau_sqrt_one_minus_alpha_cumprod = sqrt_one_minus_alpha_cumprod[self.tau_indices]
+
+        return tau_betas, tau_alphas, tau_alpha_cumprod, tau_sqrt_alpha_cumprod, tau_sqrt_one_minus_alpha_cumprod
+
+    def compute_schedule(self, betas):
+        """computes noise schedule parameters dynamically from betas."""
+        alphas = 1 - betas
+        alpha_cumprod = torch.cumprod(alphas, dim=0)
+        return betas, alphas, alpha_cumprod, torch.sqrt(alpha_cumprod), torch.sqrt(1 - alpha_cumprod)
+
+    def constrain_betas(self):
+        """constrains betas to a valid range during training."""
+        if self.trainable_beta:
+            with torch.no_grad():
+                self.betas.clamp_(min=self.beta_start, max=self.beta_end)
+
+
+
+
+class HyperParamsSDE(nn.Module):
+    def __init__(self, num_steps=1000, beta_start=1e-4, beta_end=0.02, trainable_beta=False, beta_method="linear",
+                 sigma_start=1e-3, sigma_end=10.0, start=0.0, end=1.0):
+        super().__init__()
+        self.num_steps = num_steps
+        self.beta_start = beta_start
+        self.beta_end = beta_end
+        self.trainable_beta = trainable_beta
+        self.beta_method = beta_method
+        self.sigma_start = sigma_start
+        self.sigma_end = sigma_end
+        self.start = start
+        self.end = end
+
+        if not (0 < self.beta_start < self.beta_end):
+            raise ValueError(f"beta_start ({self.beta_start}) and beta_end ({self.beta_end}) must satisfy 0 < start < end")
+        if not (0 < self.sigma_start < self.sigma_end):
+            raise ValueError(f"sigma_start ({self.sigma_start}) and sigma_end ({self.sigma_end}) must satisfy 0 < start < end")
+        if self.num_steps <= 0:
+            raise ValueError(f"num_steps ({self.num_steps}) must be positive")
+
+        beta_range = (beta_start, beta_end)
+        betas_init = self.compute_beta_schedule(beta_range, num_steps, beta_method)
+        self.time = torch.linspace(self.start, self.end, self.num_steps, dtype=torch.float32)
+        self.dt = (self.end - self.start) / self.num_steps
+
+        if trainable_beta:
+            self.betas = nn.Parameter(betas_init)
+        else:
+            self.register_buffer('betas', betas_init)
+            self.register_buffer('cum_betas', torch.cumsum(betas_init, dim=0) * self.dt)
+            self.register_buffer("sigmas", self.sigma_start * (self.sigma_end / self.sigma_start) ** self.time)
+
+    def compute_beta_schedule(self, beta_range, num_steps, method):
+        beta_min, beta_max = beta_range
+        if method == "sigmoid":
+            x = torch.linspace(-6, 6, num_steps)
+            beta = torch.sigmoid(x) * (beta_max - beta_min) + beta_min
+        elif method == "quadratic":
+            x = torch.linspace(beta_min ** 0.5, beta_max ** 0.5, num_steps)
+            beta = x ** 2
+        elif method == "constant":
+            beta = torch.full((num_steps,), beta_max)
+        elif method == "inverse_time":
+            beta = 1.0 / torch.linspace(num_steps, 1, num_steps)
+            beta = beta_min + (beta_max - beta_min) * (beta - beta.min()) / (beta.max() - beta.min())
+        elif method == "linear":
+            beta = torch.linspace(beta_min, beta_max, num_steps)
+        else:
+            raise ValueError(f"Unknown beta_method: {method}. Supported: linear, sigmoid, quadratic, constant, inverse_time")
+        beta = torch.clamp(beta, min=beta_min, max=beta_max)
+        return beta
+
+    def constrain_betas(self):
+        if self.trainable_beta:
+            with torch.no_grad():
+                self.betas.clamp_(min=self.beta_start, max=self.beta_end)
+
+    def get_variance(self, time_steps, method):
+        if method == "ve":
+            return self.sigmas[time_steps] ** 2
+        elif method == "vp":
+            return 1 - torch.exp(-self.cum_betas[time_steps])
+        elif method == "sub-vp":
+            return 1 - torch.exp(-2 * self.cum_betas[time_steps])
+        else:
+            raise ValueError(f"Unknown method: {method}")

ldm/metrics.py

@@ -0,0 +1,206 @@
+import torch
+import torch.nn.functional as F
+import lpips
+from pytorch_fid import fid_score
+import shutil
+from torchvision.utils import save_image
+import os
+
+
+class Metrics:
+    """Computes image quality metrics for evaluating diffusion models.
+
+    Supports Mean Squared Error (MSE), Peak Signal-to-Noise Ratio (PSNR), Structural
+    Similarity Index (SSIM), Fréchet Inception Distance (FID), and Learned Perceptual
+    Image Patch Similarity (LPIPS) for comparing generated and ground truth images.
+
+    Parameters
+    ----------
+    device : str, optional
+        Device for computation (e.g., 'cuda', 'cpu') (default: 'cuda').
+    fid : bool, optional
+        If True, compute FID score (default: True).
+    metrics : bool, optional
+        If True, compute MSE, PSNR, and SSIM (default: False).
+    lpips : bool, optional
+        If True, compute LPIPS using VGG backbone (default: False).
+
+    Attributes
+    ----------
+    device : str
+        Computation device.
+    fid : bool
+        Flag for FID computation.
+    metrics : bool
+        Flag for MSE, PSNR, SSIM computation.
+    lpips : bool
+        Flag for LPIPS computation.
+    lpips_model : lpips.LPIPS or None
+        LPIPS model (VGG backbone) if `lpips=True`; otherwise, None.
+    temp_dir_real : str
+        Temporary directory for real images during FID computation.
+    temp_dir_fake : str
+        Temporary directory for fake (generated) images during FID computation.
+    """
+
+    def __init__(self, device="cuda", fid=True, metrics=False, lpips=False):
+        self.device = device
+        self.fid = fid
+        self.metrics = metrics
+        self.lpips = lpips
+        self.lpips_model = lpips.LPIPS(net='vgg').to(device) if lpips else None
+        self.temp_dir_real = "temp_real"
+        self.temp_dir_fake = "temp_fake"
+
+    def compute_fid(self, real_images, fake_images):
+        """Computes the Fréchet Inception Distance (FID) between real and generated images.
+
+        Saves images to temporary directories and uses Inception V3 to compute FID,
+        cleaning up directories afterward.
+
+        Parameters
+        ----------
+        real_images : torch.Tensor
+            Real images, shape (batch_size, channels, height, width), in [-1, 1].
+        fake_images : torch.Tensor
+            Generated images, same shape, in [-1, 1].
+
+        Returns
+        -------
+        float
+            FID score, or `float('inf')` if computation fails.
+
+        Notes
+        -----
+        - Images are normalized to [0, 1] and saved as PNG files for FID computation.
+        - Uses Inception V3 with 2048-dimensional features (`dims=2048`).
+        """
+        if real_images.shape != fake_images.shape:
+            raise ValueError(f"Shape mismatch: real_images {real_images.shape}, fake_images {fake_images.shape}")
+
+        real_images = (real_images + 1) / 2
+        fake_images = (fake_images + 1) / 2
+        real_images = real_images.clamp(0, 1).cpu()
+        fake_images = fake_images.clamp(0, 1).cpu()
+
+        os.makedirs(self.temp_dir_real, exist_ok=True)
+        os.makedirs(self.temp_dir_fake, exist_ok=True)
+
+        try:
+            for i, (real, fake) in enumerate(zip(real_images, fake_images)):
+                save_image(real, f"{self.temp_dir_real}/{i}.png")
+                save_image(fake, f"{self.temp_dir_fake}/{i}.png")
+
+            fid = fid_score.calculate_fid_given_paths(
+                paths=[self.temp_dir_real, self.temp_dir_fake],
+                batch_size=50,
+                device=self.device,
+                dims=2048
+            )
+        except Exception as e:
+            print(f"Error computing FID: {e}")
+            fid = float('inf')
+        finally:
+            shutil.rmtree(self.temp_dir_real, ignore_errors=True)
+            shutil.rmtree(self.temp_dir_fake, ignore_errors=True)
+
+        return fid
+
+    def compute_metrics(self, x, x_hat):
+        """Computes MSE, PSNR, and SSIM for evaluating image quality.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Ground truth images, shape (batch_size, channels, height, width).
+        x_hat : torch.Tensor
+            Generated images, same shape as `x`.
+
+        Returns
+        -------
+        tuple
+            Tuple of (mse, psnr, ssim) as floats, where:
+            - mse: Mean squared error.
+            - psnr: Peak signal-to-noise ratio.
+            - ssim: Structural similarity index (mean over batch).
+        """
+        if x.shape != x_hat.shape:
+            raise ValueError(f"Shape mismatch: x {x.shape}, x_hat {x_hat.shape}")
+
+        mse = F.mse_loss(x_hat, x)
+        psnr = -10 * torch.log10(mse)
+        c1, c2 = (0.01 * 2) ** 2, (0.03 * 2) ** 2  # Adjusted for [-1, 1] range
+        eps = 1e-8
+        mu_x = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1)
+        mu_y = F.avg_pool2d(x_hat, kernel_size=3, stride=1, padding=1)
+        mu_xy = mu_x * mu_y
+        sigma_x_sq = F.avg_pool2d(x.pow(2), kernel_size=3, stride=1, padding=1) - mu_x.pow(2)
+        sigma_y_sq = F.avg_pool2d(x_hat.pow(2), kernel_size=3, stride=1, padding=1) - mu_y.pow(2)
+        sigma_xy = F.avg_pool2d(x * x_hat, kernel_size=3, stride=1, padding=1) - mu_xy
+        ssim = ((2 * mu_xy + c1) * (2 * sigma_xy + c2)) / (
+            (mu_x.pow(2) + mu_y.pow(2) + c1) * (sigma_x_sq + sigma_y_sq + c2) + eps
+        )
+
+        return mse.item(), psnr.item(), ssim.mean().item()
+
+    def compute_lpips(self, x, x_hat):
+        """Computes LPIPS using a pre-trained VGG network.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Ground truth images, shape (batch_size, channels, height, width), in [-1, 1].
+        x_hat : torch.Tensor
+            Generated images, same shape as `x`.
+
+        Returns
+        -------
+        float
+            Mean LPIPS score over the batch.
+
+        Raises
+        ------
+        RuntimeError
+            If `lpips=True` but `lpips_model` is not initialized.
+        """
+        if self.lpips_model is None:
+            raise RuntimeError("LPIPS model not initialized; set lpips=True in __init__")
+        if x.shape != x_hat.shape:
+            raise ValueError(f"Shape mismatch: x {x.shape}, x_hat {x_hat.shape}")
+
+        x = x.to(self.device)
+        x_hat = x_hat.to(self.device)
+        return self.lpips_model(x, x_hat).mean().item()
+
+    def forward(self, x, x_hat):
+        """Computes specified metrics for ground truth and generated images.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Ground truth images, shape (batch_size, channels, height, width), in [-1, 1].
+        x_hat : torch.Tensor
+            Generated images, same shape as `x`.
+
+        Returns
+        -------
+        tuple
+            A tuple containing:
+            - fid: FID score (float, or `float('inf')` if `fid=False` or fails).
+            - mse: Mean squared error (float, or None if `metrics=False`).
+            - psnr: Peak signal-to-noise ratio (float, or None if `metrics=False`).
+            - ssim: Structural similarity index (float, or None if `metrics=False`).
+            - lpips: LPIPS score (float, or None if `lpips=False`).
+        """
+        fid = float('inf')
+        mse, psnr, ssim = None, None, None
+        lpips_score = None
+
+        if self.metrics:
+            mse, psnr, ssim = self.compute_metrics(x, x_hat)
+        if self.fid:
+            fid = self.compute_fid(x, x_hat)
+        if self.lpips:
+            lpips_score = self.compute_lpips(x, x_hat)
+
+        return fid, mse, psnr, ssim, lpips_score