UNet_Attn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1424e95d",
   "metadata": {},
   "source": [
    "### Importing Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4900bd1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import matplotlib.pyplot as plot\n",
    "import math"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7f57a59",
   "metadata": {},
   "source": [
    "### Upsampling/Downsampling "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a2b9d73",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Upsample(dim):\n",
    "    return nn.ConvTranspose2d(dim, dim, 4, 2, 1)\n",
    "\n",
    "\n",
    "def Downsample(dim):\n",
    "    return nn.Conv2d(dim, dim, 4, 2, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c18c956",
   "metadata": {},
   "source": [
    "### Encodings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "949fbd30",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TimePositionEmbeddings(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",
    "        if len(time.size()) == 0:\n",
    "            time = torch.tensor([time.item()]).to(\"cuda\")\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",
    "        return embeddings\n",
    "\n",
    "\n",
    "class ConditionPositionEmbeddings(nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "\n",
    "        self.pe = torch.zeros(1, dim).to(\"cuda\")\n",
    "\n",
    "        i = torch.arange(dim // 2)\n",
    "        # i = i.reshape(1, -1).t()\n",
    "\n",
    "        p = -torch.arange(0, dim, 2) / dim\n",
    "        exp = torch.pow(1e4, p)\n",
    "\n",
    "        self.pe[0, 0::2] = torch.sin(i * exp)\n",
    "        self.pe[0, 1::2] = torch.cos(i * exp)\n",
    "\n",
    "    def forward(self, text):\n",
    "        N, D = text.shape\n",
    "\n",
    "        text = nn.functional.normalize(text, dim=1)\n",
    "\n",
    "        output = torch.empty(N, D).to(\"cuda\")\n",
    "\n",
    "        pos_emb = self.pe\n",
    "        output = text + pos_emb\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f63fd3df",
   "metadata": {},
   "source": [
    "### ConvOp Block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9d5a29a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConvOp(nn.Module):\n",
    "    def __init__(self, h_x, w_x, in_channels, out_channels, emb_dim):\n",
    "        super().__init__()\n",
    "        self.time_mlp = nn.Sequential(nn.GELU(), nn.Linear(emb_dim, in_channels))\n",
    "        self.cross_attn_block = Cross_Norm(h_x, w_x, in_channels)\n",
    "        self.init_conv = nn.Conv2d(\n",
    "            in_channels=in_channels,\n",
    "            out_channels=in_channels,\n",
    "            kernel_size=7,\n",
    "            padding=3,\n",
    "            groups=in_channels,\n",
    "        )\n",
    "        self.conv_block = nn.Sequential(\n",
    "            nn.GroupNorm(1, in_channels),\n",
    "            nn.Conv2d(\n",
    "                in_channels=in_channels,\n",
    "                out_channels=out_channels * 2,\n",
    "                kernel_size=3,\n",
    "                padding=(1, 1),\n",
    "            ),\n",
    "            nn.GELU(),\n",
    "            nn.GroupNorm(1, out_channels * 2),\n",
    "            nn.Conv2d(\n",
    "                in_channels=out_channels * 2,\n",
    "                out_channels=out_channels,\n",
    "                kernel_size=3,\n",
    "                padding=(1, 1),\n",
    "            ),\n",
    "        )\n",
    "\n",
    "        self.res_conv = (\n",
    "            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1)\n",
    "            if in_channels != out_channels\n",
    "            else nn.Identity()\n",
    "        )\n",
    "\n",
    "    def forward(self, x, t_emb, c_emb):\n",
    "        h = self.init_conv(x)\n",
    "\n",
    "        t_e = self.time_mlp(t_emb)\n",
    "        t_e = t_e[(...,) + (None,) * 2]\n",
    "\n",
    "        if c_emb is not None:\n",
    "            h = self.cross_attn_block(h, c_emb)\n",
    "            h = h + t_e\n",
    "        else:\n",
    "            h = h + t_e\n",
    "\n",
    "        h = self.conv_block(h)\n",
    "\n",
    "        return h + self.res_conv(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06c12648",
   "metadata": {},
   "source": [
    "### Residual Blocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "103a641e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Residual_Norm(nn.Module):\n",
    "    def __init__(self, input_dim, linear_attn=True):\n",
    "        super().__init__()\n",
    "        self.norm = nn.GroupNorm(1, input_dim)\n",
    "        self.attn = Attention(input_dim, linear_attn=linear_attn)\n",
    "        # self.attn = Attention(input_dim, linear_attn=linear_attn)\n",
    "\n",
    "    def forward(self, x):\n",
    "        input = x\n",
    "        x = self.norm(x)\n",
    "        x = self.attn(x)\n",
    "        return x + input\n",
    "\n",
    "\n",
    "class Cross_Norm(nn.Module):\n",
    "    def __init__(self, h_x, w_x, input_dim, embed_dim=256, num_heads=4):\n",
    "        super().__init__()\n",
    "        self.norm = nn.GroupNorm(1, input_dim)\n",
    "        self.attn = CrossAttention(h_x, w_x, input_dim, embed_dim=256, num_heads=4)\n",
    "        # self.attn = Attention(input_dim, linear_attn=linear_attn)\n",
    "\n",
    "    def forward(self, x, c_emb):\n",
    "        input = x\n",
    "        x = self.norm(x)\n",
    "        x = self.attn(x, c_emb)\n",
    "        return x + input"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa04563f",
   "metadata": {},
   "source": [
    "### Self Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6381b10e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Attention(nn.Module):\n",
    "    def __init__(self, input_dim, embed_dim=256, num_heads=4, linear_attn=True):\n",
    "        super().__init__()\n",
    "        self.n_head = num_heads\n",
    "        self.emb_dim = embed_dim\n",
    "        self.linear_attn = linear_attn\n",
    "        self.hidden_dim = self.emb_dim * self.n_head\n",
    "        self.key = nn.Conv2d(input_dim, self.hidden_dim, kernel_size=3, padding=(1, 1))\n",
    "        self.query = nn.Conv2d(\n",
    "            input_dim, self.hidden_dim, kernel_size=3, padding=(1, 1)\n",
    "        )\n",
    "        self.value = nn.Conv2d(\n",
    "            input_dim, self.hidden_dim, kernel_size=3, padding=(1, 1)\n",
    "        )\n",
    "        if linear_attn:\n",
    "            self.proj = nn.Sequential(\n",
    "                nn.Conv2d(self.hidden_dim, input_dim, kernel_size=3, padding=(1, 1)),\n",
    "                nn.GroupNorm(1, input_dim),\n",
    "            )\n",
    "        else:\n",
    "            self.proj = nn.Conv2d(\n",
    "                self.hidden_dim, input_dim, kernel_size=3, padding=(1, 1)\n",
    "            )\n",
    "\n",
    "    def forward(self, x):\n",
    "        N, C, H, W = x.shape\n",
    "        output = torch.empty((N, C, H, W)).to(\"cuda\")\n",
    "\n",
    "        Q = self.query(x)\n",
    "        K = self.key(x)\n",
    "        V = self.value(x)\n",
    "\n",
    "        Q = torch.reshape(Q, (N, self.n_head, self.emb_dim, H * W))\n",
    "        K = torch.reshape(K, (N, self.n_head, self.emb_dim, H * W))\n",
    "        V = torch.reshape(V, (N, self.n_head, self.emb_dim, H * W))\n",
    "\n",
    "        new_K = torch.transpose(torch.transpose(K, 1, 2), 2, 3)\n",
    "        new_Q = torch.transpose(Q, 1, 2)\n",
    "        new_V = torch.transpose(V, 1, 2)\n",
    "\n",
    "        # softmax = torch.nn.Softmax(dim=-1)\n",
    "\n",
    "        if not self.linear_attn:\n",
    "            dot_product = (torch.matmul(new_Q, new_K)) / (\n",
    "                (self.emb_dim / self.n_head) ** (0.5)\n",
    "            )\n",
    "            attention_scores = dot_product.softmax(-1)\n",
    "            Y = torch.matmul(attention_scores, new_V)\n",
    "        else:  # linear_attn\n",
    "            Q_sm = new_Q.softmax(-2)\n",
    "            K_sm = new_K.softmax(-1)\n",
    "\n",
    "            KV = torch.matmul(new_V, K_sm)\n",
    "\n",
    "            Y = torch.matmul(KV, Q_sm)\n",
    "\n",
    "        new_Y = torch.cat([Y[:, :, h, :] for h in range(self.n_head)], dim=1)\n",
    "        new_Y = torch.reshape(new_Y, (N, self.hidden_dim, H, W))\n",
    "        output = self.proj(new_Y)\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96624c69",
   "metadata": {},
   "source": [
    "### Cross Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fd385dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CrossAttention(nn.Module):\n",
    "    def __init__(self, h_x, w_x, input_dim, embed_dim=256, num_heads=4):\n",
    "        super().__init__()\n",
    "        self.n_head = num_heads\n",
    "        self.emb_dim = embed_dim\n",
    "        self.hidden_dim = self.emb_dim * self.n_head\n",
    "        self.key = nn.Conv2d(input_dim, self.hidden_dim, kernel_size=3, padding=(1, 1))\n",
    "        self.query = nn.Conv2d(\n",
    "            input_dim, self.hidden_dim, kernel_size=3, padding=(1, 1)\n",
    "        )\n",
    "        self.value = nn.Conv2d(\n",
    "            input_dim, self.hidden_dim, kernel_size=3, padding=(1, 1)\n",
    "        )\n",
    "        self.proj = nn.Conv2d(self.hidden_dim, input_dim, kernel_size=3, padding=(1, 1))\n",
    "        self.linear1 = nn.Linear(256, int(input_dim * h_x * w_x))\n",
    "\n",
    "    def forward(self, x, c_emb):\n",
    "        N, C, H, W = x.shape\n",
    "        output = torch.empty((N, C, H, W)).to(\"cuda\")\n",
    "\n",
    "        c_emb = self.linear1(c_emb)\n",
    "        c_emb = torch.nn.functional.normalize(c_emb, dim=-1)\n",
    "\n",
    "        c_emb = torch.reshape(c_emb, (N, C, H, W))\n",
    "        Q = self.query(x)\n",
    "        K = self.key(c_emb)\n",
    "        V = self.value(c_emb)\n",
    "\n",
    "        Q = torch.reshape(Q, (N, self.n_head, self.emb_dim, H * W))\n",
    "        K = torch.reshape(K, (N, self.n_head, self.emb_dim, H * W))\n",
    "        V = torch.reshape(V, (N, self.n_head, self.emb_dim, H * W))\n",
    "\n",
    "        new_K = torch.transpose(torch.transpose(K, 1, 2), 2, 3)\n",
    "        new_Q = torch.transpose(Q, 1, 2)\n",
    "        new_V = torch.transpose(V, 1, 2)\n",
    "\n",
    "        softmax = torch.nn.Softmax(dim=-1)\n",
    "\n",
    "        dot_product = (torch.matmul(new_Q, new_K)) / (\n",
    "            (self.emb_dim / self.n_head) ** (0.5)\n",
    "        )\n",
    "        attention_scores = softmax(dot_product)\n",
    "\n",
    "        Y = torch.matmul(attention_scores, new_V)\n",
    "\n",
    "        new_Y = torch.cat([Y[:, :, h, :] for h in range(self.n_head)], dim=1)\n",
    "        new_Y = torch.reshape(new_Y, (N, self.hidden_dim, H, W))\n",
    "        output = self.proj(new_Y)\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aeabf2df",
   "metadata": {},
   "source": [
    "### UNet Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18691d38",
   "metadata": {},
   "outputs": [],
   "source": [
    "class UNet(nn.Module):\n",
    "    def __init__(self, conditional_gen=False):\n",
    "        super().__init__()\n",
    "\n",
    "        im_dim = 32\n",
    "        self.img_dim = im_dim\n",
    "        text_dim = 256\n",
    "        self.text_emb_dim = text_dim\n",
    "        dim_mult = [1, 2, 4, 8]\n",
    "        self.time_emb = nn.Sequential(\n",
    "            TimePositionEmbeddings(im_dim),\n",
    "            nn.Linear(im_dim, im_dim),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(im_dim, im_dim),\n",
    "        )\n",
    "\n",
    "        # self.cond_emb = nn.Sequential(\n",
    "        #     ConditionPositionEmbeddings(text_dim),\n",
    "        #     nn.Linear(text_dim, im_dim),\n",
    "        #     nn.ReLU()\n",
    "        # )\n",
    "\n",
    "        init_dim = 32 // 3 * 2\n",
    "        self.in_conv = nn.Conv2d(3, init_dim, 7, padding=3)\n",
    "\n",
    "        dims = [init_dim, *map(lambda m: im_dim * m, dim_mult)]\n",
    "        in_out = list(zip(dims[:-1], dims[1:]))\n",
    "\n",
    "        self.downs = nn.ModuleList([])\n",
    "        idx = 0\n",
    "        for ind, (dim_in, dim_out) in enumerate(in_out):\n",
    "            is_last = ind >= (len(in_out) - 1)\n",
    "\n",
    "            self.downs.append(\n",
    "                nn.ModuleList(\n",
    "                    [\n",
    "                        ConvOp(\n",
    "                            32 / dim_mult[idx],\n",
    "                            32 / dim_mult[idx],\n",
    "                            dim_in,\n",
    "                            dim_out,\n",
    "                            emb_dim=im_dim,\n",
    "                        ),\n",
    "                        ConvOp(\n",
    "                            32 / dim_mult[idx],\n",
    "                            32 / dim_mult[idx],\n",
    "                            dim_out,\n",
    "                            dim_out,\n",
    "                            emb_dim=im_dim,\n",
    "                        ),\n",
    "                        Residual_Norm(dim_out),\n",
    "                        Downsample(dim_out) if not is_last else nn.Identity(),\n",
    "                    ]\n",
    "                )\n",
    "            )\n",
    "            idx += 1\n",
    "\n",
    "        bridge_dim = dims[-1]\n",
    "        self.bridge_block1 = ConvOp(\n",
    "            32 / dim_mult[-1], 32 / dim_mult[-1], bridge_dim, bridge_dim, im_dim\n",
    "        )\n",
    "        self.bridge_attn = Residual_Norm(dim_out, linear_attn=False)\n",
    "        self.bridge_block2 = ConvOp(\n",
    "            32 / dim_mult[-1], 32 / dim_mult[-1], bridge_dim, bridge_dim, im_dim\n",
    "        )\n",
    "\n",
    "        self.ups = nn.ModuleList([])\n",
    "        idx -= 1\n",
    "        for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):\n",
    "            is_last = ind >= (len(in_out) - 1)\n",
    "\n",
    "            self.ups.append(\n",
    "                nn.ModuleList(\n",
    "                    [\n",
    "                        ConvOp(\n",
    "                            32 / dim_mult[idx],\n",
    "                            32 / dim_mult[idx],\n",
    "                            dim_out * 2,\n",
    "                            dim_in,\n",
    "                            emb_dim=self.img_dim,\n",
    "                        ),\n",
    "                        ConvOp(\n",
    "                            32 / dim_mult[idx],\n",
    "                            32 / dim_mult[idx],\n",
    "                            dim_in,\n",
    "                            dim_in,\n",
    "                            emb_dim=self.img_dim,\n",
    "                        ),\n",
    "                        Residual_Norm(dim_in),\n",
    "                        Upsample(dim_in) if not is_last else nn.Identity(),\n",
    "                    ]\n",
    "                )\n",
    "            )\n",
    "            idx -= 1\n",
    "        self.final_ConvOp = ConvOp(32, 32, im_dim, im_dim, emb_dim=im_dim)\n",
    "        self.final_conv2d = nn.Conv2d(im_dim, 3, 1)\n",
    "\n",
    "    def forward(self, t, states):\n",
    "        if len(states) == 2:\n",
    "            x = states[0]\n",
    "            c = states[1]\n",
    "            c_emb = c\n",
    "        else:\n",
    "            x = states\n",
    "            c = None\n",
    "            c_emb = None\n",
    "\n",
    "        t_emb = self.time_emb(t)\n",
    "\n",
    "        skip_conn_vals = []\n",
    "\n",
    "        x = self.in_conv(x)\n",
    "\n",
    "        # down\n",
    "        for block1, block2, attn, downsample in self.downs:\n",
    "            x = block1(x, t_emb, c_emb)\n",
    "            x = block2(x, t_emb, c_emb)\n",
    "            x = attn(x)\n",
    "            skip_conn_vals.append(x)\n",
    "            x = downsample(x)\n",
    "\n",
    "        # bridge\n",
    "        x = self.bridge_block1(x, t_emb, c_emb)\n",
    "        x = self.bridge_attn(x)\n",
    "        x = self.bridge_block2(x, t_emb, c_emb)\n",
    "\n",
    "        # up\n",
    "        for block1, block2, attn, upsample in self.ups:\n",
    "            x = torch.cat((x, skip_conn_vals.pop()), dim=1)\n",
    "            x = block1(x, t_emb, c_emb)\n",
    "            x = block2(x, t_emb, c_emb)\n",
    "            x = attn(x)\n",
    "            x = upsample(x)\n",
    "\n",
    "        x = self.final_ConvOp(x, t_emb, c_emb)\n",
    "        if c is not None:\n",
    "            return self.final_conv2d(x), c\n",
    "        else:\n",
    "            return self.final_conv2d(x)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}