{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-23T07:34:53.034694Z",
     "start_time": "2023-02-23T07:34:50.200674Z"
    },
    "execution": {
     "iopub.execute_input": "2023-02-22T16:01:33.905754Z",
     "iopub.status.busy": "2023-02-22T16:01:33.904879Z",
     "iopub.status.idle": "2023-02-22T16:01:36.156631Z",
     "shell.execute_reply": "2023-02-22T16:01:36.155423Z",
     "shell.execute_reply.started": "2023-02-22T16:01:33.905665Z"
    }
   },
   "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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# UNet Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-23T07:34:53.142674Z",
     "start_time": "2023-02-23T07:34:53.114668Z"
    },
    "code_folding": [
     0,
     1,
     25,
     105
    ],
    "execution": {
     "iopub.execute_input": "2023-02-22T16:01:36.165052Z",
     "iopub.status.busy": "2023-02-22T16:01:36.164311Z",
     "iopub.status.idle": "2023-02-22T16:01:36.199528Z",
     "shell.execute_reply": "2023-02-22T16:01:36.198080Z",
     "shell.execute_reply.started": "2023-02-22T16:01:36.165002Z"
    }
   },
   "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\n",
    "\n",
    "        emb = math.log(10000) / (half_dim - 1)\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)\n",
    "\n",
    "\n",
    "class AttentionBlock(nn.Module):\n",
    "    def __init__(self, channels=64):\n",
    "        super().__init__()\n",
    "        self.channels = channels\n",
    "\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)  # [B, C, H, W] --> [B, C, H * W] --> [B, H*W, C]\n",
    "        h, _ = self.mhsa(h, h, h)  # [B, H*W, C]\n",
    "        h = h.swapaxes(2, 1).view(B, self.channels, H, W)  # [B, C, H*W] --> [B, C, H, W]\n",
    "        return x + h\n",
    "\n",
    "\n",
    "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",
    "        # Group 1\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",
    "        # Group 2 time embedding\n",
    "        self.dense_1 = nn.Linear(in_features=time_emb_dims, out_features=self.out_channels)\n",
    "\n",
    "        # Group 3\n",
    "        self.normlize_2 = nn.GroupNorm(num_groups=8, num_channels=self.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.match_input = nn.Identity()\n",
    "\n",
    "        if apply_attention:\n",
    "            self.attention = AttentionBlock(channels=self.out_channels)\n",
    "        else:\n",
    "            self.attention = nn.Identity()\n",
    "\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\n",
    "\n",
    "\n",
    "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)\n",
    "\n",
    "\n",
    "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)\n",
    "\n",
    "\n",
    "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",
    "        # Encoder part of the UNet. Dimension reduction.\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",
    "        # Bottleneck in between\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",
    "        # Decoder part of the UNet. Dimension restoration with skip-connections.\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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-23T07:34:53.203694Z",
     "start_time": "2023-02-23T07:34:53.177669Z"
    },
    "code_folding": [
     0,
     25,
     31,
     39,
     68
    ],
    "execution": {
     "iopub.execute_input": "2023-02-22T16:01:36.201566Z",
     "iopub.status.busy": "2023-02-22T16:01:36.201193Z",
     "iopub.status.idle": "2023-02-22T16:01:36.219110Z",
     "shell.execute_reply": "2023-02-22T16:01:36.218110Z",
     "shell.execute_reply.started": "2023-02-22T16:01:36.201528Z"
    }
   },
   "outputs": [],
   "source": [
    "def to_device(data, device):\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)\n",
    "\n",
    "class DeviceDataLoader:\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",
    "        \"\"\"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",
    "        \"\"\"Number of batches\"\"\"\n",
    "        return len(self.dl)\n",
    "\n",
    "def get_default_device():\n",
    "    return torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "def save_images(images, path, **kwargs):\n",
    "    grid = torchvision.utils.make_grid(images, **kwargs)\n",
    "    ndarr = grid.permute(1, 2, 0).to(\"cpu\").numpy()\n",
    "    im = Image.fromarray(ndarr)\n",
    "    im.save(path)\n",
    "    \n",
    "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",
    "    ele = element.gather(-1, t)\n",
    "    return ele.reshape(-1, 1, 1, 1)\n",
    "\n",
    "def setup_log_directory(config):\n",
    "    '''Log and Model checkpoint directory Setup'''\n",
    "    \n",
    "    if os.path.isdir(config.root_log_dir):\n",
    "        # Get all folders numbers in the 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",
    "        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",
    "    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",
    "    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",
    "\n",
    "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 = 0\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",
    "    for image in images:\n",
    "        video.write(image)\n",
    "\n",
    "    # Deallocating memories taken for window creation\n",
    "#     cv2.destroyAllWindows()\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}\" />'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-02T18:23:28.639407Z",
     "start_time": "2023-02-02T18:23:28.624407Z"
    }
   },
   "source": [
    "# Configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-23T07:34:54.785018Z",
     "start_time": "2023-02-23T07:34:53.639004Z"
    },
    "code_folding": [],
    "execution": {
     "iopub.execute_input": "2023-02-22T16:01:36.712080Z",
     "iopub.status.busy": "2023-02-22T16:01:36.711714Z",
     "iopub.status.idle": "2023-02-22T16:01:36.720529Z",
     "shell.execute_reply": "2023-02-22T16:01:36.719078Z",
     "shell.execute_reply.started": "2023-02-22T16:01:36.712047Z"
    }
   },
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "\n",
    "@dataclass\n",
    "class BaseConfig:\n",
    "    DEVICE = get_default_device()\n",
    "    DATASET = \"MNIST\" #  \"MNIST\", \"Cifar-10\", \"Cifar-100\", \"Flowers\"\n",
    "    \n",
    "    # For logging inferece images and saving checkpoints.\n",
    "    root_log_dir = os.path.join(\"Logs_Checkpoints\", \"Inference\")\n",
    "    root_checkpoint_dir = os.path.join(\"Logs_Checkpoints\", \"checkpoints\")\n",
    "\n",
    "    # Current log and checkpoint directory.\n",
    "    log_dir = \"version_0\"\n",
    "    checkpoint_dir = \"version_0\"\n",
    "\n",
    "@dataclass\n",
    "class TrainingConfig:\n",
    "    TIMESTEPS = 1000 # Define number of diffusion timesteps\n",
    "    IMG_SHAPE = (1, 32, 32) if BaseConfig.DATASET == \"MNIST\" else (3, 32, 32) \n",
    "    NUM_EPOCHS = 30\n",
    "    BATCH_SIZE = 128\n",
    "    LR = 2e-4\n",
    "    NUM_WORKERS = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-02T18:24:36.837306Z",
     "start_time": "2023-02-02T18:24:36.8273Z"
    }
   },
   "source": [
    "# Load Dataset & Build Dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-23T07:34:55.313081Z",
     "start_time": "2023-02-23T07:34:55.291079Z"
    },
    "code_folding": [
     0,
     21,
     37
    ],
    "execution": {
     "iopub.execute_input": "2023-02-22T16:01:37.597072Z",
     "iopub.status.busy": "2023-02-22T16:01:37.596702Z",
     "iopub.status.idle": "2023-02-22T16:01:37.608430Z",
     "shell.execute_reply": "2023-02-22T16:01:37.607135Z",
     "shell.execute_reply.started": "2023-02-22T16:01:37.597040Z"
    }
   },
   "outputs": [],
   "source": [
    "def get_dataset(dataset_name='MNIST'):\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=\"/kaggle/input/flowers-recognition/flowers\", transform=transforms)\n",
    "    \n",
    "    return dataset\n",
    "\n",
    "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\n",
    "\n",
    "def inverse_transform(tensors):\n",
    "    \"\"\"Convert tensors from [-1., 1.] to [0., 255.]\"\"\"\n",
    "    return ((tensors.clamp(-1, 1) + 1.0) / 2.0) * 255.0 "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualize Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-23T07:35:01.528613Z",
     "start_time": "2023-02-23T07:34:57.085306Z"
    },
    "code_folding": [],
    "execution": {
     "iopub.execute_input": "2023-02-22T16:01:39.192513Z",
     "iopub.status.busy": "2023-02-22T16:01:39.191545Z",
     "iopub.status.idle": "2023-02-22T16:01:39.267958Z",
     "shell.execute_reply": "2023-02-22T16:01:39.266946Z",
     "shell.execute_reply.started": "2023-02-22T16:01:39.192465Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to data/MNIST/raw/train-images-idx3-ubyte.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 9912422/9912422 [00:00<00:00, 79520059.45it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting data/MNIST/raw/train-images-idx3-ubyte.gz to data/MNIST/raw\n",
      "\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to data/MNIST/raw/train-labels-idx1-ubyte.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 28881/28881 [00:00<00:00, 61334528.52it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting data/MNIST/raw/train-labels-idx1-ubyte.gz to data/MNIST/raw\n",
      "\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to data/MNIST/raw/t10k-images-idx3-ubyte.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1648877/1648877 [00:00<00:00, 38547540.84it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting data/MNIST/raw/t10k-images-idx3-ubyte.gz to data/MNIST/raw\n",
      "\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to data/MNIST/raw/t10k-labels-idx1-ubyte.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4542/4542 [00:00<00:00, 12994903.66it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting data/MNIST/raw/t10k-labels-idx1-ubyte.gz to data/MNIST/raw\n",
      "\n"
     ]
    }
   ],
   "source": [
    "loader = get_dataloader(\n",
    "    dataset_name=BaseConfig.DATASET,\n",
    "    batch_size=128,\n",
    "    device='cpu',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-21T08:04:00.707591Z",
     "start_time": "2023-02-21T08:03:59.270574Z"
    },
    "code_folding": [],
    "execution": {
     "iopub.execute_input": "2023-02-22T16:01:39.573357Z",
     "iopub.status.busy": "2023-02-22T16:01:39.572658Z",
     "iopub.status.idle": "2023-02-22T16:01:39.920898Z",
     "shell.execute_reply": "2023-02-22T16:01:39.919984Z",
     "shell.execute_reply.started": "2023-02-22T16:01:39.573319Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 312304 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, nrow=16, padding=True, pad_value=1, normalize=True)\n",
    "    plt.imshow(grid_img.permute(1, 2, 0))\n",
    "    plt.axis(\"off\")\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Diffusion Process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-23T07:35:01.574603Z",
     "start_time": "2023-02-23T07:35:01.561607Z"
    },
    "code_folding": [
     1
    ],
    "execution": {
     "iopub.execute_input": "2023-02-22T16:01:40.386727Z",
     "iopub.status.busy": "2023-02-22T16:01:40.385687Z",
     "iopub.status.idle": "2023-02-22T16:01:40.397948Z",
     "shell.execute_reply": "2023-02-22T16:01:40.396829Z",
     "shell.execute_reply.started": "2023-02-22T16:01:40.386688Z"
    }
   },
   "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",
    "        # BETAs & ALPHAs required at different places in the Algorithm.\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",
    "        \"\"\"linear schedule, proposed in original ddpm paper\"\"\"\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": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-21T08:04:15.117858Z",
     "start_time": "2023-02-21T08:04:14.427843Z"
    },
    "execution": {
     "iopub.execute_input": "2023-02-22T16:01:41.581306Z",
     "iopub.status.busy": "2023-02-22T16:01:41.580908Z",
     "iopub.status.idle": "2023-02-22T16:01:42.154133Z",
     "shell.execute_reply": "2023-02-22T16:01:42.153007Z",
     "shell.execute_reply.started": "2023-02-22T16:01:41.581270Z"
    }
   },
   "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": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-21T08:04:15.117858Z",
     "start_time": "2023-02-21T08:04:14.427843Z"
    },
    "execution": {
     "iopub.execute_input": "2023-02-22T16:01:41.581306Z",
     "iopub.status.busy": "2023-02-22T16:01:41.580908Z",
     "iopub.status.idle": "2023-02-22T16:01:42.154133Z",
     "shell.execute_reply": "2023-02-22T16:01:42.153007Z",
     "shell.execute_reply.started": "2023-02-22T16:01:41.581270Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 560856 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": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-13T13:01:00.619395Z",
     "start_time": "2023-02-13T13:01:00.605395Z"
    },
    "code_folding": [],
    "execution": {
     "iopub.execute_input": "2023-02-22T16:01:43.160882Z",
     "iopub.status.busy": "2023-02-22T16:01:43.160156Z",
     "iopub.status.idle": "2023-02-22T16:01:43.170501Z",
     "shell.execute_reply": "2023-02-22T16:01:43.169283Z",
     "shell.execute_reply.started": "2023-02-22T16:01:43.160846Z"
    }
   },
   "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": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-23T07:36:29.926465Z",
     "start_time": "2023-02-23T07:36:29.906486Z"
    },
    "code_folding": [],
    "execution": {
     "iopub.execute_input": "2023-02-22T16:01:43.842403Z",
     "iopub.status.busy": "2023-02-22T16:01:43.841296Z",
     "iopub.status.idle": "2023-02-22T16:01:43.857743Z",
     "shell.execute_reply": "2023-02-22T16:01:43.856757Z",
     "shell.execute_reply.started": "2023-02-22T16:01:43.842351Z"
    }
   },
   "outputs": [],
   "source": [
    "# Algorithm 2: Sampling\n",
    "    \n",
    "@torch.inference_mode()\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 = torchvision.utils.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 = torchvision.utils.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": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-23T07:36:07.313353Z",
     "start_time": "2023-02-23T07:36:07.307373Z"
    },
    "code_folding": [],
    "execution": {
     "iopub.execute_input": "2023-02-22T16:01:44.742690Z",
     "iopub.status.busy": "2023-02-22T16:01:44.742000Z",
     "iopub.status.idle": "2023-02-22T16:01:44.748704Z",
     "shell.execute_reply": "2023-02-22T16:01:44.747497Z",
     "shell.execute_reply.started": "2023-02-22T16:01:44.742652Z"
    }
   },
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class ModelConfig:\n",
    "    BASE_CH = 64  # 64, 128, 256, 512\n",
    "    BASE_CH_MULT = (1, 2, 4, 8) # 32, 16, 8, 4 \n",
    "    APPLY_ATTENTION = (False, False, True, False)\n",
    "    DROPOUT_RATE = 0.1\n",
    "    TIME_EMB_MULT = 2 # 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-13T13:01:00.588388Z",
     "start_time": "2023-02-13T13:01:00.344403Z"
    },
    "code_folding": [
     0,
     13,
     23
    ],
    "hide_input": false
   },
   "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": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-13T13:01:00.603387Z",
     "start_time": "2023-02-13T13:01:00.590387Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Logging at: Logs_Checkpoints/Inference/version_0\n",
      "Model Checkpoint at: Logs_Checkpoints/checkpoints/version_0\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": 21,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-13T12:45:46.057703Z",
     "start_time": "2023-02-13T12:45:39.770695Z"
    },
    "_kg_hide-output": true,
    "hide_input": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train :: Epoch: 1/30: 100%|██████████| 469/469 [01:06<00:00,  7.08it/s, Epoch Loss: 0.0557]\n",
      "Train :: Epoch: 2/30: 100%|██████████| 469/469 [01:00<00:00,  7.73it/s, Epoch Loss: 0.0287]\n",
      "Train :: Epoch: 3/30: 100%|██████████| 469/469 [01:01<00:00,  7.64it/s, Epoch Loss: 0.0241]\n",
      "Train :: Epoch: 4/30: 100%|██████████| 469/469 [01:00<00:00,  7.73it/s, Epoch Loss: 0.0221]\n",
      "Train :: Epoch: 5/30: 100%|██████████| 469/469 [01:01<00:00,  7.68it/s, Epoch Loss: 0.0209]\n",
      "Sampling :: 100%|██████████| 999/999 [00:14<00:00, 67.20it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 16848 characters]",
      "image/png": "[encoded data, 47780 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=274x138>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train :: Epoch: 6/30: 100%|██████████| 469/469 [01:00<00:00,  7.71it/s, Epoch Loss: 0.0199]\n",
      "Train :: Epoch: 7/30: 100%|██████████| 469/469 [01:00<00:00,  7.69it/s, Epoch Loss: 0.0195]\n",
      "Train :: Epoch: 8/30: 100%|██████████| 469/469 [01:01<00:00,  7.62it/s, Epoch Loss: 0.0190]\n",
      "Train :: Epoch: 9/30: 100%|██████████| 469/469 [01:02<00:00,  7.51it/s, Epoch Loss: 0.0187]\n",
      "Train :: Epoch: 10/30: 100%|██████████| 469/469 [01:03<00:00,  7.44it/s, Epoch Loss: 0.0186]\n",
      "Sampling :: 100%|██████████| 999/999 [00:14<00:00, 66.95it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 15828 characters]",
      "image/png": "[encoded data, 36504 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=274x138>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train :: Epoch: 11/30: 100%|██████████| 469/469 [01:33<00:00,  4.99it/s, Epoch Loss: 0.0182]\n",
      "Train :: Epoch: 12/30: 100%|██████████| 469/469 [01:51<00:00,  4.22it/s, Epoch Loss: 0.0181]\n",
      "Train :: Epoch: 13/30: 100%|██████████| 469/469 [01:50<00:00,  4.23it/s, Epoch Loss: 0.0180]\n",
      "Train :: Epoch: 14/30: 100%|██████████| 469/469 [01:50<00:00,  4.25it/s, Epoch Loss: 0.0178]\n",
      "Train :: Epoch: 15/30: 100%|██████████| 469/469 [01:50<00:00,  4.25it/s, Epoch Loss: 0.0178]\n",
      "Sampling :: 100%|██████████| 999/999 [00:27<00:00, 36.96it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 16804 characters]",
      "image/png": "[encoded data, 31916 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=274x138>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train :: Epoch: 16/30: 100%|██████████| 469/469 [01:49<00:00,  4.27it/s, Epoch Loss: 0.0175]\n",
      "Train :: Epoch: 17/30: 100%|██████████| 469/469 [01:50<00:00,  4.24it/s, Epoch Loss: 0.0177]\n",
      "Train :: Epoch: 18/30: 100%|██████████| 469/469 [01:50<00:00,  4.25it/s, Epoch Loss: 0.0172]\n",
      "Train :: Epoch: 19/30: 100%|██████████| 469/469 [01:50<00:00,  4.23it/s, Epoch Loss: 0.0173]\n",
      "Train :: Epoch: 20/30: 100%|██████████| 469/469 [01:49<00:00,  4.27it/s, Epoch Loss: 0.0172]\n",
      "Sampling :: 100%|██████████| 999/999 [00:26<00:00, 37.12it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 16636 characters]",
      "image/png": "[encoded data, 37620 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=274x138>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train :: Epoch: 21/30: 100%|██████████| 469/469 [01:49<00:00,  4.28it/s, Epoch Loss: 0.0171]\n",
      "Train :: Epoch: 22/30: 100%|██████████| 469/469 [01:50<00:00,  4.25it/s, Epoch Loss: 0.0168]\n",
      "Train :: Epoch: 23/30: 100%|██████████| 469/469 [01:50<00:00,  4.25it/s, Epoch Loss: 0.0170]\n",
      "Train :: Epoch: 24/30: 100%|██████████| 469/469 [01:49<00:00,  4.29it/s, Epoch Loss: 0.0169]\n",
      "Train :: Epoch: 25/30: 100%|██████████| 469/469 [01:50<00:00,  4.26it/s, Epoch Loss: 0.0168]\n",
      "Sampling :: 100%|██████████| 999/999 [00:26<00:00, 37.03it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 16188 characters]",
      "image/png": "[encoded data, 35340 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=274x138>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train :: Epoch: 26/30: 100%|██████████| 469/469 [01:50<00:00,  4.23it/s, Epoch Loss: 0.0168]\n",
      "Train :: Epoch: 27/30: 100%|██████████| 469/469 [01:51<00:00,  4.22it/s, Epoch Loss: 0.0168]\n",
      "Train :: Epoch: 28/30: 100%|██████████| 469/469 [01:50<00:00,  4.26it/s, Epoch Loss: 0.0167]\n",
      "Train :: Epoch: 29/30: 100%|██████████| 469/469 [01:50<00:00,  4.25it/s, Epoch Loss: 0.0167]\n",
      "Train :: Epoch: 30/30: 100%|██████████| 469/469 [01:50<00:00,  4.26it/s, Epoch Loss: 0.0167]\n",
      "Sampling :: 100%|██████████| 999/999 [00:26<00:00, 37.39it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 16528 characters]",
      "image/png": "[encoded data, 39156 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=274x138>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "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 % 5 == 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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-23T07:36:13.403286Z",
     "start_time": "2023-02-23T07:36:10.026409Z"
    },
    "execution": {
     "iopub.execute_input": "2023-02-22T16:03:08.012329Z",
     "iopub.status.busy": "2023-02-22T16:03:08.011929Z",
     "iopub.status.idle": "2023-02-22T16:03:10.213564Z",
     "shell.execute_reply": "2023-02-22T16:03:10.212573Z",
     "shell.execute_reply.started": "2023-02-22T16:03:08.012295Z"
    },
    "hide_input": false
   },
   "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",
    "# checkpoint_dir = \"/kaggle/working/Logs_Checkpoints/checkpoints/version_0\"\n",
    "\n",
    "\n",
    "model.load_state_dict(torch.load(os.path.join(checkpoint_dir, \"ckpt.tar\"), map_location='cpu')['model'])\n",
    "\n",
    "model.to(BaseConfig.DEVICE)\n",
    "\n",
    "sd = SimpleDiffusion(\n",
    "    num_diffusion_timesteps = TrainingConfig.TIMESTEPS,\n",
    "    img_shape               = TrainingConfig.IMG_SHAPE,\n",
    "    device                  = BaseConfig.DEVICE,\n",
    ")\n",
    "\n",
    "log_dir = \"inference_results\"\n",
    "os.makedirs(log_dir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-23T07:42:49.677019Z",
     "start_time": "2023-02-23T07:41:04.890036Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling :: 100%|██████████| 999/999 [00:49<00:00, 20.03it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 31972 characters]",
      "image/png": "[encoded data, 77916 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=274x274>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inference_results/20240408-123303.mp4\n"
     ]
    }
   ],
   "source": [
    "generate_video = True\n",
    "\n",
    "ext = \".mp4\" if generate_video else \".png\"\n",
    "filename = f\"{datetime.now().strftime('%Y%m%d-%H%M%S')}{ext}\"\n",
    "\n",
    "save_path = os.path.join(log_dir, filename)\n",
    "\n",
    "reverse_diffusion(\n",
    "    model,\n",
    "    sd,\n",
    "    num_images=64,\n",
    "    generate_video=generate_video,\n",
    "    save_path=save_path,\n",
    "    timesteps=1000,\n",
    "    img_shape=TrainingConfig.IMG_SHAPE,\n",
    "    device=BaseConfig.DEVICE,\n",
    "    nrow=8,\n",
    ")\n",
    "print(save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-13T13:56:16.260979Z",
     "start_time": "2023-02-13T13:50:06.139878Z"
    },
    "execution": {
     "iopub.execute_input": "2023-02-22T16:03:53.237665Z",
     "iopub.status.busy": "2023-02-22T16:03:53.237279Z",
     "iopub.status.idle": "2023-02-22T16:08:06.306724Z",
     "shell.execute_reply": "2023-02-22T16:08:06.304477Z",
     "shell.execute_reply.started": "2023-02-22T16:03:53.237633Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling :: 100%|██████████| 999/999 [02:56<00:00,  5.66it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 126372 characters]",
      "image/png": "[encoded data, 307752 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=1090x274>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inference_results/20240408-123354.mp4\n"
     ]
    }
   ],
   "source": [
    "generate_video = True\n",
    "\n",
    "ext = \".mp4\" if generate_video else \".png\"\n",
    "filename = f\"{datetime.now().strftime('%Y%m%d-%H%M%S')}{ext}\"\n",
    "\n",
    "save_path = os.path.join(log_dir, filename)\n",
    "\n",
    "reverse_diffusion(\n",
    "    model,\n",
    "    sd,\n",
    "    num_images=256,\n",
    "    generate_video=generate_video,\n",
    "    save_path=save_path,\n",
    "    timesteps=1000,\n",
    "    img_shape=TrainingConfig.IMG_SHAPE,\n",
    "    device=BaseConfig.DEVICE,\n",
    "    nrow=32,\n",
    ")\n",
    "print(save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-13T14:00:18.151514Z",
     "start_time": "2023-02-13T13:57:07.698420Z"
    },
    "execution": {
     "iopub.execute_input": "2023-02-22T16:16:59.757373Z",
     "iopub.status.busy": "2023-02-22T16:16:59.756510Z",
     "iopub.status.idle": "2023-02-22T16:21:11.518330Z",
     "shell.execute_reply": "2023-02-22T16:21:11.517233Z",
     "shell.execute_reply.started": "2023-02-22T16:16:59.757332Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling :: 100%|██████████| 999/999 [02:55<00:00,  5.68it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 127564 characters]",
      "image/png": "[encoded data, 306676 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=546x546>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inference_results/20240408-123655.mp4\n"
     ]
    }
   ],
   "source": [
    "generate_video = True\n",
    "\n",
    "ext = \".mp4\" if generate_video else \".png\"\n",
    "filename = f\"{datetime.now().strftime('%Y%m%d-%H%M%S')}{ext}\"\n",
    "\n",
    "save_path = os.path.join(log_dir, filename)\n",
    "\n",
    "reverse_diffusion(\n",
    "    model,\n",
    "    sd,\n",
    "    num_images=256,\n",
    "    generate_video=generate_video,\n",
    "    save_path=save_path,\n",
    "    timesteps=1000,\n",
    "    img_shape=TrainingConfig.IMG_SHAPE,\n",
    "    device=BaseConfig.DEVICE,\n",
    "    nrow=16,\n",
    ")\n",
    "print(save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-02-22T16:30:05.578985Z",
     "iopub.status.busy": "2023-02-22T16:30:05.578388Z",
     "iopub.status.idle": "2023-02-22T16:34:16.926424Z",
     "shell.execute_reply": "2023-02-22T16:34:16.925434Z",
     "shell.execute_reply.started": "2023-02-22T16:30:05.578945Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling :: 100%|██████████| 999/999 [02:56<00:00,  5.68it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 126028 characters]",
      "image/png": "[encoded data, 305368 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=546x546>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inference_results/20240408-123956.mp4\n"
     ]
    }
   ],
   "source": [
    "generate_video = True\n",
    "\n",
    "ext = \".mp4\" if generate_video else \".png\"\n",
    "filename = f\"{datetime.now().strftime('%Y%m%d-%H%M%S')}{ext}\"\n",
    "\n",
    "save_path = os.path.join(log_dir, filename)\n",
    "\n",
    "reverse_diffusion(\n",
    "    model,\n",
    "    sd,\n",
    "    num_images=256,\n",
    "    generate_video=generate_video,\n",
    "    save_path=save_path,\n",
    "    timesteps=1000,\n",
    "    img_shape=TrainingConfig.IMG_SHAPE,\n",
    "    device=BaseConfig.DEVICE,\n",
    "    nrow=16,\n",
    ")\n",
    "print(save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-02-22T16:34:16.929146Z",
     "iopub.status.busy": "2023-02-22T16:34:16.928274Z",
     "iopub.status.idle": "2023-02-22T16:42:15.960070Z",
     "shell.execute_reply": "2023-02-22T16:42:15.959235Z",
     "shell.execute_reply.started": "2023-02-22T16:34:16.929115Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling :: 100%|██████████| 999/999 [05:40<00:00,  2.93it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 251896 characters]",
      "image/png": "[encoded data, 611460 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=1090x546>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inference_results/20240408-124257.mp4\n"
     ]
    }
   ],
   "source": [
    "generate_video = True\n",
    "\n",
    "ext = \".mp4\" if generate_video else \".png\"\n",
    "filename = f\"{datetime.now().strftime('%Y%m%d-%H%M%S')}{ext}\"\n",
    "\n",
    "save_path = os.path.join(log_dir, filename)\n",
    "\n",
    "reverse_diffusion(\n",
    "    model,\n",
    "    sd,\n",
    "    num_images=512,\n",
    "    generate_video=generate_video,\n",
    "    save_path=save_path,\n",
    "    timesteps=1000,\n",
    "    img_shape=TrainingConfig.IMG_SHAPE,\n",
    "    device=BaseConfig.DEVICE,\n",
    "    nrow=32,\n",
    ")\n",
    "print(save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-02-22T16:16:33.646107Z",
     "iopub.status.busy": "2023-02-22T16:16:33.645729Z",
     "iopub.status.idle": "2023-02-22T16:16:59.754676Z",
     "shell.execute_reply": "2023-02-22T16:16:59.753710Z",
     "shell.execute_reply.started": "2023-02-22T16:16:33.646076Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling :: 100%|██████████| 999/999 [00:19<00:00, 52.08it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 8796 characters]",
      "image/png": "[encoded data, 20864 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=138x138>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inference_results/20240408-124847.mp4\n"
     ]
    }
   ],
   "source": [
    "generate_video = True\n",
    "\n",
    "ext = \".mp4\" if generate_video else \".png\"\n",
    "filename = f\"{datetime.now().strftime('%Y%m%d-%H%M%S')}{ext}\"\n",
    "\n",
    "save_path = os.path.join(log_dir, filename)\n",
    "\n",
    "reverse_diffusion(\n",
    "    model,\n",
    "    sd,\n",
    "    num_images=16,\n",
    "    generate_video=generate_video,\n",
    "    save_path=save_path,\n",
    "    timesteps=1000,\n",
    "    img_shape=TrainingConfig.IMG_SHAPE,\n",
    "    device=BaseConfig.DEVICE,\n",
    "    nrow=4,\n",
    ")\n",
    "print(save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-13T14:06:41.407250Z",
     "start_time": "2023-02-13T14:03:27.834241Z"
    },
    "execution": {
     "iopub.execute_input": "2023-02-22T16:21:11.522425Z",
     "iopub.status.busy": "2023-02-22T16:21:11.519897Z",
     "iopub.status.idle": "2023-02-22T16:23:35.317409Z",
     "shell.execute_reply": "2023-02-22T16:23:35.316345Z",
     "shell.execute_reply.started": "2023-02-22T16:21:11.522385Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inference_results/20240408-124907.mp4\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling :: 100%|██████████| 999/999 [01:32<00:00, 10.82it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 63360 characters]",
      "image/png": "[encoded data, 155548 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=274x546>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inference_results/20240408-124907.mp4\n"
     ]
    }
   ],
   "source": [
    "generate_video = True\n",
    "\n",
    "ext = \".mp4\" if generate_video else \".png\"\n",
    "filename = f\"{datetime.now().strftime('%Y%m%d-%H%M%S')}{ext}\"\n",
    "\n",
    "save_path = os.path.join(log_dir, filename)\n",
    "print(save_path)\n",
    "\n",
    "reverse_diffusion(\n",
    "    model,\n",
    "    sd,\n",
    "    num_images=128,\n",
    "    generate_video=generate_video,\n",
    "    save_path=save_path,\n",
    "    timesteps=1000,\n",
    "    img_shape=TrainingConfig.IMG_SHAPE,\n",
    "    device=BaseConfig.DEVICE,\n",
    "    nrow=8,\n",
    ")\n",
    "print(save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-13T14:03:27.628269Z",
     "start_time": "2023-02-13T14:00:18.376509Z"
    },
    "execution": {
     "iopub.execute_input": "2023-02-22T16:23:35.323346Z",
     "iopub.status.busy": "2023-02-22T16:23:35.322602Z",
     "iopub.status.idle": "2023-02-22T16:27:47.271778Z",
     "shell.execute_reply": "2023-02-22T16:27:47.270801Z",
     "shell.execute_reply.started": "2023-02-22T16:23:35.323305Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inference_results/20240408-125042.mp4\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling :: 100%|██████████| 999/999 [02:56<00:00,  5.66it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 126424 characters]",
      "image/png": "[encoded data, 303928 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=546x546>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inference_results/20240408-125042.mp4\n"
     ]
    }
   ],
   "source": [
    "generate_video = True\n",
    "\n",
    "ext = \".mp4\" if generate_video else \".png\"\n",
    "filename = f\"{datetime.now().strftime('%Y%m%d-%H%M%S')}{ext}\"\n",
    "\n",
    "save_path = os.path.join(log_dir, filename)\n",
    "print(save_path)\n",
    "\n",
    "reverse_diffusion(\n",
    "    model,\n",
    "    sd,\n",
    "    num_images=256,\n",
    "    generate_video=generate_video,\n",
    "    save_path=save_path,\n",
    "    timesteps=1000,\n",
    "    img_shape=TrainingConfig.IMG_SHAPE,\n",
    "    device=BaseConfig.DEVICE,\n",
    "    nrow=16,\n",
    ")\n",
    "print(save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2023-02-13T14:02:13.807Z"
    },
    "execution": {
     "iopub.execute_input": "2023-02-22T16:27:47.273986Z",
     "iopub.status.busy": "2023-02-22T16:27:47.273382Z",
     "iopub.status.idle": "2023-02-22T16:30:05.576098Z",
     "shell.execute_reply": "2023-02-22T16:30:05.575242Z",
     "shell.execute_reply.started": "2023-02-22T16:27:47.273949Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling :: 100%|██████████| 999/999 [01:29<00:00, 11.10it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 63584 characters]",
      "image/png": "[encoded data, 152472 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=546x274>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inference_results/20240408-125343.png\n"
     ]
    }
   ],
   "source": [
    "generate_video = False\n",
    "\n",
    "ext = \".mp4\" if generate_video else \".png\"\n",
    "filename = f\"{datetime.now().strftime('%Y%m%d-%H%M%S')}{ext}\"\n",
    "\n",
    "save_path = os.path.join(log_dir, filename)\n",
    "\n",
    "\n",
    "reverse_diffusion(\n",
    "    model,\n",
    "    sd,\n",
    "    num_images=128,\n",
    "    generate_video=generate_video,\n",
    "    save_path=save_path,\n",
    "    timesteps=1000,\n",
    "    img_shape=TrainingConfig.IMG_SHAPE,\n",
    "    device=BaseConfig.DEVICE,\n",
    "    nrow=16,\n",
    ")\n",
    "print(save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
