all.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.nn import functional as F\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.datasets import ImageFolder\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.nn import functional as F\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.datasets import ImageFolder\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets\n",
    "import matplotlib.pyplot as plt\n",
    "from transformers import CLIPTokenizer, CLIPModel\n",
    "\n",
    "# Mapping Network\n",
    "class MappingNetwork(nn.Module):\n",
    "    def __init__(self, latent_dim, hidden_dim, num_layers):\n",
    "        super(MappingNetwork, self).__init__()\n",
    "        layers = [nn.Linear(latent_dim + 512, hidden_dim), nn.ReLU()]  # Concatenate latent vector with prompt embedding\n",
    "        for _ in range(num_layers - 1):\n",
    "            layers.extend([nn.Linear(hidden_dim, hidden_dim), nn.ReLU()])\n",
    "        self.mapping = nn.Sequential(*layers)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return self.mapping(x)\n",
    "\n",
    "# Noise Injection\n",
    "class NoiseInjection(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(NoiseInjection, self).__init__()\n",
    "        self.scale = nn.Parameter(torch.zeros(1))\n",
    "    \n",
    "    def forward(self, x):\n",
    "        noise = torch.randn_like(x)\n",
    "        return x + self.scale * noise\n",
    "\n",
    "# Adaptive Instance Normalization (AdaIN)\n",
    "class AdaIN(nn.Module):\n",
    "    def __init__(self, latent_dim, channels):\n",
    "        super(AdaIN, self).__init__()\n",
    "        self.norm = nn.InstanceNorm2d(channels)\n",
    "        self.style = nn.Linear(latent_dim, channels * 2)\n",
    "    \n",
    "    def forward(self, x, w):\n",
    "        style = self.style(w).unsqueeze(-1).unsqueeze(-1)\n",
    "        gamma, beta = style.chunk(2, 1)\n",
    "        return (1 + gamma) * self.norm(x) + beta\n",
    "\n",
    "# Self-Attention\n",
    "class SelfAttention(nn.Module):\n",
    "    def __init__(self, in_channels):\n",
    "        super(SelfAttention, self).__init__()\n",
    "        self.query = nn.Conv2d(in_channels, in_channels // 8, 1)\n",
    "        self.key = nn.Conv2d(in_channels, in_channels // 8, 1)\n",
    "        self.value = nn.Conv2d(in_channels, in_channels, 1)\n",
    "        self.gamma = nn.Parameter(torch.zeros(1))\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size, channels, height, width = x.size()\n",
    "        query = self.query(x).view(batch_size, -1, width * height).permute(0, 2, 1)\n",
    "        key = self.key(x).view(batch_size, -1, width * height)\n",
    "        attention = torch.bmm(query, key)\n",
    "        attention = F.softmax(attention, dim=-1)\n",
    "        value = self.value(x).view(batch_size, -1, width * height)\n",
    "        out = torch.bmm(value, attention.permute(0, 2, 1))\n",
    "        out = out.view(batch_size, channels, height, width)\n",
    "        return self.gamma * out + x\n",
    "\n",
    "# Blur\n",
    "class Blur(nn.Module):\n",
    "    def __init__(self, channels):\n",
    "        super(Blur, self).__init__()\n",
    "        kernel = torch.tensor([[1, 2, 1], [2, 4, 2], [1, 2, 1]], dtype=torch.float32)\n",
    "        kernel = kernel[None, None, :, :] / kernel.sum()\n",
    "        self.register_buffer('kernel', kernel.repeat(channels, 1, 1, 1))\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return F.conv2d(x, self.kernel, stride=1, padding=1, groups=x.shape[1])\n",
    "\n",
    "# Style Layer\n",
    "class StyleLayer(nn.Module):\n",
    "    def __init__(self, latent_dim, in_channels, out_channels, kernel_size=3, upsample=False, attention=False):\n",
    "        super(StyleLayer, self).__init__()\n",
    "        self.noise_injection = NoiseInjection()\n",
    "        self.adain = AdaIN(latent_dim, out_channels)\n",
    "        self.conv = nn.utils.spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size, padding=kernel_size//2))\n",
    "        self.act = nn.LeakyReLU(0.2)\n",
    "        self.upsample = upsample\n",
    "        if upsample:\n",
    "            self.blur = Blur(out_channels)\n",
    "        self.attention = SelfAttention(out_channels) if attention else None\n",
    "    \n",
    "    def forward(self, x, w):\n",
    "        if self.upsample:\n",
    "            x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)\n",
    "            x = self.blur(x)\n",
    "        x = self.conv(x)\n",
    "        x = self.noise_injection(x)\n",
    "        x = self.adain(x, w)\n",
    "        x = self.act(x)\n",
    "        if self.attention:\n",
    "            x = self.attention(x)\n",
    "        return x\n",
    "\n",
    "# Residual Block\n",
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, downsample=False):\n",
    "        super(ResidualBlock, self).__init__()\n",
    "        self.conv1 = nn.utils.spectral_norm(nn.Conv2d(in_channels, out_channels, 3, padding=1))\n",
    "        self.act = nn.LeakyReLU(0.2)\n",
    "        self.conv2 = nn.utils.spectral_norm(nn.Conv2d(out_channels, out_channels, 3, padding=1))\n",
    "        self.downsample = downsample\n",
    "        self.downsample_layer = nn.Conv2d(in_channels, out_channels, 1, stride=2) if downsample else None\n",
    "\n",
    "    def forward(self, x):\n",
    "        residual = x\n",
    "        x = self.conv1(x)\n",
    "        x = self.act(x)\n",
    "        x = self.conv2(x)\n",
    "        if self.downsample:\n",
    "            x = F.avg_pool2d(x, 2)\n",
    "            residual = self.downsample_layer(residual)\n",
    "        return x + residual\n",
    "\n",
    "# Generator with Prompt Conditioning\n",
    "class Generator(nn.Module):\n",
    "    def __init__(self, latent_dim, hidden_dim, output_channels, num_layers, clip_model):\n",
    "        super(Generator, self).__init__()\n",
    "        self.mapping = MappingNetwork(latent_dim + 512, hidden_dim, num_layers)  # Concatenate latent vector with prompt embedding\n",
    "        self.style_layers = nn.ModuleList([\n",
    "            StyleLayer(hidden_dim, hidden_dim, hidden_dim, upsample=True),\n",
    "            StyleLayer(hidden_dim, hidden_dim, hidden_dim, upsample=True, attention=True),\n",
    "            StyleLayer(hidden_dim, hidden_dim, hidden_dim, upsample=True),\n",
    "            StyleLayer(hidden_dim, hidden_dim, hidden_dim, upsample=True),\n",
    "            StyleLayer(hidden_dim, hidden_dim, hidden_dim, upsample=True),\n",
    "            StyleLayer(hidden_dim, hidden_dim, hidden_dim, upsample=True, attention=True),\n",
    "            StyleLayer(hidden_dim, hidden_dim, hidden_dim, upsample=True),\n",
    "            StyleLayer(hidden_dim, hidden_dim, hidden_dim, upsample=True),\n",
    "            StyleLayer(hidden_dim, hidden_dim, output_channels, upsample=False)\n",
    "        ])\n",
    "        self.to_rgb = nn.Sequential(\n",
    "            nn.Conv2d(output_channels, output_channels, 1),\n",
    "            nn.Tanh()\n",
    "        )\n",
    "        self.clip_model = clip_model\n",
    "    \n",
    "    def forward(self, z, prompt_embedding):\n",
    "        z = torch.cat((z, prompt_embedding), dim=1)  # Concatenate latent vector with prompt embedding\n",
    "        w = self.mapping(z)\n",
    "        w = w.unsqueeze(1).repeat(1, len(self.style_layers), 1)\n",
    "        x = torch.randn(z.shape[0], self.style_layers[0].conv.in_channels, 4, 4).to(z.device)\n",
    "        for i, style_layer in enumerate(self.style_layers):\n",
    "            x = style_layer(x, w[:, i])\n",
    "        x = self.to_rgb(x)\n",
    "        return x\n",
    "\n",
    "# Discriminator\n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, input_channels, hidden_dim, num_layers):\n",
    "        super(Discriminator, self).__init__()\n",
    "        self.layers = nn.ModuleList([\n",
    "            nn.utils.spectral_norm(nn.Conv2d(input_channels, hidden_dim, 4, stride=2, padding=1)),\n",
    "            nn.LeakyReLU(0.2)\n",
    "        ])\n",
    "        \n",
    "        for _ in range(num_layers - 1):\n",
    "            self.layers.extend([\n",
    "                ResidualBlock(hidden_dim, hidden_dim * 2, downsample=True),\n",
    "                nn.LeakyReLU(0.2)\n",
    "            ])\n",
    "            hidden_dim *= 2\n",
    "        \n",
    "        self.layers.extend([\n",
    "            nn.Conv2d(hidden_dim, 1, 4, stride=1, padding=0),\n",
    "            nn.Flatten()\n",
    "        ])\n",
    "        \n",
    "        self.model = nn.Sequential(*self.layers)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return self.model(x)\n",
    "\n",
    "# Exponential Moving Average (EMA) class for model parameters\n",
    "class ExponentialMovingAverage:\n",
    "    def __init__(self, parameters, decay):\n",
    "        self.parameters = list(parameters)\n",
    "        self.decay = decay\n",
    "        self.shadow_params = [p.clone().detach() for p in self.parameters]\n",
    "\n",
    "    def update(self, parameters):\n",
    "        for shadow_param, param in zip(self.shadow_params, parameters):\n",
    "            shadow_param.data = (1.0 - self.decay) * param.data + self.decay * shadow_param.data\n",
    "\n",
    "    def apply(self, parameters):\n",
    "        for shadow_param, param in zip(self.shadow_params, parameters):\n",
    "            param.data.copy_(shadow_param.data)\n",
    "\n",
    "# Gradient Penalty for Wasserstein GAN with Gradient Penalty (WGAN-GP)\n",
    "def compute_gradient_penalty(discriminator, real_samples, fake_samples, device):\n",
    "    alpha = torch.randn(real_samples.size(0), 1, 1, 1).to(device)\n",
    "    interpolates = (alpha * real_samples + (1 - alpha) * fake_samples).requires_grad_(True)\n",
    "    d_interpolates = discriminator(interpolates)\n",
    "    fake = torch.ones(d_interpolates.size()).to(device)\n",
    "    gradients = torch.autograd.grad(\n",
    "        outputs=d_interpolates,\n",
    "        inputs=interpolates,\n",
    "        grad_outputs=fake,\n",
    "        create_graph=True,\n",
    "        retain_graph=True,\n",
    "        only_inputs=True\n",
    "    )[0]\n",
    "    gradients = gradients.view(gradients.size(0), -1)\n",
    "    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()\n",
    "    return gradient_penalty\n",
    "\n",
    "# Placeholder for path length regularization computation\n",
    "def compute_path_length_regularization(generator, z, fake_images):\n",
    "    return 0.0\n",
    "\n",
    "# Placeholder for style mixing regularization computation\n",
    "def compute_style_mixing_regularization(generator, z):\n",
    "    return 0.0\n",
    "\n",
    "# Check if the dataset path exists\n",
    "dataset_path = 'dataset'\n",
    "if not os.path.exists(dataset_path):\n",
    "    raise FileNotFoundError(f\"Dataset path '{dataset_path}' does not exist.\")\n",
    "\n",
    "# Verify the directory structure\n",
    "for root, dirs, files in os.walk(dataset_path):\n",
    "    print(f\"Checking directory: {root}\")\n",
    "    if not dirs and not files:\n",
    "        raise ValueError(f\"No subdirectories or files found in the dataset path '{dataset_path}'. Ensure the directory structure follows the expected format for ImageFolder.\")\n",
    "    for dir_name in dirs:\n",
    "        print(f\"Found subdirectory: {dir_name}\")\n",
    "    for file_name in files:\n",
    "        print(f\"Found file: {file_name}\")\n",
    "\n",
    "# Data loading and preprocessing\n",
    "transform = transforms.Compose([\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.CenterCrop(128),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])\n",
    "])\n",
    "try:\n",
    "    dataset = datasets.ImageFolder(dataset_path, transform=transform)\n",
    "    dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)\n",
    "    print(f\"Loaded dataset with {len(dataset)} images.\")\n",
    "except Exception as e:\n",
    "    raise RuntimeError(f\"Error loading dataset: {e}\")\n",
    "\n",
    "def extract_prompt_embedding(prompt, clip_model, clip_tokenizer, device):\n",
    "    with torch.no_grad():\n",
    "        inputs = clip_tokenizer(prompt, return_tensors=\"pt\", padding=True, truncation=True).to(device)\n",
    "        text_features = clip_model.get_text_features(**inputs)\n",
    "        text_features = text_features / text_features.norm(dim=-1, keepdim=True)\n",
    "    return text_features\n",
    "\n",
    "def train(generator, discriminator, dataloader, num_epochs, latent_dim, clip_model, clip_tokenizer, device, prompts):\n",
    "    g_optim = optim.Adam(generator.parameters(), lr=0.001, betas=(0.0, 0.99))\n",
    "    d_optim = optim.Adam(discriminator.parameters(), lr=0.001, betas=(0.0, 0.99))\n",
    "    \n",
    "    g_ema = ExponentialMovingAverage(generator.parameters(), decay=0.999)\n",
    "    \n",
    "    def d_loss_fn(real_scores, fake_scores):\n",
    "        return F.relu(1.0 - real_scores).mean() + F.relu(1.0 + fake_scores).mean()\n",
    "    \n",
    "    def g_loss_fn(fake_scores):\n",
    "        return -fake_scores.mean()\n",
    "    \n",
    "    for epoch in range(num_epochs):\n",
    "        for real_images, _ in dataloader:\n",
    "            real_images = real_images.to(device)\n",
    "            batch_size = real_images.shape[0]\n",
    "            \n",
    "            # Train discriminator\n",
    "            d_optim.zero_grad()\n",
    "            \n",
    "            z = torch.randn(batch_size, latent_dim).to(device)\n",
    "            prompt = prompts[torch.randint(0, len(prompts), (1,)).item()]\n",
    "            prompt_embedding = extract_prompt_embedding(prompt, clip_model, clip_tokenizer, device)\n",
    "            fake_images = generator(z, prompt_embedding)\n",
    "            \n",
    "            real_scores = discriminator(real_images)\n",
    "            fake_scores = discriminator(fake_images.detach())\n",
    "            \n",
    "            gradient_penalty = compute_gradient_penalty(discriminator, real_images, fake_images, device)\n",
    "            \n",
    "            d_loss = d_loss_fn(real_scores, fake_scores) + 10 * gradient_penalty\n",
    "            d_loss.backward()\n",
    "            d_optim.step()\n",
    "            \n",
    "            # Train generator\n",
    "            g_optim.zero_grad()\n",
    "            \n",
    "            z = torch.randn(batch_size, latent_dim).to(device)\n",
    "            prompt = prompts[torch.randint(0, len(prompts), (1,)).item()]\n",
    "            prompt_embedding = extract_prompt_embedding(prompt, clip_model, clip_tokenizer, device)\n",
    "            fake_images = generator(z, prompt_embedding)\n",
    "            \n",
    "            fake_scores = discriminator(fake_images)\n",
    "            \n",
    "            path_length_regularization = compute_path_length_regularization(generator, z, fake_images)\n",
    "            style_mixing_regularization = compute_style_mixing_regularization(generator, z)\n",
    "            \n",
    "            g_loss = g_loss_fn(fake_scores) + 2 * path_length_regularization + 2 * style_mixing_regularization\n",
    "            g_loss.backward()\n",
    "            g_optim.step()\n",
    "            \n",
    "            g_ema.update(generator.parameters())\n",
    "        \n",
    "        print(f\"Epoch {epoch+1}/{num_epochs}, D Loss: {d_loss.item()}, G Loss: {g_loss.item()}\")\n",
    "        \n",
    "        # Save model checkpoints\n",
    "        if (epoch + 1) % 10 == 0:\n",
    "            torch.save(generator.state_dict(), f\"generator_epoch_{epoch+1}.pth\")\n",
    "            torch.save(discriminator.state_dict(), f\"discriminator_epoch_{epoch+1}.pth\")\n",
    "        \n",
    "        # Save generated images for inspection\n",
    "        if (epoch + 1) % 10 == 0:\n",
    "            with torch.no_grad():\n",
    "                z = torch.randn(16, latent_dim).to(device)\n",
    "                prompt = prompts[torch.randint(0, len(prompts), (1,)).item()]\n",
    "                prompt_embedding = extract_prompt_embedding(prompt, clip_model, clip_tokenizer, device)\n",
    "                fake_images = generator(z, prompt_embedding)\n",
    "                fake_images = (fake_images + 1) / 2  # Denormalize\n",
    "                grid = torchvision.utils.make_grid(fake_images, nrow=4)\n",
    "                plt.imshow(grid.permute(1, 2, 0).cpu().numpy())\n",
    "                plt.title(f\"Epoch {epoch+1} - Prompt: {prompt}\")\n",
    "                plt.savefig(f\"generated_images_epoch_{epoch+1}.png\")\n",
    "                plt.close()\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    # Initialize models and move to device\n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    latent_dim = 512\n",
    "    clip_model = CLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\").to(device)\n",
    "    clip_tokenizer = CLIPTokenizer.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
    "    generator = Generator(latent_dim, hidden_dim=512, output_channels=3, num_layers=8, clip_model=clip_model).to(device)\n",
    "    discriminator = Discriminator(input_channels=3, hidden_dim=64, num_layers=4).to(device)\n",
    "\n",
    "    # Define prompts for conditioning\n",
    "    prompts = [\n",
    "        \"a cat wearing sunglasses\",\n",
    "        \"a cat wearing a hat\",\n",
    "        \"a cat wearing a bowtie\",\n",
    "        \"a cat wearing a scarf\",\n",
    "        \"a cat wearing a shirt\",\n",
    "        \"a cat wearing a sweater\",\n",
    "        \"a cat wearing a jacket\",\n",
    "        \"a cat wearing a dress\",\n",
    "        \"a cat wearing a suit\",\n",
    "        \"a cat wearing a costume\"\n",
    "    ]\n",
    "\n",
    "    # Train the GAN\n",
    "    train(generator, discriminator, dataloader, num_epochs=100, latent_dim=latent_dim, clip_model=clip_model, clip_tokenizer=clip_tokenizer, device=device, prompts=prompts)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pot",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}