{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1c53094d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/6r/vdhtmmgj6hj3r0bc6c274h7r0000gq/T/ipykernel_20290/516204110.py:22: UserWarning: 'has_mps' is deprecated, please use 'torch.backends.mps.is_built()'\n",
      "  elif getattr(torch, \"has_mps\", False) and torch.has_mps:\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: mps\n",
      "Files already downloaded and verified\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 276864 characters]",
      "text/plain": [
       "<Figure size 288x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Adapted CIFAR-10 score-based model (training + sampling)\n",
    "\n",
    "# save_path = 'ckptCIFAR-OrNR.pth'\n",
    "save_path = 'ckptCIFAR-R1e-4NR.pth'\n",
    "\n",
    "\n",
    "\n",
    "import math, os, sys\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.optim import Adam\n",
    "from torchvision import transforms, utils\n",
    "from torchvision.datasets import CIFAR10\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# DEVICE: prefer CUDA, else MPS, else CPU\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "elif getattr(torch, \"has_mps\", False) and torch.has_mps:\n",
    "    device = torch.device(\"mps\")\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "print(\"Using device:\", device)\n",
    "\n",
    "# -------------------------\n",
    "# Dataset / transforms\n",
    "# -------------------------\n",
    "out_dir = \"generated\"\n",
    "os.makedirs(out_dir, exist_ok=True)\n",
    "image_size = 32\n",
    "batch_size = 128\n",
    "\n",
    "transform_train = transforms.Compose([\n",
    "    # transforms.RandomCrop(32, padding=4),\n",
    "    # transforms.RandomHorizontalFlip(),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5])\n",
    "])\n",
    "\n",
    "from torchvision.datasets import CIFAR10\n",
    "import torch\n",
    "from torch.utils.data import Subset\n",
    "\n",
    "# CIFAR-10 class names\n",
    "classes = ['airplane', 'automobile', 'bird', 'cat', 'deer',\n",
    "           'dog', 'frog', 'horse', 'ship', 'truck']\n",
    "\n",
    "# Pick the class you want\n",
    "target_class = 'automobile'\n",
    "target_idx = classes.index(target_class)  # gives 1\n",
    "\n",
    "# Load full dataset\n",
    "train_ds_full = CIFAR10(root='./data', train=True, download=True, transform=transform_train)\n",
    "\n",
    "# Find indices of only the chosen class\n",
    "car_indices = [i for i, label in enumerate(train_ds_full.targets) if label == target_idx]\n",
    "\n",
    "# Subset dataset to only include cars\n",
    "train_ds_car = Subset(train_ds_full, car_indices)\n",
    "\n",
    "# DataLoader with only car images\n",
    "data_loader = torch.utils.data.DataLoader(train_ds_car, batch_size=batch_size, shuffle=True,\n",
    "                                          num_workers=4, pin_memory=True)\n",
    "\n",
    "\n",
    "def show_images_grid(images, nrow=8, save_path=None):\n",
    "    imgs = images.clone().detach().cpu()\n",
    "    imgs = (imgs * 0.5) + 0.5  # unnormalize from [-1,1] → [0,1]\n",
    "    imgs = torch.clamp(imgs, 0, 1)\n",
    "\n",
    "    # Create grid with no padding and align pixel edges exactly\n",
    "    grid = utils.make_grid(imgs[:nrow * nrow], nrow=nrow, padding=0)\n",
    "    np_grid = np.transpose(grid.numpy(), (1, 2, 0))\n",
    "\n",
    "    # Build figure with precise pixel dimensions (no scaling, no interpolation)\n",
    "    h, w = np_grid.shape[:2]\n",
    "    fig = plt.figure(figsize=(w / 100, h / 100), dpi=100)\n",
    "    ax = plt.Axes(fig, [0., 0., 1., 1.])\n",
    "    ax.set_axis_off()\n",
    "    fig.add_axes(ax)\n",
    "\n",
    "    ax.imshow(np_grid, interpolation='nearest')  # no smoothing, no black gaps\n",
    "    fig.patch.set_facecolor('white')  # ensure no black background\n",
    "\n",
    "    if save_path:\n",
    "        plt.savefig(save_path, bbox_inches='tight', pad_inches=0)\n",
    "        plt.close(fig)\n",
    "    else:\n",
    "        plt.show()\n",
    "\n",
    "# get one batch from the dataloader\n",
    "data_iter = iter(data_loader)\n",
    "images, labels = next(data_iter)  # if your dataset has labels\n",
    "show_images_grid(images, nrow=9)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d04925e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Loaded model from ckptCIFAR-R1e-4NR.pth\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling: 100%|██████████| 1000/1000 [22:34<00:00,  1.35s/it]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 289228 characters]",
      "text/plain": [
       "<Figure size 288x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "# -------------------------\n",
    "# Time embedding utilities\n",
    "# -------------------------\n",
    "class GaussianFourierProjection(nn.Module):\n",
    "    \"\"\"Gaussian random features for time embedding.\"\"\"\n",
    "    def __init__(self, embed_dim=128, scale=30.):\n",
    "        super().__init__()\n",
    "        # fixed random weights\n",
    "        W = torch.randn(embed_dim // 2) * scale\n",
    "        self.register_buffer('W', W)  # not a parameter\n",
    "    def forward(self, t):\n",
    "        # t: (B,) in [0,1]\n",
    "        W = self.W.to(t.device)   \n",
    "        x_proj = t[:, None] * self.W[None, :] * 2 * math.pi\n",
    "        return torch.cat([torch.sin(x_proj), torch.cos(x_proj)], dim=-1)\n",
    "\n",
    "class TimeEmbedding(nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super().__init__()\n",
    "        self.prj = GaussianFourierProjection(embed_dim=dim)\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(dim, dim),\n",
    "            nn.SiLU(),\n",
    "            nn.Linear(dim, dim)\n",
    "        )\n",
    "    def forward(self, t):\n",
    "        return self.fc(self.prj(t))\n",
    "\n",
    "# -------------------------\n",
    "# Residual blocks + attention\n",
    "# -------------------------\n",
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch, time_emb_dim=None):\n",
    "        super().__init__()\n",
    "        self.time_emb_dim = time_emb_dim\n",
    "        self.norm1 = nn.GroupNorm(8, in_ch)\n",
    "        self.act1 = nn.SiLU()\n",
    "        self.conv1 = nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1)\n",
    "        self.norm2 = nn.GroupNorm(8, out_ch)\n",
    "        self.act2 = nn.SiLU()\n",
    "        self.conv2 = nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1)\n",
    "        if in_ch != out_ch:\n",
    "            self.nin_short = nn.Conv2d(in_ch, out_ch, kernel_size=1)\n",
    "        else:\n",
    "            self.nin_short = nn.Identity()\n",
    "        if time_emb_dim is not None:\n",
    "            self.time_mlp = nn.Sequential(nn.SiLU(), nn.Linear(time_emb_dim, out_ch))\n",
    "        else:\n",
    "            self.time_mlp = None\n",
    "\n",
    "    def forward(self, x, t_emb=None):\n",
    "        h = self.conv1(self.act1(self.norm1(x)))\n",
    "        if self.time_mlp is not None and t_emb is not None:\n",
    "            # t_emb: (B, dim) -> (B, out_ch, 1, 1)\n",
    "            h = h + self.time_mlp(t_emb)[:, :, None, None]\n",
    "        h = self.conv2(self.act2(self.norm2(h)))\n",
    "        return h + self.nin_short(x)\n",
    "\n",
    "class SelfAttention(nn.Module):\n",
    "    def __init__(self, ch):\n",
    "        super().__init__()\n",
    "        self.norm = nn.GroupNorm(8, ch)\n",
    "        self.q = nn.Conv2d(ch, ch, 1)\n",
    "        self.k = nn.Conv2d(ch, ch, 1)\n",
    "        self.v = nn.Conv2d(ch, ch, 1)\n",
    "        self.proj_out = nn.Conv2d(ch, ch, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, C, H, W = x.shape\n",
    "        h = self.norm(x)\n",
    "        q = self.q(h).reshape(B, C, -1)  # B,C,N\n",
    "        k = self.k(h).reshape(B, C, -1)\n",
    "        v = self.v(h).reshape(B, C, -1)\n",
    "        w = torch.bmm(q.permute(0,2,1), k)  # B,N,N\n",
    "        w = w * (C ** -0.5)\n",
    "        w = torch.softmax(w, dim=-1)\n",
    "        out = torch.bmm(v, w.permute(0,2,1))  # B,C,N\n",
    "        out = out.reshape(B, C, H, W)\n",
    "        out = self.proj_out(out)\n",
    "        return x + out\n",
    "\n",
    "# -------------------------\n",
    "# U-Net score network\n",
    "# -------------------------\n",
    "class ScoreUNet(nn.Module):\n",
    "    def __init__(self, marginal_prob_std, base_ch=64, ch_mults=(1,2,4), time_emb_dim=256):\n",
    "        super().__init__()\n",
    "        self.marginal_prob_std = marginal_prob_std\n",
    "        self.time_embed = TimeEmbedding(time_emb_dim)\n",
    "\n",
    "        self.input_conv = nn.Conv2d(3, base_ch, 3, padding=1)\n",
    "\n",
    "        # down blocks\n",
    "        in_ch = base_ch\n",
    "        self.downs = nn.ModuleList()\n",
    "        self.down_projs = nn.ModuleList()\n",
    "        for mult in ch_mults:\n",
    "            out_ch = base_ch * mult\n",
    "            self.downs.append(nn.ModuleList([\n",
    "                ResidualBlock(in_ch, out_ch, time_emb_dim=time_emb_dim),\n",
    "                ResidualBlock(out_ch, out_ch, time_emb_dim=time_emb_dim)\n",
    "            ]))\n",
    "            self.down_projs.append(nn.AvgPool2d(2))\n",
    "            in_ch = out_ch\n",
    "\n",
    "        # middle\n",
    "        self.mid1 = ResidualBlock(in_ch, in_ch, time_emb_dim=time_emb_dim)\n",
    "        self.mid_attn = SelfAttention(in_ch)\n",
    "        self.mid2 = ResidualBlock(in_ch, in_ch, time_emb_dim=time_emb_dim)\n",
    "\n",
    "        # up blocks\n",
    "        self.ups = nn.ModuleList()\n",
    "        self.up_convs = nn.ModuleList()\n",
    "        for mult in reversed(ch_mults):\n",
    "            out_ch = base_ch * mult\n",
    "            self.ups.append(nn.ModuleList([\n",
    "                ResidualBlock(in_ch + out_ch, out_ch, time_emb_dim=time_emb_dim),\n",
    "                ResidualBlock(out_ch, out_ch, time_emb_dim=time_emb_dim)\n",
    "            ]))\n",
    "            self.up_convs.append(nn.Upsample(scale_factor=2, mode='nearest'))\n",
    "            in_ch = out_ch\n",
    "\n",
    "        self.out_norm = nn.GroupNorm(8, in_ch)\n",
    "        self.out_act = nn.SiLU()\n",
    "        self.out_conv = nn.Conv2d(in_ch, 3, kernel_size=3, padding=1)\n",
    "\n",
    "    def forward(self, x, t):\n",
    "        # x: (B,3,H,W), t: (B,)\n",
    "        t = t.to(x.device)\n",
    "        t_emb = self.time_embed(t)\n",
    "\n",
    "        hs = []\n",
    "        h = self.input_conv(x)\n",
    "        for block, pool in zip(self.downs, self.down_projs):\n",
    "            h = block[0](h, t_emb)\n",
    "            h = block[1](h, t_emb)\n",
    "            hs.append(h)\n",
    "            h = pool(h)\n",
    "\n",
    "        h = self.mid1(h, t_emb)\n",
    "        h = self.mid_attn(h)\n",
    "        h = self.mid2(h, t_emb)\n",
    "\n",
    "        for (block1, block2), up in zip(self.ups, self.up_convs):\n",
    "            h = up(h)\n",
    "            # concat with skip\n",
    "            skip = hs.pop()\n",
    "            # if spatial mismatch due to odd sizes, center crop skip\n",
    "            if skip.shape[2:] != h.shape[2:]:\n",
    "                # center-crop skip to h size\n",
    "                _,_,H,W = h.shape\n",
    "                _,_,Hs,Ws = skip.shape\n",
    "                start_h = (Hs - H) // 2\n",
    "                start_w = (Ws - W) // 2\n",
    "                skip = skip[:,:,start_h:start_h+H,start_w:start_w+W]\n",
    "            h = torch.cat([h, skip], dim=1)\n",
    "            h = block1(h, t_emb)\n",
    "            h = block2(h, t_emb)\n",
    "\n",
    "        out = self.out_conv(self.out_act(self.out_norm(h)))\n",
    "        # scale by marginal std (consistent with your original design)\n",
    "        out = out / self.marginal_prob_std(t)[:, None, None, None]\n",
    "        return out\n",
    "\n",
    "# -------------------------\n",
    "# SDE helpers (same form as original)\n",
    "# -------------------------\n",
    "import functools\n",
    "def marginal_prob_std(t, sigma):\n",
    "    \"\"\"t: tensor (B,) in [0,1]\"\"\"\n",
    "    return torch.sqrt((sigma**(2 * t) - 1.) / 2. / math.log(sigma)).to(t.device)\n",
    "\n",
    "def diffusion_coeff(t, sigma):\n",
    "    return (sigma**t).to(t.device)\n",
    "\n",
    "sigma = 5.0\n",
    "marginal_prob_std_fn = functools.partial(marginal_prob_std, sigma=sigma)\n",
    "diffusion_coeff_fn = functools.partial(diffusion_coeff, sigma=sigma)\n",
    "\n",
    "# -------------------------\n",
    "# Loss function\n",
    "# -------------------------\n",
    "def loss_fn(model, x, marginal_prob_std, eps=1e-5):\n",
    "    device = x.device\n",
    "    random_t = torch.rand(x.shape[0], device=device) * (1. - eps) + eps  # (B,)\n",
    "    z = torch.randn_like(x)\n",
    "    std = marginal_prob_std(random_t)\n",
    "    perturbed_x = x + z * std[:, None, None, None]\n",
    "    score = model(perturbed_x, random_t)\n",
    "    loss = torch.mean(torch.sum((score * std[:, None, None, None] + z)**2, dim=(1,2,3)))\n",
    "    #loss=loss+ 0.0001* torch.sum(torch.abs(score))\n",
    "    return loss\n",
    "\n",
    "# -------------------------\n",
    "# Create model, optimizer\n",
    "# -------------------------\n",
    "model = ScoreUNet(marginal_prob_std=functools.partial(marginal_prob_std, sigma=sigma),\n",
    "                  base_ch=128, ch_mults=(1,2,4,8), time_emb_dim=512)\n",
    "model = model.to(device)\n",
    "# optional: DataParallel if you have >1 GPU, else single-GPU/cpu ok\n",
    "if torch.cuda.device_count() > 1:\n",
    "    model = nn.DataParallel(model)\n",
    "\n",
    "optimizer = Adam(model.parameters(), lr=5e-4, betas=(0.9,0.999), weight_decay=0.0)\n",
    "\n",
    "# -------------------------\n",
    "# Training loop\n",
    "# -------------------------\n",
    "n_epochs = 200  # adjust as you like\n",
    "print_every = 1\n",
    "\n",
    "def save_checkpoint(epoch):\n",
    "    state = {\"model\": getattr(model, \"module\", model).state_dict(),\n",
    "             \"optimizer\": optimizer.state_dict(),\n",
    "             \"epoch\": epoch}\n",
    "    torch.save(state, save_path)\n",
    "\n",
    "\n",
    "# -------------------------\n",
    "# Sampling: Predictor-Corrector (PC)\n",
    "# -------------------------\n",
    "def pc_sampler(score_model,\n",
    "               marginal_prob_std,\n",
    "               diffusion_coeff,\n",
    "               batch_size=64,\n",
    "               num_steps=500,\n",
    "               snr=0.16,\n",
    "               device=device,\n",
    "               eps=1e-3):\n",
    "\n",
    "    score_model.eval()\n",
    "    t = torch.ones(batch_size, device=device)\n",
    "    x = torch.randn(batch_size, 3, image_size, image_size, device=device) * marginal_prob_std(t)[:, None, None, None]\n",
    "    time_steps = np.linspace(1., eps, num_steps)\n",
    "    step_size = time_steps[0] - time_steps[1]\n",
    "    x_mean = None\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for time_step in tqdm(time_steps, desc=\"Sampling\"):\n",
    "            batch_time_step = torch.ones(batch_size, device=device) * float(time_step)\n",
    "            # Corrector: Langevin MCMC\n",
    "            grad = score_model(x, batch_time_step)\n",
    "            grad_norm = torch.norm(grad.reshape(grad.shape[0], -1), dim=-1).mean()\n",
    "            noise_norm = math.sqrt(np.prod(x.shape[1:]))\n",
    "            langevin_step_size = (2 * (snr * noise_norm / (grad_norm + 1e-12))**2)\n",
    "            x = x + langevin_step_size * grad\n",
    "            x = x + torch.sqrt(2 * langevin_step_size) * torch.randn_like(x)\n",
    "\n",
    "            # Predictor: Euler-Maruyama for the reverse-time SDE\n",
    "            g = diffusion_coeff(batch_time_step)\n",
    "            score = score_model(x, batch_time_step)\n",
    "            x_mean = x + (g**2)[:, None, None, None] * score * step_size\n",
    "            x = x_mean + torch.sqrt((g**2) * step_size)[:, None, None, None] * torch.randn_like(x)\n",
    "\n",
    "    return x_mean\n",
    "\n",
    "# -------------------------\n",
    "# Setup for saving generated images\n",
    "# -------------------------\n",
    "\n",
    "\n",
    "def save_generated_images(model, epoch, device, n_samples=32, num_steps=200):\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        samples = pc_sampler(\n",
    "            model,\n",
    "            marginal_prob_std_fn,\n",
    "            diffusion_coeff_fn,\n",
    "            batch_size=n_samples,\n",
    "            num_steps=num_steps,\n",
    "            snr=0.16,\n",
    "            device=device\n",
    "        )\n",
    "        # unnormalize to [0,1] range for saving\n",
    "        samples = (samples * 0.5) + 0.5\n",
    "        grid = utils.make_grid(samples, nrow=8, padding=2)\n",
    "        save_path = os.path.join(out_dir, f\"samples_epoch{epoch:03d}.png\")\n",
    "        utils.save_image(grid, save_path)\n",
    "        print(f\"Saved generated images to {save_path}\")\n",
    "\n",
    "# for epoch in range(1, n_epochs+1):\n",
    "#     model.train()\n",
    "#     pbar = tqdm(data_loader, desc=f\"Epoch {epoch}/{n_epochs}\")\n",
    "#     epoch_loss = 0.0\n",
    "#     n_items = 0\n",
    "#     for x, _ in pbar:\n",
    "#         x = x.to(device)\n",
    "#         loss = loss_fn(model, x, marginal_prob_std_fn)\n",
    "#         optimizer.zero_grad()\n",
    "#         loss.backward()\n",
    "#         optimizer.step()\n",
    "\n",
    "#         epoch_loss += loss.item() * x.shape[0]\n",
    "#         n_items += x.shape[0]\n",
    "#         pbar.set_postfix({\"loss\": f\"{(epoch_loss / n_items):.6f}\"})\n",
    "\n",
    "#     if epoch % print_every == 0:\n",
    "#         print(f\"Epoch {epoch} average loss: {epoch_loss / n_items:.6f}\")\n",
    "#         save_checkpoint(epoch)\n",
    "#     if epoch % 10 == 0:  # every 10 epochs\n",
    "#         save_generated_images(model, epoch, device, n_samples=32, num_steps=200)\n",
    "    \n",
    "\n",
    "\n",
    "# If you want to load a saved checkpoint:\n",
    "# if os.path.exists(save_path):\n",
    "#     ckpt = torch.load(save_path, map_location=device)\n",
    "#     getattr(model, \"module\", model).load_state_dict(ckpt[\"model\"])\n",
    "#     print(\"Loaded checkpoint from\", save_path)\n",
    "\n",
    "# Sampling example\n",
    "\n",
    "if os.path.exists(save_path):\n",
    "    ckpt = torch.load(save_path, map_location=device)\n",
    "    model.load_state_dict(ckpt[\"model\"])\n",
    "    print(f\"✅ Loaded model from {save_path}\")\n",
    "else:\n",
    "    raise FileNotFoundError(f\"No checkpoint found at {save_path}\")\n",
    "sample_batch_size = 81\n",
    "\n",
    "\n",
    "#Regularized  score matching images \n",
    "\n",
    "#num_steps=1000\n",
    "\n",
    "samples = pc_sampler(model, marginal_prob_std_fn, diffusion_coeff_fn,\n",
    "                     batch_size=sample_batch_size, num_steps=1000, snr=0.16, device=device)\n",
    "# # unnormalize for display\n",
    "# # samples_disp = (samples * 0.5) + 0.5\n",
    "samples_disp=samples\n",
    "show_images_grid(samples_disp, nrow=9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "34380d09",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling: 100%|██████████| 500/500 [11:05<00:00,  1.33s/it]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 289704 characters]",
      "text/plain": [
       "<Figure size 288x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#num_steps=500\n",
    "samples = pc_sampler(model, marginal_prob_std_fn, diffusion_coeff_fn,\n",
    "                     batch_size=sample_batch_size, num_steps=500, snr=0.16, device=device)\n",
    "samples_disp=samples\n",
    "show_images_grid(samples_disp, nrow=9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b4d7f1a6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Loaded model from ckptCIFAR-R5e-4NR.pth\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling: 100%|██████████| 200/200 [04:35<00:00,  1.38s/it]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 290396 characters]",
      "text/plain": [
       "<Figure size 288x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#num_steps=200\n",
    "save_path = 'ckptCIFAR-R5e-4NR.pth'\n",
    "\n",
    "if os.path.exists(save_path):\n",
    "    ckpt = torch.load(save_path, map_location=device)\n",
    "    model.load_state_dict(ckpt[\"model\"])\n",
    "    print(f\"✅ Loaded model from {save_path}\")\n",
    "else:\n",
    "    raise FileNotFoundError(f\"No checkpoint found at {save_path}\")\n",
    "samples = pc_sampler(model, marginal_prob_std_fn, diffusion_coeff_fn,\n",
    "                     batch_size=sample_batch_size, num_steps=200, snr=0.16, device=device)\n",
    "samples_disp=samples\n",
    "show_images_grid(samples_disp, nrow=9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "921952d4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Loaded model from ckptCIFAR-OrNR.pth\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling: 100%|██████████| 500/500 [10:26<00:00,  1.25s/it]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 293676 characters]",
      "text/plain": [
       "<Figure size 288x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling: 100%|██████████| 1000/1000 [22:39<00:00,  1.36s/it]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 280460 characters]",
      "text/plain": [
       "<Figure size 288x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "save_path = 'ckptCIFAR-OrNR.pth'\n",
    "#Original score matching images \n",
    "if os.path.exists(save_path):\n",
    "    ckpt = torch.load(save_path, map_location=device)\n",
    "    model.load_state_dict(ckpt[\"model\"])\n",
    "    print(f\"✅ Loaded model from {save_path}\")\n",
    "else:\n",
    "    raise FileNotFoundError(f\"No checkpoint found at {save_path}\")\n",
    "sample_batch_size = 81\n",
    "samples = pc_sampler(model, marginal_prob_std_fn, diffusion_coeff_fn,\n",
    "                     batch_size=sample_batch_size, num_steps=500, snr=0.16, device=device)\n",
    "# # unnormalize for display\n",
    "# # samples_disp = (samples * 0.5) + 0.5\n",
    "samples_disp=samples\n",
    "show_images_grid(samples_disp, nrow=9)\n",
    "\n",
    "samples = pc_sampler(model, marginal_prob_std_fn, diffusion_coeff_fn,\n",
    "                     batch_size=sample_batch_size, num_steps=1000, snr=0.16, device=device)\n",
    "\n",
    "\n",
    "#num_steps=1000\n",
    "samples_disp=samples\n",
    "show_images_grid(samples_disp, nrow=9)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "adb18336",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling: 100%|██████████| 500/500 [11:20<00:00,  1.36s/it]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 291612 characters]",
      "text/plain": [
       "<Figure size 288x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#num_steps=500\n",
    "\n",
    "samples = pc_sampler(model, marginal_prob_std_fn, diffusion_coeff_fn,\n",
    "                     batch_size=sample_batch_size, num_steps=500, snr=0.16, device=device)\n",
    "\n",
    "samples_disp=samples\n",
    "show_images_grid(samples_disp, nrow=9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "59eed2ec",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling: 100%|██████████| 200/200 [04:35<00:00,  1.38s/it]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 342676 characters]",
      "text/plain": [
       "<Figure size 288x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#num_steps=200\n",
    "\n",
    "samples = pc_sampler(model, marginal_prob_std_fn, diffusion_coeff_fn,\n",
    "                     batch_size=sample_batch_size, num_steps=200, snr=0.16, device=device)\n",
    "\n",
    "samples_disp=samples\n",
    "show_images_grid(samples_disp, nrow=9)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
