{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3YsSv7ZWExb6"
   },
   "source": [
    "# Acknowledgement\n",
    "\n",
    "Parts of this pset were inspired by\n",
    "* Berkeley CS294-158, taught by Pieter Abbeel, Wilson Yan, Kevin Frans, and Philipp Wu;\n",
    "* MIT 6.S184/6.S975, taught by Peter Holderrieth and Ezra Erives;\n",
    "* The [blog post](https://lilianweng.github.io/posts/2021-07-11-diffusion-models/) about diffusion models by Lilian Weng.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YMetH7q3Exb7"
   },
   "source": [
    "# Submission Guideline for Part 2\n",
    "\n",
    "Please include your answer to all problems, including formulas, proofs, and the figures generated in each problem, excluding code. You are required to submit the (single) pdf and all (four) notebooks (one for each problem) with your code and running outputs. Do not include code in the pdf file.\n",
    "\n",
    "Specifically, for Problem 3 in this notebook, the pdf should contain:\n",
    "- The generated figures `results/mnist_train_plot.png` and `results/image_w{w}.png` (w=0.0, 0.5, 1.0, 2.0, 4.0)\n",
    "- Answer to the short answer question about the U-Net architecture\n",
    "- Answer to the short answer question about different CFG weight $w$ in problem 3.2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pjuQYgfAnQei"
   },
   "source": [
    "# Problem 3: MNIST and Conditional Generation\n",
    "In this problem, we will write the code for conditional generation on the MNIST dataset. This part requires GPUs--you can use Google Colab for GPU access. To work on this notebook in Google Colab, copy the `pset-5` directory to your Google Drive and open this notebook. Then, start working on a GPU machine with `Runtime -> Change runtime type -> T4 GPU`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vTf_uTXrLLwP"
   },
   "source": [
    "## MNIST Dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "99dK1DZtngKq",
    "outputId": "45bdf1d7-37a8-4e9d-ef37-c2b002e8b591"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 9.91M/9.91M [00:00<00:00, 17.8MB/s]\n",
      "100%|██████████| 28.9k/28.9k [00:00<00:00, 483kB/s]\n",
      "100%|██████████| 1.65M/1.65M [00:00<00:00, 4.37MB/s]\n",
      "100%|██████████| 4.54k/4.54k [00:00<00:00, 7.02MB/s]\n"
     ]
    }
   ],
   "source": [
    "from torchvision import transforms\n",
    "from torchvision.datasets import MNIST\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "tf = transforms.Compose([transforms.ToTensor()])\n",
    "train_dataset = MNIST(\"./data\", train=True, download=True, transform=tf)\n",
    "train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)\n",
    "test_dataset = MNIST(\"./data\", train=False, download=True, transform=tf)\n",
    "test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 251
    },
    "id": "a3yr1AOrOLCK",
    "outputId": "f921e106-7789-4b25-d935-301915c4c28e"
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABKgAAAE3CAYAAACO1XkUAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQAAVehJREFUeJzt3Xm8jPX///HXcXDOse/Zd9klIkuRJbKFsvYJoU/5RpHIUrasFdJqX0JEIWVJyBJpsW/ZKkso+3IQDvP7o1/X53q9OXNm5syca2bO4367dbu9n97XXNd7zts1Z7y7rtcV4XK5XAIAAAAAAAA4JIXTAwAAAAAAAEDyxgIVAAAAAAAAHMUCFQAAAAAAABzFAhUAAAAAAAAcxQIVAAAAAAAAHMUCFQAAAAAAABzFAhUAAAAAAAAcxQIVAAAAAAAAHMUCFQAAAAAAABzFAhUAAAAAAAAcxQIVAABIlLVr10pERIRERERIt27d7rrNqVOnJHXq1BIRESGPPPKI6nvkkUckIiJCChcuLDdu3LjjtYMHD5aIiAjZvHnzHcccPXq02vbixYsybNgwKV++vGTKlEnSpUsnhQoVkmbNmsmUKVNERGTGjBnWeBP6zxyr6d+x//tfmjRpJHfu3FK3bl0ZMWKE/Pnnnx78BN0bN26czJgxI9H78bfDhw/L4MGDZfv27U4PBQAAhIGUTg8AAACEh+joaJkzZ46MGTNGoqKiVN+sWbPE5XJJypTxf/X4/fffZfz48dK9e3efjn/p0iWpVKmS/Pbbb9KiRQvp1KmTpE6dWn777TfZsGGDvPvuu/Lss89KjRo1ZNasWeq1w4cPl3379t3x5/fcc0+Cx42KirIWv27cuCF//fWXbNy4UQYOHCijRo2SSZMmSZs2bXx6TyL/LFAVLFhQnnnmGZ/3EQiHDx+WIUOGSMGCBaV8+fJODwcAAIQ4FqgAAIBfNG/eXObOnSuLFy+WVq1aqb7p06dLw4YNZfXq1Xd9bUxMjBQqVEiGDRsmnTp1kvTp03t9/MmTJ8vBgwdl3Lhxd13k+vdqpsKFC0vhwoVV35QpU2Tfvn3y9NNPe33clClT3vV1u3btkoYNG0q7du0kX758Ur16da/3DQAAkFxwix8AAPCLChUqSLly5WT69Onqz3/66SfZs2ePdOzYMd7XpkiRQkaOHClnzpyRt956y6fjHzx4UERE6tSpc9f+nDlz+rRfX5UtW1amT58ucXFxMmjQINU3b948efzxxyV//vwSFRUl2bJlk2bNmsnOnTvVdhEREXLkyBFZt26dupXw8OHDIiLyzTffSOvWraVw4cISExMjmTJlknr16sm6devuGM+ePXukZcuWkidPHomKipKcOXNKrVq1ZOnSpWq769evy4gRI6R06dISHR0tmTJlkiZNmsi2bdusbWbMmCG1atUSEZGOHTt6fEskAABAfLiCCgAA+E2nTp2kZ8+ecvz4ccmTJ4+IiEybNk1y5MghjRs3dvvaxx9/XB566CF55513pGvXrl4vKBUpUkRE/rla680333R7O2FSqVu3rhQoUEDWrVsnV65ckbRp04qIyAcffCBZs2aV5557TnLmzCm//vqrTJo0SapXry5bt26VYsWKicg/t0a+/PLLki1bNnnttdes/WbPnl1E/lkoOnfunLRv317y5s0rx48flylTpkidOnVkzZo18vDDD4uIyNmzZ6V27doiItKlSxcpUKCAnDlzRjZv3iw//vijNGrUSEREbt68KY899ph8//330q5dO+nWrZtcvHhRJk+eLNWrV5f169fLAw88IDVq1JD+/fvLiBEj5LnnnrOO48ktkQAAAHflAgAASIQ1a9a4RMT19ttvu86cOeNKnTq1a/jw4S6Xy+W6evWqK2PGjK5XXnnF5XK5XGnTpnXVrFlTvb5mzZqutGnTulwul2vjxo0uEXE9//zzVv+gQYNcIuL6+eef73rMf507d86VL18+l4i4cuTI4XryySddo0aNcn333XeuW7duuX0PNWvWdPnytcg+9vg0adLEJSKunTt3Wn8WGxt7x3Z79+51pU6d2vV///d/6s8LFChwx8/M3X7+/PNPV9asWV0NGjSw/mzx4sUuEXHNmzfP7VjHjh3rEhHX119/rf784sWLrnz58qlx/DsH06dPd7tPAAAAT3CLHwAA8JusWbPK448/bj11buHChXLx4kXp1KmTR6+vVq2aNGvWTKZOnSoHDhzw6tiZM2eWLVu2SJ8+fSRjxoyyYMEC6du3rzz88MNSpEgR+eabb7x9O36RIUMGEfmniPu//r2SyuVyyaVLl+TMmTOSPXt2KV68uPz4448e7/vf/YiIxMbGytmzZyUyMlIefPBBtZ+MGTOKiMjy5cvVOEyzZ8+WEiVKSMWKFeXMmTPWfzdu3JBHH31UNmzYINeuXfN4fAAAAJ5igQoAAPhVx44d5eDBg7JhwwaZNm2aVK5cWUqVKuXx60eOHCkul0v69evn9bGzZ88uo0aNkgMHDsiZM2fkq6++knbt2smRI0ekefPmcujQIa/3mVj/Lgj9u1AlIrJt2zZp3LixpE+fXjJmzCjZs2eX7Nmzy65du+T8+fMe7/vXX3+VNm3aSObMmSV9+vSSLVs2yZ49uyxbtkztp2bNmtK+fXuZMWOGZMuWTapXry6DBg2SvXv3qv398ssvsm/fPms89v+mTZsmt27dkjNnziTyJwIAAHAn54szAACAsFK/fn3JkyePDBkyRNasWSPjx4/36vUlSpSQjh07ypQpU7y6msiUNWtWady4sTRu3Fjy5csnI0aMkE8//VRef/11n/fpi507d0qqVKmsJwcePXpUatSoIRkyZJABAwZI8eLFJW3atBIRESE9evSQ2NhYj/YbGxsrNWrUkCtXrkiPHj2kbNmykj59eqvg/Lfffqu2//jjj6V3796yfPly+e6772TMmDEyfPhwGTdunHTr1k1E/rmiq2zZsjJ27Nh4j/tv/SsAAAB/YoEKAAD4VWRkpLRv315GjhwpMTEx0rZtW6/3MWTIEJkzZ4706dPHL0+Gq1KlioiIHD9+PNH78saqVavkyJEjUrduXet2vEWLFklsbKx8+eWX1pPw/nX27FmJiopSfxYREXHXfa9evVpOnDgh06ZNu+MJifEtwpUpU0bKlCkjvXv3lgsXLsiDDz4offv2la5du0pERIQUK1ZMTp8+LbVr15YUKdxfaB/fuAAAAHzBLX4AAMDvunTpIoMGDZIJEyaoW9s8lTt3bunevbusW7dOli1b5tFrNm3aJBcuXLhr3xdffCEi4tWthom1a9cu6dixo6RMmVLeeOMN688jIyNF5J+rlewmT54sf/755x37SZcunZw7d+6OP49vP998880dV56dO3dObt++rf4sU6ZMUqhQIbl69ar8/fffIiLSvn17+fPPP+O9guqvv/5S4/p33wAAAInFFVQAAMDv8ufPL4MHD07UPvr06SOTJk2Sn3/+2aPtP/nkE5k+fbo0atRIKleuLFmzZpWzZ8/KsmXLZM2aNVKqVCmPi7V7Iy4uTmbPni0iIjdv3pS//vpLNm7cKMuXL5c0adLIJ598IlWrVrW2b9CggaRJk0batWsn3bp1k8yZM8vGjRtl2bJlUqRIEYmLi1P7r1KlikydOlUGDBggJUuWlBQpUkiTJk3koYcekpw5c8orr7wihw8flrx588r27dtl1qxZUrZsWdm1a5e1j5kzZ8o777wjzZs3l6JFi0qqVKlk3bp1smLFCmnVqpXExMSIiEj37t1l5cqV0rt3b/n222+ldu3akiFDBjl69KisXr1aoqOjZc2aNSLyz2Jf+vTp5aOPPpI0adJIpkyZJEeOHFK7dm2//4wBAED4Y4EKAAAEpYwZM8prr70mPXv29Gj7Ll26SKZMmWTNmjUyduxYOXPmjERFRUnRokVl0KBB0rNnT/XUO3+5fv26tGvXTkREoqKiJHPmzFKqVCkZOnSodOzYUXLmzKm2L1KkiCxfvlz69+8vI0aMkMjISKlevbqsW7dOunXrJocPH1bbDx8+XM6dOycffvihXLhwQVwul/z+++9SsGBBWbFihbz66qvy/vvvS1xcnFSsWFGWLVsmU6dOVQtUjzzyiGzbtk2WLFkiJ0+elMjISClUqJCMHj3aqj8lIpIqVSpZunSpfPTRRzJr1iwZNGiQiPxzRVvlypWlQ4cO1rYxMTFWTa8ePXrI9evXpWbNmixQAQAAn0S4zOvCAQAAAAAAgCREDSoAAAAAAAA4igUqAAAAAAAAOIoFKgAAAAAAADiKBSoAAAAAAAA4igUqAAAAAAAAOIoFKgAAAAAAADiKBSoAAAAAAAA4igUqAAAAAAAAOIoFKgAAAAAAADiKBSoAAAAAAAA4igUqAAAAAAAAOIoFKgAAAAAAADiKBSoAAAAAAAA4igUqAAAAAAAAOIoFKgAAAAAAADiKBSoAAAAAAAA4igUqAAAAAAAAOIoFKgAAAAAAADiKBSoAAAAAAAA4igUqAAAAAAAAOIoFKgAAAAAAADiKBSoAAAAAAAA4igUqAAAAAAAAOIoFKgAAAAAAADiKBSoAAAAAAAA4igUqAAAAAAAAOIoFKgAAAAAAADiKBSoAAAAAAAA4igUqAAAAAAAAOIoFKgAAAAAAADiKBSoAAAAAAAA4igUqAAAAAAAAOIoFKgAAAAAAADiKBSoAAAAAAAA4igUqAAAAAAAAOIoFKgAAAAAAADiKBSoAAAAAAAA4igUqAAAAAAAAOIoFKgAAAAAAADiKBSoAAAAAAAA4igUqAAAAAAAAOIoFKgAAAAAAADgqbBeoZsyYIZkyZUr0fiIiIuSLL75I9H7gP8xteGJewxdzG56Y1/DF3IYv5jY8Ma/hi7kNT8xr/IJ2geqZZ56RZs2aOT0Mj3z44YdSsGBBiY6OlgcffFB++uknp4cU1EJlbtevXy9NmjSR3Llzh+XJ72+hMq8jR46USpUqSfr06SVHjhzSrFkz2b9/v9PDCmqhMrfjx4+XcuXKSYYMGSRDhgxStWpVWb58udPDClqhMq92o0aNkoiICOnRo4fTQwlqoTK3gwcPloiICPVfiRIlnB5WUAuVuRUROX78uDz99NOSNWtWiYmJkbJly8rmzZudHlZQCpV5LViw4B3nbEREhHTt2tXpoQWtUJnbW7duyYABA6RQoUISExMjRYoUkaFDh4rL5XJ6aEEpVOb18uXL0qNHDylQoIDExMRItWrV5Oeff3Z6WG4F7QJVqJg3b5707NlTBg0aJFu3bpX77rtP6tevL6dOnXJ6aEikK1euyH333Scffvih00OBH61bt066du0qP/zwg6xcuVJu3rwp9erVkytXrjg9NCRS3rx5ZdSoUbJlyxbZvHmz1K5dW5o2bSp79uxxemjwg59//lkmTpwo5cqVc3oo8KPSpUvLyZMnrf82bNjg9JDgB+fPn5fq1atLqlSpZPny5bJ3714ZM2aMZM6c2emhIRF+/vlndb6uXLlSRERatmzp8MiQWG+++aaMHz9ePvjgA/nll1/kzTfflLfeekvef/99p4eGRHj22Wdl5cqVMmvWLNm1a5fUq1dP6tatK8ePH3d6aPEK2QWqsWPHStmyZSVt2rSSL18+eeGFFyQ2NvaO7b744gspVqyYREdHS/369eXYsWOqf/HixVKhQgWJjo6WwoULy5AhQyQuLs6rcfz3v/+Vjh07SqlSpWTChAmSJk0amTZtWqLfY3IVLHPboEEDGTZsmDRv3jzR7wnBM69ff/21PPPMM1K6dGm57777ZMaMGXL06FHZsmVLot9jchUsc9ukSRNp2LChFCtWTO69914ZPny4pEuXTn744YdEv8fkKFjmVUQkNjZW/vOf/8jkyZP5B64fBNPcpkyZUnLmzGn9ly1btkS9t+QuWOb2zTfflHz58sn06dOlcuXKUqhQIalXr54UKVIk0e8xOQqWec2ePbs6X5csWSJFihSRmjVrJvo9JlfBMrfff/+9NG3aVBo1aiQFCxaUFi1aSL169bgzyEfBMK/Xrl2TBQsWyFtvvSU1atSQokWLyuDBg6Vo0aIyfvx4v7zPQAjZBaoUKVLIe++9J3v27JGPP/5Yvv32W3n11VfVNlevXpXhw4fLzJkzZePGjXLhwgVp06aN1f/dd99J+/btpXv37rJ3716ZOHGizJgxQ4YPH+7RGG7cuCFbtmyRunXrqnHVrVtXNm3a5J83mgwFw9zC/4J1Xi9evCgiIlmyZPF5H8ldMM7trVu35NNPP5UrV65I1apVE/X+kqtgmteuXbtKo0aN1O9b+C6Y5vbgwYOSO3duKVy4sPznP/+Ro0eP+uU9JlfBMrdffvmlPPDAA9KyZUvJkSOH3H///TJ58mS/vc/kJljm1e7GjRsye/Zs6dSpk0RERCTq/SVnwTK31apVk9WrV8uBAwdERGTHjh2yYcMGadCggX/eaDITDPMaFxcnt27dkujoaPXnMTExwX21sitIdejQwdW0aVOPt//ss89cWbNmtfL06dNdIuL64YcfrD/75ZdfXCLi+vHHH10ul8tVp04d14gRI9R+Zs2a5cqVK5eVRcS1aNGiux7z+PHjLhFxff/99+rPe/fu7apcubLHY09uQmFuTd5sm1yF4rzeunXL1ahRI1f16tU9HndyFEpzu3PnTlfatGldkZGRrowZM7qWLl3q8biTm1CZ17lz57rKlCnjunbtmsvlcrlq1qzp6t69u8fjTo5CZW6XLVvmmj9/vmvHjh2ur7/+2lW1alVX/vz5XZcuXfJ47MlNqMxtVFSUKyoqytWvXz/X1q1bXRMnTnRFR0e7ZsyY4fHYk5NQmVe7efPmuSIjI13Hjx/3eNzJUajM7a1bt1x9+vRxRUREuFKmTOmKiIi4Y5/4n1CZ16pVq7pq1qzpOn78uCsuLs41a9YsV4oUKVz33nuvx2NPaikDvwQWGKtWrZKRI0fKvn375NKlSxIXFyd///23XL16VdKkSSMi/1w2XqlSJes1JUqUkEyZMskvv/wilStXlh07dsjGjRvVKuStW7fu2A+SFnMbnoJxXrt27Sq7d+8O7v+LEAKCaW6LFy8u27dvl4sXL8rnn38uHTp0kHXr1kmpUqX8+6aTgWCY12PHjkn37t1l5cqVd/wfQPguGOZWRNT/mS9Xrpw8+OCDUqBAAZk/f7507tzZj+84+QiWub19+7Y88MADMmLECBERuf/++2X37t0yYcIE6dChg5/fdfgLlnm1mzp1qjRo0EBy587tnzeZTAXL3M6fP18++eQTmTNnjpQuXVq2b98uPXr0kNy5c3PO+iBY5nXWrFnSqVMnyZMnj0RGRkqFChWkbdu2QV3aJCQXqA4fPiyNGzeW//u//5Phw4dLlixZZMOGDdK5c2e5ceOGxx+wsbGxMmTIEHniiSfu6PPki3C2bNkkMjJS/vrrL/Xnf/31l+TMmdOzNwMlWOYW/hWM89qtWzdZsmSJrF+/XvLmzevVa/E/wTa3qVOnlqJFi4qISMWKFeXnn3+Wd999VyZOnOjxPhA887plyxY5deqUVKhQwfqzW7duyfr16+WDDz6Q69evS2RkpOdvDEEzt3eTKVMmuffee+XQoUM+vT65C6a5zZUr1x3/Y6BkyZKyYMECj16P/wmmef3XkSNHZNWqVbJw4UKvXgctmOa2d+/e0rdvX+sWs7Jly8qRI0dk5MiRLFB5KZjmtUiRIrJu3Tq5cuWKXLp0SXLlyiWtW7eWwoULe/WeklJILlBt2bJFbt++LWPGjJEUKf4pozV//vw7touLi5PNmzdL5cqVRURk//79cuHCBSlZsqSIiFSoUEH2799v/WPGW6lTp5aKFSvK6tWrrcdM3r59W1avXi3dunXzaZ/JXbDMLfwrmObV5XLJiy++KIsWLZK1a9dKoUKFfN4Xgmtu7+b27dty/fp1v+4zOQiWea1Tp47s2rVL/VnHjh2lRIkS0qdPHxanfBAsc3s3sbGx8uuvv0q7du38ts/kJJjmtnr16rJ//371ZwcOHJACBQr4vM/kKpjm9V/Tp0+XHDlySKNGjRK9r+QsmOb26tWr1hj+FRkZKbdv3/Z5n8lVMM3rv9KmTStp06aV8+fPy4oVK+Stt95K9D4DJagXqC5evCjbt29Xf5Y1a1YpWrSo3Lx5U95//31p0qSJbNy4USZMmHDH61OlSiUvvviivPfee5IyZUrp1q2bVKlSxfpLMHDgQGncuLHkz59fWrRoISlSpJAdO3bI7t27ZdiwYR6NsWfPntKhQwd54IEHpHLlyjJu3Di5cuWKdOzYMdHvP5yFwtzGxsaq/4v7+++/y/bt2yVLliySP39+3998GAuFee3atavMmTNHFi9eLOnTp5c///xTREQyZswoMTExifsBhLFQmNt+/fpJgwYNJH/+/HL58mWZM2eOrF27VlasWJHo9x+ugn1e06dPL2XKlFF/ljZtWsmaNesdfw4t2OdWRKRXr17SpEkTKVCggJw4cUIGDRokkZGR0rZt20S//3AWCnP78ssvS7Vq1WTEiBHSqlUr+emnn2TSpEkyadKkRL//cBUK8yryz//4mT59unTo0EFSpgzqf0oGjVCY2yZNmsjw4cMlf/78Urp0adm2bZuMHTtWOnXqlOj3H65CYV5XrFghLpdLihcvLocOHZLevXtLiRIlgnutwukiWPHp0KGDS0Tu+K9z584ul8vlGjt2rCtXrlyumJgYV/369V0zZ850iYjr/PnzLpfrn8JjGTNmdC1YsMBVuHBhV1RUlKtu3bquI0eOqON8/fXXrmrVqrliYmJcGTJkcFWuXNk1adIkq188KBb4/vvvu/Lnz+9KnTq1q3LlyqrYGe4UKnO7Zs2au46zQ4cO/v6RhIVQmde7jVFEXNOnT/f3jyRshMrcdurUyVWgQAFX6tSpXdmzZ3fVqVPH9c033/j95xEuQmVeTRRJT1iozG3r1q1duXLlcqVOndqVJ08eV+vWrV2HDh3y+88jnITK3LpcLtdXX33lKlOmjCsqKspVokQJ9XpooTSvK1ascImIa//+/X79GYSrUJnbS5cuubp37+7Knz+/Kzo62lW4cGHXa6+95rp+/brffybhIFTmdd68ea7ChQu7UqdO7cqZM6era9eurgsXLvj95+FPES6Xy+XtohYAAAAAAADgLykS3gQAAAAAAAAIHBaoAAAAAAAA4CgWqAAAAAAAAOAoFqgAAAAAAADgKBaoAAAAAAAA4CgWqAAAAAAAAOColJ5uGBEREchxwA9cLpdPr2Nug58vc8u8Bj/O2fDFORueOGfDF3Mbvvg8Dk+cs+GLuQ1fnswtV1ABAAAAAADAUSxQAQAAAAAAwFEsUAEAAAAAAMBRLFABAAAAAADAUSxQAQAAAAAAwFEsUAEAAAAAAMBRKZ0eAAAAAIDASJFC///ofPnyqfzxxx9b7dmzZ6u+KVOmBG5gAAAYuIIKAAAAAAAAjmKBCgAAAAAAAI5igQoAAAAAAACOogYVAAAAEKbKly+v8ubNmz1+LTWoAABJiSuoAAAAAAAA4CgWqAAAAAAAAOAoFqgAAAAAAADgKGpQ/X8lSpRQeeLEiSrXqFHDal+7dk31VahQQeV9+/b5eXQAEF7Sp09vtZs3b676pk+f7va1Y8aMsdqDBg1SfebnM4CkV7FiRZWfe+45q50pUybVV7hwYbevnT17ttUeMWKE6uP71t0VKFBA5blz57rdPi4uzmpfunQpIGMC4D+PPPKI1V6zZo1Xr127dq3VXrduneobPHhwIkYF+AdXUAEAAAAAAMBRLFABAAAAAADAUSxQAQAAAAAAwFERLpfL5dGGERGBHkuSMmtOLV++XGXz/v0rV65Y7Rs3bqi+WbNmqdyjRw8/jNB7Hk7lHcJtbr3Vt29fqz18+HDV98orr6g8bty4pBjSHXyZ2+Q+r6EgOZ2zZp2pzz//3Gr7+nMQETly5IjKDRs2VHn//v0+7zsxOGfDU3I6Z0333Xef1TZrbz7xxBMq161bV+WoqCir7e3P0P6zu3z5suozz/eNGzd6tW+7UJ/blCn/V1bWrDn15JNPun3tqlWrrHa9evX8O7AgwOdxeAr1czYxEvO9yZ1g+dkk57kNd57MLVdQAQAAAAAAwFEsUAEAAAAAAMBRyeoWvwcffNBqL1u2TPVlzpxZ5YEDB6o8depUq71gwQLVZ94OmCdPnkSN01dcDukb+y0B9r8jIiIzZ85UuVOnTkkyJlO4Xp4+dOhQlRs3bhzvtps3b1bZvN1yz549fhtXUgnnc9Z8lPyxY8dUTpMmjdX256Xqhw8fVrlatWoqnzp1ym/Hcidcz1l/+u9//6uy/fatTz/9NKmH45FwPmfNW7vefPNNle23+CXmNj1/vnbfvn0qly5d2qt924Xa3JqfsRMmTLDarVq1cvvaXbt2qWz/3Wt+VocDPo/DU6ids9545JFHVF6zZk2SHHft2rUq16pVK0mOawrnuU3uuMUPAAAAAAAAQY8FKgAAAAAAADiKBSoAAAAAAAA4KmXCm4SuKlWqqLxkyRKrbdacMh/Bu3LlSpXtj0wuXLiw6gvH+/XDmVm7yF3NiowZMwZ4NMmbWTeqWLFiKrdo0cJqly1bVvWZ9arq1Kmj8t69e/0xRHjIrFMwevRolWNiYuJ97R9//KHyxx9/rLJZN6pbt25W2/w7U7BgQZXbt2/vdlzJgXmudO7cWeXBgwervGPHjoCMw6w59cEHH6j8008/We1grUEVzsyaU+XKlXNoJJ4zv8slJ3nz5lXZXd2p48ePq9ykSROV+R7rHfNn37JlS5W//fZblS9evGi1//77b9V3zz33uD3Wo48+arXNf38k5KGHHrLa5neoEydOqNywYUOV7d+hbt686dVxkXiJqTll1pEy61m54822CLyoqCiVO3bsqHLv3r1VLlSokNXetGmT6qtevbqfRxc4XEEFAAAAAAAAR7FABQAAAAAAAEexQAUAAAAAAABHhVUNKrNe0Pz581XOkiWL1R47dqzq+/LLL1U27991d99mtmzZVJ47d67KXbt2tdrnzp2Ldz8IDLP+SsWKFVVOly6d1b5165bqM2vjwL/MOjNm7tChg9UeNWqU6suZM6fKZs2HGjVqWO0DBw4kapxIWHR0tMqXL19W2Ty37PM5bdo01XfkyBG3x/ruu++stvlZnSFDBpWLFy/udl/JwYcffqiy+TOqUKGCyoGqQZXQ34kyZcpY7eeff171mb+jT5486efRJT/33Xef2+xORESEz8f19rXutt+4caPP4wg1Zv2hRYsWxbvtjRs3VDa/8x49etR/A0uGqlWrpnJCtQ3t3/2vXbum+vLkyeO/gblx+/Ztlc3vUFu3blW5fv36VnvVqlWBGxhEJHE1p8waoOb3InfHouZU8LHXjuvfv7/qs5+XInfWD9y/f7/Vrly5suqrWbOmyuvWrUvUOAOJK6gAAAAAAADgKBaoAAAAAAAA4CgWqAAAAAAAAOCosKpBZdacyps3r8rbt2+32sOHD1d9adOmVdm8f9deI2XJkiWqr2nTpipPmjRJZXvtDLM+FQLvlVdeUblq1arxbjthwgSVu3fvHpAxwTMff/xxvH1m3SKzFlyBAgWsNjWoAm/58uVuc5UqVVT+4YcffD7Wzp07rfauXbtUn7t6gclJ+vTprXaKFCni7RO583Nu+vTpARmTWWPOrOPywgsvWO333ntP9cXFxak8depUP48OLpfLbb+9FlRC25piY2Ot9q+//qr6fv/9d5ULFSqksv2zfM+ePaqvZ8+eXo0jlL344osqFylSJN5t3377bZXfeeedgIzJZK/1KiJSsGBBlXfv3q2yWSsrVJh/hxNi/lySiv3ne/78edV3zz33uH2t/fO6du3aqs/+Oxj+4W0tKHvdqYRqTiX2WPCvmJgYlTt37qzymDFjrLb5/e2TTz5R2fz+Zv+utG3bNtWXOnVq7wfrEK6gAgAAAAAAgKNYoAIAAAAAAICjQvoWvwYNGqhsXoJqateundU2L3U1DRgwwONxmLdDdOjQQWX7pdXc4hd45i2X5qWs5qN27Y+0Nx/HjuBh3u5n3uJnst8WAucl5pY+U/bs2a22eWsn/lGsWDGrnSZNGrfb7tixI9DDuSvzNhH7Z7F5S0zJkiWTZEzJSXR0dMD2bZZCePnll632b7/95tW+8uXLZ7VPnjyp+sxbP8NJZGSkyrly5fL4tYG8rb1MmTIqV6pUyWqbtyGWL19e5X79+qn85ptv+ndwSeTs2bMqb968WeUHHnggIMdduXKlypcuXVJ53759Ki9btsxqHz58WPV99NFHKpvfnTNnzmy17eeviEjHjh09GzACxpvb+rilz1mlS5dW2X4Ln4hIvXr1VD537pzVNm//W7x4sdtj3XfffVbbvF0+lHAFFQAAAAAAABzFAhUAAAAAAAAcxQIVAAAAAAAAHBVyNagyZMhgtSdNmqT6zPv13333XZX3798fkDHdunVL5Zs3b6psPk4S/mevsdKiRQvVZ9acMh+Pffr0aat97NixAIwOgbB3716VzRo1Tz75pNWeMmVKkowJgZE1a1aVe/ToYbWLFy+exKMJDQMHDrTamTJlUn0nTpxQeezYsUkxpDusX79e5f79+1tt++96kTtroPTq1StwA0smxo8fH7B9m3U+va07ZZdcfy+bteNatWrldvuff/7Zaq9atcpv42jevLnKZt1V81x1p1u3birXrVvXapvn+O7duz3eb1Iz6znVrFlT5XTp0gXkuBcuXFA5MTXYzMfTP/bYYypHRUXFe1wkPW9qTpnWrFnj8bZDhgzx+Ti4u6FDh6ps1pwyP+vatGljtc1/6yTk6aeftto3btxQfWbtvGDGFVQAAAAAAABwFAtUAAAAAAAAcBQLVAAAAAAAAHBUyNWgGjNmjNXOkyeP6tu5c6fKgwYNUjkx92ojuFWoUMFqt23b1qvX2v9OXbt2zW9jgn898MADKhcsWNDt9hMmTAjgaCAiUrRoUavdp08f1dewYUOVR48erbK9dp9ZFy5Lliwqd+3a1W2/3aFDh1QeNmxYvNuGM3vtPfPna9b0yps3r8o7duwI3MBszNpC9tpiCxcudPvaWbNmqfzcc89ZbT7H42evfVG+fHmf9xMREeG2/5NPPvF53/hHhw4dvNreXpf15MmTPh/X25pT9jqeCX12mN/b69SpY7WXL1+u+syaSHv27HG7byf9/fffbnMwOnPmjMpmHZyKFSta7fr166u+jBkzqnzx4kU/jy75MWtMPfLII26zO95saxo8eLDPr8X/tGvXzmo3atRI9a1cuVJle80pkTtrOLqTPn16lVu3bm21zdqPW7du9Xi/TuMKKgAAAAAAADiKBSoAAAAAAAA4igUqAAAAAAAAOCroa1ClS5dO5ZYtW8a7ba9evVS+dOlSQMbkrWXLljk9hLATFRWl8quvvurxa0+cOKHygQMH/DIm+F9MTIzVHjJkiOozPxsOHjyoclLV0UnOZs+ebbXNGmGmt99+W2V7DRuzRpKvYxARGThwoMpHjx71ed+hzF19IfMz0JQrVy6rnZhaNgkxa0XZ6yWYtRLsdQZF7qxlExkZ6efRhacRI0ZYbW/PO3+ds/BMvnz5HDnuCy+8oLJZc+rUqVMqt2rVymqvX7/e7b5z5Mih8t69e622eU7ba7iIiPTt29ftvuGdmzdvqvzjjz+qbK9BVaBAAdVXrVo1lc36YfCe+R03oTpSa9asibfPmxpUZu0r+Masizt06FCrvWnTJtXXokULlS9fvqyyva6UWYuwc+fOKpufz/bP0Vq1aiUw6uDFFVQAAAAAAABwFAtUAAAAAAAAcFTQ3+I3duxYle2XspmPnDUvT00qKVPqH6P9tiQRkc8//zwph5MsmI+ZNy8Nd2fkyJEqb9++3R9Dgg/My/3N2/ZGjx5ttc3HHJu39Nkfny4icvjwYT+MEJ5K6LHzphQp/vf/R27fvu3Vazds2GC1n3nmGa9em1y8//77Vtu8vdK8FH3x4sUq22+PPXTokOozPy/NWywTw/6o8hUrVqg+8xY/M9s/O2JjY/02pnBTsmTJJDlO7969Vbafs+atnXCe/Ta9Bx98UPWdPn063m1FEr6tz868PdC8zQxJJy4uTmXzkfR20dHRKpu3/CHxzFvtzGzetufNbXwm++2EgwcP9nk/+J8uXbqobL9F+/nnn1d92bNnV9m8rdq+vbnttm3bVL7vvvtUtt9+f+HChQRGHby4ggoAAAAAAACOYoEKAAAAAAAAjmKBCgAAAAAAAI4K+hpUbdq0ibevR48eKpuPaUwq+fPnV9m8f//SpUtJOZxkwbz/3d0j1e21bkRC+57cYGB//KmISNWqVVU2Hw1trwVVuXJl1WeeO1mzZlXZrDVmZz5i9/r16/Fui8B4+umnrfZjjz3m1Wtr1qxptc36YmnSpHH7Wh5xn7B169ZZbXttJxGRjBkzun2tvaaB+dn65JNPqmx/lHJCzDpliZlHdzXoZs6c6fN+w5291sW0adMCdpw6deqo/MUXX1jtZs2aqT5qUvlHqVKlPN7W/N3at29fq23Wgpw6darK3tScMpmPV8+WLVu82+7du9fn48B7Zp0pd/i+FXj23+Eiias5Zdazou6U/+XOnVtl+781O3furPrMz8Fz586p/NZbb921fbfjHD9+XGV7DdGdO3cmMOrgxRVUAAAAAAAAcBQLVAAAAAAAAHAUC1QAAAAAAABwVNDVoKpevbrKMTEx8W4bLPdWNmrUyOkhJDsDBw5U2V0tE7NWUWLqJyRXXbp0sdrdu3dXfcWKFfN4P/6sQfPss8+qbK9BI6Lv/zZrWXz99dcqb926VeU9e/b4PK7k5NChQ1b7gw8+8Oq19u3NuVu2bJnb19prZZi/I6hn8w/732nz5/nUU095vB+zht/t27d9HpM/9/XAAw+onDZtWp/3FWrKlCljtefOnav6Spcu7fa1lSpVstr2WhUi7ms5mszP8oQ8+uijVnv27Nmqz6xrhn989NFHKvfu3dvt9u3bt4/3tZGRkSqvXLlS5Xz58lntn3/+WfX1798/4cHGo0iRIiqbtQrt4/rll19U3/z5830+LrznTc0vb+qdwTf2Op2JZdavsmezPhV8c+TIEZVv3rxptWvUqKH6xowZo/KIESNUPn/+fLzHMes7mv+OWrVqVcKDDQFcQQUAAAAAAABHsUAFAAAAAAAAR7FABQAAAAAAAEcFXQ2qDBkyqGzeN797926rfenSpSQZk8kck3lP/d9//63y0aNHAz6mcGfec1utWjWPX2vWTzh9+rRfxpScfPjhh1Y7MXWjEvL555+rbK8VZc6bWXPGPA/t9TjKli2r+lq3bq3ylStXVLbX7+jXr19Cw0Yi7dq1y6vtS5YsabWLFy+u+rZv3+6PIYWV5557TuUJEyao3K1bt3hf68+6cea+zBpIM2fOtNqZMmVSfa+//rrKrVq1Utlezypr1qyq7+zZs16PNZjlzp3bapu1YBKan40bN1pte40MT15rnz9v/x7YX9u4cWPVZ9YetY8xOTt27JjK06ZNU7lTp04q2//ev/DCC6pv+vTpKttrTpneeustlRNT12/evHkqV6hQId5tzXqu5ndpIJwNHjxYZbNulD/Za/PWqlVL9VGTyjcDBgxQ2awDaHfy5EmP92v+W8f8bDeNHj3a430HM66gAgAAAAAAgKNYoAIAAAAAAICjIlweXqft7SOFfdWgQQOVly5dGu+299xzj8pJdetW/fr1VV6+fLnK69evVzmQl2na+XrrRVLNbWLMmjVL5bZt28a77eLFi1UOh0dY+zK3/pzXy5cvW+00adJ49Vr7eWleOvzGG2+o7M1jjhNSsGBBq21eetuiRQuVzVuL7bcLmY+OHTp0qMrmbYn2n1VCwvmcTYh9frp06aL6evXq5fa1GzZssNpJ9fnqLafP2XBj//siIvLrr7+qfObMGatdu3Zt1bdnzx6/jSMYztl69epZbfP7hzfjSMxtev58rXlLQp8+fbzat78Ew9y68/TTT6tsfi+yi4uLU/nChQsqZ8uWTWX7e7f//RK587Hl5u9P+yPUW7Zs6fY4ZomMF1980WqPHz9e9dl/DycWn8cJK1y4sMpbtmyx2uZ3JPOW0WeffTZwA3Mj2M9Zb3j7Xuzfp83b9MzvRfZb+hISLD+bcJrbxMiZM6fKx48fV3nr1q0q22+Zv3HjRuAGlgiezC1XUAEAAAAAAMBRLFABAAAAAADAUSxQAQAAAAAAwFEpnR6AyXysrvkY5FSpUiXlcCy5cuWy2uajI81H4ZqPw4b3atasqXLGjBndbm+vr/Ddd98FYkjJmn0+zMeaJ8ReoykpHxt9+PBhq925c2fVN2bMGJXNR1/b68yVLl1a9dlrZojcWe+mXbt2Xo81KUVFRan82muvWe1Nmzapvs2bN6vszzp/I0eOtNpm3RLTpUuXVB43bpzfxoHQYP4OmDlzpsr2Ghxt2rRRfWYNOgSXunXrOj2EkDBnzhyVzceP22s4mbWesmbNqrK7GiALFy5U2fwebtYjMo/ljvn70z5mf9acgve8+fmXKVMmgCNJPrypoTlkyBCVBw8eHO+2Zr1XM7s7rtlnvhZJy15TSuTOGlutW7dWOVjrTnmLK6gAAAAAAADgKBaoAAAAAAAA4CgWqAAAAAAAAOCooKtBtXv3bpWvX7+usr0G1b333qv6ElMfJUUKvVbXvn17le01T9KlS6f6WrRoofKGDRt8Hgf+8fLLL6vcqFEjt9svW7bMalOfxv+2bt1613ao2rt3r9s8e/bspBxOksqSJYvK/fv3j3fbBQsWqGze625XsGBBlStVqqTyK6+8Em+/WQ8lNjZW5Y4dO6q8ePHieMeB8LRjxw6Vzb9vefPmtdqXL19OiiE55ptvvrHaZk3Mrl27un2tWb/CV4nZj/na3377LbHDSRbMGkFz585VuVChQlbb/M5k1lJ0x/yOmxjly5dX2fyOT90pwDPe1IIy60h5U+uKGlTOy507t9U2f6ebtWHD9fcnV1ABAAAAAADAUSxQAQAAAAAAwFEsUAEAAAAAAMBRES6z+Ed8G/qpboG33n33XZW7detmtc+fP6/6Xn/9dZU/++wzla9evWq1GzZsqPo6dOigcuPGjVW+ePGi1X7uuefcHscpHk7lHZyaW1PTpk2t9sKFC1Wf+d4OHz6s8mOPPWa1Dx065P/BOcyXuQ2WeUX8nDhnc+XKpfKxY8fi3fbatWsqnzhxIt5t06dPr3KOHDncjsP+HuyfzSIibdu2VXnJkiVu9xWMOGeT1tGjR632kSNHVN/DDz/st+ME2+9Ze11OEZG6deuqbNYpypAhg9X29r3Y34M/Xzt69GiV+/Tp49W+/SXY5jYxUqbUJWbvuecelZ9//vl4X5s/f36VzZqs06ZNU3nNmjVW2/w+fPPmTZV9/RknFp/HCcuZM6fK9nphmTNnVn1m3c6yZcsGbmBuhNM569S5YQqWn004za233njjDavdo0cP1Wf+jv/pp5+SYkh+5cnccgUVAAAAAAAAHMUCFQAAAAAAABwV9Lf4mbZs2WK177//fr/t1/wx2G/pExFp0qSJ1d64caPfjutPoX45ZPfu3a322LFjVZ/53sxbfjp37my1z549G4DROYvL08OTE+dsZGSkyjNnzrTarVu39nm/3mrZsqXV3rBhg+o7ffp0ko0jUDhnA6tmzZoqL1682Gqbt6K++uqrKifmltFQ+z1buHBhlbt06WK127Rpo/ry5Mnjdl/+usXvxo0bqq927doqf//9917t219CbW7hOT6PExYdHa2y/XZNsyzKlStXVO7atavKs2bN8vPo7i6czlmnbvFbu3atyrVq1XJkHKZwmtuEVKxYUWX7d+LNmzerPn+WLHAKt/gBAAAAAAAg6LFABQAAAAAAAEexQAUAAAAAAABHhVwNquzZs1vtl156SfUVLVrU7Wvtj87dtWuX6luwYIHKK1eu9HWIjgn1+3Xtj5p++eWXVZ95v3uNGjVU3r59e8DGFQyonxCeguGczZUrl9X+5ptvVF/JkiU93s/OnTtVXrp0qcpmXbnz5897vO9QxDmbtHr16mW1R40apfqOHTumcqFChXw+TjCcs/6SPn16lR955BGVn3zySZUfeughq+3tz9BeF2zGjBmqb8CAAV7tK1DCaW6h8XnsPft5OXjwYLfbmnXjkqpOTjifs2vWrFHZ/Hx2x6wrtW7dunj7zW2DRTjPralMmTIq//jjj1Z72bJlqs9evzVUUYMKAAAAAAAAQY8FKgAAAAAAADiKBSoAAAAAAAA4KqXTA/DW6dOnrXaw1C2Af3z11VdW26xBZdazCfeaU0BSOXnypNUuW7asgyMBfPfpp59a7VSpUqm+fv36qTxx4kSVn3/++cANLIhdvnxZZfvv4LtlAMlHqVKlPN52//79KqdI8b/rH27fvu23MSUntWrVcnoISCJPP/20ytHR0Vb77bffTurhBAWuoAIAAAAAAICjWKACAAAAAACAo1igAgAAAAAAgKNCrgYVwte6deusdmRkpIMjAQCEkj/++MNqz5gxQ/Vdu3ZN5fXr1yfFkAAgZJ05c8bjbVu2bKnyqFGjrPahQ4f8NiYgHL366qsq37hxw2rHxsYm9XCCAldQAQAAAAAAwFEsUAEAAAAAAMBR3OIHAADCxsmTJ1UeN26cMwMBgBC1a9euePvi4uJUHjZsmMrc1gd47rffflN50aJFVnvv3r1JPZygwBVUAAAAAAAAcBQLVAAAAAAAAHAUC1QAAAAAAABwVITL5XJ5tGFERKDHgkTycCrvwNwGP1/mlnkNfpyz4YtzNjxxzoYv5jZ88XkcnjhnwxdzG748mVuuoAIAAAAAAICjWKACAAAAAACAo1igAgAAAAAAgKM8rkEFAAAAAAAABAJXUAEAAAAAAMBRLFABAAAAAADAUSxQAQAAAAAAwFEsUAEAAAAAAMBRLFABAAAAAADAUSk93TAiIiKQ44Af+PpARuY2+Pkyt8xr8OOcDV+cs+GJczZ8Mbfhi8/j8MQ5G76Y2/DlydxyBRUAAAAAAAAcxQIVAAAAAAAAHMUCFQAAAAAAABzFAhUAAAAAAAAcxQIVAAAAAAAAHMUCFQAAAAAAABzFAhUAAAAAAAAcxQIVAAAAAAAAHMUCFQAAAAAAABzFAhUAAAAAAAAcldLpAYSKNGnSWO2DBw+qvty5c6v8/fffq1y9evXADQwAQkDatGlV7tOnj9Vu37696itQoIDKLpdL5StXrljtbt26qb5Zs2apfPv2be8HCwC4q+bNm6u8cOFCq92qVSvV99lnnyXJmAAA4YMrqAAAAAAAAOAoFqgAAAAAAADgKBaoAAAAAAAA4ChqUMUjffr0Kk+fPt1q33PPParPrHFCzZPQZdbJWbNmjcqpUqVS+f777w/4mEJN6tSpVX7qqaesds2aNVVfVFSUypUqVVLZXu/tpZdeUn2HDh1K1DjhX5GRkSoPGDBA5VdeeUVl+7l28eJF1XfixAmVU6TQ/y/F/hls/2wWEYmOjlZ54sSJ7oYNAHAjJiZG5SlTpqhs/85r1gsEAMBbXEEFAAAAAAAAR7FABQAAAAAAAEexQAUAAAAAAABHUYMqHk888YTKzZo1i3fbHTt2qDx27NhADAlJ4Msvv1S5YsWKKq9atSophxMSSpQoofL8+fNVLlOmTLyv/fPPP1U2a3w99thjVnvmzJmq78UXX1R5y5YtCQ8WAdO2bVuVBw4c6Hb7PXv2WO2OHTuqvs2bN6ts/r147rnnrPb777+v+vr27avy7NmzVb5y5YrbcQHhJF++fCofPnzYapu13RKqn3nhwgWrvXPnTtX3yCOPqLxo0SKVmzZtetf9iIgMGzZM5XfeecftOBBYZg3WefPmqZwpUyaVf/31V6v9008/BWxcAO7OrOeaI0cOlStUqKDyF198YbXNz3177VcRkaFDh6psr81r1gsF/IUrqAAAAAAAAOAoFqgAAAAAAADgqAiXh8+EjYiI8PkguXPnVtl+yJMnT/q8X3+yX+4oIlKnTh2Vzcfs2vXu3Vtlpy5P9/XxvomZ23Dw7LPPWu333ntP9Zk/m8aNG6u8evXqwA3Mxpe5DdS8Fi5cWOV169apbN6+8fbbb1vt48ePq74ff/xR5TRp0qg8fvx4q928eXPVd+bMGZXr1aun8vbt2yXYhdM5mzdvXpWXL1+usvn34pVXXrHa3t4WYr/FxH6roIhIrly5VDZ//5i3lQZKMJ2zgWTeWlC6dGmr3blzZ9W3ZMkSlc2/I+5kzpxZ5YsXL6qc0O1p/hJq56x5i9/vv/9utc0x+fre/L2vN954w20OlFCbW3+Kjo622qNGjVJ95u305vtt37691TZvqQ4WyeXzOLlJzues/Xev+e9OexmEu7G/f29/hteuXbPajz76qOr74YcfvNqXO8lpbvPnz69yt27drLa5/tCqVSuVzVuyp06darX37t2r+szbtc+dO+f9YP3Ak7nlCioAAAAAAAA4igUqAAAAAAAAOIoFKgAAAAAAADgqZVIcJBgfQ2nWnHr88cdVtj+KWUTk66+/ttr2mkUiIr169VJ54sSJKl+9etXHUSIQqlSporL93u3UqVOrPnv9JJGkqzkVzMxHk5t/v2vUqKHy+fPnPd53bGysyvZ7rc37rt9//32VV61apXLDhg2tNo++Drw//vhD5YcffljlK1euqHzz5k2fj2WvZ/XXX3+pPrMGFRLHrDFl1lx86qmnVLY/ct6sXWP+XnXHXn9ORKRly5YqFy1aVGWzxhn+cf36dZXtdT/N+mzBIkOGDE4PIdmxn9fmeWvWC+nbt6/KixcvDtzAIHny5LHaDz30kOqrWrWqx/vJmDGjyrVr11Z50aJFKm/atMlqL126VPWZ39UQeAMHDlTZ/u/W8uXLJ9k47LVia9Wqpfr8WYMqnKRMqZdb2rVrp/KYMWNUttdZNW3YsEFl87vPE088YbW7du2q+szfrW+++Wa8x3EaV1ABAAAAAADAUSxQAQAAAAAAwFEsUAEAAAAAAMBRES7z5vL4NoyICPRYAq558+ZWe+bMmarv9OnTKpv3ZdrrSh05ckT15c2bV+UePXqobNbKCRQPp/IO4TC33li4cKHKZv0xO7PWwujRowMypoT4MrdJNa/PPfecypMmTUqS4z744IMqL1u2TGV7rSyzTs7y5csDNzAvcM76pnLlylbbvB/fvNffrLPz559/Bm5gNsF8znrj9ddfV9msuWjW6Rs+fLjPx0qbNq3V3rdvn9vjzps3z+fjJEaon7P232Evv/yy6jPfm1m/yv7aXbt2BWB0/zh48KDKO3bsCNix7EJ9bt2xn1siIkOGDFG5W7du8b72pZdeUjmpfsf7UzB/Hpu1T816e++9957VdlebJiHm+/HmZzJixAiVze/CFy9e9HlciRHq52z69Omt9oABA1RfsWLFVLbXVRXR33W8/TmsX7/eag8dOlT13XvvvSp/+OGHKtt/dubvCLNG2tatW70al12oz63d5MmTVe7cubPb7e01OL/88kvVt3btWpXNOciePbvVNj/3V6xYobL579/9+/e7HZe/eDK3XEEFAAAAAAAAR7FABQAAAAAAAEexQAUAAAAAAABHhXUNKnvNKRGR6dOnW+106dKpvt69e6v8zjvvxLvfhGpQff/99yo//PDDCQ/WD8Lpfl1/MusnuJtbs45RkyZNAjImbwVz/YRgUaVKFZWXLl1qtc1z0qzx8PfffwduYG5wzvpmzpw5VrtNmzaqz7yHvkKFCipfu3YtcAOzCeVz1l4X48CBA6rvs88+U9n8fE0Me12Tnj17qr6SJUuqnFS1Ekyhds6aNWtWrlxptc1zw3xv/fv3V/mtt97y7+CCTKjNrTemTp2q8jPPPKOy/T2Y9VHs351DVTB/HteqVUvlVatW+W3fe/bsuWtbRKRVq1Y+79esS+TP3wPeCPVzNl++fFb78OHDXr3W/h7Mn8PGjRtVNutbrVu3zmpHRUW5fe3999/v8XGrV6+u8g8//HDXsXsi1Oe2fPnyVvu7775TfZGRkSqbv1vtdbAT8521aNGiKs+ePVvlK1euqFynTh2fj+UNalABAAAAAAAg6LFABQAAAAAAAEelTHiT0NG4cWOVZ86cqXJMTIzVfuGFF1Tf3LlzPT6O+dhM8xY/OCtVqlQq161bV2V3lxYm9OhPBC/zUuIff/zRajdq1Ej1PfbYYyp/8cUXARsXEq9ixYoq16tXL95t//vf/6qcVLf0hRP7Jf85c+ZUfQsWLAjYcUuXLm21T58+rfqceox5qMmdO7fK9ludRUTKlStntVOk0P+Pct++fSrPmzfPz6NDUnn22WdVbteundvt7bdszZgxIxBDQjzM26Li4uJUtj9yfsmSJarPLF9gsu/L3G9C33ftt5xlzZpV9ZmfM+Znye3bt93uG/84e/as1TbnsmrVqh7vZ+3atSqbt2/aj2N6/fXXVbbfmibi/t9M5q32Bw8edDPK5OWVV16x2mnTplV927ZtU3nw4MF+O26zZs2strkOkjKlXvZJ6PeCk7iCCgAAAAAAAI5igQoAAAAAAACOYoEKAAAAAAAAjgrpGlTmPdDmPZzmvZaTJk26a9tb69evV7lp06Yqm/diI2mZj85u2LCh2+3tj1++cOFCIIYEB9gf02rWnCpRokRSDwdeMOsJmjVRsmTJYrW/+uor1bdly5aAjSs5MutPxMbGBuxY9vO0SZMmqu/PP/8M2HHDyZQpU1QuW7asyvb5NGtONWjQQOUjR474eXQIJPsj683Hlpvfhz/44AOV7XVofH28O3wzevRole31M0VEVq5cGZDjXr161W3/E088YbXXrVun+po3b65ytmzZVD516lQiR5c82L+nFi9e3Of92OuUibivOSUi0rFjR6v9/PPPJ9lxw1n27NlVttcQu3nzpurr2bOn345bv359lefMmWO1o6OjVZ953i5evNhv4/A3VlIAAAAAAADgKBaoAAAAAAAA4CgWqAAAAAAAAOCokKtBlTVrVqtt1hrJkSOHygsXLlT5hRde8MsYGjVqpLJ5v/7t27f9chz4xl5L4W42bdqkco8ePaz2jRs3AjEkOODnn392egiwMT+f7XWkBg4cqPrMe+ozZ86s8pdffmm1n3rqKdV37dq1RI0T7tWuXVvlxNT8qlOnjsr2Og2HDx/2eb/JifkzrFatmsevXbp0qcr8zENLunTpVF6zZo3Vzpgxo+o7f/68yjNnzlT50qVLfhlT4cKFVc6UKZPKFy9eVPnXX3/1y3FD2d9//61yoGpOecus5wr/s/+7NJC13ypWrKiyvVZhYo578uRJn18bbsw6f1FRUVbb/LelWdPNG/bvziIib7zxhsr2ulNDhgxRfebv/GDGFVQAAAAAAABwFAtUAAAAAAAAcBQLVAAAAAAAAHBU0NegstecEhH5/PPPrXb27NlV38aNG1Xu0qVLQMZUrFixgOwXvps2bZrVbtiwoeoz76+212kQoWYN4A+pU6dW+cMPP1S5WbNmKpuf7d744YcfrPbVq1d93g+c9dBDD6lsr0+ze/fupB5OSIqJiVHZPA/dady4scr58uVT+b333lP56NGj8e7r2LFjHh8X/pEmTRqVCxUqFO+23bp1U3nz5s1+G0eVKlWstlnjxKwfeOHCBZUnT55stfv06eO3MSHxzM8Hu71796p85cqVQA8nLH333XdW++GHH/bqtREREVbb/D41adIklTt37hzva701Z84cq23/N3lyZ9bxO3XqlNXOli2b6jNr9f3222/x7tf++Soi8tZbb6lcqVIllc+cOWO1P/74Y9UXFxcX73GCDVdQAQAAAAAAwFEsUAEAAAAAAMBRQX+LX5s2bVR2dwnkgAEDVD579qzfxmF/DHrOnDndbjt16lS/HRd3Zz4ytWXLllY7RQq97vruu++qbD7SHs6xP4ZVRKR48eIq2x9pb15S/uOPP6p8+fJllb251QWJ16FDB5XNS8rdsV+SLCLy119/qVy6dGmV+/bta7Vnz56t+v744w+Pj4uEmbcCmOdsYuTPnz/efZuPUj537pzfjhtOlixZovK2bdtUfvDBB+N9rVmuoGjRoio/+eSTHo/D/D27Y8cOlWfOnOnxvuCZfv36qWw/V//++2/Vd/DgwYCNY/r06VY7U6ZMbrfNmDGjyk888YTV5hY/Z5mfudWrV49322+//VZlbvHzzeuvv261R44cqfpy586tcq5cuVS2/740SyqYzFInnvaJiJw8eVLloUOHut0+uTLPAfstfvfff7/qW7lypcrmrXj27885cuRQfeZ3MPO49jI3hw8fTmDUwYsrqAAAAAAAAOAoFqgAAAAAAADgKBaoAAAAAAAA4Kigq0H19NNPq2w+5tjOfJztunXr/DaOevXqqTxo0KB4t7U/clNEZMaMGX4bB/6RPn16lV977TWV7Y/avnHjhuoz7/WFc8qVK6fyBx98oLL52HlvDBkyROWEasXBvzZu3KiyWXPGfNz46dOnrfbYsWNVn/nIXbPeWNmyZa22Wftq+PDhHo4Y8bl27ZrVNmt6derUSeWJEyeqbJ/XhJg1j1atWmW1qTnlG7NujFkTyF7zxKwvZtakatSokcfH7dmzp8q3b99W2V6DY/z48arvwIEDHh8nOTNrtpnfl+21ZOznkojI5s2bfT6u/fuVyJ2/awsVKmS1zVpXZk3WUaNGqbxz506fx4XEyZ49u8oLFy5UOTo62mqfP39e9SVU8wie2bBhg9U2ayzXrFlTZXN+/FkP0s78fB49erTKR44cCchxw83zzz9vtb/++mvVV6JECZUHDx6ssr0Oq/08FLmzZpj5fTkxn/XBhCuoAAAAAAAA4CgWqAAAAAAAAOAoFqgAAAAAAADgqKCrQWWvjyBy572WZ8+etdrPPvus345bpkwZlSdPnux2HHZmLSz4n3kP9OOPPx7vtv3791eZ+XHW/PnzrXaTJk1UX2xsrMrvvvuuymvXrrXapUqVUn32+7tF3NeJM48za9as+Acc5syaNLlz57bae/fu9Xm/5mufeeYZn/dlOn78uMr2GlStWrVSfdSgSrwrV65Y7TFjxqg+s17YihUrVH7ppZestlljqmvXriqXL19e5Y8//tjrscK9CxcuqNyrV694tzVrDa1fv15lex2p7du3q74UKfT/7zTrDb744otW+z//+Y/q+/TTT1W2/x3C/xQpUkTlLFmyxLvtTz/95PNxMmTIoPJnn32mct26dVWOi4uz2v369VN9TZs2dXusRYsW+TJE+EH79u1VrlatWrzb2r/HiVA3LimYdZXN77HmdzlfvfDCCyqbdZTNur7wzNGjR622+ZmZN29et6+9evWq1d61a5fbbYcNG+bD6IIfV1ABAAAAAADAUSxQAQAAAAAAwFEsUAEAAAAAAMBRjteg+uKLL1QuVqyYymbtp4kTJ1rtU6dO+Xzce++9V+URI0aonCdPnnhf27x5c5W/+uorn8eBu8ucObPKVapUcbu9vf6NWTMFSeuhhx5SuV69elb73Llzqq9+/foq7969O979Ll68WOXx48erfPjwYZXtdTRu376t+goWLKiyWeMonKRJk0Zlsyabve5PYmpQ+VPKlPpXU0L36yNw3n//fZULFSqksr22kIjIqlWrrPby5ctV37hx41Tu2LGjH0YIf7l27ZrKZu0Re12jJ5980u2+zDqRo0aNstpm7SSzBoq9BpqIri+YnOuhmO/91q1bKts/N816Yt4w6wc++uijbre31yMyfy/nyJFDZbPm1Ny5c30YIXxhft6atV3Nf2/ZvxeZ9UHhH/Zz1qyrPGDAAJVz5sypckREhMfHOXHihMrNmjWz2lu2bPF4P/CN+fM3s2nw4MFW2zwv582bp/Lnn3+euMEFKa6gAgAAAAAAgKNYoAIAAAAAAICjHL/Fz/7Y4rvZunWryol5nKL9tr6VK1eqPvMWEvOSuqFDh1rtZcuW+TwGxC9t2rRW27z1s3Tp0m5fa95OAOe8/vrrKttvtTMfY5yYW8rMR9hHR0erbL883bxld+TIkSqbj8I+f/68z+MKNhkzZlS5atWqKh85csRqT58+XfXdvHkzcANz46233lK5TJky8W779ddfB3o4yZp5G1H37t3dZm+Yt/QndNsYAqtBgwYqm/NjvxXv5MmTbvdlL8cgom8D69mzp9vjvvrqqyrbP4eGDx+u+q5fv+52HOFk48aNKpu/a+2/13r06KH6zPn6448/VLZ/PzZLXpjfh02lSpWKt8/8XWveKmZ+vsC/unbtarXN884sfWCyn5f79+/378AgIvocNs9nk7vz0OzbtGmTyvZSGyJ33s6N4GL/XWveymn/zi6S8OdzqOIKKgAAAAAAADiKBSoAAAAAAAA4igUqAAAAAAAAOMqRGlQtW7a02mbtJ/NeS/sjq0X0Y3Zz586t+uy1bkREBg4cqHLr1q3jHZN5L/Ybb7yhsr0GFfzDXnNKRNedeeihh9y+1qxvcfjwYb+NC4ljPobe7vfff/d5v+ZnxTfffON2e3s9m4cfflj12R95LiIyduxYlc3HMYey2NhYlffs2aNymzZtrLZZq8t8DPWOHTtUPnr0qM/jsn9+jxs3TvUlVIvI/nvArJuF0HHw4EGVzcdoI7Bq1qyp8ty5c1VOly6dyk2aNLHa3j7aes2aNVZ78+bNqs+sOfnII4+o3L9/f6u9dOlS1ffjjz96NY5wMmXKFJXtj6U3v1+99957ARvHoUOHrPZTTz2l+rZt26YyNacCq0CBAioPGjTIamfLlk31uau1K5K4GqH4R758+VT+v//7P5X79Oljtf1ZS8heg1WEmlOh7PTp0ypPmDDBoZEkLa6gAgAAAAAAgKNYoAIAAAAAAICjWKACAAAAAACAoxypQWWX0D23ly9fVrlDhw5W+4MPPlB9adKkcbtve966davqGzJkiMpmnQP4n71egohIly5drLY5d2YdHbNmEJyTI0cOlc1acKtXr7ba169f92rfpUqVstpLlixRfRkzZlT52WefVfmnn366a1tEpGTJkiqbNac+/fRTq71ixQovRhx8zM/Qhg0bqvzJJ59YbbNWl5nN8/DYsWNW25wfU926dVW21yrLlCmT29fevHlT5ccff9xq79u3z+1rEbxq1arl9BCStQoVKqicPn16t9ubtaF8ZX4m1alTR+W1a9eqbP8cMseQnGtQnTt3TuWqVata7X79+qm+tm3bqmx+x7LXfzX79u/fr3LXrl1Vts/B1atXExo2/OiBBx5Qef369SpHRUVZbbMukfl3ZM6cOSqbtXnhPbPm1KuvvurQSBDMnnnmGZXt9eLMGq3JpeYyV1ABAAAAAADAUSxQAQAAAAAAwFEsUAEAAAAAAMBRjtegSsjQoUN9fu3Zs2dVnjhxotUeNmyY6rtx44bPx4FnzBpBL730UrzbmrVuzBpBhw4d8t/AkCinTp1S+dKlSyqnTPm/jxl7nQsRkSpVqqhcunRplcePHx/vazt16qTyxx9/7OGIRRYsWKCy+fcrLi7O432FGnvdKBGRJ554wmq/++67qs+sW5IuXTqV7bW8zLpeibFt2zaV7bUHRUR2797tt2MheCxfvtzpISQrkydPVrlnz54q58qVS+WcOXNa7SNHjqg+s/bNl19+6fO4ypUrp7K9JlLjxo1V35tvvunzccKN/XOxXbt2qs/8/nXy5EmV7bUjzd+15u/DNWvWJGqc8F3BggVVnjt3rsr2mlOmXr16qTxv3jy/jQv/MOfH/O5inltm9kZiXgtnFS5cWOXRo0erbP/MTa7fi7iCCgAAAAAAAI5igQoAAAAAAACOinCZz5ONb0M/XkqYJk0aq20+cnPAgAEquxueeXm6eRncDz/8oLJ5K1K48XAq75BUl4mat+2ULVs23nHs2LFD9ZmPw05ufJlbpy7/3bx5s8r2uTMfc3zPPfeonCpVKpUPHDhgtfv27av6Fi1a5PMYs2bNqrL5mPP58+f7vG9vBNs5GxkZqXJMTIzKlSpVUrlFixZW2377j8idj4M3/17Yb+OdMmWK6lu5cqXKoXjLZSids06ZOnWqyvnz57fajz76aFIPxyPBds76k3mr89ixY1VOnz59vK8135+vP6eE9rVp0ybV9/DDD/t8HFM4z63J/H5s/2w3b898+umnVb5y5UrgBhYg4fJ5bH43LlOmjNvtn3zySav9xRdfBGJIjgq2c9YsW7FhwwaPx+Hte7G/dv/+/arvtddeUzkx35edEmxz60/mv2nN78f2vzc1atRIkjElJU/mliuoAAAAAAAA4CgWqAAAAAAAAOAoFqgAAAAAAADgKEdqUCEwgv1+3YRqUF27ds1qN2jQQPUldB93uAul+gn2OjIiIv/973+ttvmoa7Nu3OrVq1Veu3at1TbrV4WDYD9n4btQOmedYtYea9q0qdXOnj17Ug/HI8npnDXrlthr9dlriYoEtgaV/bN/4MCBqm/GjBk+H8eUnOY2uQnlz+PSpUtbbbMGW9q0aVUeNmyYym+88YbVvnXrVgBG56xgO2fNOn32n7/Ind+B7Z+j5nv5448/VDZrp9pf27p1a9X3+eefezji4BVsc+tP5nefTp06qfzYY49Z7W+++SZJxpSUqEEFAAAAAACAoMcCFQAAAAAAABzFAhUAAAAAAAAcRQ2qMBLO9+smd6FcPwHx45wNX5yzCTPrMDzzzDNWu3Llyqpv69atSTGkBCXnc7ZRo0ZWu1ixYm63LVeunMrt27ePd9vFixer/N1336ls/3sSGxub4Dh9lZznNtyF8udx27Ztrfbs2bNVn1mbs0qVKiqfOHEicAMLAqF2zto/Q0VEvvzyS6u9ZMkS1bdjxw6Vo6KiVN64caPVXrlypeqz1/QNVaE2t94YPHiwymZtRft8njp1SvX16tVL5b/++su/g0sC1KACAAAAAABA0GOBCgAAAAAAAI5igQoAAAAAAACOSun0AAAAQPJj1tiIi4uz2levXk3q4SABS5cu9fm1nTp18uNIgOTjiSeeiLdv2rRpKod7zalQZ36GRkZGOjQSOGnq1KkqZ8mSReV27dpZ7VatWqm+UKw55QuuoAIAAAAAAICjWKACAAAAAACAoyJcHj7HMRQe25jchfMjOZO7UH5EMuLHORu+OGfDE+ds+GJuw1cofx4///zzVnv06NGqL3fu3Cpfvnw5ScYULDhnwxdzG748mVuuoAIAAAAAAICjWKACAAAAAACAo1igAgAAAAAAgKOoQRVGuF83fIVy/QTEj3M2fHHOhifO2fDF3IYvPo/DE+ds+GJuwxc1qAAAAAAAABD0WKACAAAAAACAo1igAgAAAAAAgKM8rkEFAAAAAAAABAJXUAEAAAAAAMBRLFABAAAAAADAUSxQAQAAAAAAwFEsUAEAAAAAAMBRLFABAAAAAADAUSxQAQAAAAAAwFEsUAEAAAAAAMBRLFABAAAAAADAUSxQAQAAAAAAwFH/DxWS3xzPZOWxAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 1200x300 with 20 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# visualize data by label\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "images_by_label = {i: [] for i in range(10)}\n",
    "\n",
    "for images, labels in train_loader:\n",
    "    for img, label in zip(images, labels):\n",
    "        if len(images_by_label[label.item()]) < 2:\n",
    "            images_by_label[label.item()].append(img.squeeze(0))\n",
    "        if all(len(images) == 2 for images in images_by_label.values()):\n",
    "            break\n",
    "    if all(len(images) == 2 for images in images_by_label.values()):\n",
    "        break\n",
    "\n",
    "# Plot the images\n",
    "fig, axes = plt.subplots(2, 10, figsize=(12, 3))\n",
    "fig.suptitle(\"MNIST Dataset\", fontsize=13, y=1.05)\n",
    "\n",
    "for label, imgs in images_by_label.items():\n",
    "    for i, img in enumerate(imgs):\n",
    "        ax = axes[i, label]\n",
    "        ax.imshow(img.numpy(), cmap=\"gray\")\n",
    "        ax.axis(\"off\")\n",
    "        if i == 0:\n",
    "            ax.set_title(f\"Label {label}\", fontsize=10)\n",
    "\n",
    "plt.tight_layout(pad=1.0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "820hB9zhPner"
   },
   "source": [
    "## 3.1 U-Net: Architecture for Image Data\n",
    "In the toy dataset, we choose MLP as the architecture of the denoising diffusion models, and use concatenation as the way to incorporate the time embedding. Although this works fine for simple synthetic distributions, it no longer suffices for complex high-dimensional distributions like images. In this problem, we will introduce the U-Net architecture specifically designed for images.\n",
    "\n",
    "Specifically, we apply [classifier-free guidance](https://arxiv.org/pdf/2207.12598) (CFG) for conditional generation of MNIST digits, conditioned on the digit label. CFG is a widely used method during diffusion model sampling to push samples towards more accurately aligning with the conditioning information (e.g. class, text caption).\n",
    "\n",
    "When applying CFG, the label embedding, together with the time embedding, is added to each hidden layer of the U-Net. A diagram of the U-Net we'll be using is shown below (we change BatchNorm to GroupNorm for better performance)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "-n6lTFLGFEoY"
   },
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import torch.nn as nn\n",
    "import math\n",
    "import torch\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "from torchvision.utils import save_image, make_grid\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch.optim as optim\n",
    "from matplotlib.pyplot import savefig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "id": "R91fDLcLFEhh"
   },
   "outputs": [],
   "source": [
    "class FourierEncoder(nn.Module):\n",
    "    \"\"\"\n",
    "    Based on https://github.com/lucidrains/denoising-diffusion-pytorch/blob/main/denoising_diffusion_pytorch/karras_unet.py#L183\n",
    "    \"\"\"\n",
    "    def __init__(self, dim: int):\n",
    "        super().__init__()\n",
    "        assert dim % 2 == 0\n",
    "        self.half_dim = dim // 2\n",
    "        self.weights = nn.Parameter(torch.randn(1, self.half_dim))\n",
    "\n",
    "    def forward(self, t: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Args:\n",
    "        - t: (bs, 1, 1, 1)\n",
    "        Returns:\n",
    "        - embeddings: (bs, dim)\n",
    "        \"\"\"\n",
    "        t = t.view(-1, 1) # (bs, 1)\n",
    "        freqs = t * self.weights * 2 * math.pi # (bs, half_dim)\n",
    "        sin_embed = torch.sin(freqs) # (bs, half_dim)\n",
    "        cos_embed = torch.cos(freqs) # (bs, half_dim)\n",
    "        return torch.cat([sin_embed, cos_embed], dim=-1) * math.sqrt(2) # (bs, dim)\n",
    "\n",
    "class ResidualLayer(nn.Module):\n",
    "    def __init__(self, channels: int, time_embed_dim: int, y_embed_dim: int):\n",
    "        super().__init__()\n",
    "        self.block1 = nn.Sequential(\n",
    "            nn.Conv2d(channels, channels, kernel_size=3, padding=1),\n",
    "            nn.GroupNorm(num_groups=8, num_channels=channels),\n",
    "            nn.SiLU(),\n",
    "        )\n",
    "        self.block2 = nn.Sequential(\n",
    "            nn.Conv2d(channels, channels, kernel_size=3, padding=1),\n",
    "            nn.GroupNorm(num_groups=8, num_channels=channels),\n",
    "            nn.SiLU(),\n",
    "        )\n",
    "        # Converts (bs, time_embed_dim) -> (bs, channels)\n",
    "        self.time_adapter = nn.Sequential(\n",
    "            nn.SiLU(),\n",
    "            nn.Linear(time_embed_dim, channels)\n",
    "        )\n",
    "        # Converts (bs, y_embed_dim) -> (bs, channels)\n",
    "        self.y_adapter = nn.Sequential(\n",
    "            nn.SiLU(),\n",
    "            nn.Linear(y_embed_dim, channels)\n",
    "        )\n",
    "\n",
    "    def forward(self, x: torch.Tensor, t_embed: torch.Tensor, y_embed: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Args:\n",
    "        - x: (bs, c, h, w)\n",
    "        - t_embed: (bs, t_embed_dim)\n",
    "        - y_embed: (bs, y_embed_dim)\n",
    "        \"\"\"\n",
    "        res = x.clone() # (bs, c, h, w)\n",
    "\n",
    "        # Initial conv block\n",
    "        x = self.block1(x) # (bs, c, h, w)\n",
    "\n",
    "        # Add time embedding\n",
    "        t_embed = self.time_adapter(t_embed).unsqueeze(-1).unsqueeze(-1) # (bs, c, 1, 1)\n",
    "        x = x + t_embed\n",
    "\n",
    "        # Add y embedding (conditional embedding)\n",
    "        y_embed = self.y_adapter(y_embed).unsqueeze(-1).unsqueeze(-1) # (bs, c, 1, 1)\n",
    "        x = x + y_embed\n",
    "\n",
    "        # Second conv block\n",
    "        x = self.block2(x) # (bs, c, h, w)\n",
    "\n",
    "        # Add back residual\n",
    "        x = x + res # (bs, c, h, w)\n",
    "\n",
    "        return x\n",
    "\n",
    "class Encoder(nn.Module):\n",
    "    def __init__(self, channels_in: int, channels_out: int, num_residual_layers: int, t_embed_dim: int, y_embed_dim: int):\n",
    "        super().__init__()\n",
    "        self.res_blocks = nn.ModuleList([\n",
    "            ResidualLayer(channels_in, t_embed_dim, y_embed_dim) for _ in range(num_residual_layers)\n",
    "        ])\n",
    "        self.downsample = nn.Conv2d(channels_in, channels_out, kernel_size=3, stride=2, padding=1)\n",
    "\n",
    "    def forward(self, x: torch.Tensor, t_embed: torch.Tensor, y_embed: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Args:\n",
    "        - x: (bs, c_in, h, w)\n",
    "        - t_embed: (bs, t_embed_dim)\n",
    "        - y_embed: (bs, y_embed_dim)\n",
    "        \"\"\"\n",
    "        # Pass through residual blocks: (bs, c_in, h, w) -> (bs, c_in, h, w)\n",
    "        for block in self.res_blocks:\n",
    "            x = block(x, t_embed, y_embed)\n",
    "\n",
    "        # Downsample: (bs, c_in, h, w) -> (bs, c_out, h // 2, w // 2)\n",
    "        x = self.downsample(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "class Midcoder(nn.Module):\n",
    "    def __init__(self, channels: int, num_residual_layers: int, t_embed_dim: int, y_embed_dim: int):\n",
    "        super().__init__()\n",
    "        self.res_blocks = nn.ModuleList([\n",
    "            ResidualLayer(channels, t_embed_dim, y_embed_dim) for _ in range(num_residual_layers)\n",
    "        ])\n",
    "\n",
    "    def forward(self, x: torch.Tensor, t_embed: torch.Tensor, y_embed: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Args:\n",
    "        - x: (bs, c, h, w)\n",
    "        - t_embed: (bs, t_embed_dim)\n",
    "        - y_embed: (bs, y_embed_dim)\n",
    "        \"\"\"\n",
    "        # Pass through residual blocks: (bs, c, h, w) -> (bs, c, h, w)\n",
    "        for block in self.res_blocks:\n",
    "            x = block(x, t_embed, y_embed)\n",
    "\n",
    "        return x\n",
    "\n",
    "class Decoder(nn.Module):\n",
    "    def __init__(self, channels_in: int, channels_out: int, num_residual_layers: int, t_embed_dim: int, y_embed_dim: int):\n",
    "        super().__init__()\n",
    "        self.upsample = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear'), nn.Conv2d(channels_in, channels_out, kernel_size=3, padding=1))\n",
    "        self.res_blocks = nn.ModuleList([\n",
    "            ResidualLayer(channels_out, t_embed_dim, y_embed_dim) for _ in range(num_residual_layers)\n",
    "        ])\n",
    "\n",
    "    def forward(self, x: torch.Tensor, t_embed: torch.Tensor, y_embed: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Args:\n",
    "        - x: (bs, c, h, w)\n",
    "        - t_embed: (bs, t_embed_dim)\n",
    "        - y_embed: (bs, y_embed_dim)\n",
    "        \"\"\"\n",
    "        # Upsample: (bs, c_in, h, w) -> (bs, c_out, 2 * h, 2 * w)\n",
    "        x = self.upsample(x)\n",
    "\n",
    "        # Pass through residual blocks: (bs, c_out, h, w) -> (bs, c_out, 2 * h, 2 * w)\n",
    "        for block in self.res_blocks:\n",
    "            x = block(x, t_embed, y_embed)\n",
    "\n",
    "        return x\n",
    "\n",
    "class MNISTUNet(nn.Module):\n",
    "    def __init__(self, channels: List[int], num_residual_layers: int, t_embed_dim: int, y_embed_dim: int):\n",
    "        super().__init__()\n",
    "        # Initial convolution: (bs, 1, 32, 32) -> (bs, c_0, 32, 32)\n",
    "        self.init_conv = nn.Sequential(\n",
    "            nn.Conv2d(1, channels[0], kernel_size=3, padding=1),\n",
    "            # nn.BatchNorm2d(channels[0]),\n",
    "            nn.GroupNorm(num_groups=8, num_channels=channels[0]),\n",
    "            nn.SiLU()\n",
    "        )\n",
    "\n",
    "        # Initialize time embedder\n",
    "        self.time_embedder = FourierEncoder(t_embed_dim)\n",
    "\n",
    "        # Initialize y embedder\n",
    "        self.y_embedder = nn.Embedding(num_embeddings = 11, embedding_dim = y_embed_dim)\n",
    "\n",
    "        # Encoders, Midcoders, and Decoders\n",
    "        encoders = []\n",
    "        decoders = []\n",
    "        for (curr_c, next_c) in zip(channels[:-1], channels[1:]):\n",
    "            encoders.append(Encoder(curr_c, next_c, num_residual_layers, t_embed_dim, y_embed_dim))\n",
    "            decoders.append(Decoder(next_c, curr_c, num_residual_layers, t_embed_dim, y_embed_dim))\n",
    "        self.encoders = nn.ModuleList(encoders)\n",
    "        self.decoders = nn.ModuleList(reversed(decoders))\n",
    "\n",
    "        self.midcoder1 = Midcoder(channels[-1], num_residual_layers, t_embed_dim, y_embed_dim)\n",
    "        self.midcoder2 = Midcoder(channels[-1], num_residual_layers, t_embed_dim, y_embed_dim)\n",
    "\n",
    "        # Final convolution\n",
    "        self.final_conv = nn.Sequential(\n",
    "            nn.GroupNorm(num_groups=8, num_channels=channels[0]),\n",
    "            nn.SiLU(),\n",
    "            nn.Conv2d(channels[0], 1, kernel_size=3, padding=1)\n",
    "        )\n",
    "\n",
    "    def forward(self, x: torch.Tensor, t: torch.Tensor, y: torch.Tensor):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "        - x: (bs, 1, 32, 32)\n",
    "        - t: (bs, 1, 1, 1)\n",
    "        - y: (bs,)\n",
    "        Returns:\n",
    "        - u_t^theta(x|y): (bs, 1, 32, 32)\n",
    "        \"\"\"\n",
    "        # Embed t and y\n",
    "        t_embed = self.time_embedder(t) # (bs, time_embed_dim)\n",
    "        y_embed = self.y_embedder(y) # (bs, y_embed_dim)\n",
    "\n",
    "        # Initial convolution\n",
    "        x = self.init_conv(x) # (bs, c_0, 32, 32)\n",
    "\n",
    "        residuals = []\n",
    "\n",
    "        # Encoders\n",
    "        for encoder in self.encoders:\n",
    "            x = encoder(x, t_embed, y_embed) # (bs, c_i, h, w) -> (bs, c_{i+1}, h // 2, w //2)\n",
    "            residuals.append(x.clone())\n",
    "\n",
    "        # Midcoder\n",
    "        x = self.midcoder1(x, t_embed, y_embed)\n",
    "        x = self.midcoder2(x, t_embed, y_embed)\n",
    "\n",
    "        # Decoders\n",
    "        for decoder in self.decoders:\n",
    "            res = residuals.pop() # (bs, c_i, h, w)\n",
    "            x = x + res\n",
    "            x = decoder(x, t_embed, y_embed) # (bs, c_i, h, w) -> (bs, c_{i-1}, 2 * h, 2 * w)\n",
    "\n",
    "        # Final convolution\n",
    "        x = self.final_conv(x) # (bs, 1, 32, 32)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Q8R9X1FHbuIy"
   },
   "source": [
    "**Please explain each components of the architecture above** (each one of `FourierEncoder`, `ResidualLayer`, `Encoder`, `Decoder`, or `Midcoder`) in your own words, (1) their role in the U-Net, (2) their inputs and outputs, and (3) a brief description of how the inputs turn into outputs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "V4XtXQBhfxFV"
   },
   "source": [
    "## 3.2 Classifier-Free Guidance\n",
    "\n",
    "Implement CFG requires a small modification to the diffusion training and sampling code.\n",
    "\n",
    "*During training*, we randomly drop out the class label with a certain probability `drop_prob=0.1`, i.e. we use a dummy embedding to replace the digit label embedding.\n",
    "\n",
    "*During sampling*, given a digit label, instead of using $\\hat{\\epsilon} = f_\\theta(x_t, t, y)$ to sample, use:\n",
    "$$\\hat{\\epsilon} = f_\\theta(x_t, t, \\varnothing) + w(f_\\theta(x_t, t, y) - f_\\theta(x_t, t, \\varnothing))$$\n",
    "where $w$ is a sampling hyperparameter that controls the strength of CFG. $\\varnothing$ indicates the unconditional model with the class label dropped out, which is supported by the dummy embedding during training. Note that $w = 1$ recovers standard sampling.\n",
    "\n",
    "Please modify the code you wrote for Problem 2 for unconditional generation to adapt for the CFG setting. You can change the model architecture from MLP to UNet. Specifically, given a data element $x$ and U-Net $f_\\theta(x, t)$, implement the following diffusion training steps similar as what we did in Problem 2, while using a different noise schedule (in step 2):\n",
    "0. Construct a class `MNISTDiffusion`\n",
    "1. Sample the diffusion timestep: $t \\sim \\text{Uniform}(0, 1)$\n",
    "2. Compute the noise-strength following a linear schedule: $\\alpha_t = 1-t, \\sigma_t = \\sqrt{1-(1-t)^2}$\n",
    "3. Sample noise $\\epsilon \\sim N(0,I)$ (same shape as $x$) and cmpute noised $x_t = \\alpha_t x + \\sigma_t \\epsilon$\n",
    "4. Estimate $\\hat{\\epsilon} = f_\\theta(x_t, t)$\n",
    "5. Optimize the loss $L = \\lVert \\epsilon - \\hat{\\epsilon} \\rVert_2^2$. Here, it suffices to just take the mean over all dimensions.\n",
    "\n",
    "*Note*: you can reuse your code from Problem 2 for functions `train`, `eval_loss`, `get_lr`, and `train_epochs`.\n",
    "\n",
    "*Hyperparameter details*\n",
    "* UNet with hidden_dims as [64, 128] and 1 blocks_per_dim\n",
    "* Train 10 epochs, batch size 128, Adam with LR 1e-3 (0 warmup steps, and `use_cosine_decay=True`)\n",
    "* Training 10 epochs takes about 7 minutes on the Google Colab T4 GPU.\n",
    "\n",
    "After training, please generate 4 images for each of the 10 digit labels using guidance strength $w$ of 0.0, 0.5, 1.0, 2.0, 4.0 respectively, and save the images as 4x10 grid of images (each column is a digit). You are required to submit the images for each guidance strength along with the training loss curve. **Comparing the results with different $w$, what can you say about its impact to the generation performance?**\n",
    "\n",
    "*Hint*: To check your answer, the final test loss is below 0.02."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "id": "KeB_8PeiffOU"
   },
   "outputs": [],
   "source": [
    "class MNISTDiffusion:\n",
    "    def __init__(self, model, data_shape, device, n_classes, drop_prob=0.1):\n",
    "        \"\"\"\n",
    "        model: neural network to estimate eps_hat (U-Net in this problem)\n",
    "        data_shape: size of the input data, (1, 28, 28) in this case\n",
    "        device: cuda\n",
    "        n_classes: number of classes for conditional generation, set to 10 in this problem\n",
    "        drop_prob: probability of dropping the condition in CFG training\n",
    "        \"\"\"\n",
    "        self.model = model.to(device)\n",
    "        self.data_shape = data_shape\n",
    "        self.device = device\n",
    "        self.drop_prob = drop_prob\n",
    "        self.n_classes = n_classes\n",
    "    \n",
    "    def loss(self, x, y):\n",
    "        \"\"\"\n",
    "        x: the input data (without adding noise) from the dataloader\n",
    "        y: the class label\n",
    "        Return:\n",
    "            The loss (as a scalar averaged over all data in the batch)\n",
    "        \"\"\"\n",
    "        raise NotImplementedError(\"Please implement this\")\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def sample(self, n, num_steps, guide_w=0.0):\n",
    "        \"\"\"\n",
    "        n: number of samples to generate (you should generate int(n/n_classes) samples for each class and thus n samples in total)\n",
    "        num_steps: number of steps in the diffusion sampling\n",
    "        guide_w: the CFG weight\n",
    "        Return:\n",
    "            The generated sample. Tensor with shape (n, *self.data_shape)\n",
    "        \"\"\"\n",
    "        raise NotImplementedError(\"Please implement this\")\n",
    "\n",
    "    def __getattr__(self, name):\n",
    "        if name in ['train', 'eval', 'parameters', 'state_dict', 'load_state_dict']:\n",
    "            return getattr(self.model, name)\n",
    "        return self.__getattribute__(name)\n",
    "\n",
    "\n",
    "# You can reuse your implementations in part 2 problem 2 for the following four functions\n",
    "\n",
    "def train(model, train_loader, optimizer, scheduler):\n",
    "    \"\"\"\n",
    "    model: model to train, the MNISTDiffusion class in this case.\n",
    "    train_loader: train_loader\n",
    "    optimizer: use torch.optim.Adam\n",
    "    scheduler: use optim.lr_scheduler.LambdaLR(\n",
    "        optimizer,\n",
    "        lr_lambda=lambda step: get_lr(step, total_steps, warmup_steps, use_cos_decay)\n",
    "    )\n",
    "    Return:\n",
    "        Tensor with train loss of each batch\n",
    "    \"\"\"\n",
    "    raise NotImplementedError(\"Please implement this\")\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def eval_loss(model, data_loader):\n",
    "    \"\"\"\n",
    "    model: model to train, the MNISTDiffusion class in this case.\n",
    "    data_loader: test_loader\n",
    "    Return:\n",
    "        Scalar with the average test loss of each batch\n",
    "    \"\"\"\n",
    "    raise NotImplementedError(\"Please implement this\")\n",
    "\n",
    "\n",
    "def get_lr(step, total_steps, warmup_steps, use_cos_decay):\n",
    "    \"\"\"\n",
    "    Function that returns the learning rate for the specific step, used in defining the scheduler:\n",
    "        scheduler = optim.lr_scheduler.LambdaLR(\n",
    "            optimizer,\n",
    "            lr_lambda=lambda step: get_lr(step, total_steps, warmup_steps, use_cos_decay)\n",
    "        )\n",
    "    \"\"\"\n",
    "    raise NotImplementedError(\"Please implement this\")\n",
    "\n",
    "\n",
    "def train_epochs(model, train_loader, test_loader, train_args):\n",
    "    \"\"\"\n",
    "    model: model to train, the MNISTDiffusion class in this case.\n",
    "    train_loader: train_loader\n",
    "    test_loader: test_loader\n",
    "    Return:\n",
    "        Two np.array for all the train losses and test losses at each step\n",
    "    \"\"\"\n",
    "    raise NotImplementedError(\"Please implement this\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "id": "qpsSdUPvffKx"
   },
   "outputs": [],
   "source": [
    "def mnist_diffusion(train_loader, test_loader):\n",
    "    \"\"\"\n",
    "    train_loader: MNIST training data\n",
    "    test_loader: MNIST test data\n",
    "\n",
    "    Returns\n",
    "    - a (# of training iterations,) numpy array of train losses evaluated every minibatch\n",
    "    - a (# of num_epochs + 1,) numpy array of test losses evaluated at the start of training and the end of every epoch\n",
    "\n",
    "    Generate 4 images for each of the 10 digits in `./results/`, using guidance strengths of 0.0, 0.5, 1.0, 2.0, 4.0.\n",
    "    Save the images as `./results/image_w{w}.png` as 4x10 grid of images (each column is a digit)\n",
    "    hint: x_gen is the output from model.sample, use the following to generate and save grids\n",
    "            grid = make_grid(x_gen*-1 + 1, nrow=10)\n",
    "            save_image(grid, save_dir + f\"image_w{w}.png\")\n",
    "    \"\"\"\n",
    "\n",
    "    n_classes = 10\n",
    "    save_dir = './results/'\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    ws_test = [0.0, 0.5, 1.0, 2.0, 4.0] # strength of generative guidance\n",
    "    num_steps = 512\n",
    "\n",
    "    model = MNISTUNet(\n",
    "        channels = [64, 128],\n",
    "        num_residual_layers = 1,\n",
    "        t_embed_dim = 64,\n",
    "        y_embed_dim = 64,\n",
    "    )\n",
    "    model = MNISTDiffusion(model, (1, 28, 28), device=\"cuda\", n_classes=n_classes, drop_prob=0.1)\n",
    "\n",
    "    raise NotImplementedError(\"Please implement this\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "id": "bTfqxUM8Exb9"
   },
   "outputs": [],
   "source": [
    "def save_training_plot(\n",
    "    train_losses: np.ndarray, test_losses: np.ndarray, title: str, fname: str\n",
    ") -> None:\n",
    "    plt.figure()\n",
    "    n_epochs = len(test_losses) - 1\n",
    "    x_train = np.linspace(0, n_epochs, len(train_losses))\n",
    "    x_test = np.arange(n_epochs + 1)\n",
    "\n",
    "    plt.plot(x_train, train_losses, label=\"train loss\")\n",
    "    plt.plot(x_test, test_losses, label=\"test loss\")\n",
    "    plt.legend()\n",
    "    plt.title(title)\n",
    "    plt.xlabel(\"Epoch\")\n",
    "    plt.ylabel(\"NLL\")\n",
    "    savefig(fname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_losses, test_losses = mnist_diffusion(train_loader, test_loader)\n",
    "print(f\"Final Test Loss: {test_losses[-1]:.4f}\")\n",
    "\n",
    "save_training_plot(\n",
    "    train_losses,\n",
    "    test_losses,\n",
    "    f\"MNIST Train Plot\",\n",
    "    f\"results/mnist_train_plot.png\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.image as mpimg\n",
    "\n",
    "ws_test = [0.0, 0.5, 1.0, 2.0, 4.0]\n",
    "for w in ws_test:\n",
    "  img_path = f'results/image_w{w}.png'\n",
    "  img = mpimg.imread(img_path)\n",
    "  plt.title(f'w={w}')\n",
    "  plt.imshow(img)\n",
    "  plt.axis('off')\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2qlmbEblExb-"
   },
   "source": [
    "# Submission Guideline for Part 2\n",
    "\n",
    "Please include your answer to all problems, including formulas, proofs, and the figures generated in each problem, excluding code. You are required to submit the (single) pdf and all (four) notebooks (one for each problem) with your code and running outputs. Do not include code in the pdf file.\n",
    "\n",
    "Specifically, for Problem 3 in this notebook, the pdf should contain:\n",
    "- The generated figures `results/mnist_train_plot.png` and `results/image_w{w}.png` (w=0.0, 0.5, 1.0, 2.0, 4.0)\n",
    "- Answer to the short answer question about the U-Net architecture\n",
    "- Answer to the short answer question about different CFG weight $w$ in problem 3.2"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}