TorchDiff 2.0.0__py3-none-any.whl

ldm/reverse_ldm.py

@@ -0,0 +1,119 @@
+import torch
+import torch.nn as nn
+
+
+
+
+class ReverseSDE(nn.Module):
+    def __init__(self, hyper_params, method):
+        super().__init__()
+        self.hyper_params = hyper_params
+        self.method = method
+
+    def forward(self, xt, noise, predicted_noise, time_steps):
+
+        dt = self.hyper_params.dt
+        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)
+        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))
+            drift = -(sigma_t ** 2 - sigma_t_prev ** 2).view(-1, 1, 1, 1) * predicted_noise * dt
+            diffusion = sigma_diff.view(-1, 1, 1, 1) * noise if noise is not None else 0
+            xt = xt + drift + diffusion
+            xt = torch.clamp(xt, -1e5, 1e5)
+
+        elif self.method == "vp":
+            drift = -0.5 * betas * xt * dt - betas * predicted_noise * dt
+            diffusion = torch.sqrt(betas * dt) * noise if noise is not None else 0
+            xt = xt + drift + diffusion
+
+        elif self.method == "sub-vp":
+            drift = -0.5 * betas * xt * dt - betas * (1 - torch.exp(-2 * cum_betas)) * predicted_noise * dt
+            diffusion = torch.sqrt(betas * (1 - torch.exp(-2 * cum_betas)) * dt) * noise if noise is not None else 0
+            xt = xt + drift + diffusion
+
+        elif self.method == "ode":
+            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)
+                drift = -0.5 * (sigma_t ** 2 - sigma_t_prev ** 2).view(-1, 1, 1, 1) * predicted_noise * dt
+            else:
+                drift = -0.5 * betas * xt * dt - 0.5 * betas * predicted_noise * dt
+            xt = xt + drift
+            xt = torch.clamp(xt, -1e5, 1e5)
+        else:
+            raise ValueError(f"Unknown method: {self.method}")
+        return xt
+
+
+
+
+class ReverseDDIM(nn.Module):
+    def __init__(self, hyper_params):
+        super().__init__()
+        self.hyper_params = hyper_params
+
+    def forward(self, xt, predicted_noise, time_steps, prev_time_steps):
+
+        if not torch.all((time_steps >= 0) & (time_steps < self.hyper_params.tau_num_steps)):
+            raise ValueError(f"time_steps must be between 0 and {self.hyper_params.tau_num_steps - 1}")
+        if not torch.all((prev_time_steps >= 0) & (prev_time_steps < self.hyper_params.tau_num_steps)):
+            raise ValueError(f"prev_time_steps must be between 0 and {self.hyper_params.tau_num_steps - 1}")
+
+        _, _, _, tau_sqrt_alpha_cumprod, tau_sqrt_one_minus_alpha_cumprod = self.hyper_params.get_tau_schedule()
+        tau_sqrt_alpha_cumprod_t = tau_sqrt_alpha_cumprod[time_steps].to(xt.device).view(-1, 1, 1, 1)
+        tau_sqrt_one_minus_alpha_cumprod_t = tau_sqrt_one_minus_alpha_cumprod[time_steps].to(xt.device).view(-1, 1, 1, 1)
+        prev_tau_sqrt_alpha_cumprod_t = tau_sqrt_alpha_cumprod[prev_time_steps].to(xt.device).view(-1, 1, 1, 1)
+        prev_tau_sqrt_one_minus_alpha_cumprod_t = tau_sqrt_one_minus_alpha_cumprod[prev_time_steps].to(xt.device).view(-1, 1, 1, 1)
+
+        eta = self.hyper_params.eta
+        x0 = (xt - tau_sqrt_one_minus_alpha_cumprod_t * predicted_noise) / tau_sqrt_alpha_cumprod_t
+        noise_coeff = eta * ((tau_sqrt_one_minus_alpha_cumprod_t / prev_tau_sqrt_alpha_cumprod_t) *
+                             prev_tau_sqrt_one_minus_alpha_cumprod_t / torch.clamp(tau_sqrt_one_minus_alpha_cumprod_t, min=1e-8))
+        direction_coeff = torch.clamp(prev_tau_sqrt_one_minus_alpha_cumprod_t ** 2 - noise_coeff ** 2, min=1e-8).sqrt()
+        xt_prev = prev_tau_sqrt_alpha_cumprod_t * x0 + noise_coeff * torch.randn_like(xt) + direction_coeff * predicted_noise
+
+        return xt_prev, x0
+
+
+
+class ReverseDDPM(nn.Module):
+    """reverse diffusion process of DDPM."""
+    def __init__(self, hyper_params):
+        super().__init__()
+        self.hyper_params = hyper_params # hyperparameters class
+
+    def forward(self, xt, predicted_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:
+            betas_t, alphas_t, alpha_bars_t, _, _ = self.hyper_params.compute_schedule(self.hyper_params.betas)
+            betas_t = betas_t[time_steps].to(xt.device)
+            alphas_t = alphas_t[time_steps].to(xt.device)
+            alpha_bars_t = alpha_bars_t[time_steps].to(xt.device)
+            alpha_bars_t_minus_1 = alpha_bars_t[time_steps - 1].to(xt.device) if time_steps.any() else None
+        else:
+            betas_t = self.hyper_params.betas[time_steps].to(xt.device)
+            alphas_t = self.hyper_params.alphas[time_steps].to(xt.device)
+            alpha_bars_t = self.hyper_params.alpha_bars[time_steps].to(xt.device)
+            alpha_bars_t_minus_1 = self.hyper_params.alpha_bars[time_steps - 1].to(xt.device) if time_steps.any() else None
+
+        sqrt_alphas_t = torch.sqrt(alphas_t).view(-1, 1, 1, 1)
+        sqrt_one_minus_alpha_bars_t = torch.sqrt(1 - alpha_bars_t).view(-1, 1, 1, 1)
+        betas_t = betas_t.view(-1, 1, 1, 1)
+
+        mu = (xt - (betas_t / sqrt_one_minus_alpha_bars_t) * predicted_noise) / sqrt_alphas_t
+
+        mask = (time_steps == 0)
+        if mask.all():
+            return mu
+
+        variance = (1 - alpha_bars_t_minus_1) / (1 - alpha_bars_t) * betas_t.squeeze()
+        std = torch.sqrt(variance).view(-1, 1, 1, 1)
+
+        z = torch.randn_like(xt).to(xt.device)
+        xt_minus_1 = mu + (~mask).float().view(-1, 1, 1, 1) * std * z
+        return xt_minus_1

ldm/sample_ldm.py

@@ -0,0 +1,254 @@
+import torch
+import torch.nn as nn
+from transformers import BertTokenizer
+
+
+class SampleLDM(nn.Module):
+    """Sampler for generating images using Latent Diffusion Models (LDM).
+
+    Generates images by iteratively denoising random noise in the latent space using a
+    reverse diffusion process, decoding the result back to the image space with a
+    pre-trained compressor, as described in Rombach et al. (2022). Supports DDPM, DDIM,
+    and SDE diffusion models, as well as conditional generation with text prompts.
+
+    Parameters
+    ----------
+    model : str
+        Diffusion model type. Supported: "ddpm", "ddim", "sde".
+    reverse_diffusion : nn.Module
+        Reverse diffusion module (e.g., ReverseDDPM, ReverseDDIM, ReverseSDE).
+    noise_predictor : nn.Module
+        Model to predict noise added during the forward diffusion process.
+    compressor_model : nn.Module
+        Pre-trained model to encode/decode between image and latent spaces (e.g., autoencoder).
+    image_shape : tuple
+        Shape of generated images as (height, width).
+    conditional_model : nn.Module, optional
+        Model for conditional generation (e.g., text embeddings), default None.
+    tokenizer : str or BertTokenizer, optional
+        Tokenizer for processing text prompts, default "bert-base-uncased".
+    batch_size : int, optional
+        Number of images to generate per batch (default: 1).
+    in_channels : int, optional
+        Number of input channels for latent representations (default: 3).
+    device : torch.device, optional
+        Device for computation (default: CUDA if available, else CPU).
+    max_length : int, optional
+        Maximum length for tokenized prompts (default: 77).
+    output_range : tuple, optional
+        Range for clamping generated images (min, max), default (-1, 1).
+
+    Attributes
+    ----------
+    device : torch.device
+        Device used for computation.
+    model : str
+        Diffusion model type ("ddpm", "ddim", "sde").
+    noise_predictor : nn.Module
+        Noise prediction model.
+    reverse : nn.Module
+        Reverse diffusion module.
+    compressor : nn.Module
+        Compressor model for latent space encoding/decoding.
+    conditional_model : nn.Module or None
+        Conditional model for text-based generation, if provided.
+    tokenizer : BertTokenizer
+        Tokenizer for text prompts.
+    in_channels : int
+        Number of input channels for latent representations.
+    image_shape : tuple
+        Shape of generated images (height, width).
+    batch_size : int
+        Batch size for generation.
+    max_length : int
+        Maximum length for tokenized prompts.
+    output_range : tuple
+        Range for clamping generated images.
+
+    Raises
+    ------
+    ValueError
+        If `image_shape` is not a tuple of two positive integers, `batch_size` is not
+        positive, `in_channels` is not positive, or `output_range` is not a tuple
+        (min, max) with min < max.
+    """
+    def __init__(self, model, reverse_diffusion, noise_predictor, compressor_model, image_shape, conditional_model=None,
+                 tokenizer="bert-base-uncased", batch_size=1, in_channels=3, device=None, max_length=77, output_range=(-1, 1)):
+        super().__init__()
+        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = model
+        self.noise_predictor = noise_predictor.to(self.device)
+        self.reverse = reverse_diffusion.to(self.device)
+        self.compressor = compressor_model.to(self.device)
+        self.conditional_model = conditional_model.to(self.device) if conditional_model else None
+        self.tokenizer = BertTokenizer.from_pretrained(tokenizer)
+        self.in_channels = in_channels
+        self.image_shape = image_shape
+        self.batch_size = batch_size
+        self.max_length = max_length
+        self.output_range = output_range
+
+        if not isinstance(image_shape, (tuple, list)) or len(image_shape) != 2 or not all(isinstance(s, int) and s > 0 for s in image_shape):
+            raise ValueError("image_shape must be a tuple of two positive integers (height, width)")
+        if batch_size <= 0:
+            raise ValueError("batch_size must be positive")
+        if in_channels <= 0:
+            raise ValueError("in_channels must be positive")
+        if not isinstance(output_range, (tuple, list)) or len(output_range) != 2 or output_range[0] >= output_range[1]:
+            raise ValueError("output_range must be a tuple (min, max) with min < max")
+
+    def tokenize(self, prompts):
+        """Tokenizes text prompts for conditional generation.
+
+        Converts input prompts into tokenized tensors using the specified tokenizer.
+
+        Parameters
+        ----------
+        prompts : str or list
+            Text prompt(s) for conditional generation. Can be a single string or a list
+            of strings.
+
+        Returns
+        -------
+        tuple
+            A tuple containing:
+            - input_ids: Tokenized input IDs (torch.Tensor, shape (batch_size, max_length)).
+            - attention_mask: Attention mask for tokenized inputs (torch.Tensor, same shape).
+
+        Raises
+        ------
+        TypeError
+            If `prompts` is not a string or a list of strings.
+        """
+        if isinstance(prompts, str):
+            prompts = [prompts]
+        elif not isinstance(prompts, list) or not all(isinstance(p, str) for p in prompts):
+            raise TypeError("prompts must be a string or list of strings")
+
+        encoded = self.tokenizer(
+            prompts,
+            padding="max_length",
+            truncation=True,
+            max_length=self.max_length,
+            return_tensors="pt"
+        )
+        return encoded["input_ids"].to(self.device), encoded["attention_mask"].to(self.device)
+
+    def forward(self, conditions=None, normalize_output=True):
+        """Generates images using the reverse diffusion process in the latent space.
+
+        Iteratively denoises random noise in the latent space using the specified reverse
+        diffusion model (DDPM, DDIM, or SDE), then decodes the result to the image space
+        with the compressor model. Supports conditional generation with text prompts.
+
+        Parameters
+        ----------
+        conditions : str or list, optional
+            Text prompt(s) for conditional generation, default None.
+        normalize_output : bool, optional
+            If True, normalizes output images to [0, 1] (default: True).
+
+        Returns
+        -------
+        torch.Tensor
+            Generated images, shape (batch_size, channels, height, width).
+            If `normalize_output` is True, images are normalized to [0, 1]; otherwise,
+            they are clamped to `output_range`.
+
+        Raises
+        ------
+        ValueError
+            If `conditions` is provided but no conditional model is specified, if a
+            conditional model is specified but `conditions` is None, or if `model` is not
+            one of "ddpm", "ddim", "sde".
+
+        Notes
+        -----
+        - Sampling is performed with `torch.no_grad()` for efficiency.
+        - The noise predictor, reverse diffusion, compressor, and conditional model
+          (if applicable) are set to evaluation mode during sampling.
+        - For DDIM, uses the subsampled tau schedule (`tau_num_steps`); for DDPM/SDE,
+          uses the full number of steps (`num_steps`).
+        - The compressor model is assumed to have `encode` and `decode` methods for
+          latent space conversion.
+        """
+        if conditions is not None and self.conditional_model is None:
+            raise ValueError("Conditions provided but no conditional model specified")
+        if conditions is None and self.conditional_model is not None:
+            raise ValueError("Conditions must be provided for conditional model")
+
+        noisy_samples = torch.randn(self.batch_size, self.in_channels, self.image_shape[0], self.image_shape[1]).to(self.device)
+
+        self.noise_predictor.eval()
+        self.compressor.eval()
+        self.reverse.eval()
+        if self.conditional_model:
+            self.conditional_model.eval()
+
+        with torch.no_grad():
+            xt = noisy_samples
+            xt, _ = self.compressor.encode(xt)
+
+            if self.model == "ddim":
+                num_steps = self.reverse.hyper_params.tau_num_steps
+            elif self.model == "ddpm" or self.model == "sde":
+                num_steps = self.reverse.hyper_params.num_steps
+            else:
+                raise ValueError(f"Unknown model: {self.model}. Supported: ddpm, ddim, sde")
+
+            for t in reversed(range(num_steps)):
+                time_steps = torch.full((self.batch_size,), t, device=self.device, dtype=torch.long)
+                prev_time_steps = torch.full((self.batch_size,), max(t - 1, 0), device=self.device, dtype=torch.long)
+
+                if self.model == "sde":
+                    noise = torch.randn_like(xt) if getattr(self.reverse, "method", None) != "ode" else None
+
+                if self.conditional_model is not None and conditions is not None:
+                    input_ids, attention_masks = self.tokenize(conditions)
+                    key_padding_mask = (attention_masks == 0)
+                    y = self.conditional_model(input_ids, key_padding_mask)
+                    predicted_noise = self.noise_predictor(xt, time_steps, y)
+                else:
+                    predicted_noise = self.noise_predictor(xt, time_steps)
+
+                if self.model == "sde":
+                    xt = self.reverse(xt, noise, predicted_noise, time_steps)
+                elif self.model == "ddim":
+                    xt, _ = self.reverse(xt, predicted_noise, time_steps, prev_time_steps)
+                elif self.model == "ddpm":
+                    xt = self.reverse(xt, predicted_noise, time_steps)
+                else:
+                    raise ValueError(f"Unknown model: {self.model}. Supported: ddpm, ddim, sde")
+
+            x = self.compressor.decode(xt)
+            generated_imgs = torch.clamp(x, min=self.output_range[0], max=self.output_range[1])
+            if normalize_output:
+                generated_imgs = (generated_imgs - self.output_range[0]) / (self.output_range[1] - self.output_range[0])
+
+        return generated_imgs
+
+    def to(self, device):
+        """Moves the module and its components to the specified device.
+
+        Parameters
+        ----------
+        device : torch.device
+            Target device for computation.
+
+        Returns
+        -------
+        self
+            The module moved to the specified device.
+
+        Notes
+        -----
+        - Moves `noise_predictor`, `reverse`, `compressor`, and `conditional_model`
+          (if applicable) to the specified device.
+        """
+        self.device = device
+        self.noise_predictor.to(device)
+        self.reverse.to(device)
+        self.compressor.to(device)
+        if self.conditional_model:
+            self.conditional_model.to(device)
+        return super().to(device)