{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "import os\n",
    "import cv2\n",
    "import math\n",
    "import base64\n",
    "import random\n",
    "import numpy as np\n",
    "from PIL import Image \n",
    "from tqdm import tqdm\n",
    "from datetime import datetime\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.cuda import amp\n",
    "import torch.nn.functional as F\n",
    "from torch.optim import Adam, AdamW\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "import torchvision\n",
    "import torchvision.transforms as TF\n",
    "import torchvision.datasets as datasets\n",
    "from torchvision.utils import make_grid\n",
    "\n",
    "from torchmetrics import MeanMetric\n",
    "\n",
    "from IPython.display import display, HTML, clear_output\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_device(data, device):\n",
    "    \"\"\"将张量移动到选择的设备\"\"\"\n",
    "    \"\"\"Move tensor(s) to chosen device\"\"\"\n",
    "    if isinstance(data, (list, tuple)):\n",
    "        return [to_device(x, device) for x in data]\n",
    "    return data.to(device, non_blocking=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DeviceDataLoader:\n",
    "    \"\"\"包装一个数据加载器，来把数据移动到另一个设备上\"\"\"\n",
    "    \"\"\"Wrap a dataloader to move data to a device\"\"\"\n",
    "\n",
    "    def __init__(self, dl, device):\n",
    "        self.dl = dl\n",
    "        self.device = device\n",
    "\n",
    "    def __iter__(self):\n",
    "        \"\"\"在移动到设备后生成一个批次的数据\"\"\"\n",
    "        \"\"\"Yield a batch of data after moving it to device\"\"\"\n",
    "        for b in self.dl:\n",
    "            yield to_device(b, self.device)\n",
    "\n",
    "    def __len__(self):\n",
    "        \"\"\"批次的数量\"\"\"\n",
    "        \"\"\"Number of batches\"\"\"\n",
    "        return len(self.dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_default_device():\n",
    "    return torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_images(images, path, **kwargs):\n",
    "    grid = make_grid(images, **kwargs)\n",
    "    ndarr = grid.permute(1,2,0).to(\"cpu\").numpy()\n",
    "    im = Image.fromarray(ndarr)\n",
    "    im.save(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get(element: torch.Tensor, t: torch.Tensor):\n",
    "    \"\"\"\n",
    "    Get value at index position \"t\" in \"element\" and \n",
    "        reshape it to have the same dimension as a batch of images\n",
    "\n",
    "    获得在\"element\"中位置\"t\"并且reshape，以和一组照片有相同的维度\n",
    "    \"\"\"\n",
    "    ele = element.gather(-1, t)\n",
    "    return ele.reshape(-1, 1, 1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "element = torch.tensor([[1,2,3,4,5],\n",
    "                        [2,3,4,5,6],\n",
    "                        [3,4,5,6,7]])\n",
    "t = torch.tensor([1,2,0]).unsqueeze(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1, 2, 3, 4, 5],\n",
      "        [2, 3, 4, 5, 6],\n",
      "        [3, 4, 5, 6, 7]])\n",
      "tensor([[1],\n",
      "        [2],\n",
      "        [0]])\n"
     ]
    }
   ],
   "source": [
    "print(element)\n",
    "print(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[[2]]],\n",
      "\n",
      "\n",
      "        [[[4]]],\n",
      "\n",
      "\n",
      "        [[[3]]]])\n"
     ]
    }
   ],
   "source": [
    "extracted_scores = get(element, t)\n",
    "print(extracted_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup_log_directory(config):\n",
    "    \"\"\"\n",
    "    Log and Model checkpoint directory Setup\n",
    "    记录并且建模目录准备\n",
    "    \"\"\"\n",
    "\n",
    "    if os.path.isdir(config.root_log_dir):\n",
    "        # Get all folders numbers in the root_log_dir\n",
    "        # 在root_log_dir下获得所有文件夹数目\n",
    "        folder_numbers = [int(folder.replace(\"version_\", \"\")) for folder in os.listdir(config.root_log_dir)]\n",
    "\n",
    "        # Find the latest version number present in the log_dir\n",
    "        # 找到在log_dir下的最新版本数字\n",
    "        last_version_number = max(folder_numbers)\n",
    "\n",
    "        # New version name\n",
    "        version_name = f\"version{last_version_number + 1}\"\n",
    "\n",
    "    else:\n",
    "        version_name = config.log_dir\n",
    "\n",
    "    # Update the training config default directory\n",
    "    # 更新训练config默认目录\n",
    "    log_dir         = os.path.join(config.root_log_dir,         version_name)\n",
    "    checkpoint_dir  = os.path.join(config.root_checkpoint_dir,  version_name) \n",
    "\n",
    "    # Create new directory for saving new experiment version\n",
    "    # 创建一个新目录来保存新的实验版本\n",
    "    os.makedirs(log_dir,        exist_ok=True)\n",
    "    os.makedirs(checkpoint_dir, exist_ok=True)\n",
    "\n",
    "    print(f\"Logging at: {log_dir}\")\n",
    "    print(f\"Model Checkpoint at: {checkpoint_dir}\")\n",
    "\n",
    "    return log_dir, checkpoint_dir\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "def frames2vid(images, save_path):\n",
    "\n",
    "    WIDTH = images[0].shape[1]\n",
    "    HEIGHT = images[0].shape[0]\n",
    "\n",
    "    # fourcc = cv2.VideoWriter_fourcc(*'XVID')\n",
    "    fourcc = cv2.VideoWriter_fourcc(*'mp4v')\n",
    "    video = cv2.VideoWriter(save_path, fourcc, 25, (WIDTH, HEIGHT))\n",
    "\n",
    "    # Appending the images to the video one by one\n",
    "    # 一个接一个的将照片追加到视频\n",
    "    for image in images:\n",
    "        video.write(image)\n",
    "    \n",
    "    # Deallocating memories taken for window creation\n",
    "    # 释放创建window占用的内存\n",
    "        \n",
    "    video.release()\n",
    "    return\n",
    "\n",
    "def display_gif(gif_path):\n",
    "    b64 = base64.b64encode(open(gif_path,'rb').read()).decode('ascii')\n",
    "    display(HTML(f'<img src=\"data:image/gif;base64,{b64}\" />'))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "\n",
    "@dataclass\n",
    "class BaseConfig:\n",
    "    DEVICE = get_default_device()\n",
    "    DATASET = \"Flowers\" #MNIST \"cifar-10\" \"Flowers\"\n",
    "\n",
    "    # 记录推断日志信息并保存存档点\n",
    "    root_log_dir = os.path.join(\"Logs_Checkpoints\", \"Inference\")\n",
    "    root_checkpoint_dir = os.path.join(\"Logs_Checkpoints\",\"checkpoints\")\n",
    "\n",
    "    #目前的日志和存档点目录\n",
    "    log_dir = \"version_0\"\n",
    "    checkpoint_dir = \"version_0\"\n",
    "\n",
    "@dataclass\n",
    "class TrainingConfig:\n",
    "    TIMESTEPS = 1000\n",
    "    IMG_SHAPE = (1,32,32) if BaseConfig.DATASET == \"MNIST\" else (3,32,32)\n",
    "    NUM_EPOCHS = 800\n",
    "    BATCH_SIZE = 32\n",
    "    LR = 2e-4\n",
    "    NUM_WORKERS = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Dataset & Build Dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dataset(dataset_name='MNIST'):\n",
    "    \"\"\"\n",
    "    Returns the dataset class object that will be passed to the Dataloader\n",
    "    Three preprocessing transforms,\n",
    "    and one augmentation are applied to every image in the dataset\n",
    "    返回数据集的类对象\n",
    "    这个类对象将会被传递给DataLoader\n",
    "    数据集中的每个图象将会应用三个预处理转换\n",
    "    和一个增强\n",
    "    \"\"\"\n",
    "    transforms = TF.Compose(\n",
    "        [\n",
    "            TF.ToTensor(),\n",
    "            TF.Resize((32,32),\n",
    "                      interpolation=TF.InterpolationMode.BICUBIC,\n",
    "                      antialias=True),\n",
    "            TF.RandomHorizontalFlip(),\n",
    "            TF.Lambda(lambda t: (t * 2) - 1) # scale between [-1, 1]\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    if dataset_name.upper() == \"MNIST\":\n",
    "        dataset = datasets.MNIST(root=\"data\", train=True, download=True, transform=transforms)\n",
    "    elif dataset_name == \"Cifar-10\":\n",
    "        dataset = datasets.CIFAR10(root=\"data\", train=True,download=True, transform=transforms)\n",
    "    elif dataset_name == \"Cifar-100\":\n",
    "        dataset = datasets.CIFAR10(root=\"data\", train=True,download=True, transform=transforms)\n",
    "    elif dataset_name == \"Flowers\":\n",
    "        dataset = datasets.ImageFolder(root=\"data/flowers\", transform=transforms)\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dataloader(dataset_name='MNIST',\n",
    "                   batch_size=32,\n",
    "                   pin_memory=False,\n",
    "                   shuffle=True,\n",
    "                   num_workers=0,\n",
    "                   device=\"cpu\"\n",
    "                   ):\n",
    "    dataset     = get_dataset(dataset_name=dataset_name)\n",
    "    dataLoader  = DataLoader(dataset, batch_size=batch_size,\n",
    "                             pin_memory=pin_memory,\n",
    "                             num_workers=num_workers,\n",
    "                             shuffle=shuffle\n",
    "                             )\n",
    "    device_dataloader = DeviceDataLoader(dataLoader, device)\n",
    "    return device_dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [],
   "source": [
    "def inverse_transform(tensors):\n",
    "    \"\"\"\n",
    "    Convert tensors from [-1., 1.] to [0., 255.]\n",
    "    \"\"\"\n",
    "    return ((tensors.clamp(-1, 1) + 1.0) / 2.0) * 255.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset ImageFolder\n",
      "    Number of datapoints: 4317\n",
      "    Root location: data/flowers\n",
      "    StandardTransform\n",
      "Transform: Compose(\n",
      "               ToTensor()\n",
      "               Resize(size=(32, 32), interpolation=bicubic, max_size=None, antialias=True)\n",
      "               RandomHorizontalFlip(p=0.5)\n",
      "               Lambda()\n",
      "           )\n"
     ]
    }
   ],
   "source": [
    "dataset = get_dataset(dataset_name='Flowers')\n",
    "print(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [],
   "source": [
    "loader = get_dataloader(\n",
    "    dataset_name=BaseConfig.DATASET,\n",
    "    batch_size=128,\n",
    "    device='cpu'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1155728 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(12, 6), facecolor='white')\n",
    "\n",
    "for b_image, _ in loader:\n",
    "    b_image = inverse_transform(b_image).cpu()\n",
    "    grid_img = make_grid(b_image / 255.0, \n",
    "                         nrow = 16, \n",
    "                         padding=True,\n",
    "                         pad_value=1,\n",
    "                         normalize=True \n",
    "                        )\n",
    "    plt.imshow(grid_img.permute(1, 2, 0))\n",
    "    plt.axis(\"off\")\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## UNet Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Architecture"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Architecture\n",
    "\n",
    "Comprises 5 components\n",
    "\n",
    "1. Encoder\n",
    "2. Bottleneck\n",
    "3. Decoder\n",
    "4. Self-attention\n",
    "5. Sinusoidal time embeddings\n",
    "\n",
    "### Details\n",
    "\n",
    "1. There are four levels in the encoder and decoder path with bottleneck blocks between them\n",
    "2. Each encoder stage comprises two residual blocks with convolutional downsampling except the last level\n",
    "3. Each corresponding decoder stage comprises three residual blocks and uses 2x nearest neighbors with convolutions to upsample the input from the previous level.\n",
    "4. Each stage in the encoder path is connected to the decoder path with the help of skip connections\n",
    "5. The model uses \"Self-Attention\" modules at a single feature map resolution\n",
    "6. Every residual block in the model gets the inputs from the previous layer (and others in the decoder path) and the embedding of the current timestep. The timestep embedding informs the model of the input's current position in the Markov chain.\n",
    "\n",
    "![UNet-architecture](https://learnopencv.com/wp-content/uploads/2023/02/denoising-diffusion-probabilistic-models_UNet_model_architecture-768x307.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SinusoidalPositionEmbeddings(nn.Module):\n",
    "    def __init__(self, total_time_steps=1000, time_emb_dims=128, time_emb_dims_exp=512):\n",
    "        super().__init__()\n",
    "\n",
    "        half_dim = time_emb_dims // 2 # half_dim=64\n",
    "\n",
    "        emb = math.log(10000) / (half_dim - 1) # log_e(10000) = 9.21 / 63 =0.14619 000\n",
    "        emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)\n",
    "\n",
    "        ts = torch.arange(total_time_steps, dtype=torch.float32) \n",
    "\n",
    "        emb = torch.unsqueeze(ts, dim=-1) * torch.unsqueeze(emb, dim=0)\n",
    "        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)\n",
    "\n",
    "        self.time_blocks = nn.Sequential(\n",
    "            nn.Embedding.from_pretrained(emb),\n",
    "            nn.Linear(in_features=time_emb_dims, out_features=time_emb_dims_exp),\n",
    "            nn.SiLU(),\n",
    "            nn.Linear(in_features=time_emb_dims_exp, out_features=time_emb_dims_exp),\n",
    "        )\n",
    "\n",
    "    def forward(self, time):\n",
    "        return self.time_blocks(time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttentionBlock(nn.Module):\n",
    "    def __init__(self, channels=64):\n",
    "        super().__init__()\n",
    "        self.channels = channels\n",
    "        self.group_norm = nn.GroupNorm(num_groups=8, num_channels=channels)\n",
    "        self.mhsa = nn.MultiheadAttention(embed_dim=self.channels, num_heads=4, batch_first=True)\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, _, H, W = x.shape\n",
    "        h = self.group_norm(x)\n",
    "        h = h.reshape(B, self.channels, H * W).swapaxes(1, 2)\n",
    "        h, _ = self.mhsa(h,h,h)\n",
    "        h = h.swapaxes(2, 1).view(B, self.channels, H, W)\n",
    "        return x + h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResnetBlock(nn.Module):\n",
    "    def __init__(self, *, in_channels, out_channels, dropout_rate=0.1, time_emb_dims=512, apply_attention=False):\n",
    "        super().__init__()\n",
    "        self.in_channels = in_channels\n",
    "        self.out_channels = out_channels\n",
    "\n",
    "        self.act_fn = nn.SiLU()\n",
    "\n",
    "        # Group1\n",
    "        self.normlize_1 = nn.GroupNorm(num_groups=8, num_channels=self.in_channels)\n",
    "        self.conv_1 = nn.Conv2d(in_channels=self.in_channels, out_channels=self.out_channels, kernel_size=3, stride=1, padding=\"same\")\n",
    "\n",
    "        # Group2 time embedding\n",
    "        self.dense_1 = nn.Linear(in_features=time_emb_dims, out_features=self.out_channels)\n",
    "\n",
    "        # Group3 \n",
    "        self.normlize_2 = nn.GroupNorm(num_groups=8, num_channels=out_channels)\n",
    "        self.dropout = nn.Dropout2d(p=dropout_rate)\n",
    "        self.conv_2 = nn.Conv2d(in_channels=self.out_channels, out_channels=self.out_channels, kernel_size=3, stride=1, padding=\"same\")\n",
    "\n",
    "        if self.in_channels != self.out_channels:\n",
    "            self.match_input = nn.Conv2d(in_channels=self.in_channels,out_channels=self.out_channels, kernel_size=1, stride=1)\n",
    "        else:\n",
    "            self.attention = nn.Identity()\n",
    "        \n",
    "        if apply_attention:\n",
    "            self.attention = AttentionBlock(channels=self.out_channels)\n",
    "        else:\n",
    "            self.attention = nn.Identity()\n",
    "    def forward(self, x, t):\n",
    "        #group 1\n",
    "        h = self.act_fn(self.normlize_1(x))\n",
    "        h = self.conv_1(h)\n",
    "\n",
    "        #group 2\n",
    "        # add in timestep embedding\n",
    "        h += self.dense_1(self.act_fn(t))[:, :, None, None]\n",
    "\n",
    "        #group 3\n",
    "        h = self.act_fn(self.normlize_2(h))\n",
    "        h = self.dropout(h)\n",
    "        h = self.conv_2(h)\n",
    "\n",
    "        # Residual and attention\n",
    "        h = h + self.match_input(x)\n",
    "        h = self.attention(h)\n",
    "\n",
    "        return h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DownSample(nn.Module):\n",
    "    def __init__(self, channels):\n",
    "        super().__init__()\n",
    "        self.downsample = nn.Conv2d(in_channels=channels, out_channels=channels, kernel_size=3, stride=2, padding=1)\n",
    "\n",
    "    def forward(self, x, *args):\n",
    "        return self.downsample(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [],
   "source": [
    "class UpSample(nn.Module):\n",
    "    def __init__(self, in_channels):\n",
    "        super().__init__()\n",
    "\n",
    "        self.upsample = nn.Sequential(\n",
    "            nn.Upsample(scale_factor=2, mode=\"nearest\"),\n",
    "            nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=3, stride=1, padding=1),\n",
    "        )\n",
    "\n",
    "    def forward(self, x, *args):\n",
    "        return self.upsample(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [],
   "source": [
    "class UNet(nn.Module):\n",
    "    def __init__(\n",
    "            self,\n",
    "            input_channels=3,\n",
    "            output_channels=3,\n",
    "            num_res_blocks=2,\n",
    "            base_channels=128,\n",
    "            base_channels_multiples=(1,2,4,8),\n",
    "            apply_attention=(False,False,True,False),\n",
    "            dropout_rate=0.1,\n",
    "            time_multiple=4,\n",
    "    ):\n",
    "        super().__init__()\n",
    "\n",
    "        time_emb_dims_exp = base_channels * time_multiple\n",
    "        self.time_embeddings = SinusoidalPositionEmbeddings(time_emb_dims=base_channels, time_emb_dims_exp=time_emb_dims_exp)\n",
    "\n",
    "        self.first = nn.Conv2d(in_channels=input_channels, out_channels=base_channels, kernel_size=3, stride=1, padding=\"same\")\n",
    "\n",
    "        num_resolutions = len(base_channels_multiples)\n",
    "\n",
    "        self.encoder_blocks = nn.ModuleList()\n",
    "        curr_channels = [base_channels]\n",
    "        in_channels = base_channels\n",
    "\n",
    "        for level in range(num_resolutions):\n",
    "            out_channels = base_channels * base_channels_multiples[level]\n",
    "\n",
    "            for _ in range(num_res_blocks):\n",
    "\n",
    "                block = ResnetBlock(\n",
    "                    in_channels = in_channels,\n",
    "                    out_channels=out_channels,\n",
    "                    dropout_rate=dropout_rate,\n",
    "                    time_emb_dims=time_emb_dims_exp,\n",
    "                    apply_attention=apply_attention[level],\n",
    "                )\n",
    "                self.encoder_blocks.append(block)\n",
    "\n",
    "                in_channels = out_channels\n",
    "                curr_channels.append(in_channels)\n",
    "\n",
    "            if level != (num_resolutions - 1):\n",
    "                self.encoder_blocks.append(DownSample(channels=in_channels))\n",
    "                curr_channels.append(in_channels)\n",
    "\n",
    "        self.bottleneck_blocks = nn.ModuleList(\n",
    "            (\n",
    "                ResnetBlock(\n",
    "                    in_channels=in_channels,\n",
    "                    out_channels=in_channels,\n",
    "                    dropout_rate=dropout_rate,\n",
    "                    time_emb_dims=time_emb_dims_exp,\n",
    "                    apply_attention=True,\n",
    "                ),\n",
    "                ResnetBlock(\n",
    "                    in_channels=in_channels,\n",
    "                    out_channels=in_channels,\n",
    "                    dropout_rate=dropout_rate,\n",
    "                    time_emb_dims=time_emb_dims_exp,\n",
    "                    apply_attention=False,\n",
    "                ),\n",
    "            )\n",
    "        )\n",
    "\n",
    "        self.decoder_blocks = nn.ModuleList()\n",
    "\n",
    "        for level in reversed(range(num_resolutions)):\n",
    "            out_channels = base_channels * base_channels_multiples[level]\n",
    "\n",
    "            for _ in range(num_res_blocks + 1):\n",
    "                encoder_in_channels = curr_channels.pop()\n",
    "                block = ResnetBlock(\n",
    "                    in_channels=encoder_in_channels + in_channels,\n",
    "                    out_channels=out_channels,\n",
    "                    dropout_rate=dropout_rate,\n",
    "                    time_emb_dims=time_emb_dims_exp,\n",
    "                    apply_attention=apply_attention[level]\n",
    "                )\n",
    "\n",
    "                in_channels = out_channels\n",
    "                self.decoder_blocks.append(block)\n",
    "            \n",
    "            if level != 0:\n",
    "                self.decoder_blocks.append(UpSample(in_channels))\n",
    "\n",
    "        self.final = nn.Sequential(\n",
    "            nn.GroupNorm(num_groups=8, num_channels=in_channels),\n",
    "            nn.SiLU(),\n",
    "            nn.Conv2d(in_channels=in_channels, out_channels=output_channels, kernel_size=3, stride=1, padding=\"same\"),\n",
    "        )\n",
    "\n",
    "    def forward(self, x, t):\n",
    "\n",
    "        time_emb = self.time_embeddings(t)\n",
    "\n",
    "        h = self.first(x)\n",
    "        outs = [h]\n",
    "\n",
    "        for layer in self.encoder_blocks:\n",
    "            h = layer(h, time_emb)\n",
    "            outs.append(h)\n",
    "\n",
    "        for layer in self.bottleneck_blocks:\n",
    "            h = layer(h, time_emb)\n",
    "\n",
    "        for layer in self.decoder_blocks:\n",
    "            if isinstance(layer, ResnetBlock):\n",
    "                out = outs.pop()\n",
    "                h = torch.cat([h, out], dim=1)\n",
    "            h = layer(h, time_emb)\n",
    "\n",
    "        h = self.final(h)\n",
    "\n",
    "        return h\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Diffusion Process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleDiffusion:\n",
    "    def __init__(\n",
    "            self,\n",
    "            num_diffusion_timesteps=1000,\n",
    "            img_shape=(3, 64, 64),\n",
    "            device=\"cpu\"\n",
    "    ):\n",
    "        self.num_diffusion_timesteps = num_diffusion_timesteps\n",
    "        self.img_shape = img_shape\n",
    "        self.device=device\n",
    "\n",
    "        self.initialize()\n",
    "\n",
    "    def initialize(self):\n",
    "        self.beta = self.get_betas()\n",
    "        self.alpha = 1 - self.beta\n",
    "\n",
    "        self.sqrt_beta                       = torch.sqrt(self.beta)\n",
    "        self.alpha_cumulative                = torch.cumprod(self.alpha, dim=0)\n",
    "        self.sqrt_alpha_cumulative           = torch.sqrt(self.alpha_cumulative)\n",
    "        self.one_by_sqrt_alpha               = 1. / torch.sqrt(self.alpha)\n",
    "        self.sqrt_one_minus_alpha_cumulative = torch.sqrt(1-self.alpha_cumulative)\n",
    "\n",
    "    def get_betas(self):\n",
    "        scale = 1000 / self.num_diffusion_timesteps\n",
    "        beta_start = scale * 1e-4\n",
    "        beta_end = scale * 0.02\n",
    "        return torch.linspace(\n",
    "            beta_start,\n",
    "            beta_end,\n",
    "            self.num_diffusion_timesteps,\n",
    "            dtype=torch.float32,\n",
    "            device=self.device,\n",
    "        )\n",
    "\n",
    "def forward_diffusion(sd: SimpleDiffusion, x0: torch.Tensor, timesteps: torch.Tensor):\n",
    "    eps = torch.randn_like(x0) #Noise\n",
    "    mean = get(sd.sqrt_alpha_cumulative, t=timesteps) * x0 # Image scaled\n",
    "    std_dev = get(sd.sqrt_one_minus_alpha_cumulative, t=timesteps) # Noise scaled\n",
    "    sample = mean + std_dev * eps # scaled inputs * scaled noise\n",
    "\n",
    "    return sample, eps # return ... gt noise --> model predicts this)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample Forward Diffusion Process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [],
   "source": [
    "sd = SimpleDiffusion(num_diffusion_timesteps=TrainingConfig.TIMESTEPS, device=\"cpu\")\n",
    "\n",
    "loader = iter(  # converting dataloader into an iterator for now.\n",
    "    get_dataloader(\n",
    "        dataset_name=BaseConfig.DATASET,\n",
    "        batch_size=6,\n",
    "        device=\"cpu\",\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 830760 characters]",
      "text/plain": [
       "<Figure size 1000x500 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "x0s, _ = next(loader)\n",
    "\n",
    "noisy_images = []\n",
    "specific_timesteps = [0, 10, 50, 100, 150, 200, 250, 300, 400, 600, 800, 999]\n",
    "\n",
    "for timestep in specific_timesteps:\n",
    "    timestep = torch.as_tensor(timestep, dtype=torch.long)\n",
    "\n",
    "    xts, _ = forward_diffusion(sd, x0s, timestep)\n",
    "    xts = inverse_transform(xts) / 255.0\n",
    "    xts = make_grid(xts, nrow=1, padding=1)\n",
    "\n",
    "    noisy_images.append(xts)\n",
    "\n",
    "# Plot and see samples at different timesteps\n",
    "\n",
    "_, ax = plt.subplots(1, len(noisy_images), figsize=(10, 5), facecolor=\"white\")\n",
    "\n",
    "for i, (timestep, noisy_sample) in enumerate(zip(specific_timesteps, noisy_images)):\n",
    "    ax[i].imshow(noisy_sample.squeeze(0).permute(1, 2, 0))\n",
    "    ax[i].set_title(f\"t={timestep}\", fontsize=8)\n",
    "    ax[i].axis(\"off\")\n",
    "    ax[i].grid(False)\n",
    "\n",
    "plt.suptitle(\"Forward Diffusion Process\", y=0.9)\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class ModelConfig:\n",
    "    BASE_CH = 64  # 64, 128, 256, 256\n",
    "    BASE_CH_MULT = (1, 2, 4, 4) # 32, 16, 8, 8 \n",
    "    APPLY_ATTENTION = (False, True, True, False)\n",
    "    DROPOUT_RATE = 0.1\n",
    "    TIME_EMB_MULT = 4 # 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = UNet(\n",
    "    input_channels          = TrainingConfig.IMG_SHAPE[0],\n",
    "    output_channels         = TrainingConfig.IMG_SHAPE[0],\n",
    "    base_channels           = ModelConfig.BASE_CH,\n",
    "    base_channels_multiples = ModelConfig.BASE_CH_MULT,\n",
    "    apply_attention         = ModelConfig.APPLY_ATTENTION,\n",
    "    dropout_rate            = ModelConfig.DROPOUT_RATE,\n",
    "    time_multiple           = ModelConfig.TIME_EMB_MULT,\n",
    ")\n",
    "model.to(BaseConfig.DEVICE)\n",
    "\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=TrainingConfig.LR)\n",
    "\n",
    "dataloader = get_dataloader(\n",
    "    dataset_name  = BaseConfig.DATASET,\n",
    "    batch_size    = TrainingConfig.BATCH_SIZE,\n",
    "    device        = BaseConfig.DEVICE,\n",
    "    pin_memory    = True,\n",
    "    num_workers   = TrainingConfig.NUM_WORKERS,\n",
    ")\n",
    "\n",
    "loss_fn = nn.MSELoss()\n",
    "\n",
    "sd = SimpleDiffusion(\n",
    "    num_diffusion_timesteps = TrainingConfig.TIMESTEPS,\n",
    "    img_shape               = TrainingConfig.IMG_SHAPE,\n",
    "    device                  = BaseConfig.DEVICE,\n",
    ")\n",
    "\n",
    "scaler = amp.GradScaler()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class ModelConfig:\n",
    "    BASE_CH = 64  # 64, 128, 256, 256\n",
    "    BASE_CH_MULT = (1, 2, 4, 4) # 32, 16, 8, 8 \n",
    "    APPLY_ATTENTION = (False, True, True, False)\n",
    "    DROPOUT_RATE = 0.1\n",
    "    TIME_EMB_MULT = 4 # 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = UNet(\n",
    "    input_channels          = TrainingConfig.IMG_SHAPE[0],\n",
    "    output_channels         = TrainingConfig.IMG_SHAPE[0],\n",
    "    base_channels           = ModelConfig.BASE_CH,\n",
    "    base_channels_multiples = ModelConfig.BASE_CH_MULT,\n",
    "    apply_attention         = ModelConfig.APPLY_ATTENTION,\n",
    "    dropout_rate            = ModelConfig.DROPOUT_RATE,\n",
    "    time_multiple           = ModelConfig.TIME_EMB_MULT,\n",
    ")\n",
    "model.to(BaseConfig.DEVICE)\n",
    "\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=TrainingConfig.LR)\n",
    "\n",
    "dataloader = get_dataloader(\n",
    "    dataset_name  = BaseConfig.DATASET,\n",
    "    batch_size    = TrainingConfig.BATCH_SIZE,\n",
    "    device        = BaseConfig.DEVICE,\n",
    "    pin_memory    = True,\n",
    "    num_workers   = TrainingConfig.NUM_WORKERS,\n",
    ")\n",
    "\n",
    "loss_fn = nn.MSELoss()\n",
    "\n",
    "sd = SimpleDiffusion(\n",
    "    num_diffusion_timesteps = TrainingConfig.TIMESTEPS,\n",
    "    img_shape               = TrainingConfig.IMG_SHAPE,\n",
    "    device                  = BaseConfig.DEVICE,\n",
    ")\n",
    "\n",
    "scaler = amp.GradScaler()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Logging at: Logs_Checkpoints\\Inference\\version1\n",
      "Model Checkpoint at: Logs_Checkpoints\\checkpoints\\version1\n"
     ]
    }
   ],
   "source": [
    "total_epochs = TrainingConfig.NUM_EPOCHS + 1\n",
    "log_dir, checkpoint_dir = setup_log_directory(config=BaseConfig())\n",
    "\n",
    "generate_video = False\n",
    "ext = \".mp4\" if generate_video else \".png\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Algorithm 1: Training\n",
    "\n",
    "def train_one_epoch(model, sd, loader, optimizer, scaler, loss_fn, epoch=800, \n",
    "                   base_config=BaseConfig(), training_config=TrainingConfig()):\n",
    "    \n",
    "    loss_record = MeanMetric()\n",
    "    model.train()\n",
    "\n",
    "    with tqdm(total=len(loader), dynamic_ncols=True) as tq:\n",
    "        tq.set_description(f\"Train :: Epoch: {epoch}/{training_config.NUM_EPOCHS}\")\n",
    "         \n",
    "        for x0s, _ in loader:\n",
    "            tq.update(1)\n",
    "            \n",
    "            ts = torch.randint(low=1, high=training_config.TIMESTEPS, size=(x0s.shape[0],), device=base_config.DEVICE)\n",
    "            xts, gt_noise = forward_diffusion(sd, x0s, ts)\n",
    "\n",
    "            with amp.autocast():\n",
    "                pred_noise = model(xts, ts)\n",
    "                loss = loss_fn(gt_noise, pred_noise)\n",
    "\n",
    "            optimizer.zero_grad(set_to_none=True)\n",
    "            scaler.scale(loss).backward()\n",
    "\n",
    "            # scaler.unscale_(optimizer)\n",
    "            # torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "\n",
    "            scaler.step(optimizer)\n",
    "            scaler.update()\n",
    "\n",
    "            loss_value = loss.detach().item()\n",
    "            loss_record.update(loss_value)\n",
    "\n",
    "            tq.set_postfix_str(s=f\"Loss: {loss_value:.4f}\")\n",
    "\n",
    "        mean_loss = loss_record.compute().item()\n",
    "    \n",
    "        tq.set_postfix_str(s=f\"Epoch Loss: {mean_loss:.4f}\")\n",
    "    \n",
    "    return mean_loss "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def reverse_diffusion(model, sd, timesteps=1000, img_shape=(3, 64, 64), \n",
    "                      num_images=5, nrow=8, device=\"cpu\", **kwargs):\n",
    "\n",
    "    x = torch.randn((num_images, *img_shape), device=device)\n",
    "    model.eval()\n",
    "\n",
    "    if kwargs.get(\"generate_video\", False):\n",
    "        outs = []\n",
    "\n",
    "    for time_step in tqdm(iterable=reversed(range(1, timesteps)), \n",
    "                          total=timesteps-1, dynamic_ncols=False, \n",
    "                          desc=\"Sampling :: \", position=0):\n",
    "\n",
    "        ts = torch.ones(num_images, dtype=torch.long, device=device) * time_step\n",
    "        z = torch.randn_like(x) if time_step > 1 else torch.zeros_like(x)\n",
    "\n",
    "        predicted_noise = model(x, ts)\n",
    "\n",
    "        beta_t                            = get(sd.beta, ts)\n",
    "        one_by_sqrt_alpha_t               = get(sd.one_by_sqrt_alpha, ts)\n",
    "        sqrt_one_minus_alpha_cumulative_t = get(sd.sqrt_one_minus_alpha_cumulative, ts) \n",
    "\n",
    "        x = (\n",
    "            one_by_sqrt_alpha_t\n",
    "            * (x - (beta_t / sqrt_one_minus_alpha_cumulative_t) * predicted_noise)\n",
    "            + torch.sqrt(beta_t) * z\n",
    "        )\n",
    "\n",
    "        if kwargs.get(\"generate_video\", False):\n",
    "            x_inv = inverse_transform(x).type(torch.uint8)\n",
    "            grid = make_grid(x_inv, nrow=nrow, pad_value=255.0).to(\"cpu\")\n",
    "            ndarr = torch.permute(grid, (1, 2, 0)).numpy()[:, :, ::-1]\n",
    "            outs.append(ndarr)\n",
    "\n",
    "    if kwargs.get(\"generate_video\", False): # Generate and save video of the entire reverse process. \n",
    "        frames2vid(outs, kwargs['save_path'])\n",
    "        display(Image.fromarray(outs[-1][:, :, ::-1])) # Display the image at the final timestep of the reverse process.\n",
    "        return None\n",
    "\n",
    "    else: # Display and save the image at the final timestep of the reverse process. \n",
    "        x = inverse_transform(x).type(torch.uint8)\n",
    "        grid = make_grid(x, nrow=nrow, pad_value=255.0).to(\"cpu\")\n",
    "        pil_image = TF.functional.to_pil_image(grid)\n",
    "        pil_image.save(kwargs['save_path'], format=save_path[-3:].upper())\n",
    "        display(pil_image)\n",
    "        return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train :: Epoch: 1/800:   0%|          | 0/135 [00:00<?, ?it/s]\n"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "Can't pickle local object 'get_dataset.<locals>.<lambda>'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[106], line 6\u001b[0m\n\u001b[0;32m      3\u001b[0m gc\u001b[38;5;241m.\u001b[39mcollect()\n\u001b[0;32m      5\u001b[0m \u001b[38;5;66;03m# Algorithm 1: Training\u001b[39;00m\n\u001b[1;32m----> 6\u001b[0m \u001b[43mtrain_one_epoch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msd\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdataloader\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mscaler\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mloss_fn\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mepoch\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mepoch\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m      8\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m epoch \u001b[38;5;241m%\u001b[39m \u001b[38;5;241m20\u001b[39m \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[0;32m      9\u001b[0m     save_path \u001b[38;5;241m=\u001b[39m os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mjoin(log_dir, \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mepoch\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;132;01m{\u001b[39;00mext\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n",
      "Cell \u001b[1;32mIn[104], line 12\u001b[0m, in \u001b[0;36mtrain_one_epoch\u001b[1;34m(model, sd, loader, optimizer, scaler, loss_fn, epoch, base_config, training_config)\u001b[0m\n\u001b[0;32m      9\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m tqdm(total\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mlen\u001b[39m(loader), dynamic_ncols\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m) \u001b[38;5;28;01mas\u001b[39;00m tq:\n\u001b[0;32m     10\u001b[0m     tq\u001b[38;5;241m.\u001b[39mset_description(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTrain :: Epoch: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mepoch\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m/\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtraining_config\u001b[38;5;241m.\u001b[39mNUM_EPOCHS\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m---> 12\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mx0s\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m_\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mloader\u001b[49m\u001b[43m:\u001b[49m\n\u001b[0;32m     13\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtq\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mupdate\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[0;32m     15\u001b[0m \u001b[43m        \u001b[49m\u001b[43mts\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrandint\u001b[49m\u001b[43m(\u001b[49m\u001b[43mlow\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mhigh\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtraining_config\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mTIMESTEPS\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msize\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mx0s\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mshape\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbase_config\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mDEVICE\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[1;32mIn[73], line 12\u001b[0m, in \u001b[0;36mDeviceDataLoader.__iter__\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m     10\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"在移动到设备后生成一个批次的数据\"\"\"\u001b[39;00m\n\u001b[0;32m     11\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Yield a batch of data after moving it to device\"\"\"\u001b[39;00m\n\u001b[1;32m---> 12\u001b[0m \u001b[43m\u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mb\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdl\u001b[49m\u001b[43m:\u001b[49m\n\u001b[0;32m     13\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01myield\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mto_device\u001b[49m\u001b[43m(\u001b[49m\u001b[43mb\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\Cxyoz\\.conda\\py-133-pytorch\\Lib\\site-packages\\torch\\utils\\data\\dataloader.py:439\u001b[0m, in \u001b[0;36mDataLoader.__iter__\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    437\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_iterator\n\u001b[0;32m    438\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m--> 439\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_get_iterator\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\Cxyoz\\.conda\\py-133-pytorch\\Lib\\site-packages\\torch\\utils\\data\\dataloader.py:387\u001b[0m, in \u001b[0;36mDataLoader._get_iterator\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    385\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m    386\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcheck_worker_number_rationality()\n\u001b[1;32m--> 387\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_MultiProcessingDataLoaderIter\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\Cxyoz\\.conda\\py-133-pytorch\\Lib\\site-packages\\torch\\utils\\data\\dataloader.py:1040\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter.__init__\u001b[1;34m(self, loader)\u001b[0m\n\u001b[0;32m   1033\u001b[0m w\u001b[38;5;241m.\u001b[39mdaemon \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[0;32m   1034\u001b[0m \u001b[38;5;66;03m# NB: Process.start() actually take some time as it needs to\u001b[39;00m\n\u001b[0;32m   1035\u001b[0m \u001b[38;5;66;03m#     start a process and pass the arguments over via a pipe.\u001b[39;00m\n\u001b[0;32m   1036\u001b[0m \u001b[38;5;66;03m#     Therefore, we only add a worker to self._workers list after\u001b[39;00m\n\u001b[0;32m   1037\u001b[0m \u001b[38;5;66;03m#     it started, so that we do not call .join() if program dies\u001b[39;00m\n\u001b[0;32m   1038\u001b[0m \u001b[38;5;66;03m#     before it starts, and __del__ tries to join but will get:\u001b[39;00m\n\u001b[0;32m   1039\u001b[0m \u001b[38;5;66;03m#     AssertionError: can only join a started process.\u001b[39;00m\n\u001b[1;32m-> 1040\u001b[0m \u001b[43mw\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstart\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   1041\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_index_queues\u001b[38;5;241m.\u001b[39mappend(index_queue)\n\u001b[0;32m   1042\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_workers\u001b[38;5;241m.\u001b[39mappend(w)\n",
      "File \u001b[1;32mc:\\Users\\Cxyoz\\.conda\\py-133-pytorch\\Lib\\multiprocessing\\process.py:121\u001b[0m, in \u001b[0;36mBaseProcess.start\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    118\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m _current_process\u001b[38;5;241m.\u001b[39m_config\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdaemon\u001b[39m\u001b[38;5;124m'\u001b[39m), \\\n\u001b[0;32m    119\u001b[0m        \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdaemonic processes are not allowed to have children\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[0;32m    120\u001b[0m _cleanup()\n\u001b[1;32m--> 121\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_popen \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_Popen\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[0;32m    122\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_sentinel \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_popen\u001b[38;5;241m.\u001b[39msentinel\n\u001b[0;32m    123\u001b[0m \u001b[38;5;66;03m# Avoid a refcycle if the target function holds an indirect\u001b[39;00m\n\u001b[0;32m    124\u001b[0m \u001b[38;5;66;03m# reference to the process object (see bpo-30775)\u001b[39;00m\n",
      "File \u001b[1;32mc:\\Users\\Cxyoz\\.conda\\py-133-pytorch\\Lib\\multiprocessing\\context.py:224\u001b[0m, in \u001b[0;36mProcess._Popen\u001b[1;34m(process_obj)\u001b[0m\n\u001b[0;32m    222\u001b[0m \u001b[38;5;129m@staticmethod\u001b[39m\n\u001b[0;32m    223\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_Popen\u001b[39m(process_obj):\n\u001b[1;32m--> 224\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_default_context\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_context\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mProcess\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_Popen\u001b[49m\u001b[43m(\u001b[49m\u001b[43mprocess_obj\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\Cxyoz\\.conda\\py-133-pytorch\\Lib\\multiprocessing\\context.py:336\u001b[0m, in \u001b[0;36mSpawnProcess._Popen\u001b[1;34m(process_obj)\u001b[0m\n\u001b[0;32m    333\u001b[0m \u001b[38;5;129m@staticmethod\u001b[39m\n\u001b[0;32m    334\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_Popen\u001b[39m(process_obj):\n\u001b[0;32m    335\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpopen_spawn_win32\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Popen\n\u001b[1;32m--> 336\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mPopen\u001b[49m\u001b[43m(\u001b[49m\u001b[43mprocess_obj\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\Cxyoz\\.conda\\py-133-pytorch\\Lib\\multiprocessing\\popen_spawn_win32.py:95\u001b[0m, in \u001b[0;36mPopen.__init__\u001b[1;34m(self, process_obj)\u001b[0m\n\u001b[0;32m     93\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m     94\u001b[0m     reduction\u001b[38;5;241m.\u001b[39mdump(prep_data, to_child)\n\u001b[1;32m---> 95\u001b[0m     \u001b[43mreduction\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdump\u001b[49m\u001b[43m(\u001b[49m\u001b[43mprocess_obj\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mto_child\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     96\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[0;32m     97\u001b[0m     set_spawning_popen(\u001b[38;5;28;01mNone\u001b[39;00m)\n",
      "File \u001b[1;32mc:\\Users\\Cxyoz\\.conda\\py-133-pytorch\\Lib\\multiprocessing\\reduction.py:60\u001b[0m, in \u001b[0;36mdump\u001b[1;34m(obj, file, protocol)\u001b[0m\n\u001b[0;32m     58\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdump\u001b[39m(obj, file, protocol\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[0;32m     59\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m'''Replacement for pickle.dump() using ForkingPickler.'''\u001b[39;00m\n\u001b[1;32m---> 60\u001b[0m     \u001b[43mForkingPickler\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfile\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mprotocol\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdump\u001b[49m\u001b[43m(\u001b[49m\u001b[43mobj\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[1;31mAttributeError\u001b[0m: Can't pickle local object 'get_dataset.<locals>.<lambda>'"
     ]
    }
   ],
   "source": [
    "for epoch in range(1, total_epochs):\n",
    "    torch.cuda.empty_cache()\n",
    "    gc.collect()\n",
    "    \n",
    "    # Algorithm 1: Training\n",
    "    train_one_epoch(model, sd, dataloader, optimizer, scaler, loss_fn, epoch=epoch)\n",
    "\n",
    "    if epoch % 20 == 0:\n",
    "        save_path = os.path.join(log_dir, f\"{epoch}{ext}\")\n",
    "        \n",
    "        # Algorithm 2: Sampling\n",
    "        reverse_diffusion(model, sd, timesteps=TrainingConfig.TIMESTEPS, num_images=32, generate_video=generate_video,\n",
    "            save_path=save_path, img_shape=TrainingConfig.IMG_SHAPE, device=BaseConfig.DEVICE,\n",
    "        )\n",
    "\n",
    "        # clear_output()\n",
    "        checkpoint_dict = {\n",
    "            \"opt\": optimizer.state_dict(),\n",
    "            \"scaler\": scaler.state_dict(),\n",
    "            \"model\": model.state_dict()\n",
    "        }\n",
    "        torch.save(checkpoint_dict, os.path.join(checkpoint_dir, \"ckpt.tar\"))\n",
    "        del checkpoint_dict"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "DLML",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
