diff options
Diffstat (limited to 'models')
| -rw-r--r-- | models/__init__.py | 6 | ||||
| -rw-r--r-- | models/autoencoder.py | 103 | ||||
| -rw-r--r-- | models/unet.py | 140 |
3 files changed, 249 insertions, 0 deletions
diff --git a/models/__init__.py b/models/__init__.py new file mode 100644 index 0000000..db05aa9 --- /dev/null +++ b/models/__init__.py @@ -0,0 +1,6 @@ +""" +Model definitions for MNIST Diffusion. +""" + +from .unet import UNetMNIST +from .autoencoder import AEModule, Autoencoder diff --git a/models/autoencoder.py b/models/autoencoder.py new file mode 100644 index 0000000..353ec89 --- /dev/null +++ b/models/autoencoder.py @@ -0,0 +1,103 @@ +""" +Autoencoder model for MNIST. +Can be used for latent diffusion. +""" + +import torch +import torch.nn as nn + + +class AEModule(nn.Module): + """Base class for autoencoder modules (encoder and decoder).""" + + def __init__(self): + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + raise NotImplementedError("Subclasses must implement forward method.") + + def encode(self, x: torch.Tensor) -> torch.Tensor: + """Encode input to latent space.""" + raise NotImplementedError("Subclasses must implement encode method.") + + def decode(self, z: torch.Tensor) -> torch.Tensor: + """Decode from latent space to image space.""" + raise NotImplementedError("Subclasses must implement decode method.") + + +class Autoencoder(AEModule): + """ + Convolutional Autoencoder for MNIST images. + + Args: + input_channels: Number of input image channels (1 for MNIST) + latent_channels: Number of channels in latent space + """ + + def __init__(self, input_channels: int = 1, latent_channels: int = 1): + super().__init__() + self.input_channels = input_channels + self.latent_channels = latent_channels + + # Encoder: 28x28 -> 8x8 + self.encoder = nn.Sequential( + nn.Conv2d(input_channels, 16, kernel_size=3, padding=1), + nn.ReLU(), + nn.Conv2d(16, 16, kernel_size=3, padding=1), + nn.ReLU(), + nn.MaxPool2d(kernel_size=2), # 14x14 + + nn.Conv2d(16, 32, kernel_size=3, padding=1), + nn.ReLU(), + nn.Conv2d(32, 32, kernel_size=3, padding=1), + nn.ReLU(), + nn.MaxPool2d(kernel_size=2), # 7x7 + + nn.Conv2d(32, 64, kernel_size=3, padding=1), + nn.ReLU(), + nn.Conv2d(64, 64, kernel_size=3, padding=1), + nn.ReLU(), + nn.Conv2d(64, latent_channels, kernel_size=3, padding=1), + nn.ReLU(), + + # 7x7 -> 8x8 (no activation - latent space should be unconstrained) + nn.Conv2d(latent_channels, latent_channels, kernel_size=2, padding=1), + ) + + # Decoder: 8x8 -> 28x28 + self.decoder = nn.Sequential( + # 8x8 -> 7x7 + nn.Conv2d(latent_channels, latent_channels, kernel_size=2, padding=0), + nn.ReLU(), + + nn.Conv2d(latent_channels, 64, kernel_size=3, padding=1), + nn.ReLU(), + nn.Conv2d(64, 64, kernel_size=3, padding=1), + nn.ReLU(), + nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2), # 14x14 + nn.ReLU(), + + nn.Conv2d(32, 32, kernel_size=3, padding=1), + nn.ReLU(), + nn.ConvTranspose2d(32, 16, kernel_size=2, stride=2), # 28x28 + nn.ReLU(), + + nn.Conv2d(16, 16, kernel_size=3, padding=1), + nn.ReLU(), + nn.Conv2d(16, input_channels, kernel_size=3, padding=1), + nn.Sigmoid(), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Full autoencoder forward pass.""" + z = self.encoder(x) + x_recon = self.decoder(z) + return x_recon + + def encode(self, x: torch.Tensor) -> torch.Tensor: + """Encode input to latent space.""" + return self.encoder(x) + + def decode(self, z: torch.Tensor) -> torch.Tensor: + """Decode from latent space to image space.""" + return self.decoder(z) diff --git a/models/unet.py b/models/unet.py new file mode 100644 index 0000000..f21859a --- /dev/null +++ b/models/unet.py @@ -0,0 +1,140 @@ +""" +UNet model for MNIST diffusion. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import sys +sys.path.append("..") +from src.config import IMG_SIZE, NB_CHANNEL, NB_LABEL, DIFFU_STEPS + + +class SelfAttention(nn.Module): + """Self-attention module for UNet.""" + + def __init__(self, nb_channels: int, nb_heads: int): + super().__init__() + self.attention = nn.MultiheadAttention(nb_channels, nb_heads) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + _, c, w, h = x.shape + x = x.view(-1, c, w * h) + x = x.permute(2, 0, 1) + x, _ = self.attention(x, x, x) + x = x.permute(1, 2, 0) + x = x.view(-1, c, w, h) + return x + + +class UNetMNIST(nn.Module): + """ + UNet architecture for MNIST diffusion model. + + Inputs: + xt: image at step t (NB_CHANNEL x IMG_SIZE x IMG_SIZE) + t: step number (1) + vec: one-hot vector of the label (NB_LABEL) + + Output: + Predicted noise or denoised image (NB_CHANNEL x IMG_SIZE x IMG_SIZE) + """ + + def __init__(self, use_attention: bool = False): + super().__init__() + self.use_attention = use_attention + + # Encoder for timestep t + self.encodet = nn.Linear(1, IMG_SIZE * IMG_SIZE) + + # Encoder for label vector + self.encodevec = nn.Linear(NB_LABEL, IMG_SIZE * IMG_SIZE) + + # UNet encoder (2 extra channels for t and vec) + self.conv1 = nn.Conv2d(NB_CHANNEL + 2, 64, 3, padding=1) + self.conv2 = nn.Conv2d(64, 64, 3, padding=1) + self.maxpool1 = nn.MaxPool2d(2, 2) + + self.conv3 = nn.Conv2d(64, 128, 3, padding=1) + self.conv4 = nn.Conv2d(128, 128, 3, padding=1) + self.maxpool2 = nn.MaxPool2d(2, 2) + + # Bottleneck + self.conv5 = nn.Conv2d(128, 256, 3, padding=1) + self.conv6 = nn.Conv2d(256, 256, 3, padding=1) + + # Optional attention layers + if use_attention: + self.att1 = SelfAttention(128, 8) + self.att2 = SelfAttention(256, 8) + self.att3 = SelfAttention(256, 8) + + # UNet decoder + self.upconv1 = nn.ConvTranspose2d(256, 128, 2, stride=2) + self.conv7 = nn.Conv2d(256, 128, 3, padding=1) + self.conv8 = nn.Conv2d(128, 128, 3, padding=1) + + self.upconv2 = nn.ConvTranspose2d(128, 64, 2, stride=2) + self.conv9 = nn.Conv2d(128, 64, 3, padding=1) + self.conv10 = nn.Conv2d(64, 64, 3, padding=1) + + # Output layer + self.conv11 = nn.Conv2d(64, NB_CHANNEL, 3, padding=1) + + def forward(self, xt: torch.Tensor, t: torch.Tensor, vec: torch.Tensor) -> torch.Tensor: + """ + Forward pass of UNet. + + Args: + xt: Noisy image at timestep t [B, C, H, W] + t: Timestep [B, 1] + vec: Label one-hot vector [B, NB_LABEL] + + Returns: + Predicted noise or denoised image [B, C, H, W] + """ + # Encode timestep and label + t_enc = F.relu(self.encodet(t / DIFFU_STEPS)) + t_enc = t_enc.view(-1, 1, IMG_SIZE, IMG_SIZE) + + vec_enc = F.relu(self.encodevec(vec)) + vec_enc = vec_enc.view(-1, 1, IMG_SIZE, IMG_SIZE) + + # Concatenate input with embeddings + x = torch.cat((xt, t_enc, vec_enc), dim=1) + + # Encoder path + x1 = F.relu(self.conv1(x)) + x1 = F.relu(self.conv2(x1)) + + x2 = self.maxpool1(x1) + x2 = F.relu(self.conv3(x2)) + x2 = F.relu(self.conv4(x2)) + + # Bottleneck + x3 = self.maxpool2(x2) + if self.use_attention: + x3 = self.att1(x3) + x3 = F.relu(self.conv5(x3)) + if self.use_attention: + x3 = self.att2(x3) + x3 = F.relu(self.conv6(x3)) + if self.use_attention: + x3 = self.att3(x3) + + # Decoder path with skip connections + x4 = self.upconv1(x3) + x4 = torch.cat((x4, x2), dim=1) + x4 = F.relu(self.conv7(x4)) + x4 = F.relu(self.conv8(x4)) + + x5 = self.upconv2(x4) + x5 = torch.cat((x5, x1), dim=1) + x5 = F.relu(self.conv9(x5)) + x5 = F.relu(self.conv10(x5)) + + # Output + out = self.conv11(x5) + + return out |