imagen_train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['KMP_DUPLICATE_LIB_OK']='True'\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "import nibabel as nib\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.optim import Adam\n",
    "\n",
    "import glob\n",
    "import time\n",
    "import yaml\n",
    "from yaml.loader import SafeLoader\n",
    "from torch.utils.data.dataset import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "import pandas as pd\n",
    "from skimage.metrics import structural_similarity as ssim\n",
    "from tqdm import tqdm\n",
    "\n",
    "import cv2\n",
    "from imagen_pytorch3D import Unet, NullUnet, Imagen, SRUnet256, alpha_cosine_log_snr\n",
    "# from imagen_pytorch import load_imagen_from_checkpoint, ImagenTrainer\n",
    "\n",
    "from trainer import ImagenTrainer\n",
    "import torchvision.transforms.functional as TF\n",
    "import torchvision.transforms as T\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from PIL import Image \n",
    "import PIL \n",
    "import pandas as pd\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 118,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randn(1,1,32,32,32)\n",
    "y = torch.randn(1,1,32,32,32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD5CAYAAADhukOtAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/MnkTPAAAACXBIWXMAAAsTAAALEwEAmpwYAAAaV0lEQVR4nO2de6yV5ZXGn8VNlDtyKXK1SEEKAnpqtEOtYrXa1IjNxLZJW9KSYqc1mSb1D+MkUyeZP9rJ2KZ/TDqho5FOOrVOL9ZOTAevGC9VoVIEL2gRBDyAgFyKFT2HNX/szczB+Z7nHPY5Z2/q+/wSwj7vOmt/736/b51v7/fZa63ITBhj3v8MaPUEjDHNwcFuTCE42I0pBAe7MYXgYDemEBzsxhTCoN44R8RVAH4AYCCAf8vM76jfHzlyZE6YMKHS1tnZqY5TOd7R0UF9Bg4cSG3qWMeOHaO2wYMHV44fPXqU+igGDeLL/6c//YnahgwZQm3stannU/NX0uyAAfxewWzqnKnnU+dFoa6DvkbNUb02NseJEydSnzFjxlSOb9++Hfv3768MmIaDPSIGAvgXAFcA2AHgmYi4NzOfZz4TJkzAbbfdVmk7ePAgPRZbqP3791OfkSNHUpu68N966y1qmzRpUuX4li1bqI/izDPPpLbHH3+c2qZNm0ZtBw4cqBx/7LHHqI+a/9tvv01tp59+OrUNHz68cvyNN96gPmeccQa1qXOmGDVqVOV4o4HJbjwAcOTIEWo77bTTqG306NGV4zfddBP1+cxnPlM5fvXVV1Of3ryNvxDAK5m5JTPfAXAXgGt78XzGmH6kN8E+GcD2Lj/vqI8ZY05B+n2DLiJWRMTaiFh76NCh/j6cMYbQm2DfCWBql5+n1MdOIDNXZmZbZrapz9HGmP6lN8H+DIBZEXF2RAwB8DkA9/bNtIwxfU3Du/GZ2RERNwL4b9Sktzsyc5PyOXbsGN3tVnIS22FmO62AltfY7iegd31ff/31yvH169dTn0WLFlHbu+++S23qOdXrHjt2bOX4Rz7yEeqza9cuams0K5K9NqZoAMCf//xnalMSmlIF2DwOHz5MfYYNG9bQPNSOu7oeP/axj1WOX3755dSHxYtSC3qls2fmfQDu681zGGOag79BZ0whONiNKQQHuzGF4GA3phAc7MYUQq9240+Wjo4OmgjBJCOAJ4wwSQ7QSTJKnti2bRu1vfrqqyftc/HFF1ObynpTtnXr1lEbk41mz55Nfa688kpq27dvH7U9++yz1MYy6RrJ/gKAqVOnUpvyY9fb+PHjqY+6PtR1peZx7rnnUtsnP/nJynGW2QbwBCWZ4EMtxpj3FQ52YwrBwW5MITjYjSkEB7sxhdDU3fgDBw7gnnvuqbSpEk1sJ3nnzv+XUfu/qHJKu3fvpjaW7ALwnVilJGzaxHODrrnmGmqbNWsWtT333HPUtmPHjspxlYihdn3b29upTSUUsd34PXv2UJ/p06ef9PMB+pwxVUOpDKo8ltrtVn4qIeq8886rHFeJQex8ypqB1GKMeV/hYDemEBzsxhSCg92YQnCwG1MIDnZjCqGp0tvhw4exZs2aSptKPpgxY0bleCMJEAAwYsQIalOJGizJZOjQodRHdVtRtd+WLl1Kbeq1sc4v6lhKOlRrrGroqSQlhpJS1bGUrMhQEpp6PtYCDNDrqGoAsrp8ah6qhh7Dd3ZjCsHBbkwhONiNKQQHuzGF4GA3phAc7MYUQq+kt4jYCuAwgE4AHZnZJg82aBDNsFKtc5gsp+qIqTZOqtVUIzYlMynJ6KGHHqK2L3zhC9S2fPlyamPZcg888AD1UVmAqhWSOmcf+MAHKsfVeiiZT2WU7d27l9qYxDZ8+HDqozImlRz2iU98gtpYiyeAt69SWW9Kqmb0hc5+WWby1TbGnBL4bbwxhdDbYE8AqyNiXUSs6IsJGWP6h96+jV+cmTsjYgKA+yPixcx8tOsv1P8IrAD0V1GNMf1Lr6IvM3fW/98D4FcALqz4nZWZ2ZaZbQ52Y1pHw9EXEcMiYsTxxwCuBLCxryZmjOlbevM2fiKAX9UlgEEA/iMzfysPNmgQJkyYUGlTcse8efMqx5WMw44D6Cw1lXnF5DxVLFNJeUryuvPOO6lt2bJl1PaVr3ylcvzjH/849fnd735HbY888gi1qYy+d955p3JcFalUvPXWWyd9LICvvypgyaQwAFi8eDG1XX/99dQ2efJkauvo6KgcV9cpm7+S5BoO9szcAmBBo/7GmObiD9HGFIKD3ZhCcLAbUwgOdmMKwcFuTCE0teDk4MGDaXG9zZs3Uz9WLHHUqFHUR/UNU5KG+uLPwYMHK8eVXPfqq69Sm5LsXnrpJWpTktcll1xSOT5//nzq8+Uvf5narr76ampTkt1rr71WOd5IL73u/FQRSJa1N2fOHOrDCpwCwGc/+1lqO+ecc6hNZWGy/mwqM4/JdQrf2Y0pBAe7MYXgYDemEBzsxhSCg92YQmjqbvzAgQPpDuPMmTOpH6vFddZZZ1GfadOmUZuqGacSNdatW1c5rhJyVJshlbSgkjE2bdpEbVu3bq0cV4kYF1xwAbWxJCRA72h/9KMfrRxXteRUssu+ffuo7dChQ9TGEmHGjRtHfVTttw9/+MPUppQcdY0wVIsnthvPdvYB39mNKQYHuzGF4GA3phAc7MYUgoPdmEJwsBtTCE2V3t5++22a8KJqtS1cuLBynLUYArSspWQXlbDA2k2ppASWxANomU/N4/zzz6e2bdu2VY5v376d+qgkJNWSSSXyXH755ZXjS5YsoT5z586lNpXYpK6DI0eOVI6rmodK8lLSoWoNpaRUJsup1lts7VVLLt/ZjSkEB7sxheBgN6YQHOzGFIKD3ZhCcLAbUwjdSm8RcQeATwPYk5nz6mNjAfwMwAwAWwFcn5lvdvdc77zzDpWAZs2aRf1YJo/KKFOZbcqvvb2d2phEoqQ3dSzVtohJRoCWythcBg3ip7qR1kSAzlL77W+rO4GtXr2a+qhrYMEC3nyI1d0DuHw1ZcoU6jNmzBhqU+dMSWUqG62Rc8bWXl1vPbmz3wngqveM3QzgwcycBeDB+s/GmFOYboO93m/9vWU/rwWwqv54FYClfTstY0xf0+hn9omZefz97i7UOroaY05hev112czMiKAfSCJiBYAV9ce9PZwxpkEavbPvjohJAFD/fw/7xcxcmZltmdnmYDemdTQa7PcCWFZ/vAzAr/tmOsaY/qIn0ttPAVwKYFxE7ADwbQDfAXB3RCwHsA3A9T052MCBA2mml2p1w1oJ7d27l/qwNlMAsHv3bmpTbZeYHDZixAjqo7Kk1Dsd1dKItaECeAaYap909OhRalNZVCpzjMlGykcVlVSS17nnnkttjUhU5513HrWpYpQTJkxoyK8RH5aZp6TSboM9Mz9PTNU5jMaYUxJ/g86YQnCwG1MIDnZjCsHBbkwhONiNKYSmFpwcMmQI7c/WiIyjUFLNyJEjqU31Pdu5c2fluMr+UqiikipLaujQodS2f/970xhqqDVUPfNYkU1AF8xctGhR5biSp1TWmJJSVTHKPXuqv++leq+9+SZP4FQSseo5p+RN1iNOyYMTJ1Z/Q11Jtr6zG1MIDnZjCsHBbkwhONiNKQQHuzGF4GA3phCaKr0BPNNLZYcxVG8wJScpyYtJVwCXB5m8AwBvvPEGtan+X8pPvW5WmHH+/PnURxV6VH3UVFYW60Wm+qgp6UqxY8eOk/ZR/f7UNaD626kioUouZdeV6rPHshiVpOg7uzGF4GA3phAc7MYUgoPdmEJwsBtTCE3djVc16FRyCtuZVvW21A652vVVO6oMtSutdmFVcsrXvvY1alu+fDm1sd1i9bpU4sf69eup7fnnn6e2Xbt2VY6z9l+ATigaMmQItamEEabyqJZXSslZunQptSmVRK0/20FnCTLKphKofGc3phAc7MYUgoPdmEJwsBtTCA52YwrBwW5MIfSk/dMdAD4NYE9mzquP3QrgqwCOa2K3ZOZ93T3X4MGDaQ2yzs5O6tfe3l45rmqWKduYMWOoTUkkW7ZsqRxX9fOUFPKlL32J2m644QZqU4kaDz/8cOX4Aw88QH02b95Mbew1A1ryYudTnWf1fCppSK0xqxm3bt066qNkPtVyTEmiM2bMoDaWbKQStlhbLinXUcv/cSeAqyrGv5+ZC+v/ug10Y0xr6TbYM/NRAPxWYoz5i6A3n9lvjIgNEXFHRPD3xcaYU4JGg/2HAGYCWAigHcBt7BcjYkVErI2IteqzrTGmf2ko2DNzd2Z2ZuYxAD8CcKH43ZWZ2ZaZbep74saY/qWhYI+Iru05rgOwsW+mY4zpL3oivf0UwKUAxkXEDgDfBnBpRCwEkAC2AuA6UReOHDlCJQ9V24vVLXv22Wepz4gRI6hNZdix2l4Al0JU9t0111xDbYsXL6a2VatWUZt63U8++WTl+O7du6mPQkmYstUQkYCUhKaeT9ZWE3ITk+XUNXDw4EFqu+eee6hNvXNVUirLBFUtzJhMqeTLboM9Mz9fMXx7d37GmFMLf4POmEJwsBtTCA52YwrBwW5MITjYjSmEphacHDZsGC644IJKm5JPzj777Mrx1157jfq8/PLL1KYKLCqJimXLTZ8+nfpcccUV1PbII49QG5PQAN0WiBW/VAUWjx49Sm2qCKSSN5lMqY6lWl6prDH1zUwmpY4fP576NNKKDNDn7EMf+hC1jR07tnJcZfMxuU5dG76zG1MIDnZjCsHBbkwhONiNKQQHuzGF4GA3phCaKr2NHj0a1113XaVNFSJkBRYXLFhAfVShRCXxrFmzhtpY/7LZs2dTnw0bNlDbvn37qG3OnDnUpvrYLVmypHKcFSgEgFdeeYXalFR2zjnnUBuTNzdu5NnQKsNOyVAqW47Jpep6iwhqUxKxypZ7/PHHqe3SSy+tHGeSM8CzAN3rzRjjYDemFBzsxhSCg92YQnCwG1MITd2N7+zsxIEDByptareYJWOoXViVeLBo0SJqUzvMrH5eo1Vzp02bRm0qyUS9bra+KlmEtR8CdM04VReO7UyrJCTVdkkdS702trOudtUVahdfqTyq3dTq1asrx7/+9a+f9LHU/HxnN6YQHOzGFIKD3ZhCcLAbUwgOdmMKwcFuTCH0pP3TVAA/BjARtXZPKzPzBxExFsDPAMxArQXU9ZnJdRXUWtMwGY3VLAO4/DNoEJ++kiBGjRpFbUp6YxKbahml2vEo+Ue1/lGvbdOmTZXj7e3t1EetB6tpB+g5snPTqEypJEBVd40l8qiWXYcOHaI2Vi8OqLU3Y6jkpYceeqhy/KKLLqI+8+fPrxzvbQ26DgDfysy5AC4C8I2ImAvgZgAPZuYsAA/WfzbGnKJ0G+yZ2Z6Zv68/PgzgBQCTAVwL4Hj3wVUAlvbTHI0xfcBJfWaPiBkAFgF4CsDEzDz+3nAXam/zjTGnKD0O9ogYDuAXAL6ZmSd8qMlaxnxl1nxErIiItRGxVn0uN8b0Lz0K9ogYjFqg/yQzf1kf3h0Rk+r2SQAqdyAyc2VmtmVm2/Dhw/tizsaYBug22KOWSXA7gBcy83tdTPcCWFZ/vAzAr/t+esaYvqInWW9/BeCLAJ6LiPX1sVsAfAfA3RGxHMA2ANd3e7BBgzBxYvVHe5XJxeQTJXUoCUJ9nFByGGu5o+rMKRln7ty51DZy5EhqU6+bSX1KQtu7d29D81BZaqxuoMpuVPKakq7UOWPXgZIv2XkGtNyoMhVVZh67fn7zm99Qn5kzZ1aOK6m322DPzMcAsAp8l3fnb4w5NfA36IwpBAe7MYXgYDemEBzsxhSCg92YQmhqwcmIoDKJajPEiuupIoSNFmxUX/w544wzKsenT59OfTZv3kxtSpZT82CZgwCwdevWynG1VuPGjaM21ZJJzZ9JW0oCVG2X1DlTsHZIKvtOyWRK2lJzVJmFTKZ8+umnqc+TTz5ZOa5kWd/ZjSkEB7sxheBgN6YQHOzGFIKD3ZhCcLAbUwhNld4ALq+oHmAsG2rq1KnUh/UaA3g/NEBLMkwOW7BgAfVRctKuXbuobceOHdSm5CuWpaZes+pRpjLKlGTHstsaLWCijqWuHbZWSupV0qbK9Gu0KOb48eMrx7dv3059nnjiicpxmdFJLcaY9xUOdmMKwcFuTCE42I0pBAe7MYXQ1N34jo4OunOqdhHZLvi+ffuoj9r9nDRpErUpDh8+XDmuklbmzJlDbap9lXptKlGD7RarBI5GUXXh2PlUySLqnCl1Qu2ss0QepbooBUJdp2o9VELRmDFjKsfVeWZJMk6EMcY42I0pBQe7MYXgYDemEBzsxhSCg92YQuhWeouIqQB+jFpL5gSwMjN/EBG3AvgqgONZFLdk5n3quQYMGEATJJQcxuQaJZGoBBTVgkjJP0x6Uy2SpkyZQm1KsnvmmWeoTb021p5I1eRrtNWUkryYtKXkRjVHldikElDYc6pkItXiSaHOi5I+mayorsUXX3yxclzJdT3R2TsAfCszfx8RIwCsi4j767bvZ+Y/9+A5jDEtpie93toBtNcfH46IFwBM7u+JGWP6lpP6zB4RMwAsAvBUfejGiNgQEXdERPXXgIwxpwQ9DvaIGA7gFwC+mZmHAPwQwEwAC1G7899G/FZExNqIWMs+8xpj+p8eBXtEDEYt0H+Smb8EgMzcnZmdmXkMwI8AXFjlm5krM7MtM9vUBowxpn/pNtijtsV4O4AXMvN7Xca7bp9fB2Bj30/PGNNX9GQ3/q8AfBHAcxGxvj52C4DPR8RC1OS4rQBu6O6JBgwYQGUSJYcx2UK1NGq0TY+SO0aPHl05ruQOhWobpeQwlbHFpLLXX3+d+qhMLtbyCtDnjPl1dHRQHyV5qY+ASs5j14hqXaVQMp+ysTZUAL9W1eti0qG67nuyG/8YgCoBUWrqxphTC3+DzphCcLAbUwgOdmMKwcFuTCE42I0phFOm/VMjmVxKJlMSj8rWUvIJk3GUZKTmobL2zjzzTGq77LLLqG3nzp2V4yxLCtDrqNZK2VTBTIaSFNV5Ue2amIQ5YcIE6qPWQ72u/fv3U5s610zCVDHBCouqc+I7uzGF4GA3phAc7MYUgoPdmEJwsBtTCA52YwqhqdKbynpT2VVMglBZb6x/VneorCE2R5Wn32jfMCXjnHXWWdS2ePHiyvE1a9ZQn/b2dmpTktfkybw6GZu/es1qrcaOHUttKvuOZSQ22utNZaIpCVZdV8ymsjMbwXd2YwrBwW5MITjYjSkEB7sxheBgN6YQHOzGFEJTpbfMRGdnZ6XtzTffpH5MYlP9uhRK0lDPyaQ3JeMolNzIspoALV+xbK7Zs2dTHyXzqUKJBw4coDaW7aeeT8laquecynpjWWDKp1HJa9iwYdSm+voxeVMVMm2kH53v7MYUgoPdmEJwsBtTCA52YwrBwW5MIXS7Gx8RQwE8CuC0+u//PDO/HRFnA7gLwJkA1gH4YmbKLcJjx47RnWS1W8kSHUaNGkV9VDJDo+2a2E69SoRRSSuqBdGuXbuoTe0ks934efPmUZ8nnniC2tRaqfp67Nwo5UIli6jXrBJh1A45gylGgF4PpaCotWKJSKoWnlI1GD25sx8FsCQzF6DWnvmqiLgIwHcBfD8zzwHwJoDlJ310Y0zT6DbYs8bx2/Hg+r8EsATAz+vjqwAs7Y8JGmP6hp72Zx9Y7+C6B8D9AP4I4EBmHn9vsgMAT242xrScHgV7ZnZm5kIAUwBcCGBOTw8QESsiYm1ErG20Ta4xpvec1G58Zh4A8DCAiwGMjojjG3xTAFR2J8jMlZnZlpltque4MaZ/6TbYI2J8RIyuPz4dwBUAXkAt6P+6/mvLAPy6n+ZojOkDepIIMwnAqogYiNofh7sz878i4nkAd0XEPwJ4FsDt3T1RRFB5QiWFMGlFfSw4ePAgtSlJQ9VcY35K5lOSkWolpJIxGqnHtnDhQuoza9Ysatu6dSu1qVZDLFFDSWFKTlLHaiR5SclkjdSLA/T1qFo5MVlOXQPsOlV1GbsN9szcAGBRxfgW1D6/G2P+AvA36IwpBAe7MYXgYDemEBzsxhSCg92YQohGsmcaPljEGwC21X8cB2Bv0w7O8TxOxPM4kb+0eUzPzPFVhqYG+wkHjlibmW0tObjn4XkUOA+/jTemEBzsxhRCK4N9ZQuP3RXP40Q8jxN538yjZZ/ZjTHNxW/jjSmElgR7RFwVES9FxCsRcXMr5lCfx9aIeC4i1kfE2iYe946I2BMRG7uMjY2I+yPi5fr/Y1o0j1sjYmd9TdZHxKeaMI+pEfFwRDwfEZsi4m/r401dEzGPpq5JRAyNiKcj4g/1efxDffzsiHiqHjc/iwieuldFZjb1H4CBqJW1+iCAIQD+AGBus+dRn8tWAONacNxLAJwPYGOXsX8CcHP98c0AvtuiedwK4KYmr8ckAOfXH48AsBnA3GaviZhHU9cEQAAYXn88GMBTAC4CcDeAz9XH/xXA35zM87bizn4hgFcyc0vWSk/fBeDaFsyjZWTmowDe203xWtQKdwJNKuBJ5tF0MrM9M39ff3wYteIok9HkNRHzaCpZo8+LvLYi2CcD2N7l51YWq0wAqyNiXUSsaNEcjjMxM48XEN8FYGIL53JjRGyov83v948TXYmIGajVT3gKLVyT98wDaPKa9EeR19I36BZn5vkArgbwjYi4pNUTAmp/2VH7Q9QKfghgJmo9AtoB3NasA0fEcAC/APDNzDyh7Esz16RiHk1fk+xFkVdGK4J9J4CpXX6mxSr7m8zcWf9/D4BfobWVd3ZHxCQAqP+/pxWTyMzd9QvtGIAfoUlrEhGDUQuwn2TmL+vDTV+Tqnm0ak3qxz6AkyzyymhFsD8DYFZ9Z3EIgM8BuLfZk4iIYREx4vhjAFcC2Ki9+pV7USvcCbSwgOfx4KpzHZqwJlEr0HY7gBcy83tdTE1dEzaPZq9JvxV5bdYO43t2Gz+F2k7nHwH8XYvm8EHUlIA/ANjUzHkA+ClqbwffRe2z13LUeuY9COBlAA8AGNuiefw7gOcAbEAt2CY1YR6LUXuLvgHA+vq/TzV7TcQ8mromAM5DrYjrBtT+sPx9l2v2aQCvAPhPAKedzPP6G3TGFELpG3TGFIOD3ZhCcLAbUwgOdmMKwcFuTCE42I0pBAe7MYXgYDemEP4HKIdgaxwqzwUAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD5CAYAAADhukOtAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/MnkTPAAAACXBIWXMAAAsTAAALEwEAmpwYAAAZwklEQVR4nO2dfYyV5ZnGr9sBFGH4HBhHOjggFMUviiPRqq1L1bJNGzAxVpMa/zClJTXZJt0/jJts3WTTtJttm/6x6Yauprbp1rr9ULsxq6y1tR+WdsBhAEHkSxm+hmGQQUFG4N4/zksyuO99zcw757yH9rl+CWHmuec55znPvNecc57r3Pdt7g4hxF8/F9R7AUKIcpDYhUgEiV2IRJDYhUgEiV2IRJDYhUiEMaOZbGbLAHwHQAOA/3D3r7Ofb2xs9BkzZuTGiliAbA6LnTlzJoxdcEH898/Mcsfff//9cE5DQ0OhdZw4cSKMjRkT/9qi2zx+/Hih+2L7GO0HwPcxgu0HW0eR+2IUfVxFr7kLL7wwdzzSCgBMmTIld3zPnj3o6+vLfQCFxW5mDQD+DcAdALoB/NnMnnX316I5M2bMwNe+9rXcGBNMtIkDAwPhnFOnToWxd955J4xNnDgxjI0dOzZ3/ODBg+GcCRMmhDEmsk2bNoUxdhEcO3Ysd3zDhg3hnK6urjDG9njcuHFhjO1jBPuDdPLkyTB20UUXhbFInEy07I9pJEyAX3Psdz137tzc8S9+8YvhnBUrVuSOL1u2LJwzmj+JSwBsd/ed7j4A4EkAy0dxe0KIGjIasc8CsGfQ993ZmBDiPKTmB3RmttLMOsysI3qJKYSoPaMR+14ArYO+/1A2dg7uvtrd2929vbGxcRR3J4QYDaMR+58BzDezOWY2DsC9AJ6tzrKEENWm8Gm8u58ys4cAPI+K9fa4u28eal5kQTBrIjqpj07HgWIntAA/Ud279/+9cAEAvPZaaECgvb09jLGT+j179oQxdgo+bdq03PH58+eHc95+++0w1tvbG8ZOnz4dxiL7avz48eEc9jsr+hYwsj6ZTcauAXadMstu8uTJYez666/PHb/11lvDOS0tLbnjTBOj8tnd/TkAz43mNoQQ5aBP0AmRCBK7EIkgsQuRCBK7EIkgsQuRCKM6jR8pAwMD6O7uzo21trbmjgOxXcMSWlhSxbvvvhvGmOW1devW3PH+/v5wzqJFi8LYpEmTwlhfX1+h2KxZ+Z9Ybm5uDufccMMNYSyyGwFuyx06dCh3nCXWsOSZonZplNTC5rCkLJaQw5gzZ04YW7p0ae44+53t27cvd5ytXc/sQiSCxC5EIkjsQiSCxC5EIkjsQiRCqafx/f39WLNmTW7s0ksvDedNnTo1d/zo0aPhnAMHDoQxdlLPTp8PHz6cO97U1BTO2bVrVxhjj/mqq64KY2vXrg1jmzfn5yKxvYr2F+COB0uEiRJGWAIKq9fH5jF3JTp1Z0krLCGHPWaWhHLllVeGMebYRESJQWx9emYXIhEkdiESQWIXIhEkdiESQWIXIhEkdiESoVTr7fjx43j11VdzY9E4EHdAYckMrK4aS7hgsahzB0vuYPXprrvuujB2xx13hDFmNb3wwgu546xrTVS3biiYzRMlBxVt1cR+10UsQFbHjyWTsN/17Nmzw9jixYvDWJTwwqzIKDmM7a+e2YVIBIldiESQ2IVIBIldiESQ2IVIBIldiEQYlfVmZrsBHANwGsApd497HaFiC0TN7JllMH369Nzxoi18WFYTu83IxmEWyXvvvRfGmC330Y9+NIzdd999YSyy0X7/+9+Hc6J6ZkBcww0ALr744jA2c+bM3HGWbcb2kdlh7HddJOuN2basfdWNN94Yxlidv8juZZZiEeutGj7737h7XHlQCHFeoJfxQiTCaMXuAF4ws3VmtrIaCxJC1IbRvoy/xd33mtlMAGvMbKu7vzz4B7I/AisB/p5MCFFbRvXM7u57s/97APwCwJKcn1nt7u3u3l70c9FCiNFTWH1mNsHMGs9+DeBOAJuqtTAhRHUZzcv4ZgC/yCyMMQD+093/h01oaGgIWx6xLKSWlpYRL47ZFsx2idoWAcCRI0dyx1nWGLP51q1bV2gdd911VxhbtWpV7jjLovvDH/4Qxl555ZUwtn379jAWWZjMrmO2FrPXTpw4EcaiV5PMUmS/z/b22F1esWJFGLv88svDWGQrRlYvwK/hiMJid/edAOIcTSHEeYXeRAuRCBK7EIkgsQuRCBK7EIkgsQuRCMayvKpNY2OjR32tWKbRlClTcsdZjzJmTUSZdwDv9RbFTp48Gc5hMfYhI1bYcMGCBWHsM5/5TO44K27Jes6x/ejs7AxjO3bsGPHtRb30AF4wk2XERTbaJZdcEs5hNtmdd94Zxm699dYw1tjYGMaifnrsE6dRX7nbb78dnZ2duRe/ntmFSASJXYhEkNiFSASJXYhEkNiFSIRS2z+NGzcOra2tuTF2WhmdTDc1NYVzopZRAD/ljBJ1gDhZZ9u2beEcdhof1WkDeBLE1q1bw9hbb72VO85aE11//fVhbN68eWFszpw5YeyWW27JHZ8wYUI4hyW0sNP448ePh7EIlpDDrsX58+eHMZbMFbXDAuLHzfYqWqPaPwkhJHYhUkFiFyIRJHYhEkFiFyIRJHYhEqFU621gYAB79uzJjbFabZFFxWwhluzCLB5Wfyyqg8ZaPLHbY1YNs+xY/bRo3u7du8M5W7ZsCWPMAmxubg5jH//4x3PHb7vttnAOS9b58Ic/HMai9klAnGTCHhe7PqIEFIDXPWTXd3SbkydPDudE1w5LANMzuxCJILELkQgSuxCJILELkQgSuxCJILELkQhD1qAzs8cBfBpAj7tfnY1NA/ATAG0AdgO4x93zeyMNYty4cR5lo7GsrLa2ttzxWbNmhXOK1nc7cODAiGOsft67774bxpiFdvTo0TBW7XZHrO0Ss6iKZJux3wurhXfDDTeEsU9+8pNhLLrepk+fHs6Jah4CvA1V0T2ONMgsxWjOzTffjPXr1xeuQfd9AMs+MPYwgBfdfT6AF7PvhRDnMUOKPeu33veB4eUAnsi+fgLAiuouSwhRbYq+Z2929/3Z1wdQ6egqhDiPGfXHZd3dzSx8429mKwGsBHiFGCFEbSn6zH7QzFoAIPu/J/pBd1/t7u3u3s4OZ4QQtaWo+p4F8ED29QMAnqnOcoQQtWLIl/Fm9mMAtwFoMrNuAF8F8HUAT5nZgwDeBHDPcO8weinPXuIfOZLv6jHrhxWOZNlJUVYeAPT29uaOs0wo1pqIWTUsI47ZP9FaouwvgFuRzJplmXlR9hXLECyaPchaSkXWZ3RNAcCVV14Zxtjvhe0VszCjfWTXTvR7ZvczpNjd/b4g9Imh5gohzh/0JlqIRJDYhUgEiV2IRJDYhUgEiV2IRCi911uUqcZsqMg+YTYI65PFrDJmn0T9wZitVdSOYXYYi0WwgofMpmT99KZOnRrGrr766tzxlpaWcA7rz8cy4i655JIwdvjw4dxxlo147NixMMauK/ahMWZTRtcIK5oaXYsqOCmEkNiFSAWJXYhEkNiFSASJXYhEkNiFSIRSrbcxY8aEVg6zOyILgvV6Y3YMs11YRlmUhcQKTjIiWwjg+8F6okWFGVlBz7lz54YxZpUxuzSyN9njYtmIzG584403wliUTdnY2BjOoZlj5DGzzE1mvUV7VTTDLkLP7EIkgsQuRCJI7EIkgsQuRCJI7EIkQqmn8Q0NDWFrneiD/UCxZAZ2Qs7aJ7GkluikniVAsDp5zc1xuf277747jN17771hbMGCBbnjbD/27dsXxtavXx/Gtm7dGsYOHjyYO97d3R3OYXvFTuoZUbIOax3G3Inbb789jEXJP0CxOoXMGYquU7ZPemYXIhEkdiESQWIXIhEkdiESQWIXIhEkdiESYTjtnx4H8GkAPe5+dTb2KIDPAziU/dgj7v7cULfl7mGLH2ZNRDYDa9VUpHUOwBMWIvvqwIED4Rxmy91///1h7LOf/WwYY4/tV7/6Ve74mjVrwjkbN24MY2+++WYYY+2aooQRtnaWSFK0buChQ4dyxzdv3hzOYfX62DX3uc99Loxdc801YWz8+PG542x/mS0XMZxn9u8DWJYz/m13X5T9G1LoQoj6MqTY3f1lAH0lrEUIUUNG8579ITPrMrPHzSyuKSyEOC8oKvbvArgcwCIA+wF8M/pBM1tpZh1m1sHeDwshakshsbv7QXc/7e5nAHwPwBLys6vdvd3d21nReyFEbSkkdjMbXKvoLgCbqrMcIUStGI719mMAtwFoMrNuAF8FcJuZLQLgAHYD+MJw7uzkyZPYtWtXbozZHUePHs0dP3LkSDiHWSSsjhizNCLrjbXcWb58eRhbuHBhGPvlL38ZxiJ7DQA6Ojpyx5k9yF5xMTusSD02ZkWya6BI1hibx+b09PSEsaeffjqMsYyzVatWhTGWLRcR1adj1+KQYnf3+3KGHxv2qoQQ5wX6BJ0QiSCxC5EIErsQiSCxC5EIErsQiVBqwcmLL74Y7e3tYSwiahnEMrL2798fxljxRVbEMrJrrr322nAOs9f++Mc/hrGXXnopjLGimJHl1draGs5hNhQrAsmssgi2v6ztErMA2fqj+2OtlVjrMNa+imXSrVu3LoxdeumlueMTJ04M50QWG7Pe9MwuRCJI7EIkgsQuRCJI7EIkgsQuRCJI7EIkQqnW26RJk7B06dLcGLNWoqw3ZmuxzCWWAcZ6m+3evTt3nFlGUZYfwC1AVqBwYGAgjPX15VcQY5ltUVFGgNtaLS0tYSzKSIz2EOA2X9FilBGsSCWD7WNvb28Yi7IRAWDJkvxyENOnTw/nRHvFHpee2YVIBIldiESQ2IVIBIldiESQ2IVIhFJP44H4VJUld0SnzyxRYNKkSWHsqquuCmOzZ88OY9FJPUviYVx22WVhjCVqvPXWW2Gsv78/d5ydZrNTX1Yzjt1mlDASOSsAb73FTtxZ7bcIVmuQrYP9Xk6cOBHGOjs7w1h0XbHkJbb3EXpmFyIRJHYhEkFiFyIRJHYhEkFiFyIRJHYhEmE47Z9aAfwAQDMq7Z5Wu/t3zGwagJ8AaEOlBdQ97h73Y8qILA9WOyuqdcYsF9aaqLGxMYyxVjxREgRL7mBrZI+ZJbsw26i7uzt3nNl1RVsyscfGbKgIZqUy640lIkWJISxhhN0ei7H9YPUSn3/++dzxBQsWhHMWL16cOz7aGnSnAHzF3RcCuBHAl8xsIYCHAbzo7vMBvJh9L4Q4TxlS7O6+393XZ18fA7AFwCwAywE8kf3YEwBW1GiNQogqMKL37GbWBuAjANYCaHb3s/WaD6DyMl8IcZ4ybLGb2UQAPwPwZXc/5zOZXnkDlPsmyMxWmlmHmXVEH+UUQtSeYYndzMaiIvQfufvPs+GDZtaSxVsA5JaGcffV7t7u7u3s8+pCiNoypNitcrz3GIAt7v6tQaFnATyQff0AgGeqvzwhRLUYTtbbzQDuB7DRzDqzsUcAfB3AU2b2IIA3Adwz1A2ZGcaPH58bY9ZKZJMwe4fZIKyFD6sx1tTUlDu+bdu2cM7rr78exubMmRPGWPYd4/Dhw7njrO0Ss6FYBhirTxdZVMzKYxYgsxvZ+qMYyxpj62Aw65BldUZZb7/5zW/COW1tbbnj1I4OIxnu/jsAkXn3iaHmCyHOD/QJOiESQWIXIhEkdiESQWIXIhEkdiESodSCk2YW2hpFspBYZlstiKy3efPmhXM2bNgQxtgnCpnlFbV4AmKLh1mKzK5hVhkrtBm1J2JZY8xeY4Uei1hlbM5777034tsDENrKAM+0jH5nv/3tb8M5V1xxRe44s1j1zC5EIkjsQiSCxC5EIkjsQiSCxC5EIkjsQiRCqd7VmTNnQmuAZVAVKRrICu8xG4rNi/LxFy1aFM5ha9y7d28YO3ToUBhj9lWUeRVlwwHc5mOZXMz6jLIY2X0xC7CoZVfk2mHFPtk6mD3IbMrI3mTXR0dHR+64rDchhMQuRCpI7EIkgsQuRCJI7EIkQqmn8e4enroXOVlnc1iNMXaKzFyBKMaq5rJ2Uuzkf+fOnWHs6NGjYSx6bOwUmdX/YyfdUbILu78irZoAnpxSpA1YEScB4G4Cq23IiE7xWd26rq6u3HFWl1HP7EIkgsQuRCJI7EIkgsQuRCJI7EIkgsQuRCIMab2ZWSuAH6DSktkBrHb375jZowA+D+BsxsYj7v7cELcV1ulitktkTRSxXACelMCSIKI1MjuptbU1jLE1sv1gSTLR+ll9NGYBsj1mNmV0f2zv2T6yGCN6bMxSZPXpWIztI7uuohiz0Xbs2JE7zqzB4fjspwB8xd3Xm1kjgHVmtiaLfdvd/3UYtyGEqDPD6fW2H8D+7OtjZrYFwKxaL0wIUV1G9J7dzNoAfATA2mzoITPrMrPHzWxqtRcnhKgewxa7mU0E8DMAX3b3fgDfBXA5gEWoPPN/M5i30sw6zKyD1UkXQtSWYYndzMaiIvQfufvPAcDdD7r7aXc/A+B7AJbkzXX31e7e7u7t7DPkQojaMqTYrXLE+BiALe7+rUHjLYN+7C4Am6q/PCFEtRjOafzNAO4HsNHMOrOxRwDcZ2aLULHjdgP4wlA3dMEFF4TWSxEbjdknzBYq2mYoskhY9hdrn8RsucmTJ4cxZvH09vbmju/atSucwyxAVq+P2XlFYPYUs5SYTRlZZWwP2e2xjLii7ciia5XdV/SWmOpoqIW4++8A5O0M9dSFEOcX+gSdEIkgsQuRCBK7EIkgsQuRCBK7EIlQasFJILYGWDZRVDySWWiMohlUkQ3FrBrWjoetf8qUKWHspptuCmNRy6Dt27cXWkdjY2OheVHxRWavMVgBUbaOyPos2gKMWYDMgmXXd2T3TpgwYcTrYGvXM7sQiSCxC5EIErsQiSCxC5EIErsQiSCxC5EIpVpvZhZm8rBsswhmeTHbgtkgzBqaOHFi7jizp5gdw3q2MXuwpaUljC1ZkltWAL/+9a/DOfv27QtjRSwjIM6IY9mIrBgl+30yezNaf9EMNWZtsX50LPsxWmORYp80AzCMCCH+qpDYhUgEiV2IRJDYhUgEiV2IRJDYhUiE0rPeWPZSRGRp1CKDilllka3BLChm8bCCgiy7qq+vL4xNmzYtd/yaa64J5/T09ISxKHsN4PsYWYdsr1iM3Rezm6Jrh/VRK9qDj9mDLBbZecx6Y5ZoOGfEM4QQf5FI7EIkgsQuRCJI7EIkgsQuRCIMeRpvZhcBeBnAhdnP/9Tdv2pmcwA8CWA6gHUA7nf3IY/Ho9NMdsoZndIWPTVlJ5msNlk0jyWtzJgxI4yx9kmHDx8OY+yEfObMmbnj7e3t4ZxXXnkljLGT+igxCIiTWliyC3Mg2GNmNeii+2PXADsFZ/fFrh3mJkSJPEVaoo22Bt1JAEvd/TpU2jMvM7MbAXwDwLfdfR6AIwAeHMZtCSHqxJBi9wrvZN+Ozf45gKUAfpqNPwFgRS0WKISoDsPtz96QdXDtAbAGwA4Ab7v72dc03QBm1WSFQoiqMCyxu/tpd18E4EMAlgC4Yrh3YGYrzazDzDpYsQYhRG0Z0Wm8u78N4CUANwGYYmZnTwk+BCC3O4G7r3b3dndvZz3HhRC1ZUixm9kMM5uSfT0ewB0AtqAi+ruzH3sAwDM1WqMQogoMJxGmBcATZtaAyh+Hp9z9v83sNQBPmtk/A3gVwGPDucPIpmI2WmSTMHuN2TjMdmF2WGSfMOuNJVz09/eHsaKtoZqbm3PHZ8+eHc5pa2sLYywxiBGtv2jtN2ZDFanVxpKQ2N6zOnMsMatIIg97XMzmixhy5929C8BHcsZ3ovL+XQjxF4A+QSdEIkjsQiSCxC5EIkjsQiSCxC5EIhizr6p+Z2aHALyZfdsEoLe0O4/ROs5F6ziXv7R1XObuuamWpYr9nDs263D3OO9S69A6tI6qrkMv44VIBIldiESop9hX1/G+B6N1nIvWcS5/Neuo23t2IUS56GW8EIlQF7Gb2TIze93MtpvZw/VYQ7aO3Wa20cw6zayjxPt93Mx6zGzToLFpZrbGzN7I/p9ap3U8amZ7sz3pNLNPlbCOVjN7ycxeM7PNZvZ32Xipe0LWUeqemNlFZvYnM9uQreOfsvE5ZrY2081PzCzul5WHu5f6D0ADKmWt5gIYB2ADgIVlryNby24ATXW4348BWAxg06CxfwHwcPb1wwC+Uad1PArg70vejxYAi7OvGwFsA7Cw7D0h6yh1TwAYgInZ12MBrAVwI4CnANybjf87gFUjud16PLMvAbDd3Xd6pfT0kwCW12EddcPdXwbwwe6My1Ep3AmUVMAzWEfpuPt+d1+ffX0MleIos1DynpB1lIpXqHqR13qIfRaAPYO+r2exSgfwgpmtM7OVdVrDWZrdfX/29QEA+VUoyuEhM+vKXubX/O3EYMysDZX6CWtRxz35wDqAkvekFkVeUz+gu8XdFwP4WwBfMrOP1XtBQOUvOyp/iOrBdwFcjkqPgP0AvlnWHZvZRAA/A/Bldz+njE+Ze5KzjtL3xEdR5DWiHmLfC6B10Pdhscpa4+57s/97APwC9a28c9DMWgAg+z9uxVJD3P1gdqGdAfA9lLQnZjYWFYH9yN1/ng2Xvid566jXnmT3/TZGWOQ1oh5i/zOA+dnJ4jgA9wJ4tuxFmNkEM2s8+zWAOwFs4rNqyrOoFO4E6ljA86y4Mu5CCXtilZ5FjwHY4u7fGhQqdU+idZS9JzUr8lrWCeMHThs/hcpJ5w4A/1CnNcxFxQnYAGBzmesA8GNUXg6+j8p7rwdR6Zn3IoA3APwvgGl1WscPAWwE0IWK2FpKWMctqLxE7wLQmf37VNl7QtZR6p4AuBaVIq5dqPxh+cdB1+yfAGwH8F8ALhzJ7eoTdEIkQuoHdEIkg8QuRCJI7EIkgsQuRCJI7EIkgsQuRCJI7EIkgsQuRCL8Hw6nemXYMBvUAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "hr = np.load('./results/target.npy')\n",
    "plt.imshow(hr[0][0][:,16,:],cmap='gray')\n",
    "plt.show()\n",
    "\n",
    "hr = np.load('./results/pred.npy')\n",
    "plt.imshow(hr[0][0][:,16,:],cmap='gray')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['./HCP3D/valid\\\\109830\\\\T1w\\\\T1w_acpc_dc_restore_brain_sim036T_4x_groundtruth.nii.gz',\n",
       " './HCP3D/valid\\\\112516\\\\T1w\\\\T1w_acpc_dc_restore_brain_sim036T_4x_groundtruth.nii.gz']"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lr_files = glob.glob('./HCP3D/valid/*/T1w/lr_norm.nii.gz')\n",
    "hr_files = glob.glob('./HCP3D/valid/*/T1w/T1w_acpc_dc_restore_brain_sim036T_4x_groundtruth.nii.gz')\n",
    "hr_files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(6, 1, 32, 32, 32)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = np.load('./results/figures/conditional_iqt_0_lr.npy')\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(256, 256, 256) (256, 256, 256)\n",
      "(256, 256, 256) (256, 256, 256)\n"
     ]
    }
   ],
   "source": [
    "def cube(data):\n",
    "\n",
    "    hyp_norm = data\n",
    "\n",
    "    if len(hyp_norm.shape)>3:\n",
    "        hyp_norm = hyp_norm[:,:, 2:258, 27:283]\n",
    "    else:\n",
    "        hyp_norm = hyp_norm[2:258, 27:283]\n",
    "\n",
    "    return hyp_norm\n",
    "\n",
    "def save2patches(lr_files, hr_files, patch_size=32, save_path = './HCP3D/valid_patch/', ratio=0.7):\n",
    "    for lr_file in lr_files:\n",
    "        lr = nib.load(lr_file).get_fdata().astype(np.float32)\n",
    "        hr = nib.load(lr_file.replace('lr_norm', 'T1w_acpc_dc_restore_brain_sim036T_4x_groundtruth')).get_fdata().astype(np.float32)\n",
    "        fname = lr_file.split('\\\\')[1]\n",
    "        if not os.path.isdir(save_path+fname):\n",
    "            os.makedirs(save_path+fname)\n",
    "        lr = cube(lr)\n",
    "        hr = cube(hr)\n",
    "        print(lr.shape, hr.shape)\n",
    "\n",
    "        x,y,z = lr.shape\n",
    "        total_voxel = patch_size**3\n",
    "        for i in range(0,x-patch_size,patch_size):\n",
    "            for j in range(0,y-patch_size,patch_size):\n",
    "                for k in range(0,z-patch_size,patch_size):\n",
    "                    lr_patch = lr[i:i+patch_size, j: j+patch_size, k: k+patch_size]\n",
    "                    hr_patch = hr[i:i+patch_size, j: j+patch_size, k: k+patch_size]\n",
    "                    non_zero = np.count_nonzero(lr_patch)\n",
    "                    non_zero_proportion = (non_zero/total_voxel)\n",
    "                    if (non_zero_proportion > ratio):\n",
    "                        with open(save_path+fname+'/lr_patch_'+str(patch_size)+'_'+str(i)+'_'+str(j)+'_'+str(k)+'.npy', 'wb') as f:\n",
    "                            np.save(f, lr_patch)\n",
    "                        with open(save_path+fname+'/hr_patch_'+str(patch_size)+'_'+str(i)+'_'+str(j)+'_'+str(k)+'.npy', 'wb') as f:\n",
    "                            np.save(f, hr_patch)\n",
    "save2patches(lr_files, hr_files)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pred(images, imagen, iter):\n",
    "    lowres_images = images.cuda()  # starting un-resoluted images [1, 3, 64, 64]\n",
    "\n",
    "    pred = imagen.sample(\n",
    "        start_at_unet_number = 2,              # start at unet number 2\n",
    "        start_image_or_video = lowres_images,  # pass in low resolution images to be resoluted\n",
    "        cond_scale = 1)\n",
    "\n",
    "    pred.shape # (1, 3, 256, 256)\n",
    "    # define a transform to convert a tensor to PIL image\n",
    "    transform = T.ToPILImage()\n",
    "\n",
    "    # convert the tensor to PIL image using above transform\n",
    "    img = transform(pred[0])\n",
    "    img.save(save_path.split('.')[0]+'_iter_{}.png'.format(iter))\n",
    "    \n",
    "def evaluate(imagen):\n",
    "    # do the above for many many many many steps\n",
    "    # now you can sample an image based on the text embeddings as well as low resolution images\n",
    "\n",
    "    lowres_images = torch.randn(3, 3, 64, 64).cuda()  # starting un-resoluted images\n",
    "\n",
    "    images = imagen.sample(\n",
    "        texts = [\n",
    "            'a whale breaching from afar',\n",
    "            'young girl blowing out candles on her birthday cake',\n",
    "            'fireworks with blue and green sparkles'\n",
    "        ],\n",
    "        start_at_unet_number = 2,              # start at unet number 2\n",
    "        start_image_or_video = lowres_images,  # pass in low resolution images to be resoluted\n",
    "        cond_scale = 3.)\n",
    "\n",
    "    images.shape # (3, 3, 256, 256)\n",
    "    \n",
    "def np2tensor(x, length, mode='2d'):\n",
    "    x = torch.tensor(x)\n",
    "    if mode == '2d':\n",
    "        if length == 2:\n",
    "            x = torch.unsqueeze(x,0)\n",
    "        elif length == 3:\n",
    "            x = torch.unsqueeze(x,0)\n",
    "    else:\n",
    "        if length == 3:\n",
    "            x = torch.unsqueeze(x,0)\n",
    "        elif length == 4:\n",
    "            x = torch.unsqueeze(x,0)\n",
    "    return x\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dataset(Dataset):\n",
    "    def __init__(\n",
    "        self,\n",
    "        hr_files,\n",
    "        lr_files,\n",
    "        fake = False\n",
    "    ):\n",
    "        self.hrfiles = hr_files\n",
    "        self.lrfiles = lr_files\n",
    "        self.fake = fake\n",
    "        \n",
    "        assert len(self.hrfiles) == len(self.hrfiles), \"Length should be same\"\n",
    "    \n",
    "    def transform(self, img, size=(256,256)):\n",
    "        return TF.resize(img, size)\n",
    "        \n",
    "    def normalize(self, img):\n",
    "        img = (img-img.min())/(img.max()-img.min())\n",
    "        return img\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.hrfiles)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "\n",
    "        if not self.fake:\n",
    "            hrfile = self.hrfiles[idx]\n",
    "            lrfile = self.hrfiles[idx].replace('groundtruth_', 'lr_')\n",
    "            \n",
    "            hrimg = np.load(hrfile).astype(np.float32)\n",
    "            hrimg = np2tensor(hrimg, len(hrimg.shape))\n",
    "            hrimg = self.transform(hrimg)\n",
    "            hrimg = self.normalize(hrimg)\n",
    "            \n",
    "            lrimg = np.load(lrfile).astype(np.float32)\n",
    "            lrimg = np2tensor(lrimg, len(lrimg.shape))\n",
    "    #         lrimg = self.transform(lrimg,size=(64,64))\n",
    "            lrimg = self.transform(lrimg)\n",
    "            lrimg = self.normalize(lrimg)\n",
    "        \n",
    "        else:\n",
    "            hrimg = torch.randn(1,32,32,32)\n",
    "            lrimg = torch.randn(1,32,32,32)\n",
    "\n",
    "        return hrimg, lrimg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "146 146\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 1, 32, 32, 32])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hr_files = glob.glob('./HCP2D/train/*/groundtruth_*.npy')\n",
    "lr_files = glob.glob('./HCP2D/train/*/lr_*.npy')\n",
    "batch_size = 2\n",
    "\n",
    "print(len(hr_files), len(lr_files))\n",
    "train_dataset = Dataset(hr_files, lr_files, fake=True)\n",
    "train_loader =  DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=False)\n",
    "data = next(iter(train_loader))\n",
    "len(train_loader)\n",
    "data[0].shape\n",
    "# cnt = 10\n",
    "# for i,data in enumerate(train_loader):\n",
    "#     print(i)\n",
    "#     plt.imshow(data[0][0][0],cmap='gray')\n",
    "#     plt.show()\n",
    "#     plt.imshow(data[1][0][0],cmap='gray')\n",
    "#     plt.show()\n",
    "#     if i == cnt:\n",
    "#         break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "146 146\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 1, 32, 32, 32])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hr_files_test = glob.glob('./HCP2D/test/*/groundtruth_*.npy')\n",
    "lr_files_test = glob.glob('./HCP2D/test/*/lr_*.npy')\n",
    "batch_size = 1\n",
    "\n",
    "print(len(hr_files), len(lr_files))\n",
    "valid_dataset = Dataset(hr_files_test, lr_files_test, fake=True)\n",
    "valid_loader =  DataLoader(valid_dataset, batch_size=batch_size, shuffle=False, drop_last=False)\n",
    "data = next(iter(valid_loader))\n",
    "len(valid_loader)\n",
    "data[0].shape\n",
    "# cnt = 10\n",
    "# for i,data in enumerate(valid_loader):\n",
    "#     print(i)\n",
    "#     plt.imshow(data[0][0][0],cmap='gray')\n",
    "#     plt.show()\n",
    "#     plt.imshow(data[1][0][0],cmap='gray')\n",
    "#     plt.show()\n",
    "#     if i == cnt:\n",
    "#         break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The base dimension of your u-net should ideally be no smaller than 128, as recommended by a professional DDPM trainer https://nonint.com/2022/05/04/friends-dont-let-friends-train-small-diffusion-models/\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "SRUnet256(\n",
       "  (init_conv): CrossEmbedLayer(\n",
       "    (convs): ModuleList(\n",
       "      (0): Conv3d(1, 16, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "      (1): Conv3d(1, 8, kernel_size=(7, 7, 7), stride=(1, 1, 1), padding=(3, 3, 3))\n",
       "      (2): Conv3d(1, 8, kernel_size=(15, 15, 15), stride=(1, 1, 1), padding=(7, 7, 7))\n",
       "    )\n",
       "  )\n",
       "  (to_time_hiddens): Sequential(\n",
       "    (0): LearnedSinusoidalPosEmb()\n",
       "    (1): Linear(in_features=17, out_features=128, bias=True)\n",
       "    (2): SiLU()\n",
       "  )\n",
       "  (to_time_cond): Sequential(\n",
       "    (0): Linear(in_features=128, out_features=128, bias=True)\n",
       "  )\n",
       "  (to_time_tokens): Sequential(\n",
       "    (0): Linear(in_features=128, out_features=64, bias=True)\n",
       "    (1): Rearrange('b (r d) -> b r d', r=2)\n",
       "  )\n",
       "  (norm_cond): LayerNorm((32,), eps=1e-05, elementwise_affine=True)\n",
       "  (init_resnet_block): ResnetBlock(\n",
       "    (time_mlp): Sequential(\n",
       "      (0): SiLU()\n",
       "      (1): Linear(in_features=128, out_features=64, bias=True)\n",
       "    )\n",
       "    (block1): Block(\n",
       "      (groupnorm): GroupNorm(8, 32, eps=1e-05, affine=True)\n",
       "      (activation): SiLU()\n",
       "      (project): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "    )\n",
       "    (block2): Block(\n",
       "      (groupnorm): GroupNorm(8, 32, eps=1e-05, affine=True)\n",
       "      (activation): SiLU()\n",
       "      (project): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "    )\n",
       "    (gca): GlobalContext(\n",
       "      (to_k): Conv3d(32, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "      (net): Sequential(\n",
       "        (0): Conv3d(32, 16, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "        (1): SiLU()\n",
       "        (2): Conv3d(16, 32, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "        (3): Sigmoid()\n",
       "      )\n",
       "    )\n",
       "    (res_conv): Identity()\n",
       "  )\n",
       "  (downs): ModuleList(\n",
       "    (0): ModuleList(\n",
       "      (0): Sequential(\n",
       "        (0): Rearrange('b c (h s1) (w s2) (d s3) -> b (c s1 s2 s3) h w d', s1=2, s2=2, s3=2)\n",
       "        (1): Conv3d(256, 32, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "      )\n",
       "      (1): ResnetBlock(\n",
       "        (time_mlp): Sequential(\n",
       "          (0): SiLU()\n",
       "          (1): Linear(in_features=128, out_features=64, bias=True)\n",
       "        )\n",
       "        (block1): Block(\n",
       "          (groupnorm): GroupNorm(8, 32, eps=1e-05, affine=True)\n",
       "          (activation): SiLU()\n",
       "          (project): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "        )\n",
       "        (block2): Block(\n",
       "          (groupnorm): GroupNorm(8, 32, eps=1e-05, affine=True)\n",
       "          (activation): SiLU()\n",
       "          (project): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "        )\n",
       "        (res_conv): Identity()\n",
       "      )\n",
       "      (2): ModuleList(\n",
       "        (0): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=64, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 32, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 32, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(32, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(32, 16, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(16, 32, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Identity()\n",
       "        )\n",
       "        (1): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=64, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 32, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 32, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(32, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(32, 16, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(16, 32, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Identity()\n",
       "        )\n",
       "      )\n",
       "      (3): Identity()\n",
       "      (4): None\n",
       "    )\n",
       "    (1): ModuleList(\n",
       "      (0): Sequential(\n",
       "        (0): Rearrange('b c (h s1) (w s2) (d s3) -> b (c s1 s2 s3) h w d', s1=2, s2=2, s3=2)\n",
       "        (1): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "      )\n",
       "      (1): ResnetBlock(\n",
       "        (time_mlp): Sequential(\n",
       "          (0): SiLU()\n",
       "          (1): Linear(in_features=128, out_features=256, bias=True)\n",
       "        )\n",
       "        (block1): Block(\n",
       "          (groupnorm): GroupNorm(8, 128, eps=1e-05, affine=True)\n",
       "          (activation): SiLU()\n",
       "          (project): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "        )\n",
       "        (block2): Block(\n",
       "          (groupnorm): GroupNorm(8, 128, eps=1e-05, affine=True)\n",
       "          (activation): SiLU()\n",
       "          (project): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "        )\n",
       "        (res_conv): Identity()\n",
       "      )\n",
       "      (2): ModuleList(\n",
       "        (0): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=256, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 128, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 128, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(128, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(128, 64, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(64, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Identity()\n",
       "        )\n",
       "        (1): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=256, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 128, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 128, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(128, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(128, 64, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(64, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Identity()\n",
       "        )\n",
       "        (2): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=256, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 128, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 128, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(128, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(128, 64, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(64, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Identity()\n",
       "        )\n",
       "        (3): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=256, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 128, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 128, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(128, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(128, 64, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(64, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Identity()\n",
       "        )\n",
       "      )\n",
       "      (3): Identity()\n",
       "      (4): None\n",
       "    )\n",
       "    (2): ModuleList(\n",
       "      (0): Sequential(\n",
       "        (0): Rearrange('b c (h s1) (w s2) (d s3) -> b (c s1 s2 s3) h w d', s1=2, s2=2, s3=2)\n",
       "        (1): Conv3d(1024, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "      )\n",
       "      (1): ResnetBlock(\n",
       "        (time_mlp): Sequential(\n",
       "          (0): SiLU()\n",
       "          (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "        )\n",
       "        (block1): Block(\n",
       "          (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "          (activation): SiLU()\n",
       "          (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "        )\n",
       "        (block2): Block(\n",
       "          (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "          (activation): SiLU()\n",
       "          (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "        )\n",
       "        (res_conv): Identity()\n",
       "      )\n",
       "      (2): ModuleList(\n",
       "        (0): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(256, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Identity()\n",
       "        )\n",
       "        (1): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(256, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Identity()\n",
       "        )\n",
       "        (2): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(256, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Identity()\n",
       "        )\n",
       "        (3): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(256, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Identity()\n",
       "        )\n",
       "        (4): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(256, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Identity()\n",
       "        )\n",
       "        (5): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(256, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Identity()\n",
       "        )\n",
       "        (6): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(256, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Identity()\n",
       "        )\n",
       "        (7): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(256, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Identity()\n",
       "        )\n",
       "      )\n",
       "      (3): TransformerBlock(\n",
       "        (layers): ModuleList(\n",
       "          (0): ModuleList(\n",
       "            (0): Attention(\n",
       "              (norm): LayerNorm()\n",
       "              (to_q): Linear(in_features=256, out_features=512, bias=False)\n",
       "              (to_kv): Linear(in_features=256, out_features=128, bias=False)\n",
       "              (to_context): Sequential(\n",
       "                (0): LayerNorm((32,), eps=1e-05, elementwise_affine=True)\n",
       "                (1): Linear(in_features=32, out_features=128, bias=True)\n",
       "              )\n",
       "              (to_out): Sequential(\n",
       "                (0): Linear(in_features=512, out_features=256, bias=False)\n",
       "                (1): LayerNorm()\n",
       "              )\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): LayerNorm()\n",
       "              (1): Linear(in_features=256, out_features=512, bias=False)\n",
       "              (2): GELU(approximate=none)\n",
       "              (3): LayerNorm()\n",
       "              (4): Linear(in_features=512, out_features=256, bias=False)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (4): None\n",
       "    )\n",
       "  )\n",
       "  (ups): ModuleList(\n",
       "    (0): ModuleList(\n",
       "      (0): ResnetBlock(\n",
       "        (time_mlp): Sequential(\n",
       "          (0): SiLU()\n",
       "          (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "        )\n",
       "        (block1): Block(\n",
       "          (groupnorm): GroupNorm(8, 512, eps=1e-05, affine=True)\n",
       "          (activation): SiLU()\n",
       "          (project): Conv3d(512, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "        )\n",
       "        (block2): Block(\n",
       "          (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "          (activation): SiLU()\n",
       "          (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "        )\n",
       "        (res_conv): Conv3d(512, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "      )\n",
       "      (1): ModuleList(\n",
       "        (0): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 512, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(512, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(256, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Conv3d(512, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "        )\n",
       "        (1): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 512, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(512, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(256, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Conv3d(512, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "        )\n",
       "        (2): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 512, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(512, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(256, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Conv3d(512, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "        )\n",
       "        (3): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 512, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(512, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(256, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Conv3d(512, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "        )\n",
       "        (4): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 512, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(512, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(256, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Conv3d(512, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "        )\n",
       "        (5): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 512, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(512, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(256, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Conv3d(512, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "        )\n",
       "        (6): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 512, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(512, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(256, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Conv3d(512, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "        )\n",
       "        (7): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 512, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(512, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(256, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Conv3d(512, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "        )\n",
       "      )\n",
       "      (2): TransformerBlock(\n",
       "        (layers): ModuleList(\n",
       "          (0): ModuleList(\n",
       "            (0): Attention(\n",
       "              (norm): LayerNorm()\n",
       "              (to_q): Linear(in_features=256, out_features=512, bias=False)\n",
       "              (to_kv): Linear(in_features=256, out_features=128, bias=False)\n",
       "              (to_context): Sequential(\n",
       "                (0): LayerNorm((32,), eps=1e-05, elementwise_affine=True)\n",
       "                (1): Linear(in_features=32, out_features=128, bias=True)\n",
       "              )\n",
       "              (to_out): Sequential(\n",
       "                (0): Linear(in_features=512, out_features=256, bias=False)\n",
       "                (1): LayerNorm()\n",
       "              )\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): LayerNorm()\n",
       "              (1): Linear(in_features=256, out_features=512, bias=False)\n",
       "              (2): GELU(approximate=none)\n",
       "              (3): LayerNorm()\n",
       "              (4): Linear(in_features=512, out_features=256, bias=False)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (3): PixelShuffleUpsample(\n",
       "        (net): Sequential(\n",
       "          (0): Conv3d(256, 1024, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "          (1): SiLU()\n",
       "          (2): PixelShuffle3D()\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (1): ModuleList(\n",
       "      (0): ResnetBlock(\n",
       "        (time_mlp): Sequential(\n",
       "          (0): SiLU()\n",
       "          (1): Linear(in_features=128, out_features=256, bias=True)\n",
       "        )\n",
       "        (block1): Block(\n",
       "          (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "          (activation): SiLU()\n",
       "          (project): Conv3d(256, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "        )\n",
       "        (block2): Block(\n",
       "          (groupnorm): GroupNorm(8, 128, eps=1e-05, affine=True)\n",
       "          (activation): SiLU()\n",
       "          (project): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "        )\n",
       "        (res_conv): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "      )\n",
       "      (1): ModuleList(\n",
       "        (0): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=256, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 128, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(128, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(128, 64, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(64, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "        )\n",
       "        (1): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=256, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 128, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(128, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(128, 64, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(64, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "        )\n",
       "        (2): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=256, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 128, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(128, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(128, 64, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(64, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "        )\n",
       "        (3): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=256, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(256, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 128, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(128, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(128, 64, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(64, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Conv3d(256, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "        )\n",
       "      )\n",
       "      (2): Identity()\n",
       "      (3): PixelShuffleUpsample(\n",
       "        (net): Sequential(\n",
       "          (0): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "          (1): SiLU()\n",
       "          (2): PixelShuffle3D()\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (2): ModuleList(\n",
       "      (0): ResnetBlock(\n",
       "        (time_mlp): Sequential(\n",
       "          (0): SiLU()\n",
       "          (1): Linear(in_features=128, out_features=64, bias=True)\n",
       "        )\n",
       "        (block1): Block(\n",
       "          (groupnorm): GroupNorm(8, 64, eps=1e-05, affine=True)\n",
       "          (activation): SiLU()\n",
       "          (project): Conv3d(64, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "        )\n",
       "        (block2): Block(\n",
       "          (groupnorm): GroupNorm(8, 32, eps=1e-05, affine=True)\n",
       "          (activation): SiLU()\n",
       "          (project): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "        )\n",
       "        (res_conv): Conv3d(64, 32, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "      )\n",
       "      (1): ModuleList(\n",
       "        (0): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=64, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 64, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(64, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 32, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(32, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(32, 16, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(16, 32, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Conv3d(64, 32, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "        )\n",
       "        (1): ResnetBlock(\n",
       "          (time_mlp): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=128, out_features=64, bias=True)\n",
       "          )\n",
       "          (block1): Block(\n",
       "            (groupnorm): GroupNorm(8, 64, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(64, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (block2): Block(\n",
       "            (groupnorm): GroupNorm(8, 32, eps=1e-05, affine=True)\n",
       "            (activation): SiLU()\n",
       "            (project): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "          )\n",
       "          (gca): GlobalContext(\n",
       "            (to_k): Conv3d(32, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "            (net): Sequential(\n",
       "              (0): Conv3d(32, 16, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (1): SiLU()\n",
       "              (2): Conv3d(16, 32, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "              (3): Sigmoid()\n",
       "            )\n",
       "          )\n",
       "          (res_conv): Conv3d(64, 32, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "        )\n",
       "      )\n",
       "      (2): Identity()\n",
       "      (3): PixelShuffleUpsample(\n",
       "        (net): Sequential(\n",
       "          (0): Conv3d(32, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "          (1): SiLU()\n",
       "          (2): PixelShuffle3D()\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (mid_block1): ResnetBlock(\n",
       "    (time_mlp): Sequential(\n",
       "      (0): SiLU()\n",
       "      (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "    )\n",
       "    (block1): Block(\n",
       "      (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "      (activation): SiLU()\n",
       "      (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "    )\n",
       "    (block2): Block(\n",
       "      (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "      (activation): SiLU()\n",
       "      (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "    )\n",
       "    (res_conv): Identity()\n",
       "  )\n",
       "  (mid_attn): TransformerBlock(\n",
       "    (layers): ModuleList(\n",
       "      (0): ModuleList(\n",
       "        (0): Attention(\n",
       "          (norm): LayerNorm()\n",
       "          (to_q): Linear(in_features=256, out_features=512, bias=False)\n",
       "          (to_kv): Linear(in_features=256, out_features=128, bias=False)\n",
       "          (to_out): Sequential(\n",
       "            (0): Linear(in_features=512, out_features=256, bias=False)\n",
       "            (1): LayerNorm()\n",
       "          )\n",
       "        )\n",
       "        (1): Sequential(\n",
       "          (0): LayerNorm()\n",
       "          (1): Linear(in_features=256, out_features=512, bias=False)\n",
       "          (2): GELU(approximate=none)\n",
       "          (3): LayerNorm()\n",
       "          (4): Linear(in_features=512, out_features=256, bias=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (mid_block2): ResnetBlock(\n",
       "    (time_mlp): Sequential(\n",
       "      (0): SiLU()\n",
       "      (1): Linear(in_features=128, out_features=512, bias=True)\n",
       "    )\n",
       "    (block1): Block(\n",
       "      (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "      (activation): SiLU()\n",
       "      (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "    )\n",
       "    (block2): Block(\n",
       "      (groupnorm): GroupNorm(8, 256, eps=1e-05, affine=True)\n",
       "      (activation): SiLU()\n",
       "      (project): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "    )\n",
       "    (res_conv): Identity()\n",
       "  )\n",
       "  (upsample_combiner): UpsampleCombiner()\n",
       "  (final_res_block): ResnetBlock(\n",
       "    (time_mlp): Sequential(\n",
       "      (0): SiLU()\n",
       "      (1): Linear(in_features=128, out_features=64, bias=True)\n",
       "    )\n",
       "    (block1): Block(\n",
       "      (groupnorm): GroupNorm(8, 32, eps=1e-05, affine=True)\n",
       "      (activation): SiLU()\n",
       "      (project): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "    )\n",
       "    (block2): Block(\n",
       "      (groupnorm): GroupNorm(8, 32, eps=1e-05, affine=True)\n",
       "      (activation): SiLU()\n",
       "      (project): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       "    )\n",
       "    (gca): GlobalContext(\n",
       "      (to_k): Conv3d(32, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "      (net): Sequential(\n",
       "        (0): Conv3d(32, 16, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "        (1): SiLU()\n",
       "        (2): Conv3d(16, 32, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n",
       "        (3): Sigmoid()\n",
       "      )\n",
       "    )\n",
       "    (res_conv): Identity()\n",
       "  )\n",
       "  (final_conv): Conv3d(32, 1, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))\n",
       ")"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# unet for imagen\n",
    "unet1 = NullUnet()\n",
    "\n",
    "unet2 = SRUnet256(\n",
    "    dim = 32,\n",
    "    dim_mults = (1, 4, 8), #(1, 2, 4, 8),\n",
    "    channels=1,\n",
    "    num_resnet_blocks = (2, 4, 8),#(2, 4, 8, 8),\n",
    "    layer_attns = (False, False, True),\n",
    "    layer_cross_attns = False,\n",
    "    cond_on_text =False\n",
    ")\n",
    "\n",
    "imagen = Imagen(\n",
    "    condition_on_text = False, \n",
    "    unets = (unet1, unet2),\n",
    "    image_sizes = (8, 32),\n",
    "    channels=1,\n",
    "    timesteps = 500,\n",
    "    cond_drop_prob = 0.0\n",
    ").to(device)\n",
    "\n",
    "trainer = ImagenTrainer(\n",
    "    imagen = imagen,\n",
    "    #cosine_decay_max_steps = len(train_loader)*10,\n",
    "    split_valid_from_train = False # whether to split the validation dataset from the training\n",
    ")\n",
    "\n",
    "trainer.add_train_dataset(train_dataset, batch_size = batch_size)\n",
    "trainer.add_valid_dataset(valid_dataset, batch_size = 1)\n",
    "\n",
    "unet2.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "unet 2 has not been trained\n",
      "when sampling, you can pass stop_at_unet_number to stop early in the cascade, so it does not try to generate with untrained unets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2it [00:22, 11.09s/it]\n"
     ]
    }
   ],
   "source": [
    "test_img = torch.randn(1,1,32,32,32)\n",
    "outputs = trainer.sample(batch_size = 1, return_all_unet_outputs = True, return_pil_images = False, start_image_or_video = test_img.to(device), start_at_unet_number = 2) # returns List[Image]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2it [00:25, 12.69s/it]\n",
      "2it [00:24, 12.28s/it]\n",
      "2it [00:24, 12.42s/it]\n",
      "2it [00:26, 13.46s/it]\n",
      "2it [00:25, 12.78s/it]\n",
      "2it [00:26, 13.29s/it]\n",
      "2it [00:26, 13.02s/it]\n",
      "2it [00:25, 12.91s/it]\n",
      "2it [00:25, 12.88s/it]\n",
      "2it [00:25, 12.54s/it]\n",
      "2it [00:25, 12.60s/it]\n",
      "2it [00:25, 12.62s/it]\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "too many values to unpack (expected 2)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[1;32m~\\AppData\\Local\\Temp/ipykernel_7112/2720244002.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      6\u001b[0m     \u001b[0mtrainer\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mupdate\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0munet_number\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;36m2\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      7\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 8\u001b[1;33m \u001b[0mvalid_loss\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mpreds\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtrainer\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalid_step\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0munet_number\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;36m2\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mmax_batch_size\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;36m2\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      9\u001b[0m \u001b[0mvalid_loss\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmean\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mvalid_loss\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     10\u001b[0m \u001b[0mtrain_ls\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmean\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0marray\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mtrain_ls\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mValueError\u001b[0m: too many values to unpack (expected 2)"
     ]
    }
   ],
   "source": [
    "train_ls = []\n",
    "for i in range(2):\n",
    "    print(i)\n",
    "    loss = trainer.train_step(unet_number = 2, max_batch_size = 2)\n",
    "    train_ls.append(loss)\n",
    "    trainer.update(unet_number = 2)\n",
    "    \n",
    "valid_loss, preds, data = trainer.valid_step(unet_number = 2, max_batch_size = 2)\n",
    "valid_loss = np.mean(valid_loss)\n",
    "train_ls = np.mean(np.array(train_ls))\n",
    "\n",
    "print(\"Loss: {} {}\".format(train_ls, valid_loss))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# working training loop\n",
    "train_ls = []\n",
    "valid_ls = []\n",
    "lst_best = 0\n",
    "best = 100\n",
    "for i in range(20000):\n",
    "    loss = trainer.train_step(unet_number = 2, max_batch_size = 2)\n",
    "    train_ls.append(loss)\n",
    "    train_loss_save = pd.DataFrame({'loss': train_ls}).to_csv(os.getcwd()+('/train_loss.csv'), index=False)\n",
    "    trainer.update(unet_number = 2)\n",
    "    if not (i % 50):\n",
    "        print(f'unet: 2, Epoch: {i}, loss: {loss}')\n",
    "        valid_loss, preds = trainer.valid_step(unet_number = 2, max_batch_size = 2)\n",
    "        valid_loss = np.mean(valid_loss)\n",
    "        valid_ls.append(valid_loss)\n",
    "        valid_loss_save = pd.DataFrame({'loss': valid_ls}).to_csv(os.getcwd()+('/valid_loss.csv'), index=False)\n",
    "        print(f'valid loss: {valid_loss}')\n",
    "        if best > valid_loss:\n",
    "            print(\"Best model!\")\n",
    "            best = valid_loss\n",
    "            for j in range(len(preds)):\n",
    "                preds[j][0][0].save('./results/figures/'+f'conditional_iqt_{j}_pred.png')\n",
    "            trainer.save('./results/model/checkpoint.pt')\n",
    "trainer.save('./results/model/last_checkpoint.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(train_ls, label='train')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(valid_ls, label='valid')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.load('./results/model/checkpoint.pt');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images, outputs, lst = trainer.sample(batch_size = 1, return_all_unet_outputs = True, return_pil_images = True, start_image_or_video = test_img.to(device), start_at_unet_number = 2) # returns List[Image]\n",
    "ls = F.l1_loss(test_img2, outputs[0].cpu())\n",
    "print(\"Test loss: \",ls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lst2 = []\n",
    "for i in range(len(lst)):\n",
    "    if (i % 20 == 0) or (i == len(lst)-1):\n",
    "        lst2.append(lst[i][0,0].cpu().numpy())\n",
    "        \n",
    "plt.imshow(test_img[0,0].cpu().numpy(), cmap='gray')\n",
    "plt.title(\"Input\".format(i))\n",
    "plt.show()\n",
    "plt.imshow(outputs[-1][0,0].cpu().numpy(), cmap='gray')\n",
    "plt.title(\"Time Step: {}\".format(i))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# idx = torch.where(test_img2==0)\n",
    "# img = lst[-1].cpu()+1\n",
    "# img[idx] = 0\n",
    "# img.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "# ims is a list of lists, each row is a list of artists to draw in the\n",
    "# current frame; here we are just animating one artist, the image, in\n",
    "# each frame\n",
    "ims = []\n",
    "for i in range(len(lst2)):\n",
    "    im = ax.imshow(lst2[i],cmap='gray', animated=True)\n",
    "    title = ax.text(70, -10, f'Timestep: {i*20}', fontsize=15, bbox={'facecolor': 'blue',\n",
    "                                       'alpha': 0.5, 'pad': 5})\n",
    "    if i == 0:\n",
    "        ax.imshow(lst2[i],cmap='gray')  # show an initial one first\n",
    "    ims.append([im,title])\n",
    "\n",
    "ani = animation.ArtistAnimation(fig, ims, interval=100, blit=True,\n",
    "                                repeat_delay=5000)\n",
    "\n",
    "# To save the animation, use e.g.\n",
    "#\n",
    "ani.save(\"movie.mp4\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(test_img[0][0],cmap='gray')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "iqt",
   "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.8.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "0315dcab0c71d3b3d40117005573438bfb357873381d8be12f58ac795438fa13"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}