examples/ddpm.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "if 'has_changed_dir' not in globals():\n",
    "    repo_path = os.path.abspath(os.path.join('..'))\n",
    "    \n",
    "    if repo_path not in sys.path:\n",
    "        sys.path.append(repo_path)\n",
    "    \n",
    "    os.chdir(repo_path)\n",
    "    \n",
    "    globals()['has_changed_dir'] = True\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from typing import Tuple\n",
    "\n",
    "from PIL import Image\n",
    "from tqdm import tqdm\n",
    "\n",
    "import neunet as nnet\n",
    "import neunet.nn as nn\n",
    "from data_loader import load_mnist, load_utkface  # noqa F401\n",
    "from neunet import Tensor\n",
    "from neunet.optim import Adam\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://arxiv.org/abs/2006.11239\n",
    "# https://arxiv.org/abs/2102.09672\n",
    "\n",
    "# https://huggingface.co/blog/annotated-diffusion\n",
    "# https://lilianweng.github.io/posts/2021-07-11-diffusion-models/\n",
    "# https://nn.labml.ai/diffusion/ddpm/index.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def linear_schedule(start, end, timesteps):\n",
    "    return nnet.tensor(\n",
    "        np.linspace(start, end, timesteps, dtype=np.float32),\n",
    "        requires_grad=False,\n",
    "        device=device,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class Diffusion:\n",
    "    def __init__(self, timesteps: int, beta_start: float, beta_end: float, criterion, model=None):\n",
    "        self.model = model\n",
    "\n",
    "        self.beta_start = beta_start\n",
    "        self.beta_end = beta_end\n",
    "        self.timesteps = timesteps\n",
    "\n",
    "        self.betas = linear_schedule(beta_start, beta_end, timesteps)\n",
    "        self.sqrt_betas = nnet.tensor(\n",
    "            np.sqrt(self.betas.cpu().numpy(), dtype=np.float32),\n",
    "            requires_grad=False,\n",
    "            device=device,\n",
    "        )\n",
    "\n",
    "        self.alphas = 1 - self.betas\n",
    "        self.inv_sqrt_alphas = nnet.tensor(\n",
    "            1 / np.sqrt(self.alphas.cpu().numpy(), dtype=np.float32),\n",
    "            requires_grad=False,\n",
    "            device=device,\n",
    "        )\n",
    "\n",
    "        self.alphas_cumprod = nnet.tensor(\n",
    "            np.cumprod(self.alphas.cpu().numpy(), axis=0, dtype=np.float32),\n",
    "            requires_grad=False,\n",
    "            device=device,\n",
    "        )\n",
    "        self.sqrt_alphas_cumprod = nnet.tensor(\n",
    "            np.sqrt(self.alphas_cumprod.cpu().numpy(), dtype=np.float32),\n",
    "            requires_grad=False,\n",
    "            device=device,\n",
    "        )\n",
    "        self.sqrt_one_minus_alphas_cumprod = nnet.tensor(\n",
    "            np.sqrt(1 - self.alphas_cumprod.cpu().numpy(), dtype=np.float32),\n",
    "            requires_grad=False,\n",
    "            device=device,\n",
    "        )\n",
    "\n",
    "        self.scaled_alphas = (1 - self.alphas) / self.sqrt_one_minus_alphas_cumprod\n",
    "\n",
    "        # self.sqrt_recip_alphas = np.sqrt(1.0 / self.alphas)\n",
    "        # self.alphas_cumprod_prev = np.concatenate([np.array([1]), self.alphas_cumprod[:-1]]) #np.pad(alphas_cumprod[:-1], (1, 0), value=1.0)\n",
    "        # self.posterior_variance = self.betas * (1. - self.alphas_cumprod_prev) / (1. - self.alphas_cumprod)\n",
    "\n",
    "        self.criterion = criterion\n",
    "        self.optimizer = Adam(self.model.parameters(), lr=2e-4)\n",
    "\n",
    "    def forward(self, x: Tensor, t=None):\n",
    "        \"\"\"\n",
    "        https://arxiv.org/abs/2006.11239\n",
    "\n",
    "        Algorithm 1: Training; according to the paper\n",
    "        \"\"\"\n",
    "        timesteps_selection = np.random.randint(1, self.timesteps, (x.shape[0],), dtype=np.int32)\n",
    "        noise = nnet.tensor(np.random.normal(size=x.shape), requires_grad=False, device=x.device)\n",
    "\n",
    "        x_t = (\n",
    "            self.sqrt_alphas_cumprod[timesteps_selection, None, None, None] * x\n",
    "            + self.sqrt_one_minus_alphas_cumprod[timesteps_selection, None, None, None] * noise\n",
    "        )\n",
    "        \n",
    "        x = self.model.forward(\n",
    "            nnet.tensor(x_t, requires_grad=False, device=x.device),\n",
    "            timesteps_selection / self.timesteps,\n",
    "        )\n",
    "\n",
    "        return x, noise\n",
    "\n",
    "    def ddpm_denoise_sample(\n",
    "        self,\n",
    "        sample_num=None,\n",
    "        image_size=None,\n",
    "        states_step_size=1,\n",
    "        x_t=None,\n",
    "        orig_x=None,\n",
    "        mask=None,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        https://arxiv.org/abs/2006.11239\n",
    "\n",
    "        Algorithm 2: Sampling; according to the paper\n",
    "        \"\"\"\n",
    "\n",
    "        if mask is not None:\n",
    "            if orig_x is None:\n",
    "                raise ValueError(\"orig_x must not be None\")\n",
    "            if orig_x.shape != mask.shape:\n",
    "                raise ValueError(\"orig_x shape must be equal to mask shape\")\n",
    "\n",
    "        if x_t is None:\n",
    "            if orig_x is None:\n",
    "                x_t = np.random.normal(size=(sample_num, *image_size))\n",
    "            else:\n",
    "                x_t = np.random.normal(size=orig_x.shape)\n",
    "\n",
    "        x_t = nnet.tensor(x_t, device=device)\n",
    "        x_ts = []\n",
    "        for t in tqdm(\n",
    "            reversed(range(0, self.timesteps)),\n",
    "            desc=\"ddpm denoisinig samples\",\n",
    "            total=self.timesteps,\n",
    "        ):\n",
    "            noise = (\n",
    "                nnet.tensor(np.random.normal(size=x_t.shape), device=device)\n",
    "                if t > 1\n",
    "                else 0\n",
    "            )\n",
    "            eps = (\n",
    "                self.model.forward(x_t, np.array([t]) / self.timesteps).reshape(x_t.shape).detach()\n",
    "            )\n",
    "\n",
    "            x_t = (\n",
    "                self.inv_sqrt_alphas[t] * (x_t - eps * self.scaled_alphas[t])\n",
    "                + self.sqrt_betas[t] * noise\n",
    "            )\n",
    "            # x_t = self.sqrt_recip_alphas[t] * (x_t - self.betas[t] * eps / self.sqrt_one_minus_alphas_cumprod[t]) + np.sqrt(self.posterior_variance[t]) * noise\n",
    "\n",
    "            if mask is not None:\n",
    "                orig_x_noise = nnet.tensor(\n",
    "                    np.random.normal(size=orig_x.shape),\n",
    "                    device=device,\n",
    "                )\n",
    "\n",
    "                orig_x_t = (\n",
    "                    self.sqrt_alphas_cumprod[t] * orig_x\n",
    "                    + self.sqrt_one_minus_alphas_cumprod[t] * orig_x_noise\n",
    "                )\n",
    "                x_t = orig_x_t * mask + x_t * (1 - mask)\n",
    "\n",
    "            if t % states_step_size == 0:\n",
    "                x_ts.append(x_t.cpu().detach().numpy())\n",
    "\n",
    "        return x_t.to(\"cpu\").detach().numpy(), x_ts\n",
    "\n",
    "    def ddim_denoise_sample(\n",
    "        self,\n",
    "        sample_num=None,\n",
    "        image_size=None,\n",
    "        states_step_size=1,\n",
    "        eta=1.0,\n",
    "        perform_steps=100,\n",
    "        x_t=None,\n",
    "        orig_x=None,\n",
    "        mask=None,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        https://arxiv.org/abs/2010.02502\n",
    "\n",
    "        Denoising Diffusion Implicit Models (DDIM) sampling; according to the paper\n",
    "        \"\"\"\n",
    "        # TODO: FIX LOOP with appropriate tensors types\n",
    "        if mask is not None:\n",
    "            if orig_x is None:\n",
    "                raise ValueError(\"orig_x must not be None\")\n",
    "            if orig_x.shape != mask.shape:\n",
    "                raise ValueError(\"orig_x shape must be equal to mask shape\")\n",
    "\n",
    "        if x_t is None:\n",
    "            if orig_x is None:\n",
    "                x_t = np.random.normal(size=(sample_num, *image_size))\n",
    "            else:\n",
    "                x_t = np.random.normal(size=orig_x.shape)\n",
    "\n",
    "        x_t = nnet.tensor(x_t, device=device)\n",
    "        x_ts = []\n",
    "        for t in tqdm(\n",
    "            reversed(range(1, self.timesteps)[:perform_steps]),\n",
    "            desc=\"ddim denoisinig samples\",\n",
    "            total=perform_steps,\n",
    "        ):\n",
    "            noise = (\n",
    "                nnet.tensor(np.random.normal(size=x_t.shape), device=device)\n",
    "                if t > 1\n",
    "                else 0\n",
    "            )\n",
    "            eps = self.model.forward(x_t, np.array([t]) / self.timesteps).reshape(\n",
    "                x_t.shape\n",
    "            ).detach()\n",
    "\n",
    "            x0_t = (x_t - eps * nnet.sqrt(1 - self.alphas_cumprod[t])) / nnet.sqrt(\n",
    "                self.alphas_cumprod[t]\n",
    "            )\n",
    "\n",
    "            sigma = eta * nnet.sqrt(\n",
    "                (1 - self.alphas_cumprod[t - 1])\n",
    "                / (1 - self.alphas_cumprod[t])\n",
    "                * (1 - self.alphas_cumprod[t] / self.alphas_cumprod[t - 1])\n",
    "            )\n",
    "            c = nnet.sqrt((1 - self.alphas_cumprod[t - 1]) - sigma**2)\n",
    "\n",
    "            x_t = nnet.sqrt(self.alphas_cumprod[t - 1]) * x0_t - c * eps + sigma * noise\n",
    "\n",
    "            if mask is not None:\n",
    "                orig_x_noise = nnet.tensor(\n",
    "                    np.random.normal(size=orig_x.shape),\n",
    "                    device=device,\n",
    "                )\n",
    "\n",
    "                orig_x_t = (\n",
    "                    self.sqrt_alphas_cumprod[t] * orig_x\n",
    "                    + self.sqrt_one_minus_alphas_cumprod[t] * orig_x_noise\n",
    "                )\n",
    "                x_t = orig_x_t * mask + x_t * (1 - mask)\n",
    "\n",
    "            if t % states_step_size == 0:\n",
    "                x_ts.append(x_t.cpu().detach().numpy())\n",
    "\n",
    "        return x_t.to(\"cpu\").detach().numpy(), x_ts\n",
    "\n",
    "    def get_images_set(\n",
    "        self,\n",
    "        x_num: int,\n",
    "        y_num: int,\n",
    "        margin: int,\n",
    "        images: np.ndarray,\n",
    "        image_size: Tuple[int, int, int],\n",
    "    ):\n",
    "        def denormalize(x):\n",
    "            return (x - np.min(x)) / (np.max(x) - np.min(x)) * 255\n",
    "\n",
    "        channels, H_size, W_size = image_size\n",
    "\n",
    "        images_array = np.full(\n",
    "            (y_num * (margin + H_size), x_num * (margin + W_size), channels),\n",
    "            255,\n",
    "            dtype=np.uint8,\n",
    "        )\n",
    "        num = 0\n",
    "        for i in range(y_num):\n",
    "            for j in range(x_num):\n",
    "                y = i * (margin + H_size)\n",
    "                x = j * (margin + W_size)\n",
    "\n",
    "                images_array[y : y + H_size, x : x + W_size] = denormalize(\n",
    "                    images[num].transpose(1, 2, 0)\n",
    "                )\n",
    "\n",
    "                num += 1\n",
    "\n",
    "        images_array = images_array[\n",
    "            : (y_num - 1) * (H_size + margin) + H_size,\n",
    "            : (x_num - 1) * (W_size + margin) + W_size,\n",
    "        ]\n",
    "\n",
    "        if channels == 1:\n",
    "            return Image.fromarray(images_array.squeeze(axis=2)).convert(\"L\")\n",
    "        else:\n",
    "            return Image.fromarray(images_array)\n",
    "\n",
    "    def train(self, dataset, epochs, batch_size, image_path, image_size, save_path, save_every_epochs=1):\n",
    "        channels, H_size, W_size = image_size\n",
    "\n",
    "        data_batches = np.array_split(dataset, np.arange(batch_size, len(dataset), batch_size))\n",
    "\n",
    "        loss_history = []\n",
    "        for epoch in range(epochs):\n",
    "            tqdm_range = tqdm(enumerate(data_batches), total=len(data_batches))\n",
    "\n",
    "            losses = []\n",
    "            for batch_num, (batch) in tqdm_range:\n",
    "                batch = batch.reshape(-1, channels, H_size, W_size)\n",
    "                # print(batch.shape)\n",
    "                output, noise = self.forward(\n",
    "                    nnet.tensor(batch, requires_grad=True, device=device, dtype=nnet.float32)\n",
    "                )\n",
    "                loss = self.criterion(output, noise)\n",
    "                losses.append(loss.item())\n",
    "                self.optimizer.zero_grad()\n",
    "                loss.backward()\n",
    "                self.optimizer.step()\n",
    "\n",
    "                tqdm_range.set_description(f\"loss: {losses[-1]:.7f} | epoch {epoch + 1}/{epochs}\")\n",
    "\n",
    "                if batch_num == (len(data_batches) - 1):\n",
    "                    epoch_loss = nnet.tensor(losses, device=device).mean().detach().cpu().numpy()\n",
    "\n",
    "                    tqdm_range.set_description(\n",
    "                        f\"loss: {epoch_loss:.7f} | epoch {epoch + 1}/{epochs}\"\n",
    "                    )\n",
    "\n",
    "            if (epoch + 1) % save_every_epochs == 0:\n",
    "                margin = 10\n",
    "                x_num, y_num = 5, 5\n",
    "\n",
    "                samples, samples_in_time = self.ddpm_denoise_sample(\n",
    "                    x_num * y_num, (channels, H_size, W_size)\n",
    "                )\n",
    "                images_grid = self.get_images_set(\n",
    "                    x_num, y_num, margin, samples, (channels, H_size, W_size)\n",
    "                )\n",
    "                images_grid.save(f\"{image_path}/np_ddpm_{epoch + 1}.png\")\n",
    "\n",
    "                images_grid_in_time = []\n",
    "                for sample in samples_in_time:\n",
    "                    images_grid_in_time.append(\n",
    "                        self.get_images_set(\n",
    "                            x_num, y_num, margin, sample, (channels, H_size, W_size)\n",
    "                        )\n",
    "                    )\n",
    "\n",
    "                images_grid_in_time[0].save(\n",
    "                    f\"{image_path}/np_ddpm_in_time.gif\",\n",
    "                    save_all=True,\n",
    "                    append_images=images_grid_in_time[1:],\n",
    "                    duration=50,\n",
    "                    loop=0,\n",
    "                )\n",
    "\n",
    "                if not Path(save_path).exists():\n",
    "                    Path(save_path).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "                nnet.save(self.model.state_dict(), f\"{save_path}/ddpm_{epoch + 1}.nt\")\n",
    "\n",
    "            loss_history.append(epoch_loss)\n",
    "\n",
    "        return loss_history\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class ResBlock(nn.Module):\n",
    "    def __init__(self, input_channels, output_channels, time_emb_dim, up=False):\n",
    "        super().__init__()\n",
    "        self.time_embedding = nn.Linear(time_emb_dim, output_channels)\n",
    "        self.input_channels = input_channels\n",
    "        self.output_channels = output_channels\n",
    "        if up:\n",
    "            self.conv1 = nn.Conv2d(\n",
    "                2 * input_channels, output_channels, kernel_size=(3, 3), padding=(1, 1)\n",
    "            )\n",
    "            self.transform = nn.ConvTranspose2d(\n",
    "                output_channels,\n",
    "                output_channels,\n",
    "                kernel_size=(4, 4),\n",
    "                stride=(2, 2),\n",
    "                padding=(1, 1),\n",
    "            )\n",
    "\n",
    "        else:\n",
    "            self.conv1 = nn.Conv2d(\n",
    "                input_channels, output_channels, kernel_size=(3, 3), padding=(1, 1)\n",
    "            )\n",
    "            self.transform = nn.Conv2d(\n",
    "                output_channels,\n",
    "                output_channels,\n",
    "                kernel_size=(4, 4),\n",
    "                stride=(2, 2),\n",
    "                padding=(1, 1),\n",
    "            )\n",
    "\n",
    "        self.conv2 = nn.Conv2d(output_channels, output_channels, kernel_size=(3, 3), padding=(1, 1))\n",
    "        self.relu1 = nn.LeakyReLU(alpha=0.01)\n",
    "        self.relu2 = nn.LeakyReLU(alpha=0.01)\n",
    "        self.relu3 = nn.LeakyReLU(alpha=0.01)\n",
    "\n",
    "        self.bnorm1 = nn.BatchNorm2d(output_channels, momentum=0.1, eps=1e-5)\n",
    "        self.bnorm2 = nn.BatchNorm2d(output_channels, momentum=0.1, eps=1e-5)\n",
    "\n",
    "    def forward(self, x, t):\n",
    "        x = self.conv1.forward(x)\n",
    "        h = self.relu1.forward(x)\n",
    "        h = self.bnorm1.forward(h)\n",
    "\n",
    "        t = self.time_embedding.forward(t)\n",
    "\n",
    "        time_emb = self.relu2.forward(t)\n",
    "\n",
    "        time_emb = time_emb[(...,) + (None,) * 2]\n",
    "\n",
    "        h = h + time_emb\n",
    "\n",
    "        h = self.conv2.forward(h)\n",
    "        h = self.relu3.forward(h)\n",
    "        h = self.bnorm2.forward(h)\n",
    "\n",
    "        return self.transform.forward(h)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "    \"\"\"Implements the sinusoidal positional encoding.\"\"\"\n",
    "\n",
    "    def __init__(self, max_len, d_model, dropout_rate=0.1, data_type=np.float32):\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "        self.d_model = d_model\n",
    "        self.dropout_rate = dropout_rate\n",
    "        self.max_len = max_len\n",
    "\n",
    "        self.data_type = data_type\n",
    "\n",
    "        pe = np.zeros((max_len, d_model))  # (max_len, d_model)\n",
    "        position = np.arange(0, max_len)[:, np.newaxis]  # (max_len, 1)\n",
    "        div_term = np.exp(np.arange(0, d_model, 2) * (-np.log(10000.0) / d_model))  # (d_model,)\n",
    "\n",
    "        pe[:, 0::2] = np.sin(position * div_term)\n",
    "        pe[:, 1::2] = np.cos(position * div_term)\n",
    "\n",
    "        self.pe = nnet.tensor(\n",
    "            pe[:, np.newaxis, :].astype(self.data_type),\n",
    "            requires_grad=False,\n",
    "            device=device,\n",
    "        )  # (max_len, 1, d_model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"x: (batch_size, seq_len, d_model)\"\"\"\n",
    "        x = x + self.pe[: x.shape[0], :]  # (batch_size, seq_len, d_model)\n",
    "\n",
    "        # x: (batch_size, seq_len, d_model)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleUNet(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        image_channels,\n",
    "        image_size,\n",
    "        down_channels=(32, 64, 128, 256, 512),\n",
    "        up_channels=(512, 256, 128, 64, 32),\n",
    "    ):\n",
    "        noise_channels = image_channels\n",
    "        time_emb_dim = 32\n",
    "\n",
    "        self.time_embedding = nn.Sequential(\n",
    "            PositionalEncoding(max_len=1000, d_model=time_emb_dim),\n",
    "            nn.Linear(time_emb_dim, time_emb_dim),\n",
    "            nn.LeakyReLU(),\n",
    "        )\n",
    "\n",
    "        if image_size & (image_size - 1) != 0:\n",
    "            self.input_conv = nn.ConvTranspose2d(\n",
    "                image_channels, down_channels[0], kernel_size=(5, 5)\n",
    "            )\n",
    "            self.output_conv = nn.Conv2d(up_channels[-1], noise_channels, kernel_size=(5, 5))\n",
    "        else:\n",
    "            self.input_conv = nn.Conv2d(\n",
    "                image_channels, down_channels[0], kernel_size=(3, 3), padding=(1, 1)\n",
    "            )\n",
    "            self.output_conv = nn.ConvTranspose2d(\n",
    "                up_channels[-1], noise_channels, kernel_size=(3, 3), padding=(1, 1)\n",
    "            )\n",
    "\n",
    "        self.down_layers = nn.ModuleList(\n",
    "            [\n",
    "                ResBlock(down_channels[i], down_channels[i + 1], time_emb_dim).to(device)\n",
    "                for i in range(len(down_channels) - 1)\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        self.up_layers = nn.ModuleList(\n",
    "            [\n",
    "                ResBlock(up_channels[i], up_channels[i + 1], time_emb_dim, up=True).to(device)\n",
    "                for i in range(len(up_channels) - 1)\n",
    "            ]\n",
    "        )\n",
    "\n",
    "    def forward(self, x, t):\n",
    "        t = nnet.tensor(\n",
    "            np.asarray(t[:, None, None], dtype=np.float32),\n",
    "            requires_grad=False,\n",
    "            device=x.device,\n",
    "        )\n",
    "\n",
    "        # for layer in self.time_embedding:\n",
    "        #     # t = t.reshape(t.shape[0], -1)\n",
    "        #     t = layer.forward(t)\n",
    "\n",
    "        t = self.time_embedding.forward(t)\n",
    "        t = t.reshape(t.shape[0], -1)\n",
    "\n",
    "        x = self.input_conv.forward(x)\n",
    "\n",
    "        residual_inputs = []\n",
    "        for down_layer in self.down_layers:\n",
    "            x = down_layer.forward(x, t)\n",
    "            residual_inputs.append(x)\n",
    "\n",
    "        for up_layer in self.up_layers:\n",
    "            residual_x = residual_inputs.pop()\n",
    "\n",
    "            x = nnet.concatenate(*(x, residual_x), axis=1)\n",
    "            x = up_layer.forward(x, t)\n",
    "        return self.output_conv.forward(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_size = (3, 32, 32)\n",
    "# image_size = (1, 28, 28) # for mnist\n",
    "\n",
    "training_data = load_utkface(image_size=(3, 32, 32))\n",
    "# training_data, _, _, _ = load_mnist() # for mnist\n",
    "\n",
    "diffusion = Diffusion(\n",
    "    model=SimpleUNet(\n",
    "        image_channels=image_size[0],\n",
    "        image_size = image_size[2],\n",
    "        down_channels=(128, 256, 512, 1024), # (32, 64, 128) for mnist\n",
    "        up_channels=(1024, 512, 256, 128), # (128, 64, 32) for mnist\n",
    "    ).to(device),\n",
    "    timesteps=300,\n",
    "    beta_start=0.0001,\n",
    "    beta_end=0.02,\n",
    "    criterion=nn.MSELoss(),\n",
    ")\n",
    "\n",
    "# diffusion.model.load_state_dict(nnet.load(\"saved models/utkface_model/ddpm_3.nt\")) # load saved model example if it exists\n",
    "\n",
    "diffusion.train(\n",
    "    training_data,\n",
    "    epochs=3,\n",
    "    batch_size=5,\n",
    "    image_path=\"generated images\",\n",
    "    save_path = \"saved models/utkface_model\", # for \"saved models/mnist_model\" mnist\n",
    "    image_size=image_size,\n",
    ")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}