TorchDiff 2.0.0__py3-none-any.whl

sde/sample_sde.py

@@ -0,0 +1,216 @@
+import torch
+import torch.nn as nn
+from transformers import BertTokenizer
+
+
+
+
+class SampleSDE(nn.Module):
+    """Sampler for generating images using SDE-based generative models.
+
+    Generates images by iteratively denoising random noise using the reverse SDE process
+    and a trained noise predictor, as described in Song et al. (2021). Supports both
+    unconditional and conditional generation with text prompts.
+
+    Parameters
+    ----------
+    reverse_diffusion : ReverseSDE
+        Reverse SDE diffusion module for denoising.
+    noise_predictor : nn.Module
+        Model to predict noise added during the forward SDE process.
+    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".
+    max_length : int, optional
+        Maximum length for tokenized prompts (default: 77).
+    batch_size : int, optional
+        Number of images to generate per batch (default: 1).
+    in_channels : int, optional
+        Number of input channels for generated images (default: 3).
+    device : torch.device, optional
+        Device for computation (default: CUDA if available, else CPU).
+    output_range : tuple, optional
+        Range for clamping generated images (min, max), default (-1, 1).
+
+    Attributes
+    ----------
+    device : torch.device
+        Device used for computation.
+    reverse : ReverseSDE
+        Reverse SDE diffusion module.
+    noise_predictor : nn.Module
+        Noise prediction model.
+    conditional_model : nn.Module or None
+        Conditional model for text-based generation, if provided.
+    tokenizer : BertTokenizer
+        Tokenizer for text prompts.
+    max_length : int
+        Maximum length for tokenized prompts.
+    in_channels : int
+        Number of input channels.
+    image_shape : tuple
+        Shape of generated images (height, width).
+    batch_size : int
+        Batch size for generation.
+    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, or `output_range` is not a tuple (min, max) with min < max.
+    """
+    def __init__(self, reverse_diffusion, noise_predictor, image_shape, conditional_model=None,
+                 tokenizer="bert-base-uncased", max_length=77, batch_size=1, in_channels=3, device=None, output_range=(-1, 1)):
+        super().__init__()
+        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.reverse = reverse_diffusion.to(self.device)
+        self.noise_predictor = noise_predictor.to(self.device)
+        self.conditional_model = conditional_model.to(self.device) if conditional_model else None
+        self.tokenizer = BertTokenizer.from_pretrained(tokenizer)
+        self.max_length = max_length
+        self.in_channels = in_channels
+        self.image_shape = image_shape
+        self.batch_size = batch_size
+        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 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 SDE sampling process.
+
+        Iteratively denoises random noise to generate images using the reverse SDE process
+        and noise predictor. 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, in_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, or if
+            a conditional model is specified but `conditions` is None.
+
+        Notes
+        -----
+        - Sampling is performed with `torch.no_grad()` for efficiency.
+        - The noise predictor, reverse SDE, and conditional model (if applicable) are set
+          to evaluation mode during sampling.
+        """
+        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.reverse.eval()
+        if self.conditional_model:
+            self.conditional_model.eval()
+
+        with torch.no_grad():
+            xt = noisy_samples
+            for t in reversed(range(self.reverse.hyper_params.num_steps)):
+                noise = torch.randn_like(xt) if self.reverse.method != "ode" else None
+                time_steps = torch.full((self.batch_size,), t, device=self.device, dtype=torch.long)
+
+                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)
+
+                xt = self.reverse(xt, noise, predicted_noise, time_steps)
+
+            generated_imgs = torch.clamp(xt, 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`, and `conditional_model` (if applicable) to
+          the specified device.
+        - The `compressor` attribute is not defined in this implementation and should be
+          removed or implemented if intended.
+        """
+        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)

sde/text_encoder.py

@@ -0,0 +1,152 @@
+import torch
+import torch.nn as nn
+import math
+from transformers import BertModel, DistilBertModel
+
+
+
+
+class TextEncoder(torch.nn.Module):
+    def __init__(
+            self,
+            use_pretrained_model=True,
+            model_name="bert-base-uncased",
+            vocabulary_size=30522,
+            num_layers=6,
+            input_dimension=768,
+            output_dimension=768,
+            num_heads=8,
+            context_length=77,
+            dropout_rate=0.1,
+            qkv_bias=False,
+            scaling_value=4,
+            epsilon=1e-5
+    ):
+        super().__init__()
+        self.use_pretrained_model = use_pretrained_model
+        if self.use_pretrained_model:
+            # self.bert = DistilBertModel.from_pretrained("distilbert-base-uncased")
+            self.bert = BertModel.from_pretrained(model_name)
+            for param in self.bert.parameters():
+                param.requires_grad = False
+            self.projection = nn.Linear(self.bert.config.hidden_size, output_dimension)
+        else:
+            self.embedding = Embedding(
+                vocabulary_size=vocabulary_size,
+                embedding_dimension=input_dimension,
+                context_length=context_length
+            )
+            self.layers = torch.nn.ModuleList([
+                EncoderLayer(
+                    input_dimension=input_dimension,
+                    output_dimension=output_dimension,
+                    num_heads=num_heads,
+                    dropout_rate=dropout_rate,
+                    qkv_bias=qkv_bias,
+                    scaling_value=scaling_value,
+                    epsilon=epsilon
+                )
+                for _ in range(num_layers)
+            ])
+    def forward(self, x, attention_mask=None):
+        if self.use_pretrained_model:
+            x = self.bert(input_ids=x, attention_mask=attention_mask)
+            x = x.last_hidden_state
+            x = self.projection(x)
+        else:
+            x = self.embedding(x)
+            for layer in self.layers:
+                x = layer(x, attention_mask=attention_mask)
+        return x
+#-----------------------------------------------------------------------------------------------------------------------
+class EncoderLayer(torch.nn.Module):
+    def __init__(
+            self,
+            input_dimension,
+            output_dimension,
+            num_heads,
+            dropout_rate,
+            qkv_bias,
+            scaling_value,
+            epsilon=1e-5
+    ):
+        super().__init__()
+        self.attention = nn.MultiheadAttention(
+            embed_dim=input_dimension,
+            num_heads=num_heads,
+            dropout=dropout_rate,
+            bias=qkv_bias,
+            batch_first=True
+        )
+        self.output_projection = nn.Linear(input_dimension, output_dimension) if input_dimension != output_dimension else nn.Identity()
+        self.norm1 = nn.LayerNorm(normalized_shape=input_dimension, eps=epsilon)
+        self.dropout1 = nn.Dropout(dropout_rate)
+        self.feedforward = FeedForward(
+            embedding_dimension=input_dimension,
+            scaling_value=scaling_value,
+            dropout_rate=dropout_rate
+        )
+        self.norm2 = nn.LayerNorm(normalized_shape=output_dimension, eps=epsilon)
+        self.dropout2 = nn.Dropout(dropout_rate)
+    def forward(self, x, attention_mask=None):
+        attn_output, _ = self.attention(x, x, x, key_padding_mask=attention_mask)
+        attn_output = self.output_projection(attn_output)
+        x = self.norm1(x + self.dropout1(attn_output))
+        ff_output = self.feedforward(x)
+        x = self.norm2(x + self.dropout2(ff_output))
+        return x
+#-----------------------------------------------------------------------------------------------------------------------
+class FeedForward(torch.nn.Module):
+    def __init__(self, embedding_dimension, scaling_value, dropout_rate=0.1):
+        super().__init__()
+        self.layers = torch.nn.Sequential(
+            torch.nn.Linear(
+                in_features=embedding_dimension,
+                out_features=embedding_dimension * scaling_value,
+                bias=True
+            ),
+            torch.nn.GELU(),
+            torch.nn.Dropout(dropout_rate),
+            torch.nn.Linear(
+                in_features=embedding_dimension * scaling_value,
+                out_features=embedding_dimension,
+                bias=True
+            )
+        )
+    def forward(self, x):
+        return self.layers(x)
+#-----------------------------------------------------------------------------------------------------------------------
+class Embedding(torch.nn.Module):
+    def __init__(
+        self,
+        vocabulary_size,
+        embedding_dimension=768,
+        context_length=77
+    ):
+        super().__init__()
+        self.token_embedding = nn.Embedding(
+            num_embeddings=vocabulary_size,
+            embedding_dim=embedding_dimension
+        )
+        self.embedding_dimension = embedding_dimension
+        self.context_length = context_length
+        self.register_buffer("positional_encoding", self._generate_positional_encoding(context_length))
+
+    def _generate_positional_encoding(self, seq_len):
+        position = torch.arange(seq_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, self.embedding_dimension, 2, dtype=torch.float) *
+                             -(math.log(10000.0) / self.embedding_dimension))
+        pos_enc = torch.zeros((seq_len, self.embedding_dimension), device=position.device)
+        pos_enc[:, 0::2] = torch.sin(position * div_term)
+        pos_enc[:, 1::2] = torch.cos(position * div_term)
+        return pos_enc.unsqueeze(0)
+
+    def forward(self, token_ids):
+        assert token_ids.dim() == 2, "Input token_ids should be of shape (batch_size, seq_len)"
+        token_embedded = self.token_embedding(token_ids)
+        seq_len = token_ids.size(1)
+        if seq_len > self.context_length:
+            position_encoded = self._generate_positional_encoding(seq_len).to(token_embedded.device)
+        else:
+            position_encoded = self.positional_encoding[:, :seq_len, :].to(token_embedded.device)
+        return token_embedded + position_encoded