{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exercise Session 14: Diffusion Models\n",
    "\n",
    "Recently, diffusion models have been broadly used in generative AI:\n",
    "- MidJourney (images),\n",
    "- Suno AI (music),\n",
    "- Sora (video),\n",
    "\n",
    "and many others!\n",
    "\n",
    "Despite being powerful and realistic, diffusion models remain conceptually simple to analyze. In this tutorial, we are going to code our own diffusion model.\n",
    "\n",
    "Let us consider a situation where you have a dataset made of many images, for example flowers. As humans, we know that certain characteristics make an image look like a flower: the shape of the petals, the stem, typical colors, textures, etc.\n",
    "Thus, generating a random flower requires respecting many constraints. Mathematically, it can be seen as drawing random images from a complicated probability distribution, which we call $P_{\\rm real}$, representing the distribution of real flowers.\n",
    "\n",
    "In the context of this practical session, using real flower images would require too much computational power.  \n",
    "Therefore, we are going to focus on a simpler dataset, that is **MNIST**.  \n",
    "Here, the distribution $P_{\\rm real}$ corresponds to the distribution of handwritten digits that we want to sample from.\n",
    "\n",
    "We will follow the usual procedure of diffusion models that we will explain on the way. Let us start with some introductory work."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 237
    },
    "executionInfo": {
     "elapsed": 6928,
     "status": "ok",
     "timestamp": 1764328617283,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "DkAMRPLjA_Se",
    "outputId": "a97e629d-8b57-4b5e-b29f-0a0067a3c5ee"
   },
   "outputs": [],
   "source": [
    "from keras.datasets import mnist\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import math\n",
    "\n",
    "#We load the data\n",
    "(x_train, _), (x_test, _) = mnist.load_data()\n",
    "\n",
    "#We renormalize the data to a float in [0,1]\n",
    "x_train = x_train.astype('float32') / 255.\n",
    "x_test = x_test.astype('float32') / 255.\n",
    "\n",
    "#We reshape the images to be 2d matrices of dimension n times d\n",
    "x_train = x_train.reshape((len(x_train), np.prod(x_train.shape[1:])))\n",
    "x_test = x_test.reshape((len(x_test), np.prod(x_test.shape[1:])))\n",
    "print(x_train.shape)\n",
    "print(x_test.shape)\n",
    "\n",
    "#We plot the first ten images\n",
    "%matplotlib inline\n",
    "\n",
    "n = 10\n",
    "plt.figure(figsize=(20, 2))\n",
    "for i in range(n):\n",
    "    ax = plt.subplot(1, n, i+1)\n",
    "    plt.imshow(x_test[i].reshape(28, 28))\n",
    "    plt.gray()\n",
    "    ax.get_xaxis().set_visible(False)\n",
    "    ax.get_yaxis().set_visible(False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Diffusions models use CNN architecture. Let's transform the data in an adapted shape for a CNN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 74,
     "status": "ok",
     "timestamp": 1764328617357,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "lTn3j0XdF-Vb",
    "outputId": "b7a59765-f29d-4b21-af4e-694436c6260a"
   },
   "outputs": [],
   "source": [
    "# Reshaping the data for CNN\n",
    "x_train_img = x_train.reshape(-1, 28, 28, 1).astype(\"float32\")\n",
    "x_test_img  = x_test.reshape(-1, 28, 28, 1).astype(\"float32\")\n",
    "\n",
    "x_train_img.shape, x_test_img.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Denoising with CNN\n",
    "\n",
    "Before coding the diffusion model, let us first code a CNN denoiser.\n",
    "\n",
    "We note $\\{ \\mathbf{x}_{{\\rm real},\\, i} \\}_{i \\in [n]}$ the set of $n$ images of dimensions $d \\times d$ contained in the dataset. We will artificially noise each of them for given values of noise $\\sigma_i$ and train a CNN to denoise them. Let us highlight that the CNN is aware of the noise level $\\sigma_i$ for each picture and will take this parameter in consideration during training.\n",
    "\n",
    "Let us specify the setting. For each image $\\mathbf{x}_{{\\rm real},\\, i} \\in \\mathbb{R}^{d \\times d}$ we will add noise of magnitude $\\sigma_i \\in \\mathbb{R}_+$. Then the noisy image is\n",
    "$$\\mathbf{x}_{{\\rm noisy},\\, i} = \\mathbf{x}_{{\\rm real},\\, i} + \\sigma_i \\boldsymbol{\\varepsilon}_i \\, ,$$\n",
    "where $\\boldsymbol{\\varepsilon}_i$ is an image of i.i.d. Gaussian pixels (each pixel distributed as $N(0,1)$).\n",
    "\n",
    "Our goal is to train a CNN $\\mathbf{f}_{\\mathbf{w}}$ that recovers the clean image from the noisy one, while also receiving the noise level $\\sigma_i$ as an additional input.\n",
    "The denoising loss for image $i$ is therefore:\n",
    "$$ L_i(\\mathbf{w}) = \\left\\|\\, \\mathbf{f}_{\\mathbf{w}}(\\mathbf{x}_{{\\rm noisy},\\, i}, \\sigma_i) - \\mathbf{x}_{{\\rm real},\\, i} \\,\\right\\|^2$$\n",
    "In practice, training is done in mini-batches that will have a size $B$.\n",
    "\n",
    "Additionally, for computational purpose, we will define an effective input $\\mathbf{x}_{{\\rm in},\\, i}$ as the concatenation of $\\mathbf{x}_{{\\rm noisy},\\, i}$ and $\\sigma_i$. Note that $\\sigma_i$ must be transformed in an adapted shape. Another viewpoint, is to consider that the noise is another channel. For instance, for color images there are 3 channels (RGB). Here there will be two (gray level, noise).  \n",
    "\n",
    "For the range of noise values, we will define a set of possible values for the noise and we will choose them uniformly.\n",
    "\n",
    "This will conclude the training part."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 8,
     "status": "ok",
     "timestamp": 1764328617365,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "Fm-auC0FGHw5",
    "outputId": "aabc3641-7e07-4cb7-e3ea-6328bfd052bc"
   },
   "outputs": [],
   "source": [
    "# Let's define the possible values for the noise.\n",
    "SIGMA_SET = np.array([0.00, 0.05, 0.10, 0.20, 0.35, 0.50], dtype=\"float32\")\n",
    "\n",
    "# We create a function to pick them uniformly\n",
    "def sample_sigma(batch_size):\n",
    "    return np.random.choice(SIGMA_SET, size=(batch_size, 1, 1, 1)).astype(\"float32\")\n",
    "\n",
    "# For a given value of noise and a given image, compute the noisy image.\n",
    "def add_noise(x, sigma):\n",
    "    \"\"\"\n",
    "    x: (B,28,28,1) in [0,1]\n",
    "    sigma: (B,1,1,1)\n",
    "    -> x_noisy: (B,28,28,1)\n",
    "    \"\"\"\n",
    "    eps = np.random.randn(*x.shape).astype(\"float32\")\n",
    "\n",
    "    # your code\n",
    "    # x_noisy = \n",
    "\n",
    "    return np.clip(x_noisy, 0.0, 1.0)\n",
    "\n",
    "# Just a bit of formatting for giving to the network\n",
    "def with_sigma_channel(x_noisy, sigma):\n",
    "    \"\"\"\n",
    "    transform the scalar sigma as a constant image channel.\n",
    "    -> (B,28,28,2)  [image_noisy, sigma_map]\n",
    "    \"\"\"\n",
    "    sigma_map = np.tile(sigma, (1, 28, 28, 1))  # (B,28,28,1)\n",
    "    return np.concatenate([x_noisy, sigma_map], axis=-1)\n",
    "\n",
    "# Sanity check\n",
    "_b = 8\n",
    "_sigma = sample_sigma(_b)\n",
    "_xn = add_noise(x_train_img[:_b], _sigma)\n",
    "_xin = with_sigma_channel(_xn, _sigma)\n",
    "_xin.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 1099,
     "status": "ok",
     "timestamp": 1764328618485,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "tz2_4lqkGSK0"
   },
   "outputs": [],
   "source": [
    "BATCH = 128\n",
    "\n",
    "# Let us define the function that creates the batches of data for training\n",
    "def train_generator():\n",
    "    N = x_train_img.shape[0]\n",
    "    while True:\n",
    "        idx = np.random.randint(0, N, size=(BATCH,))\n",
    "        x_clean = x_train_img[idx]                                  # (B,28,28,1)\n",
    "        sigma   = sample_sigma(BATCH)                                # (B,1,1,1)\n",
    "        x_noisy = add_noise(x_clean, sigma)                          # (B,28,28,1)\n",
    "        x_in    = with_sigma_channel(x_noisy, sigma)                 # (B,28,28,2)\n",
    "        yield x_in, x_clean\n",
    "\n",
    "# The same for test\n",
    "def val_generator():\n",
    "    N = x_test_img.shape[0]\n",
    "    while True:\n",
    "        idx = np.random.randint(0, N, size=(BATCH,))\n",
    "        x_clean = x_test_img[idx]\n",
    "        sigma   = sample_sigma(BATCH)\n",
    "        x_noisy = add_noise(x_clean, sigma)\n",
    "        x_in    = with_sigma_channel(x_noisy, sigma)\n",
    "        yield x_in, x_clean"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that almost all the pipeline is done, we just have to choose the architecture. Since the noise level is staying low (${\\rm max} \\, \\sigma_i = 0.5$), we do not need a big architecture, the following will do."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 257
    },
    "executionInfo": {
     "elapsed": 1692,
     "status": "ok",
     "timestamp": 1764328620179,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "OpHQKMY7GhG9",
    "outputId": "e0a7057a-6eb7-4ccc-a1d8-f4ed7a11b62b"
   },
   "outputs": [],
   "source": [
    "from keras import layers, models, optimizers\n",
    "\n",
    "model = models.Sequential([\n",
    "    layers.Input(shape=(28, 28, 2)),                 # 1 image channel + 1 σ channel\n",
    "    layers.Conv2D(32, 3, padding=\"same\", activation=\"relu\"),\n",
    "    layers.Conv2D(64, 3, padding=\"same\", activation=\"relu\"),\n",
    "    layers.Conv2D(32, 3, padding=\"same\", activation=\"relu\"),\n",
    "    layers.Conv2D(1,  3, padding=\"same\", activation=\"sigmoid\")  # output in [0,1]\n",
    "])\n",
    "\n",
    "model.compile(optimizer=optimizers.Adam(2e-4), loss=\"mse\")\n",
    "model.summary()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EQ3kP2-6pYnT"
   },
   "source": [
    "Let's train the model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 37630,
     "status": "ok",
     "timestamp": 1764328657809,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "aL0NQ5n0Gk7X",
    "outputId": "563bc8f6-8a49-4ae5-be50-4778b1a1021c"
   },
   "outputs": [],
   "source": [
    "steps_per_epoch = 60000 // BATCH\n",
    "val_steps = 10000 // BATCH\n",
    "\n",
    "hist = model.fit(\n",
    "    train_generator(),\n",
    "    steps_per_epoch=steps_per_epoch,\n",
    "    validation_data=val_generator(),\n",
    "    validation_steps=val_steps,\n",
    "    epochs=5\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "osr6tdpEpgf3"
   },
   "source": [
    "Now that the model is trained, we would like to see if on a test image, the CNN is able to denoise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 3,
     "status": "ok",
     "timestamp": 1764328657814,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "kmFojZF9qn1m"
   },
   "outputs": [],
   "source": [
    "# Denoising\n",
    "def denoise(x_clean_or_noisy, sigma_value):\n",
    "    \"\"\"\n",
    "    x_clean_or_noisy: (B,28,28,1) en [0,1]\n",
    "    sigma_value: float (ex: 0.2) or array shape (B,1,1,1)\n",
    "    -> denoised image\n",
    "    \"\"\"\n",
    "    if np.isscalar(sigma_value):                          # Just to be sure that sigma is not already reshaped.\n",
    "        sigma = np.full((x_clean_or_noisy.shape[0],1,1,1), float(sigma_value), dtype=\"float32\")\n",
    "    else:\n",
    "        sigma = sigma_value.astype(\"float32\")\n",
    "    \n",
    "    # your code \n",
    "    # x_noisy =        \n",
    "    # x_in    = \n",
    "    \n",
    "    return model.predict(x_in, verbose=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 556
    },
    "executionInfo": {
     "elapsed": 1049,
     "status": "ok",
     "timestamp": 1764328658876,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "HP75SLImG4gc",
    "outputId": "1eb57a26-d6ed-4226-efa6-7f2ca1135cb3"
   },
   "outputs": [],
   "source": [
    "# Plotting for several images\n",
    "n = 10\n",
    "x0 = x_test_img[:n]\n",
    "sigma_demo = 0.2\n",
    "\n",
    "\n",
    "x_noisy = add_noise(x0, np.full((n,1,1,1), sigma_demo, dtype=\"float32\"))\n",
    "x_hat   = denoise(x0, sigma_demo)\n",
    "\n",
    "plt.figure(figsize=(20, 6))\n",
    "for i in range(n):\n",
    "    ax = plt.subplot(3, n, i+1);   plt.imshow(x0[i].reshape(28,28));      plt.gray(); ax.set_axis_off()\n",
    "    ax = plt.subplot(3, n, n+i+1); plt.imshow(x_noisy[i].reshape(28,28));  plt.gray(); ax.set_axis_off()\n",
    "    ax = plt.subplot(3, n, 2*n+i+1); plt.imshow(x_hat[i].reshape(28,28)); plt.gray(); ax.set_axis_off()\n",
    "plt.suptitle(f\"Clean / Noisy (σ={sigma_demo}) / Denoised\", y=0.98, fontsize=14)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "executionInfo": {
     "elapsed": 1100,
     "status": "ok",
     "timestamp": 1764328659977,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "AlMprcxSYmFd",
    "outputId": "ff744ad5-621f-495d-cfde-c39142bcc7da"
   },
   "outputs": [],
   "source": [
    "# Plotting for several values of sigma\n",
    "idx = 0\n",
    "x0 = x_test_img[idx:idx+1]  # (1,28,28,1)\n",
    "sigmas = np.array([0.00, 0.05, 0.10, 0.20, 0.35, 0.50], dtype=\"float32\")\n",
    "\n",
    "imgs = []\n",
    "titles = []\n",
    "for s in sigmas:\n",
    "    sigma = np.full((1,1,1,1), s, dtype=\"float32\")\n",
    "    x_noisy = add_noise(x0, sigma)\n",
    "    x_hat   = denoise(x0, s)\n",
    "    imgs += [x_noisy[0,...,0], x_hat[0,...,0]]\n",
    "    titles += [f\"Noisy σ={s:.2f}\", f\"Denoised σ={s:.2f}\"]\n",
    "\n",
    "plt.figure(figsize=(12, 2*len(sigmas)))\n",
    "for i, (im, tl) in enumerate(zip(imgs, titles), 1):\n",
    "    ax = plt.subplot(len(sigmas), 2, i); plt.imshow(im, cmap=\"gray\"); ax.set_axis_off(); ax.set_title(tl)\n",
    "plt.suptitle(\"Balayage multi-σ pour une image\", y=0.99); plt.tight_layout(); plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 307
    },
    "executionInfo": {
     "elapsed": 1798,
     "status": "ok",
     "timestamp": 1764328661776,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "Bk7gVex-G5SV",
    "outputId": "3d9a7a6a-2967-46d5-92b3-d3c5cb98611c"
   },
   "outputs": [],
   "source": [
    "# Let's plot the MSE vs the noise level\n",
    "def eval_curve(sigmas_eval, N=500):\n",
    "    idx = np.random.choice(x_test_img.shape[0], size=N, replace=False)\n",
    "    X = x_test_img[idx]\n",
    "    mse_list, psnr_list = [], []\n",
    "    for s in sigmas_eval:\n",
    "        sigma = np.full((N,1,1,1), s, dtype=\"float32\")\n",
    "        X_noisy = add_noise(X, sigma)\n",
    "        X_hat   = denoise(X, s)\n",
    "        mse_list.append(np.mean((X_hat - X)**2))\n",
    "    return np.array(mse_list)\n",
    "\n",
    "sigmas_eval = np.linspace(0.0, 0.8, 9, dtype=\"float32\")\n",
    "mse_v = eval_curve(sigmas_eval, N=500)\n",
    "plt.figure(figsize=(10,3))\n",
    "plt.plot(sigmas_eval, mse_v, marker='o'); plt.xlabel(\"σ\"); plt.ylabel(\"MSE\"); plt.title(\"MSE vs σ\")\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deterministic Model\n",
    "\n",
    "To build our first diffusion model, we now incorporate a notion of time into both the noising and the denoising process.  \n",
    "More precisely, we define a continuous interpolation between the real image and the pure Gaussian noise:\n",
    "\n",
    "$$\n",
    "\\mathbf{x}_{{\\rm noisy},\\, i}(t)\n",
    "= a(t)\\,\\mathbf{x}_{{\\rm real},\\, i}\n",
    "+ b(t)\\,\\boldsymbol{\\varepsilon}_i \\, ,\n",
    "$$\n",
    "\n",
    "where $a(t)$ and $b(t)$ satisfy  \n",
    "$$\n",
    "{\\rm at}\\; t=0,\\; \\mathbf{x}_{{\\rm noisy},\\, i} = \\mathbf{x}_{{\\rm real},\\, i}\\, , \\quad {\\rm and} \\;\n",
    "{\\rm at}\\; t=T,\\; \\mathbf{x}_{{\\rm noisy},\\, i} = \\boldsymbol{\\varepsilon}_i\\, .\n",
    "$$\n",
    "In other words, the parameter $t$ smoothly controls the noise intensity.  \n",
    "During training, the role previously played by the noise level $\\sigma_i$ is now replaced by the time index $t$.\n",
    "\n",
    "This defines the ideal behaviour we want from our forward noising process.  \n",
    "However, to implement diffusion properly, first we need to describe the iterations. The forward dynamics are defined between successive timesteps by:\n",
    "$$\n",
    "\\mathbf{x}_t\n",
    "= \\sqrt{\\alpha_t}\\,\\mathbf{x}_{t-1}\n",
    "\\;+\\;\n",
    "\\sqrt{1-\\alpha_t}\\,\\boldsymbol{\\varepsilon}_t,\n",
    "\\qquad\n",
    "\\boldsymbol{\\varepsilon}_t \\sim \\mathcal{N}(0,I).\n",
    "$$\n",
    "From these local updates, one can show, by expanding the recursion, that $\\mathbf{x}_t$ always admits the closed form\n",
    "$$\n",
    "\\mathbf{x}_t\n",
    "= \\sqrt{\\bar{\\alpha}_t}\\,\\mathbf{x}_0\n",
    "\\;+\\;\n",
    "\\sqrt{1-\\bar{\\alpha}_t}\\,\\boldsymbol{\\varepsilon},\n",
    "$$\n",
    "where\n",
    "$$\n",
    "\\bar{\\alpha}_t \\;=\\; \\alpha_1 \\alpha_2 \\cdots \\alpha_t.\n",
    "$$\n",
    "This equation connects with our first approach through:\n",
    "$$\n",
    "a(t) = \\sqrt{\\bar{\\alpha}_t}\\,,\n",
    "\\quad\n",
    "b(t) = \\sqrt{1 - \\bar{\\alpha}_t}.\n",
    "$$\n",
    "One can also note that:\n",
    "$$\n",
    "\\alpha_t = \\frac{\\bar\\alpha_t}{\\bar\\alpha_{t-1}} \\, , \\quad \\beta_t = 1- \\alpha_t\\, .\n",
    "$$\n",
    "This more technical approach is essential to code the iteration of the forward and backward processes.\n",
    "\n",
    "Secondly, passing $t$ directly as a raw scalar to a CNN is not ideal.  \n",
    "Convolutional networks operate through spatial filters, and they react better when the conditioning signal (time, here) is expressed in a richer form rather than as a single value.\n",
    "This is why diffusion models typically use a time embedding based on simple Fourier features: instead of feeding $t$ alone, we feed several sinusoids of $t$ at different frequencies. This is nothing more than a feature map $\\phi$ which also allows to increase the amount of parameters and improve the precision of the model.\n",
    "\n",
    "More precisely, after normalizing $t_i$ into $[0,1]$, we define:\n",
    "$$\n",
    "\\phi(t_i) =\n",
    "\\big(\\sin(2\\pi\\,\\omega_1 t_i),\\; \\cos(2\\pi\\,\\omega_1 t_i),\\; \\dots,\\;\n",
    "      \\sin(2\\pi\\,\\omega_K t_i),\\; \\cos(2\\pi\\,\\omega_K t_i) \\big),\n",
    "$$\n",
    "where the frequencies are chosen as $\\omega_k = 2^k$.  \n",
    "\n",
    "They give the network a multi-scale description of the time index, which makes the conditioning on $t_i$ much easier to learn.  \n",
    "\n",
    "Once the vector $\\phi(t_i)$ is computed, we simply broadcast it across the spatial dimensions and concatenate it to the noisy image as additional channels.  \n",
    "Thus, the network receives an input tensor of shape\n",
    "$(d, d,\\, 1 + 2K)$:\n",
    "- $1$ channel for pixels values,\n",
    "- $2K$ channels for the time embedding.\n",
    "\n",
    "With this input, the denoising objective becomes:\n",
    "$$\n",
    "L_i(\\mathbf{w}) =\n",
    "\\left\\|\\,\n",
    "\\mathbf{f}_{\\mathbf{w}}\\big(\\mathbf{x}_{{\\rm noisy},\\, i}(t_i),\\, \\phi(t_i)\\big)\n",
    "- \\mathbf{x}_{{\\rm real},\\, i}\n",
    "\\right\\|^2.\n",
    "$$\n",
    "\n",
    "This time embedding mechanism is a crucial ingredient of modern diffusion models, and it ensures that the CNN can effectively learn how to denoise images at different stages of the diffusion process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 2,
     "status": "ok",
     "timestamp": 1764328661776,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "L6zquykoiMQf"
   },
   "outputs": [],
   "source": [
    "# Rescaling [0,1] -> [-1,1]\n",
    "def to_symmetric(x):      # x in [0,1]\n",
    "    return x * 2.0 - 1.0  # in [-1,1]\n",
    "\n",
    "def from_symmetric(x):    # x in [-1,1]\n",
    "    return np.clip((x + 1.0) * 0.5, 0.0, 1.0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 1,
     "status": "ok",
     "timestamp": 1764328661776,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "OuLByf7vY0AF"
   },
   "outputs": [],
   "source": [
    "# The function used for time interpolation are cosines.\n",
    "T = 200\n",
    "t = (np.arange(T) + 0.5)/T\n",
    "s = 0.008\n",
    "f = lambda u: np.cos((np.pi/2)*(u+s)/(1+s))**2\n",
    "alpha_bar = (f(t) / f(0)).astype(\"float32\")\n",
    "alpha = np.concatenate([[alpha_bar[0]], alpha_bar[1:]/alpha_bar[:-1]]).astype(\"float32\")\n",
    "beta  = (1.0 - alpha).astype(\"float32\")\n",
    "\n",
    "# We compute the noisy image at a given time\n",
    "def q_sample(x0_sym, t_idx, eps):\n",
    "    \"\"\"x_t = sqrt(\\bar{α}_t) x0 + sqrt(1-\\bar{α}_t) ε\"\"\"\n",
    "    a = np.sqrt(alpha_bar[t_idx]).reshape(-1,1,1,1).astype(\"float32\")\n",
    "    b = np.sqrt(1.0 - alpha_bar[t_idx]).reshape(-1,1,1,1).astype(\"float32\")\n",
    "    return a * x0_sym + b * eps\n",
    "\n",
    "# Temporal embedding through Fourier feature (2K channels)\n",
    "K = 8\n",
    "def t_embed_map(t_idx, K=K):\n",
    "    tt = (t_idx.astype(\"float32\") / (T-1)).reshape(-1,1,1,1)  # in [0,1]\n",
    "\n",
    "    # your code\n",
    "    # freqs =               # [1,2,4,...]\n",
    "    # ang =                 # (B,1,1,K)\n",
    "    # sin =                 # (B,1,1,K)\n",
    "    # emb =                 # (B,1,1,2K)\n",
    "    \n",
    "    return np.tile(emb, (1,28,28,1)).astype(\"float32\")        # (B,28,28,2K)\n",
    "\n",
    "def make_input(x_like, t_idx):  # x_like: (B,28,28,1) en [-1,1]\n",
    "    return np.concatenate([x_like, t_embed_map(t_idx, K=K)], axis=-1).astype(\"float32\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DpURSg9KDXm5"
   },
   "source": [
    "Architecture:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 689
    },
    "executionInfo": {
     "elapsed": 98,
     "status": "ok",
     "timestamp": 1764328661873,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "bpuSdvsVijNj",
    "outputId": "0cac9f40-7a65-4e7e-d0d2-37d04ecb94e9"
   },
   "outputs": [],
   "source": [
    "from keras import layers, models, optimizers, Input\n",
    "\n",
    "CIN = 1 + 2*K  # 1 channel x_t + 2K time-embedding channels\n",
    "\n",
    "inp = Input(shape=(28, 28, CIN))                       # 28x28\n",
    "\n",
    "# --- Encoding level 1 (28x28) ---\n",
    "e1 = layers.Conv2D(64, 3, padding=\"same\", activation=\"swish\")(inp)\n",
    "e1 = layers.Conv2D(64, 3, padding=\"same\", activation=\"swish\")(e1)    # 28x28\n",
    "\n",
    "# --- Encoding level 2 (14x14) ---\n",
    "e2 = layers.Conv2D(128, 3, strides=2, padding=\"same\", activation=\"swish\")(e1)  # 14x14\n",
    "e2 = layers.Conv2D(128, 3, padding=\"same\", activation=\"swish\")(e2)             # 14x14\n",
    "\n",
    "# --- Bottleneck (14x14) ---\n",
    "b  = layers.Conv2D(128, 3, padding=\"same\", activation=\"swish\")(e2)\n",
    "b  = layers.Conv2D(128, 3, padding=\"same\", activation=\"swish\")(b)\n",
    "\n",
    "# --- Decoder level 2 -> 1 (14x14 -> 28x28) ---\n",
    "d1 = layers.Conv2DTranspose(64, 3, strides=2, padding=\"same\", activation=\"swish\")(b)  # 28x28\n",
    "d1 = layers.Concatenate()([d1, e1])                                                   # skip connection\n",
    "d1 = layers.Conv2D(64, 3, padding=\"same\", activation=\"swish\")(d1)\n",
    "d1 = layers.Conv2D(64, 3, padding=\"same\", activation=\"swish\")(d1)\n",
    "\n",
    "# --- Output: x0_hat in [-1,1] ---\n",
    "out = layers.Conv2D(1, 3, padding=\"same\", activation=None)(d1)\n",
    "\n",
    "model_diff = models.Model(inp, out)\n",
    "model_diff.compile(optimizer=optimizers.Adam(1e-3), loss=\"mse\")\n",
    "model_diff.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 7,
     "status": "ok",
     "timestamp": 1764328661881,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "HiOa8G43ilq_"
   },
   "outputs": [],
   "source": [
    "BATCH = 256\n",
    "\n",
    "def diffusion_train_generator():\n",
    "    N = x_train_img.shape[0]\n",
    "    while True:\n",
    "        idx = np.random.randint(0, N, size=(BATCH,))\n",
    "        x0 = x_train_img[idx]                                # [0,1]\n",
    "        x0_sym = to_symmetric(x0)                            # [-1,1]\n",
    "\n",
    "        t_idx = np.random.randint(0, T, size=(BATCH,), dtype=np.int64)\n",
    "        eps = np.random.randn(*x0.shape).astype(\"float32\")   # ε ~ N(0,1)\n",
    "\n",
    "        x_t = q_sample(x0_sym, t_idx, eps)                   # [-1,1]\n",
    "        x_in = make_input(x_t, t_idx)\n",
    "\n",
    "        yield x_in, x0_sym"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "n8owKULPDQmP"
   },
   "source": [
    "Training!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 659231,
     "status": "ok",
     "timestamp": 1764329321112,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "2WGArG8jiou-",
    "outputId": "b0aaa04f-46a8-436a-996a-49e0fb3b8485"
   },
   "outputs": [],
   "source": [
    "steps_per_epoch = 60000 // BATCH\n",
    "hist = model_diff.fit(diffusion_train_generator(),\n",
    "                      steps_per_epoch=steps_per_epoch,\n",
    "                      epochs=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deterministic Reverse Process\n",
    "\n",
    "Once the denoiser is trained, we can use it to run the diffusion process backwards.  \n",
    "Recall that the forward noising step was defined as\n",
    "$$\n",
    "\\mathbf{x}_t\n",
    "= \\sqrt{\\bar{\\alpha}_t}\\,\\mathbf{x}_0\n",
    "\\;+\\;\n",
    "\\sqrt{1-\\bar{\\alpha}_t}\\,\\boldsymbol{\\varepsilon}.\n",
    "$$\n",
    "\n",
    "In the deterministic reverse process, we start from pure noise at $t=T$ and move backwards to $t=0$.  \n",
    "At each step, the model receives the current noisy image $\\mathbf{x}_t$ together with its time embedding $\\phi(t)$, and predicts an estimate of the clean image:\n",
    "$$\n",
    "\\widehat{\\mathbf{x}}_0 = \\mathbf{f}_{\\mathbf{w}}(\\mathbf{x}_t , \\phi(t)).\n",
    "$$\n",
    "\n",
    "With this estimate, we recover first the effective noise:\n",
    "$$\n",
    "\\widehat{\\boldsymbol{\\varepsilon}}_t\n",
    "=\n",
    "\\frac{\\mathbf{x}_t - \\sqrt{\\bar{\\alpha}_t}\\,\\widehat{\\mathbf{x}}_0}\n",
    "     {\\sqrt{\\,1-\\bar{\\alpha}_t\\,}}.\n",
    "$$\n",
    "\n",
    "Secondly, we use the estimates $\\widehat{\\mathbf{x}}_0$ and $\\widehat{\\boldsymbol{\\varepsilon}}_t$ but evaluated at $t-1$:\n",
    "$$\n",
    "\\mathbf{x}_{t-1}\n",
    "=\n",
    "\\sqrt{\\bar{\\alpha}_{t-1}}\\,\\widehat{\\mathbf{x}}_0\n",
    "\\;+\\;\n",
    "\\sqrt{1-\\bar{\\alpha}_{t-1}}\\,\\widehat{\\boldsymbol{\\varepsilon}}_t.\n",
    "$$\n",
    "\n",
    "No randomness is added in this process, so starting from the same initial noise always produces the same reconstructed image.  \n",
    "This is the deterministic version of diffusion sampling.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 609
    },
    "executionInfo": {
     "elapsed": 15195,
     "status": "ok",
     "timestamp": 1764329336320,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "KnZcTYUdiq0r",
    "outputId": "c4259c57-e990-41db-d81c-a5dc72d9ee05"
   },
   "outputs": [],
   "source": [
    "# (We want to repeat the process 16 times to have a great plot!)\n",
    "\n",
    "def sample_ddim(n=16, steps=100):\n",
    "    t_seq = np.linspace(T-1, 0, steps, dtype=int)\n",
    "    x = np.random.randn(n,28,28,1).astype(\"float32\")  # x_T ~ N(0,1)\n",
    "\n",
    "    for t in t_seq:\n",
    "        t_arr = np.full((n,), t, dtype=np.int64)\n",
    "        x_in = make_input(x, t_arr)\n",
    "\n",
    "        ab_t   = float(alpha_bar[t])\n",
    "        ab_prev= float(alpha_bar[t-1]) if t > 0 else 1.0\n",
    "\n",
    "        # your code\n",
    "        # x0_hat =          # x0_hat in [-1,1]\n",
    "        # eps_hat =         # reconstituted ε̂\n",
    "        # x = \n",
    "\n",
    "    return from_symmetric(x)  # -> [0,1]\n",
    "\n",
    "samples = sample_ddim(n=16, steps=200)\n",
    "\n",
    "plt.figure(figsize=(6,6))\n",
    "for i in range(min(16, samples.shape[0])):\n",
    "    ax = plt.subplot(4,4,i+1)\n",
    "    plt.imshow(samples[i,...,0], cmap=\"gray\", vmin=0.0, vmax=1.0)\n",
    "    ax.set_axis_off()\n",
    "plt.suptitle(\"Samples from deterministic process\"); plt.tight_layout(); plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jOJZS30gPwzu"
   },
   "source": [
    "A really fancy type of videos that we see often try to show the forward and the backward pass. So that from a test sample we can generate another one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 578
    },
    "executionInfo": {
     "elapsed": 15660,
     "status": "ok",
     "timestamp": 1764329351981,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "Y6bRLooeruRD",
    "outputId": "e25a116b-e550-4927-c18e-8d43495ca189"
   },
   "outputs": [],
   "source": [
    "from matplotlib.animation import FuncAnimation\n",
    "from IPython.display import HTML\n",
    "\n",
    "\n",
    "def forward_noising_frames(x0_img01, n_frames=30):\n",
    "    \"\"\"x0_img01 shape (1,28,28,1) in [0,1]. Return frames with progressive noise.\"\"\"\n",
    "    x0s = to_symmetric(x0_img01)\n",
    "    frames = [x0_img01[0,...,0]]\n",
    "    t_seq = np.linspace(0, T-1, n_frames, dtype=int)\n",
    "    for t in t_seq:\n",
    "        eps = np.random.randn(*x0s.shape).astype(\"float32\")\n",
    "        x_t = q_sample(x0s, np.array([t]), eps)\n",
    "        frames.append(from_symmetric(x_t[0,...,0]))\n",
    "    return frames, int(t_seq[-1]), x_t\n",
    "\n",
    "def reverse_from_xt_frames(x_t, t_start, n_frames=30, steps_per_frame=None):\n",
    "    \"\"\"\n",
    "    Denoising from a given x_t (x_t in [-1,1]), until t=0.\n",
    "    If steps_per_frame is None, we pick it to cover the whole interval.\n",
    "    \"\"\"\n",
    "    frames = [from_symmetric(x_t[0,...,0])]\n",
    "    t_cur = int(t_start)\n",
    "\n",
    "    if steps_per_frame is None:\n",
    "        steps_per_frame = max(1, math.ceil((t_cur + 1) / n_frames))\n",
    "\n",
    "    for _ in range(n_frames):\n",
    "        for _ in range(steps_per_frame):\n",
    "            if t_cur < 0: break\n",
    "            t = t_cur\n",
    "            t_arr = np.full((1,), t, dtype=np.int64)\n",
    "            x_in  = make_input(x_t, t_arr)\n",
    "\n",
    "            # your code\n",
    "            # x0_hat = \n",
    "            # ab_t    = \n",
    "            # ab_prev = \n",
    "            # eps_hat = \n",
    "            # x_t     = \n",
    "            \n",
    "            t_cur  -= 1\n",
    "        frames.append(from_symmetric(x_t[0,...,0]))\n",
    "        if t_cur < 0:  # if already done, we replicate\n",
    "            break\n",
    "\n",
    "    # if you want exactly n+1 frames, you can pad:\n",
    "    while len(frames) < n_frames + 1:\n",
    "        frames.append(frames[-1])\n",
    "\n",
    "    return frames\n",
    "\n",
    "\n",
    "# Choice of a starting sample\n",
    "idx = np.random.randint(0, x_test_img.shape[0])\n",
    "x0_img = x_test_img[idx:idx+1]  # (1,28,28,1) en [0,1]\n",
    "\n",
    "# Forward\n",
    "fwd_frames, t_star, x_t_star = forward_noising_frames(x0_img, n_frames=25)\n",
    "\n",
    "# Backward\n",
    "rev_frames = reverse_from_xt_frames(x_t_star, t_star, n_frames=35, steps_per_frame=None)\n",
    "\n",
    "frames = fwd_frames + rev_frames\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "im = ax.imshow(frames[0], cmap=\"gray\", vmin=0, vmax=1); ax.axis(\"off\")\n",
    "def update(i): im.set_data(frames[i]); return [im]\n",
    "ani = FuncAnimation(fig, update, frames=len(frames), interval=120)\n",
    "plt.close()\n",
    "HTML(ani.to_jshtml())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stochastic Reverse Process\n",
    "\n",
    "The deterministic reverse process always produces the same final image when starting from the same initial noise.  \n",
    "However, real diffusion models are probabilistic: at each step of the reverse trajectory, a small amount of randomness is injected, allowing the model to explore multiple plausible reconstructions.\n",
    "\n",
    "The key idea is the following.  \n",
    "Even if the forward noising process was defined as\n",
    "$$\n",
    "\\mathbf{x}_t = a(t)\\,\\mathbf{x}_0 + b(t)\\,\\boldsymbol{\\varepsilon},\n",
    "$$\n",
    "the reverse direction is not unique: the mapping from $\\mathbf{x}_t$ back to $\\mathbf{x}_{t-1}$ is not a single point but a whole probability distribution.  \n",
    "Adding a stochastic component in the reverse step lets us sample from this distribution instead of collapsing to a single trajectory.\n",
    "\n",
    "In practice, we proceed exactly as in the deterministic case: the model predicts an estimate $\\widehat{\\mathbf{x}}_0$ from $\\mathbf{x}_t$ and $\\phi(t)$. We can still recover the effective noise $\\widehat{\\boldsymbol{\\varepsilon}}_t$ through the identity\n",
    "$$\n",
    "\\mathbf{x}_t = a(t)\\,\\widehat{\\mathbf{x}}_0 + b(t)\\,\\widehat{\\boldsymbol{\\varepsilon}}_t.\n",
    "$$\n",
    "\n",
    "The difference comes in the update rule. Instead of jumping deterministically to\n",
    "$$\n",
    "\\mathbf{x}_{t-1}\n",
    "= a(t-1)\\,\\widehat{\\mathbf{x}}_0 + b(t-1)\\,\\widehat{\\boldsymbol{\\varepsilon}}_t,\n",
    "$$\n",
    "we add a controlled noise term:\n",
    "$$\n",
    "\\mathbf{x}_{t-1}\n",
    "= a(t-1)\\,\\widehat{\\mathbf{x}}_0\n",
    "\\;+\\; b(t-1)\\,\\widehat{\\boldsymbol{\\varepsilon}}_t\n",
    "\\;+\\; \\eta(t)\\,\\mathbf{z}_t,\n",
    "$$\n",
    "where $\\mathbf{z}_t$ is a Gaussian noise and $\\eta(t)$ controls the amount of randomness injected at step $t$. But what becomes the updates of the backward process?\n",
    "\n",
    "With this new Gaussian noise, one can apply a central limit theorem. Namely, the next image $\\mathbf{x}_{t-1}$ for a given $\\mathbf{x}_{t}$ is gaussian:\n",
    "$$\n",
    "\\mathbf{x}_{t-1} \\,\\big|\\, \\mathbf{x}_t \\;\\sim\\; N \\big(\\,\\boldsymbol{\\mu}_t,\\; c_t\\,\\mathbf{I}\\big)\\, ,\n",
    "$$\n",
    "$$\n",
    "\\mathbf{x}_{t-1} = \\boldsymbol{\\mu}_t + \\sqrt{c_t}\\,\\mathbf{g}_t,\n",
    "\\qquad \\mathbf{g}_t\\sim\\mathcal{N}(0,I).\n",
    "$$\n",
    "It remains to compute the mean and the variance. One can show that:\n",
    "\n",
    "- **Mean**\n",
    "$$\n",
    "\\boldsymbol{\\mu}_t\n",
    "=\n",
    "\\frac{\\sqrt{\\alpha_t}(1-\\bar{\\alpha}_{t-1})}{1-\\bar{\\alpha}_t}\\,\\mathbf{x}_t\n",
    "\\;+\\;\n",
    "\\frac{\\sqrt{\\bar{\\alpha}_{t-1}}\\,\\beta_t}{1-\\bar{\\alpha}_t}\\,\n",
    "\\widehat{\\mathbf{x}}_0 ,\n",
    "$$\n",
    "\n",
    "- **Variance**\n",
    "$$\n",
    "c_t\n",
    "=\n",
    "\\frac{1-\\bar{\\alpha}_{t-1}}{1-\\bar{\\alpha}_t}\\,\\beta_t.\n",
    "$$\n",
    "\n",
    "This stochastic variant allows the reverse process to generate different valid images from the same initial noise, illustrating the generative power of diffusion models.\n",
    "\n",
    "The good news is that the difference from the deterministic model is in the reverse process, we do not need to re-train the model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 2,
     "status": "ok",
     "timestamp": 1764329351985,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "gwaXHgnCujV5"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def sample_ddpm(n=16):\n",
    "    \"\"\"\n",
    "    Stochastic reverse diffusion:\n",
    "    - start from x_T ~ N(0, I)\n",
    "    - at each step t, predict x0_hat from (x_t, t)\n",
    "    - compute the mean/variance\n",
    "    - sample x_{t-1} from this Gaussian\n",
    "    Returns samples in [0,1].\n",
    "    \"\"\"\n",
    "    # x_T ~ N(0, I)\n",
    "    x = np.random.randn(n, 28, 28, 1).astype(\"float32\")\n",
    "\n",
    "    for t in reversed(range(T)):\n",
    "        t_arr = np.full((n,), t, dtype=np.int64)\n",
    "\n",
    "        # predict x0_hat in [-1,1]\n",
    "        x_in   = make_input(x, t_arr)\n",
    "        x0_hat = model_diff.predict(x_in, verbose=0).astype(\"float32\")\n",
    "\n",
    "        # your code\n",
    "\n",
    "        # diffusion coefficients at timestep t\n",
    "        # ab_t    = \n",
    "        # ab_prev = \n",
    "        # a_t     = \n",
    "        # b_t     = \n",
    "\n",
    "        # reverse mean\n",
    "        # coef_x0 = \n",
    "        # coef_xt = \n",
    "        \n",
    "        # mean    = \n",
    "\n",
    "        # reverse variance\n",
    "        # var    = \n",
    "\n",
    "        if t > 0:\n",
    "            # stochastic step: add noise\n",
    "            noise = np.random.randn(*x.shape).astype(\"float32\")\n",
    "            x = mean + np.sqrt(var) * noise\n",
    "        else:\n",
    "            # last step: no noise\n",
    "            x = mean\n",
    "\n",
    "    # convert [-1,1] → [0,1]\n",
    "    return from_symmetric(x)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 609
    },
    "executionInfo": {
     "elapsed": 13579,
     "status": "ok",
     "timestamp": 1764329365564,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "V3P_OMdxGNYF",
    "outputId": "b13518d9-099a-40fd-f3ea-2baf39f59996"
   },
   "outputs": [],
   "source": [
    "samples_ddpm = sample_ddpm(n=16)\n",
    "\n",
    "plt.figure(figsize=(6,6))\n",
    "for i in range(16):\n",
    "    ax = plt.subplot(4,4,i+1)\n",
    "    plt.imshow(samples_ddpm[i,...,0], cmap=\"gray\", vmin=0.0, vmax=1.0)\n",
    "    ax.axis(\"off\")\n",
    "plt.suptitle(\"Samples from stochastic process\"); plt.tight_layout(); plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nhSMeZcXX24E"
   },
   "source": [
    "We generate two deterministic runs and two stochastic runs starting from the exact same noised image.  \n",
    "The deterministic ones allow us to check that the reverse process always gives the same result, while the stochastic ones illustrate that randomness can lead to different outcomes.  \n",
    "Having all four side-by-side makes the contrast immediately visible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 598
    },
    "executionInfo": {
     "elapsed": 47359,
     "status": "ok",
     "timestamp": 1764329460511,
     "user": {
      "displayName": "Hugo Tabanelli",
      "userId": "14976810322055414022"
     },
     "user_tz": -60
    },
    "id": "rXhARdyMXNo7",
    "outputId": "f35b8041-fd3c-4eef-c7b3-052f4f69f899"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.animation import FuncAnimation\n",
    "from IPython.display import HTML\n",
    "\n",
    "# -------- 1) Forward noising process --------\n",
    "def forward_noising_frames(x0_img01, n_frames=20):\n",
    "    \"\"\"\n",
    "    Generates a sequence of progressively noised versions of a clean image.\n",
    "    x0_img01 : (1,28,28,1) in [0,1]\n",
    "    Returns: list of frames, final timestep t_star, and the noised x_t at t_star.\n",
    "    \"\"\"\n",
    "    x0s = to_symmetric(x0_img01)          # map to [-1,1]\n",
    "    frames = [x0_img01[0,...,0]]          # first frame = clean image\n",
    "    t_seq = np.linspace(0, int(0.95*T), n_frames, dtype=int)\n",
    "\n",
    "    x_t = None\n",
    "    for t in t_seq:\n",
    "        eps = np.random.randn(*x0s.shape).astype(\"float32\")\n",
    "        x_t = q_sample(x0s, np.array([t]), eps)            # forward process\n",
    "        frames.append(from_symmetric(x_t[0,...,0]))\n",
    "    return frames, int(t_seq[-1]), x_t\n",
    "\n",
    "\n",
    "# -------- 2) Reverse deterministic --------\n",
    "def reverse_ddim_from_xt_frames(x_t, t_start, n_frames=30):\n",
    "    \"\"\"\n",
    "    Deterministic denoising starting from x_t at timestep t_start.\n",
    "    Returns a list of intermediate frames for visualization.\n",
    "    \"\"\"\n",
    "    frames = [from_symmetric(x_t[0,...,0])]\n",
    "    t_cur = int(t_start)\n",
    "    steps_per_frame = max(1, (t_cur+1) // n_frames)\n",
    "\n",
    "    for _ in range(n_frames):\n",
    "        for _ in range(steps_per_frame):\n",
    "            if t_cur < 0:\n",
    "                break\n",
    "            t = t_cur\n",
    "            t_arr = np.full((1,), t, dtype=np.int64)\n",
    "\n",
    "            # deterministic x0-prediction model\n",
    "            x_in   = make_input(x_t, t_arr)\n",
    "            x0_hat = model_diff.predict(x_in, verbose=0).astype(\"float32\")\n",
    "\n",
    "            # your code\n",
    "\n",
    "            # ab_t    = \n",
    "            # ab_prev = \n",
    "\n",
    "            # deterministic update\n",
    "            # eps_hat = \n",
    "            # x_t     = \n",
    "\n",
    "            t_cur -= 1\n",
    "\n",
    "        frames.append(from_symmetric(x_t[0,...,0]))\n",
    "        if t_cur < 0:\n",
    "            break\n",
    "\n",
    "    return frames\n",
    "\n",
    "\n",
    "# -------- 3) Reverse stochastic --------\n",
    "def reverse_ddpm_from_xt_frames(x_t, t_start, n_frames=30):\n",
    "    \"\"\"\n",
    "    Stochastic reverse process:\n",
    "    adds Gaussian noise at every step except t=0.\n",
    "    Returns intermediate frames for visualization.\n",
    "    \"\"\"\n",
    "    frames = [from_symmetric(x_t[0,...,0])]\n",
    "    t_cur = int(t_start)\n",
    "    steps_per_frame = max(1, (t_cur+1) // n_frames)\n",
    "\n",
    "    for _ in range(n_frames):\n",
    "        for _ in range(steps_per_frame):\n",
    "            if t_cur < 0:\n",
    "                break\n",
    "            t = t_cur\n",
    "            t_arr = np.full((1,), t, dtype=np.int64)\n",
    "\n",
    "            # x0-prediction model\n",
    "            x_in   = make_input(x_t, t_arr)\n",
    "            x0_hat = model_diff.predict(x_in, verbose=0).astype(\"float32\")\n",
    "\n",
    "            # your code\n",
    "\n",
    "            # ab_t    = \n",
    "            # ab_prev = \n",
    "            # a_t     = \n",
    "            # b_t     = \n",
    "\n",
    "            # stochastic mean\n",
    "            # coef_x0 = \n",
    "            # coef_xt = \n",
    "            # mean    = \n",
    "\n",
    "            # stochastic variance\n",
    "            # var = \n",
    "\n",
    "            if t > 0:\n",
    "                noise = np.random.randn(*x_t.shape).astype(\"float32\")\n",
    "                x_t   = mean + np.sqrt(var) * noise\n",
    "            else:\n",
    "                x_t   = mean\n",
    "\n",
    "            t_cur -= 1\n",
    "\n",
    "        frames.append(from_symmetric(x_t[0,...,0]))\n",
    "        if t_cur < 0:\n",
    "            break\n",
    "\n",
    "    return frames\n",
    "\n",
    "\n",
    "# -------- 4) 4-way animation: two deterministic runs + two stochastic runs --------\n",
    "def four_run_diffusion_animation(n_fwd=20, n_back=30):\n",
    "    \"\"\"\n",
    "    Produces an animation showing:\n",
    "    - forward noising\n",
    "    - two deterministic reverse runs\n",
    "    - two stochastic reverse runs\n",
    "    \"\"\"\n",
    "    # 1) sample a real test image\n",
    "    idx = np.random.randint(0, x_test_img.shape[0])\n",
    "    x0  = x_test_img[idx:idx+1]\n",
    "\n",
    "    # 2) forward trajectory (shared)\n",
    "    fwd_frames, t_star, x_t_star = forward_noising_frames(x0, n_frames=n_fwd)\n",
    "\n",
    "    # 3) four reverse trajectories starting from the same x_t*\n",
    "    ddim1 = reverse_ddim_from_xt_frames(x_t_star.copy(), t_star, n_frames=n_back)\n",
    "    ddim2 = reverse_ddim_from_xt_frames(x_t_star.copy(), t_star, n_frames=n_back)\n",
    "    ddpm1 = reverse_ddpm_from_xt_frames(x_t_star.copy(), t_star, n_frames=n_back)\n",
    "    ddpm2 = reverse_ddpm_from_xt_frames(x_t_star.copy(), t_star, n_frames=n_back)\n",
    "\n",
    "    # equalize sequence lengths\n",
    "    L = max(len(ddim1), len(ddim2), len(ddpm1), len(ddpm2))\n",
    "    for seq in (ddim1, ddim2, ddpm1, ddpm2):\n",
    "        while len(seq) < L:\n",
    "            seq.append(seq[-1])\n",
    "\n",
    "    # 4) assemble composite frames\n",
    "    frames = []\n",
    "\n",
    "    # forward phase (same image in all 4 quadrants)\n",
    "    for fr in fwd_frames:\n",
    "        top    = np.concatenate([fr, fr], axis=1)\n",
    "        bottom = np.concatenate([fr, fr], axis=1)\n",
    "        tile   = np.concatenate([top, bottom], axis=0)\n",
    "        frames.append(tile)\n",
    "\n",
    "    # reverse trajectories\n",
    "    for k in range(L):\n",
    "        top    = np.concatenate([ddim1[k], ddim2[k]], axis=1)\n",
    "        bottom = np.concatenate([ddpm1[k], ddpm2[k]], axis=1)\n",
    "        tile   = np.concatenate([top, bottom], axis=0)\n",
    "        frames.append(tile)\n",
    "\n",
    "    # 5) animation rendering\n",
    "    fig, ax = plt.subplots(figsize=(5,5))\n",
    "    im = ax.imshow(frames[0], cmap=\"gray\", vmin=0, vmax=1)\n",
    "    ax.axis(\"off\")\n",
    "\n",
    "    ax_text_ddim = ax.text(-0.05, 0.75, \"Deterministic\", transform=ax.transAxes,\n",
    "                           va=\"center\", ha=\"right\", fontsize=12, color=\"white\")\n",
    "    ax_text_ddpm = ax.text(-0.05, 0.25, \"Stochastic\", transform=ax.transAxes,\n",
    "                           va=\"center\", ha=\"right\", fontsize=12, color=\"white\")\n",
    "\n",
    "    def update(i):\n",
    "        im.set_data(frames[i])\n",
    "        return [im, ax_text_ddim, ax_text_ddpm]\n",
    "\n",
    "    ani = FuncAnimation(fig, update, frames=len(frames), interval=120)\n",
    "    plt.close()\n",
    "    return HTML(ani.to_jshtml())\n",
    "\n",
    "\n",
    "# run the animation\n",
    "four_run_diffusion_animation(n_fwd=20, n_back=30)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "authorship_tag": "ABX9TyMf/8RzFk66VuDZGVXEmnm6",
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}