TorchDiff 2.0.0__py3-none-any.whl

ldm/train_ldm.py

@@ -0,0 +1,412 @@
+import torch
+from tqdm import tqdm
+import torch.nn as nn
+from torch.cuda.amp import GradScaler, autocast
+from torch.amp import GradScaler, autocast
+from torch.optim.lr_scheduler import LambdaLR
+from transformers import BertTokenizer
+import warnings
+
+
+
+
+class TrainLDM(nn.Module):
+    """Trainer for Latent Diffusion Models (LDM).
+
+    Manages the training process for LDMs, optimizing a noise predictor to learn the noise
+    added by the forward diffusion process in the latent space, as described in Rombach
+    et al. (2022). Uses a pre-trained compressor model to encode images into a latent
+    space, supports conditional training with text prompts, mixed precision, learning rate
+    scheduling, early stopping, and checkpointing.
+
+    Parameters
+    ----------
+    forward_model : nn.Module
+        Forward diffusion module (e.g., ForwardDDPM, ForwardSDE) to add noise in the latent space.
+    hyper_params_model : nn.Module
+        Hyperparameter module (e.g., HyperParamsDDPM, HyperParamsSDE) defining the noise schedule.
+    noise_predictor : nn.Module
+        Model to predict noise added during the forward diffusion process.
+    compressor_model : nn.Module
+        Pre-trained model to encode images into the latent space and decode back (e.g., autoencoder).
+    optimizer : torch.optim.Optimizer
+        Optimizer for training the noise predictor and conditional model (if applicable).
+    objective : callable
+        Loss function to compute the difference between predicted and actual noise.
+    data_loader : torch.utils.data.DataLoader
+        DataLoader for training data.
+    conditional_model : nn.Module, optional
+        Model for conditional generation (e.g., text embeddings), default None.
+    val_loader : torch.utils.data.DataLoader, optional
+        DataLoader for validation data, default None.
+    max_epoch : int, optional
+        Maximum number of training epochs (default: 1000).
+    device : torch.device, optional
+        Device for computation (default: CUDA if available, else CPU).
+    store_path : str, optional
+        Path to save model checkpoints (default: "ldm_model.pth").
+    patience : int, optional
+        Number of epochs to wait for improvement before early stopping (default: 10).
+    warmup_epochs : int, optional
+        Number of epochs for learning rate warmup (default: 100).
+    tokenizer : BertTokenizer, optional
+        Tokenizer for processing text prompts, default None (loads "bert-base-uncased").
+    max_length : int, optional
+        Maximum length for tokenized prompts (default: 77).
+    val_frequency : int, optional
+        Frequency (in epochs) for validation (default: 10).
+
+    Attributes
+    ----------
+    device : torch.device
+        Device used for computation.
+    forward_diffusion : nn.Module
+        Forward diffusion module.
+    hyper_params_model : nn.Module
+        Hyperparameter module for the noise schedule.
+    noise_predictor : nn.Module
+        Noise prediction model.
+    compressor_model : nn.Module
+        Compressor model for latent space encoding/decoding.
+    conditional_model : nn.Module or None
+        Conditional model for text-based training, if provided.
+    optimizer : torch.optim.Optimizer
+        Optimizer for training.
+    objective : callable
+        Loss function for training.
+    data_loader : torch.utils.data.DataLoader
+        Training data loader.
+    val_loader : torch.utils.data.DataLoader or None
+        Validation data loader, if provided.
+    max_epoch : int
+        Maximum training epochs.
+    store_path : str
+        Path for saving checkpoints.
+    max_length : int
+        Maximum length for tokenized prompts.
+    patience : int
+        Patience for early stopping.
+    scheduler : torch.optim.lr_scheduler.ReduceLROnPlateau
+        Learning rate scheduler based on validation or training loss.
+    warmup_lr_scheduler : torch.optim.lr_scheduler.LambdaLR
+        Learning rate scheduler for warmup.
+    tokenizer : BertTokenizer
+        Tokenizer for text prompts.
+    val_frequency : int
+        Frequency for validation.
+
+    Raises
+    ------
+    ValueError
+        If the default tokenizer ("bert-base-uncased") fails to load and no tokenizer is provided.
+    """
+    def __init__(self, forward_model, hyper_params_model, noise_predictor, compressor_model, optimizer, objective, data_loader,
+                 conditional_model=None, val_loader=None, max_epoch=1000, device=None, store_path=None,
+                 patience=10, warmup_epochs=100, tokenizer=None, max_length=77, val_frequency=10):
+        super().__init__()
+        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.forward_diffusion = forward_model.to(device)
+        self.hyper_params_model = hyper_params_model.to(device)
+        self.noise_predictor = noise_predictor
+        self.compressor_model = compressor_model
+        self.conditional_model = conditional_model
+        self.optimizer = optimizer
+        self.objective = objective
+        self.data_loader = data_loader
+        self.val_loader = val_loader
+        self.max_epoch = max_epoch
+        self.store_path = store_path or "ldm_model.pth"
+        self.max_length = max_length
+        self.patience = patience
+        self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=self.patience, factor=0.5)
+        self.warmup_lr_scheduler = self.warmup_scheduler(self.optimizer, warmup_epochs)
+        self.val_frequency = val_frequency
+        if tokenizer is None:
+            try:
+                self.tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+            except Exception as e:
+                raise ValueError(f"Failed to load default tokenizer: {e}. Please provide a tokenizer.")
+
+    def load_checkpoint(self, checkpoint_path):
+        """Loads a training checkpoint to resume training.
+
+        Restores the state of the noise predictor, conditional model (if applicable),
+        and optimizer from a saved checkpoint.
+
+        Parameters
+        ----------
+        checkpoint_path : str
+            Path to the checkpoint file.
+
+        Returns
+        -------
+        tuple
+            A tuple containing:
+            - epoch: The epoch at which the checkpoint was saved (int).
+            - loss: The loss at the checkpoint (float).
+
+        Raises
+        ------
+        FileNotFoundError
+            If the checkpoint file is not found.
+        KeyError
+            If the checkpoint is missing required keys ('model_state_dict_noise_predictor'
+            or 'optimizer_state_dict').
+
+        Warns
+        -----
+        warnings.warn
+            If the optimizer state cannot be loaded, if the checkpoint contains a
+            conditional model state but none is defined, or if no conditional model
+            state is provided when expected.
+        """
+        try:
+            checkpoint = torch.load(checkpoint_path, map_location=self.device)
+        except FileNotFoundError:
+            raise FileNotFoundError(f"Checkpoint file not found at {checkpoint_path}")
+
+        if 'model_state_dict_noise_predictor' not in checkpoint:
+            raise KeyError("Checkpoint missing 'model_state_dict_noise_predictor' key")
+        self.noise_predictor.load_state_dict(checkpoint['model_state_dict_noise_predictor'])
+
+        if self.conditional_model is not None:
+            if 'model_state_dict_conditional' in checkpoint and checkpoint['model_state_dict_conditional'] is not None:
+                self.conditional_model.load_state_dict(checkpoint['model_state_dict_conditional'])
+            else:
+                warnings.warn(
+                    "Checkpoint contains no 'model_state_dict_conditional' or it is None, skipping conditional model loading")
+        elif 'model_state_dict_conditional' in checkpoint and checkpoint['model_state_dict_conditional'] is not None:
+            warnings.warn(
+                "Checkpoint contains conditional model state, but no conditional model is defined in this instance")
+
+        if 'optimizer_state_dict' not in checkpoint:
+            raise KeyError("Checkpoint missing 'optimizer_state_dict' key")
+        try:
+            self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        except ValueError as e:
+            warnings.warn(f"Optimizer state loading failed: {e}. Continuing without optimizer state.")
+
+        epoch = checkpoint.get('epoch', -1)
+        loss = checkpoint.get('loss', float('inf'))
+
+        self.noise_predictor.to(self.device)
+        if self.conditional_model is not None:
+            self.conditional_model.to(self.device)
+
+        print(f"Loaded checkpoint from {checkpoint_path} at epoch {epoch} with loss {loss:.4f}")
+        return epoch, loss
+
+    @staticmethod
+    def warmup_scheduler(optimizer, warmup_epochs=10):
+        """Creates a learning rate scheduler for warmup.
+
+        Generates a scheduler that linearly increases the learning rate from 0 to the
+        optimizer's initial value over the specified warmup epochs, then maintains it.
+
+        Parameters
+        ----------
+        optimizer : torch.optim.Optimizer
+            Optimizer to apply the scheduler to.
+        warmup_epochs : int, optional
+            Number of epochs for the warmup phase (default: 10).
+
+        Returns
+        -------
+        torch.optim.lr_scheduler.LambdaLR
+            Learning rate scheduler for warmup.
+        """
+        def lr_lambda(epoch):
+            if epoch < warmup_epochs:
+                return epoch / warmup_epochs
+            return 1.0
+
+        return LambdaLR(optimizer, lr_lambda)
+
+    def forward(self):
+        """Trains the LDM to predict noise added by the forward diffusion process in the latent space.
+
+        Executes the training loop, optimizing the noise predictor and conditional model
+        (if applicable) using mixed precision, gradient clipping, and learning rate
+        scheduling. Uses a pre-trained compressor model to encode images into the latent
+        space. Supports validation, early stopping, and checkpointing.
+
+        Returns
+        -------
+        tuple
+            A tuple containing:
+            - train_losses: List of mean training losses per epoch (list of float).
+            - best_val_loss: Best validation or training loss achieved (float).
+
+        Notes
+        -----
+        - Training uses mixed precision via `torch.cuda.amp` for efficiency.
+        - The compressor model is assumed pre-trained and set to evaluation mode.
+        - Checkpoints are saved when the validation (or training) loss improves, and on
+          early stopping.
+        - Early stopping is triggered if no improvement occurs for `patience` epochs.
+        """
+
+        self.noise_predictor.train()
+        self.noise_predictor.to(self.device)
+        if self.conditional_model is not None:
+            self.conditional_model.train()
+            self.conditional_model.to(self.device)
+        if self.compressor_model is not None:
+            self.compressor_model.eval() # the model is already trained
+            self.compressor_model.to(self.device)
+
+        scaler = GradScaler()
+        train_losses = []
+        best_val_loss = float("inf")
+        wait = 0
+        for epoch in range(self.max_epoch):
+            train_losses_ = []
+            for x, y in tqdm(self.data_loader):
+                x = x.to(self.device)
+                with torch.no_grad():
+                    x, _ = self.compressor_model.encode(x) # using compressor model bring x to latent space
+                if self.conditional_model is not None:
+                    y_list = y.cpu().numpy().tolist() if isinstance(y, torch.Tensor) else y
+                    y_list = [str(item) for item in y_list]
+                    y_encoded = self.tokenizer(
+                        y_list,
+                        padding="max_length",
+                        truncation=True,
+                        max_length=self.max_length,
+                        return_tensors="pt"
+                    ).to(self.device)
+                    input_ids = y_encoded["input_ids"]
+                    attention_mask = y_encoded["attention_mask"]
+                    y_encoded = self.conditional_model(input_ids, attention_mask)
+                else:
+                    y_encoded = None
+
+                self.optimizer.zero_grad()
+                with autocast(device_type='cuda' if self.device.type == 'cuda' else 'cpu'):
+                    noise = torch.randn_like(x).to(self.device)
+                    t = torch.randint(0, self.hyper_params_model.num_steps, (x.shape[0],)).to(self.device)
+                    assert x.device == noise.device == t.device, "Device mismatch detected"
+                    assert t.shape[0] == x.shape[0], "Timestep batch size mismatch"
+                    noisy_x = self.forward_diffusion(x, noise, t)
+                    p_noise = self.noise_predictor(noisy_x, t, y_encoded)
+                    loss = self.objective(p_noise, noise)
+                scaler.scale(loss).backward()
+
+                nn.utils.clip_grad_norm_(self.noise_predictor.parameters(), max_norm=1.0)
+                if self.conditional_model is not None:
+                    nn.utils.clip_grad_norm_(self.conditional_model.parameters(), max_norm=1.0)
+                scaler.step(self.optimizer)
+                scaler.update()
+                self.warmup_lr_scheduler.step()
+                train_losses_.append(loss.item())
+
+            mean_train_loss = torch.mean(torch.tensor(train_losses_)).item()
+            train_losses.append(mean_train_loss)
+            print(f"\nEpoch: {epoch + 1} | Train Loss: {mean_train_loss:.4f}", end="")
+
+            if self.val_loader is not None and (epoch + 1) % self.val_frequency == 0:
+                val_loss = self.validate()
+                print(f" | Val Loss: {val_loss:.4f}")
+                current_best = val_loss
+                self.scheduler.step(val_loss)
+            else:
+                print()
+                current_best = mean_train_loss
+                self.scheduler.step(mean_train_loss)
+
+            if current_best < best_val_loss:
+                best_val_loss = current_best
+                wait = 0
+                try:
+                    torch.save({
+                        'epoch': epoch + 1,
+                        'model_state_dict_noise_predictor': self.noise_predictor.state_dict(),
+                        'model_state_dict_conditional': self.conditional_model.state_dict() if self.conditional_model is not None else None,
+                        'optimizer_state_dict': self.optimizer.state_dict(),
+                        'loss': best_val_loss,
+                        'hyper_params_model': self.hyper_params_model,
+                        'max_epoch': self.max_epoch,
+                    }, self.store_path)
+                    print(f"Model saved at epoch {epoch + 1}")
+                except Exception as e:
+                    print(f"Failed to save model: {e}")
+            else:
+                wait += 1
+                if wait >= self.patience:
+                    print("Early stopping triggered")
+                    try:
+                        torch.save({
+                            'epoch': epoch + 1,
+                            'model_state_dict_noise_predictor': self.noise_predictor.state_dict(),
+                            'model_state_dict_conditional': self.conditional_model.state_dict() if self.conditional_model is not None else None,
+                            'optimizer_state_dict': self.optimizer.state_dict(),
+                            'loss': best_val_loss,
+                            'hyper_params_model': self.hyper_params_model,
+                            'max_epoch': self.max_epoch,
+                        }, self.store_path + "_early_stop.pth")
+                        print(f"Final model saved at {self.store_path}_early_stop.pth")
+                    except Exception as e:
+                        print(f"Failed to save final model: {e}")
+                    break
+
+        return train_losses, best_val_loss
+
+    def validate(self):
+        """Validates the LDM on the validation dataset.
+
+        Computes the validation loss using the noise predictor and forward diffusion
+        process in the latent space, with optional conditional inputs.
+
+        Returns
+        -------
+        float
+            Mean validation loss across the validation dataset.
+
+        Notes
+        -----
+        - Validation is performed with `torch.no_grad()` for efficiency.
+        - The compressor model is used to encode validation data into the latent space.
+        - The noise predictor and conditional model (if applicable) are set to evaluation
+          mode during validation and restored to training mode afterward.
+        """
+        self.noise_predictor.eval()
+        if self.conditional_model is not None:
+            self.conditional_model.eval()
+
+        val_losses = []
+        with torch.no_grad():
+            for x, y in self.val_loader:
+                x = x.to(self.device)
+                with torch.no_grad():
+                    x, _ = self.compressor_model.encode(x)
+
+                if self.conditional_model is not None:
+                    y_list = y.cpu().numpy().tolist() if isinstance(y, torch.Tensor) else y
+                    y_list = [str(item) for item in y_list]
+                    y_encoded = self.tokenizer(
+                        y_list,
+                        padding="max_length",
+                        truncation=True,
+                        max_length=self.max_length,
+                        return_tensors="pt"
+                    ).to(self.device)
+                    input_ids = y_encoded["input_ids"]
+                    attention_mask = y_encoded["attention_mask"]
+                    y_encoded = self.conditional_model(input_ids, attention_mask)
+                else:
+                    y_encoded = None
+
+                noise = torch.randn_like(x).to(self.device)
+                t = torch.randint(0, self.hyper_params_model.num_steps, (x.shape[0],)).to(self.device)
+                assert x.device == noise.device == t.device, "Device mismatch detected"
+                assert t.shape[0] == x.shape[0], "Timestep batch size mismatch"
+                noisy_x = self.forward_diffusion(x, noise, t)
+                p_noise = self.noise_predictor(noisy_x, t, y_encoded)
+                loss = self.objective(p_noise, noise)
+                val_losses.append(loss.item())
+
+        mean_val_loss = torch.mean(torch.tensor(val_losses)).item()
+        self.noise_predictor.train()
+        if self.conditional_model is not None:
+            self.conditional_model.train()
+        return mean_val_loss

sde/forward_sde.py

@@ -0,0 +1,98 @@
+import torch
+import torch.nn as nn
+
+
+
+class ForwardSDE(nn.Module):
+    """Forward diffusion process for SDE-based generative models.
+
+    Implements the forward diffusion process for score-based generative models using
+    Stochastic Differential Equations (SDEs), supporting Variance Exploding (VE),
+    Variance Preserving (VP), sub-Variance Preserving (sub-VP), and ODE methods, as
+    described in Song et al. (2021).
+
+    Parameters
+    ----------
+    hyper_params : object
+        Hyperparameter object containing SDE-specific parameters. Expected to have
+        attributes:
+        - `dt`: Time step size for SDE integration (float).
+        - `sigmas`: Sigma values for VE method (torch.Tensor, optional).
+        - `betas`: Beta values for VP, sub-VP, or ODE methods (torch.Tensor).
+        - `cum_betas`: Cumulative beta values for sub-VP method (torch.Tensor, optional).
+    method : str
+        SDE method to use. Supported methods: "ve", "vp", "sub-vp", "ode".
+
+    Attributes
+    ----------
+    hyper_params : object
+        Stores the provided hyperparameter object.
+    method : str
+        Selected SDE method.
+
+    Raises
+    ------
+    ValueError
+        If `method` is not one of the supported methods ("ve", "vp", "sub-vp", "ode").
+    """
+    def __init__(self, hyper_params, method):
+        super().__init__()
+        self.hyper_params = hyper_params
+        self.method = method
+
+    def forward(self, x0, noise, time_steps):
+        """Applies the forward SDE diffusion process to the input data.
+
+        Perturbs the input data `x0` by adding noise according to the specified SDE
+        method at given time steps, incorporating drift and diffusion terms as applicable.
+
+        Parameters
+        ----------
+        x0 : torch.Tensor
+            Input data tensor, shape (batch_size, channels, height, width).
+        noise : torch.Tensor
+            Gaussian noise tensor, same shape as `x0`.
+        time_steps : torch.Tensor
+            Tensor of time step indices (long), shape (batch_size,), where each value
+            is in the range [0, hyper_params.num_steps - 1].
+
+        Returns
+        -------
+        torch.Tensor
+            Noisy data tensor at the specified time steps, same shape as `x0`.
+
+        Raises
+        ------
+        ValueError
+            If `method` is not one of the supported methods ("ve", "vp", "sub-vp", "ode").
+        """
+        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