---
license: apache-2.0
---

# GAN for Comic Faces Paired Generation

## Model Overview

This model implements a **Generative Adversarial Network (GAN)** with a **UNet generator** and a **PatchGAN discriminator**. The network is designed to generate paired images of comic faces based on a synthetic dataset of comic faces. The model aims to generate high-quality image pairs where the first image is transformed into a second target image (e.g., photo-to-cartoon or cartoon-to-photo transformations).

- **Dataset:** [Comic Faces Paired Synthetic Dataset](https://www.kaggle.com/datasets/defileroff/comic-faces-paired-synthetic)
- **Batch Size:** 32  
- **Input Shape:** (3, 256, 256) (RGB Images)  
- **Output Shape:** (3, 256, 256)  

## Model Architecture

### Generator: **UNet**

The generator uses a **UNet architecture**, which is designed for image-to-image translation tasks. It has an encoder-decoder structure with skip connections, allowing for high-resolution output. The architecture includes the following layers:

- **Encoder Path (Contracting Path):**  
  The encoder consists of **DoubleConv** layers that progressively downsample the input image to extract features. It uses **MaxPool2d** to reduce spatial dimensions.

- **Bottleneck:**  
  The deepest layer of the network (with 1024 feature channels) processes the smallest version of the image.

- **Decoder Path (Expanding Path):**  
  The decoder uses **Upsample** layers to progressively increase the spatial dimensions and **DoubleConv** layers to refine the output. Skip connections are used to combine features from the encoder path.

- **Final Convolution:**  
  The final layer outputs the transformed image using a **1x1 convolution**.

### Discriminator: **PatchGANDiscriminator**

The discriminator uses a **PatchGAN** architecture, which classifies patches of the image as real or fake. The discriminator works by processing the **input image and output image pair** (3 channels for the input image + 3 channels for the generated output). It progressively reduces the spatial dimensions using **Conv2d** and **LeakyReLU** activations, while normalizing each layer with **InstanceNorm2d**. The final output is a probability score indicating whether the patch is real or fake.

---

### Generator Code (UNet):

```python
class UNet(nn.Module, PyTorchModelHubMixin):
    def __init__(self, in_channels, out_channels):
        super(UNet, self).__init__()

        # Contracting Path (Encoder)
        self.down_conv1 = DoubleConv(in_channels, 64)
        self.down_conv2 = DoubleConv(64, 128)
        self.down_conv3 = DoubleConv(128, 256)
        self.down_conv4 = DoubleConv(256, 512)
        self.down_conv5 = DoubleConv(512, 1024)

        # Downsampling
        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)

        # Upsampling layers using nn.Upsample
        self.upsample = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)

        # Decoder (Expanding Path)
        self.up_conv1 = DoubleConv(1024 + 512, 512)
        self.up_conv2 = DoubleConv(512 + 256, 256)
        self.up_conv3 = DoubleConv(256 + 128, 128)
        self.up_conv4 = DoubleConv(128 + 64, 64)

        # Final 1x1 convolution to get desired number of output channels
        self.final_conv = nn.Conv2d(64, out_channels, kernel_size=1)

    def forward(self, x):
        x1 = self.down_conv1(x)
        x2 = self.down_conv2(self.maxpool(x1))
        x3 = self.down_conv3(self.maxpool(x2))
        x4 = self.down_conv4(self.maxpool(x3))
        x5 = self.down_conv5(self.maxpool(x4))

        x = self.upsample(x5)
        x = torch.cat([x4, x], dim=1)
        x = self.up_conv1(x)

        x = self.upsample(x)
        x = torch.cat([x3, x], dim=1)
        x = self.up_conv2(x)

        x = self.upsample(x)
        x = torch.cat([x2, x], dim=1)
        x = self.up_conv3(x)

        x = self.upsample(x)
        x = torch.cat([x1, x], dim=1)
        x = self.up_conv4(x)

        return self.final_conv(x)


class PatchGANDiscriminator(nn.Module, PyTorchModelHubMixin):
    def __init__(self, in_channels=6):
        super().__init__()

        self.layers = nn.Sequential(
            nn.Conv2d(in_channels, 64, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
            nn.InstanceNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
            nn.InstanceNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(256, 512, kernel_size=4, stride=1, padding=1),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(512, 1, kernel_size=4, stride=1, padding=1),
        )

    def forward(self, x):
        return self.layers(x)
```