diffusion_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "import torchvision\n",
    "from torchvision import transforms\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import warnings \n",
    "import math\n",
    "from torch.optim import Adam\n",
    "import wandb\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"mps\")\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_and_transform_data(IMG_SIZE, BATCH_SIZE):\n",
    "    data_transforms = torchvision.transforms.Compose([\n",
    "        transforms.Resize((IMG_SIZE, IMG_SIZE)),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Lambda(lambda x: 2*x - 1) # normalize to [-1, 1]\n",
    "    ])\n",
    "\n",
    "    train = torchvision.datasets.Flowers102(root='./flower', download=False, split='train', transform=data_transforms)\n",
    "    test = torchvision.datasets.Flowers102(root='./flower', download=False, split='test', transform=data_transforms)\n",
    "\n",
    "    # Since we are training a generative model, we don't need\n",
    "    # splitted train and test datasets. We can just concatenate them\n",
    "    data = torch.utils.data.ConcatDataset([train, test])\n",
    "    data_loader = torch.utils.data.DataLoader(data, batch_size=BATCH_SIZE, shuffle=True, drop_last=True)\n",
    "\n",
    "    return data_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reverse_transforms(img):\n",
    "    '''\n",
    "    Reverse the transforms we applied while loading the data\n",
    "    in order to visualize the images\n",
    "    '''\n",
    "    reverse_tr = transforms.Compose([\n",
    "        transforms.Lambda(lambda x: (x + 1) / 2), # denormalize to [0, 1]\n",
    "        transforms.ToPILImage()\n",
    "    ])\n",
    "    \n",
    "    return reverse_tr(img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Because the model was trained on my local machine, for the sake of \n",
    "# speed I will use a smaller image size and batch size\n",
    "IMG_SIZE = 64\n",
    "BATCH_SIZE = 128\n",
    "\n",
    "data_loader = load_and_transform_data(IMG_SIZE, BATCH_SIZE)\n",
    "len(data_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def linear_beta_schedule(beta_0, beta_T, T):\n",
    "    '''Linear schedule for beta from beta_0 to beta_T in T steps'''\n",
    "    return torch.linspace(beta_0, beta_T, T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Take the original image and time step, and return the image after the diffusion process\n",
    "def forward_diffusion(x_0, t, alpha_cumprods):\n",
    "    alpha_cumprod = alpha_cumprods[t]\n",
    "    if len(alpha_cumprod.shape) == 0:\n",
    "        alpha_cumprod = alpha_cumprod.unsqueeze(0)        \n",
    "    alpha_cumprod = alpha_cumprod[:, None, None, None]\n",
    "    \n",
    "    noise = torch.randn_like(x_0)\n",
    "    x_t = x_0*torch.sqrt(alpha_cumprod) + torch.sqrt(1-alpha_cumprod)*noise\n",
    "\n",
    "    return x_t.squeeze(), noise\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T = 350 # number of times the diffusion process is applied\n",
    "betas = linear_beta_schedule(1e-4, 2e-2, T)\n",
    "alphas = 1 - betas\n",
    "alpha_cumprods = torch.cumprod(alphas, dim=0).to(device) # pre-compute cumulative product of alphas\n",
    "betas = betas.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = next(iter(data_loader))[0][0].to(device)\n",
    "\n",
    "plt.figure(figsize=(20,15))\n",
    "plt.axis('off')\n",
    "\n",
    "num_images = 8\n",
    "step_size = int(T/num_images)\n",
    "\n",
    "for idx in range(0, T, step_size):\n",
    "    t = torch.Tensor([idx]).type(torch.int64).to(device)\n",
    "    plt.subplot(1, num_images+1, int(idx/step_size) + 1)\n",
    "    \n",
    "    noisy_image, noise = forward_diffusion(image, idx, alpha_cumprods)\n",
    "    plt.imshow(reverse_transforms(noisy_image))\n",
    "    plt.title(f't = {idx}')\n",
    "    plt.axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This U-Net architecture has been adapted from https://colab.research.google.com/drive/1sjy9odlSSy0RBVgMTgP7s99NXsqglsUL?usp=sharing\n",
    "class Block(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch, time_emb_dim, up=False):\n",
    "        super().__init__()\n",
    "        self.time_mlp =  nn.Linear(time_emb_dim, out_ch)\n",
    "        if up:\n",
    "            self.conv1 = nn.Conv2d(2*in_ch, out_ch, 3, padding=1)\n",
    "            self.transform = nn.ConvTranspose2d(out_ch, out_ch, 4, 2, 1)\n",
    "        else:\n",
    "            self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)\n",
    "            self.transform = nn.Conv2d(out_ch, out_ch, 4, 2, 1)\n",
    "        \n",
    "        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)\n",
    "        self.bnorm1 = nn.BatchNorm2d(out_ch)\n",
    "        self.bnorm2 = nn.BatchNorm2d(out_ch)\n",
    "        self.relu  = nn.ReLU()\n",
    "        \n",
    "    def forward(self, x, t, ):\n",
    "        # First Conv\n",
    "        h = self.bnorm1(self.relu(self.conv1(x)))\n",
    "        # Time embedding\n",
    "        time_emb = self.relu(self.time_mlp(t))\n",
    "        # Extend last 2 dimensions\n",
    "        time_emb = time_emb[(..., ) + (None, ) * 2]\n",
    "        # Add time channel\n",
    "        h = h + time_emb\n",
    "        # Second Conv\n",
    "        h = self.bnorm2(self.relu(self.conv2(h)))\n",
    "        # Down or Upsample\n",
    "        return self.transform(h)\n",
    "\n",
    "\n",
    "class SinusoidalPositionEmbeddings(nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "\n",
    "    def forward(self, time):\n",
    "        device = time.device\n",
    "        half_dim = self.dim // 2\n",
    "        embeddings = math.log(10000) / (half_dim - 1)\n",
    "        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)\n",
    "        embeddings = time[:, None] * embeddings[None, :]\n",
    "        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)\n",
    "        \n",
    "        return embeddings\n",
    "\n",
    "\n",
    "class SimpleUnet(nn.Module):\n",
    "    \"\"\"\n",
    "    A simplified variant of the U-Net architecture.\n",
    "    \"\"\"\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        image_channels = 3\n",
    "        down_channels = (64, 128, 256, 512, 1024)\n",
    "        up_channels = (1024, 512, 256, 128, 64)\n",
    "        out_dim = 1 \n",
    "        time_emb_dim = 32\n",
    "\n",
    "        # Time embedding\n",
    "        self.time_mlp = nn.Sequential(\n",
    "                SinusoidalPositionEmbeddings(time_emb_dim),\n",
    "                nn.Linear(time_emb_dim, time_emb_dim),\n",
    "                nn.ReLU()\n",
    "            )\n",
    "        \n",
    "        # Initial projection\n",
    "        self.conv0 = nn.Conv2d(image_channels, down_channels[0], 3, padding=1)\n",
    "\n",
    "        # Downsample\n",
    "        self.downs = nn.ModuleList([Block(down_channels[i], down_channels[i+1], \\\n",
    "                                    time_emb_dim) \\\n",
    "                    for i in range(len(down_channels)-1)])\n",
    "        # Upsample\n",
    "        self.ups = nn.ModuleList([Block(up_channels[i], up_channels[i+1], \\\n",
    "                                        time_emb_dim, up=True) \\\n",
    "                    for i in range(len(up_channels)-1)])\n",
    "\n",
    "        self.output = nn.Conv2d(up_channels[-1], 3, out_dim)\n",
    "\n",
    "    def forward(self, x, timestep):\n",
    "        # Embedd time\n",
    "        t = self.time_mlp(timestep)\n",
    "        # Initial conv\n",
    "        x = self.conv0(x)\n",
    "        # Unet\n",
    "        residual_inputs = []\n",
    "        for down in self.downs:\n",
    "            x = down(x, t)\n",
    "            residual_inputs.append(x)\n",
    "        for up in self.ups:\n",
    "            residual_x = residual_inputs.pop()\n",
    "            # Add residual x as additional channels\n",
    "            x = torch.cat((x, residual_x), dim=1)           \n",
    "            x = up(x, t)\n",
    "        return self.output(x)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SimpleUnet()\n",
    "print(\"Number of parameters: \", sum(p.numel() for p in model.parameters()))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will compute the loss as the authors did in the original paper, which is a simple MSE loss:\n",
    "<p align=\"center\">\n",
    "    <img src=\"./images/loss_equation.png\" width=35%> \n",
    "</p>\n",
    "<p align=\"center\">\n",
    "    Source: <i>Denoising Diffusion Probabilistic Models</i> paper\n",
    "</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_loss(model, x_0, t):\n",
    "    x_t, noise = forward_diffusion(x_0, t, alpha_cumprods)\n",
    "    pred_noise = model(x_t, t)\n",
    "\n",
    "    return torch.mean((pred_noise - noise)**2)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following block, the sampling procedure is implemented as discussed in the original paper.\n",
    "<p align=\"center\">\n",
    "    <img src=\"./images/sampling.png\" width=35%> \n",
    "</p>\n",
    "<p align=\"center\">\n",
    "    Source: <i>Denoising Diffusion Probabilistic Models</i> paper\n",
    "</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def sample_timestep(x_t, t, alpha_cumprods):\n",
    "    if t == 0:\n",
    "        alpha_t = alpha_cumprods[0]\n",
    "        z = 0\n",
    "    else:\n",
    "        alpha_t = alpha_cumprods[t] / alpha_cumprods[t-1]\n",
    "        z = torch.randn_like(x_t, device=device) # Sample noise from a normal distribution\n",
    "\n",
    "    alphas_cumprod_prev = torch.nn.functional.pad(alpha_cumprods[:-1], (1, 0), value=1.0)\n",
    "    posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alpha_cumprods)\n",
    "    posterior_variance_t = posterior_variance[t].to(device)\n",
    "    pred_noise = model(x_t, t)\n",
    "    prev_x = (x_t - (1-alpha_t)/torch.sqrt(1-alpha_cumprods[t])*pred_noise)/torch.sqrt(alpha_t) + z*torch.sqrt(posterior_variance_t)\n",
    "\n",
    "    return prev_x\n",
    "    \n",
    "\n",
    "@torch.no_grad()\n",
    "def sample_plot_image(return_final=False):\n",
    "    imgs = [] # store images during denoising\n",
    "    img_size = IMG_SIZE\n",
    "    img = torch.randn((1, 3, img_size, img_size), device=device)\n",
    "    \n",
    "    num_images = 8 # number of images to plot\n",
    "    stepsize = int(T/num_images)\n",
    "\n",
    "    for i in range(0,T)[::-1]:\n",
    "        t = torch.full((1,), i, device=device, dtype=torch.long)\n",
    "        img = sample_timestep(img, t, alpha_cumprods)\n",
    "        if return_final and i == 0: # we already got x_0\n",
    "            return reverse_transforms(img.detach().cpu().squeeze())\n",
    "        if i % stepsize == 0:\n",
    "            imgs.append(reverse_transforms(img.detach().cpu().squeeze()))\n",
    "\n",
    "    return imgs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.init(\n",
    "    project=\"Diffusion model\", \n",
    "    config={\n",
    "        \"learning_rate\": 0.001,\n",
    "        \"batch_size\": BATCH_SIZE,\n",
    "        \"epochs\": 280,\n",
    "    }\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.to(device)\n",
    "optimizer = Adam(model.parameters(), lr=wandb.config['learning_rate'])\n",
    "epochs = wandb.config['epochs']\n",
    "\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    model.train()\n",
    "    epoch_loss = 0\n",
    "    for step, batch in enumerate(data_loader):\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        batch = batch[0].to(device)\n",
    "        t = torch.randint(0, T, (wandb.config['batch_size'],), device=device).long().to(device)\n",
    "        loss = get_loss(model, batch, t)\n",
    "        epoch_loss += loss.detach().cpu().item()\n",
    "        \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    \n",
    "    print(f\"Epoch {epoch} | Loss: {epoch_loss/len(data_loader)} \", end=\"\\r\")\n",
    "    \n",
    "    model.eval()\n",
    "    imgs = sample_plot_image()\n",
    "    wandb.log({\n",
    "        \"epoch_loss\": epoch_loss/len(data_loader),\n",
    "        \"images\": [wandb.Image(img) for img in imgs]\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "imgs = sample_plot_image()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot images\n",
    "fig, axs = plt.subplots(1, len(imgs), figsize=(20, 20))\n",
    "for i, img in enumerate(imgs):\n",
    "    axs[i].imshow(img)\n",
    "    axs[i].axis('off')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), \"model.pt\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.9 ('base')",
   "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.10.9"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "892acf4471d4e46ccf7e8a968c938b95284af0f7e87e91c50416527849cb4ab4"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}