tutorial notebooks/edm tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/AshishKumar4/FlaxDiff/blob/main/tutorial%20notebooks/edm%20tutorial.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>\n",
    "\n",
    "# Building better diffusion models with EDM\n",
    "\n",
    "In this notebookk, we would discuss the paper [\"Elucidating the design space of Diffusion based models\" by Karras et al](https://arxiv.org/pdf/2206.00364), along with ideas from other subsequent papers after DDIM/DDPM, and look at some more generalized way of looking at the noise schedules and diffusion models in general. \n",
    "\n",
    "**This is part 2 in the diffusion series, and so its expected and strongly recommended that the reader go through the [Part 1: Simple Diffusion](./simple%20diffusion%20flax.ipynb) first to understand the basics and get familiar with the way we implement the ideas**\n",
    "\n",
    "In the previous notebook, we discussed how to build, train and sample from a denoising model built as specified by the DDPM and DDIM papers, and also touched upon the idea the forward and reverse diffusion processes can be treated as ODEs/SDEs instead of discrete markov chains, letting us use ODE/SDE solvers to integrate the denoising model's output to generate images. We would carry on from these ideas in this notebook. \n",
    "\n",
    "The generated images in the previous notebook weren't really of great. There are several reasons for that such as low training epochs, basic model architecture, etc. Subsequent papers after DDPM and DDIM introduced several improvements and ideas to increase the quality, such as better noise schedulers, better ways to model the gradient of log likelihood, etc. The EDM paper by Karras et al takes these ideas and also introduces some new generalizations.\n",
    "\n",
    "## New Ideas\n",
    "\n",
    "### 1. Noise Schedules\n",
    "\n",
    "The noise schedules in the DDPM and DDIM papers were simple linear schedules. In the previous notebook, we used the cosine noise schedule to train our model. We discussed how its a 'Variance preserving noise schedule'. By that we meant that if we formulate our forward diffusion process of adding noise as \n",
    "\n",
    "$x_t = \\alpha_t x_0 + \\sigma_t \\epsilon_0$\n",
    "\n",
    "where \n",
    "- $x_t$ is the data sample at time $t$\n",
    "- $x_0$ is the initial data sample\n",
    "- $\\epsilon$ is the Gaussian noise\n",
    "- $\\alpha$ $\\sigma_t$ are the signal and noise rates at time $t$ respectively\n",
    "\n",
    "Then the variance of the data sample at time $t$ is given by\n",
    "\n",
    "$Var(x_t) = \\alpha_t^2 Var(x_0) + \\sigma_t^2$\n",
    "\n",
    "Assuming the $Var(x_0) = 1$, we have \n",
    "\n",
    "$Var(x_t) = \\alpha_t^2 + \\sigma_t^2$\n",
    "\n",
    "than, the noise schedule is variance preserving if \n",
    "\n",
    "$\\alpha_t^2 + \\sigma_t^2 = 1$\n",
    "\n",
    "The awesome [Score based generative modeling through stochastic differential equations](https://arxiv.org/pdf/2011.13456) paper also introduced the variance exploding noise schedule, where the variance increases with time instead of being constant\n",
    "\n",
    "$\\alpha_t^2 + \\sigma_t^2 > 1$\n",
    "\n",
    "When dealing with a variance exploding noise schedule, one can overlook $alpha_t$ by setting it to 1 and just talk about the noise schedule as $\\sigma_t$\n",
    "\n",
    "$x_t = x_0 + \\sigma_t \\epsilon_0$\n",
    "\n",
    "**In the notebook from now on, we shall deal with the noise schedule as $\\sigma_t$ only**\n",
    "\n",
    "Now, ofcourse a model shouldn't be given an input sample that has arbitrarily high variance, so we scale the input samples by the variance of the noise schedule\n",
    "\n",
    "$Var(x_t) = 1 + \\sigma_t^2$\n",
    "\n",
    "thus the scaling factor of $x_t$ when given to the model is $\\frac{1}{\\sqrt{1 + \\sigma_t^2}}$\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install jax[cuda12]==0.4.28 flax[all] orbax grain-nightly augmax clu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm\n",
    "from flax import linen as nn\n",
    "import jax\n",
    "from typing import Dict, Callable, Sequence, Any, Union\n",
    "from dataclasses import field\n",
    "import jax.numpy as jnp\n",
    "import tensorflow_datasets as tfds\n",
    "import grain.python as pygrain\n",
    "# import tensorflow as tf\n",
    "import numpy as np\n",
    "import augmax\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from clu import metrics\n",
    "from flax.training import train_state  # Useful dataclass to keep train state\n",
    "import optax\n",
    "from flax import struct                # Flax dataclasses\n",
    "import time\n",
    "import os\n",
    "from datetime import datetime\n",
    "from flax.training import orbax_utils\n",
    "import functools\n",
    "from tensorflow_datasets.core.utils import gcs_utils\n",
    "gcs_utils._is_gcs_disabled = True\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Some Important Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "normalizeImage = lambda x: jax.nn.standardize(x, mean=[127.5], std=[127.5])\n",
    "denormalizeImage = lambda x: (x + 1.0) * 127.5\n",
    "\n",
    "\n",
    "def plotImages(imgs, fig_size=(8, 8), dpi=100):\n",
    "    fig = plt.figure(figsize=fig_size, dpi=dpi)\n",
    "    imglen = imgs.shape[0]\n",
    "    for i in range(imglen):\n",
    "        plt.subplot(fig_size[0], fig_size[1], i + 1)\n",
    "        plt.imshow(tf.cast(denormalizeImage(imgs[i, :, :, :]), tf.uint8))\n",
    "        plt.axis(\"off\")\n",
    "    plt.show()\n",
    "\n",
    "class RandomClass():\n",
    "    def __init__(self, rng: jax.random.PRNGKey):\n",
    "        self.rng = rng\n",
    "\n",
    "    def get_random_key(self):\n",
    "        self.rng, subkey = jax.random.split(self.rng)\n",
    "        return subkey\n",
    "    \n",
    "    def get_sigmas(self, steps):\n",
    "        return jnp.tan(self.theta_min + steps * (self.theta_max - self.theta_min)) / self.kappa\n",
    "\n",
    "    def reset_random_key(self):\n",
    "        self.rng = jax.random.PRNGKey(42)\n",
    "\n",
    "class MarkovState(struct.PyTreeNode):\n",
    "    pass\n",
    "\n",
    "class RandomMarkovState(MarkovState):\n",
    "    rng: jax.random.PRNGKey\n",
    "\n",
    "    def get_random_key(self):\n",
    "        rng, subkey = jax.random.split(self.rng)\n",
    "        return RandomMarkovState(rng), subkey"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Pipeline\n",
    "\n",
    "For simplicity, we will use the oxford flowers dataset for this notebook. We will use a newer pipeline for this notebook, based on the `google/grain` library, which helps us avoid any tensorflow dependencies apart from having to use it the very first time to download the dataset.\n",
    "\n",
    "**Tensorflow is required to download the dataset in the very first run. You can install tensorflow cpu version if you are having issues with the cuda stuff.**\n",
    "\n",
    "**If you have previously downloaded the dataset using the TFDS pipeline in the previous notebook, you might need to run clean the previous dataset stored at `~/tensorflow_datasets/`, otherwise the pipeline will throw an error**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_labels_oxford_flowers102(path):\n",
    "    def load_labels():\n",
    "        with open(path, \"r\") as f:\n",
    "            textlabels = [i.strip() for i in f.readlines()]\n",
    "        return textlabels\n",
    "    return load_labels\n",
    "\n",
    "# Configure the following for your datasets\n",
    "dataToLabelGenMap = {\n",
    "    \"oxford_flowers102\": load_labels_oxford_flowers102(\"~/tensorflow_datasets/oxford_flowers102/2.1.1/label.labels.txt\"),   # Change this if required!\n",
    "}\n",
    "\n",
    "def get_dataset(data_name=\"oxford_flowers102\", batch_size=64, image_scale=256, method=jax.image.ResizeMethod.LANCZOS3):\n",
    "    data_source = tfds.data_source(data_name, split=\"all\", try_gcs=False)\n",
    "    \n",
    "    gpu_device = jax.devices(\"gpu\")[0]\n",
    "    cpu_device = jax.devices(\"cpu\")[0]\n",
    "    \n",
    "    print(f\"Gpu Device: {gpu_device}, Cpu Device: {cpu_device}\")\n",
    "        \n",
    "    def preprocess(image):\n",
    "        # image = jax.device_put(image, device=jax.devices(\"cpu\")[0])\n",
    "        image = (image - 127.5) / 127.5\n",
    "        image = jax.image.resize(image, (image_scale, image_scale, 3), method=method)\n",
    "        image = jnp.clip(image, -1.0, 1.0)\n",
    "        image = jax.device_put(image, device=jax.devices(\"gpu\")[0]) \n",
    "        return  image\n",
    "    \n",
    "    preprocess = jax.jit(preprocess, backend=\"cpu\")\n",
    "\n",
    "    augments = augmax.Chain(\n",
    "        augmax.HorizontalFlip(0.5),\n",
    "        augmax.RandomContrast((-0.05, 0.05), 1.),\n",
    "        augmax.RandomBrightness((-0.2, 0.2), 1.)\n",
    "    )\n",
    "\n",
    "    augments = jax.jit(augments, backend=\"cpu\")\n",
    "    \n",
    "    if os.path.exists(f\"./datacache/{data_name}_labels.pkl\"):\n",
    "        print(\"Loading labels from cache\")\n",
    "        with open(f\"./datacache/{data_name}_labels.pkl\", \"rb\") as f:\n",
    "            import pickle\n",
    "            embed = pickle.load(f)\n",
    "            embed_labels = embed[\"embed_labels\"]\n",
    "            embed_labels_full = embed[\"embed_labels_full\"]\n",
    "            null_labels = embed[\"null_labels\"]\n",
    "            null_labels_full = embed[\"null_labels_full\"]\n",
    "    else:\n",
    "        print(\"No cache found, generating labels\")\n",
    "        textlabels = dataToLabelGenMap[data_name]()\n",
    "        \n",
    "        model, tokenizer = defaultTextEncodeModel()\n",
    "\n",
    "        embed_labels, embed_labels_full = encodePrompts(textlabels, model, tokenizer)\n",
    "        embed_labels = embed_labels.tolist()\n",
    "        embed_labels_full = embed_labels_full.tolist()\n",
    "        \n",
    "        null_labels, null_labels_full = encodePrompts([\"\"], model, tokenizer)\n",
    "        null_labels = null_labels.tolist()[0]\n",
    "        null_labels_full = null_labels_full.tolist()[0]\n",
    "        \n",
    "        os.makedirs(\"./datacache\", exist_ok=True)\n",
    "        with open(f\"./datacache/{data_name}_labels.pkl\", \"wb\") as f:\n",
    "            import pickle\n",
    "            pickle.dump({\n",
    "                \"embed_labels\": embed_labels,\n",
    "                \"embed_labels_full\": embed_labels_full,\n",
    "                \"null_labels\": null_labels,\n",
    "                \"null_labels_full\": null_labels_full\n",
    "                }, f)\n",
    "        \n",
    "    embed_labels = [np.array(i, dtype=np.float16) for i in embed_labels]\n",
    "    embed_labels_full = [np.array(i, dtype=np.float16) for i in embed_labels_full]\n",
    "    null_labels = np.array(null_labels, dtype=np.float16)\n",
    "    null_labels_full  = np.array(null_labels_full, dtype=np.float16)\n",
    "    \n",
    "    def labelizer(labelidx:int) -> jnp.array:\n",
    "        label_pooled = embed_labels[labelidx]\n",
    "        label_seq = embed_labels_full[labelidx]\n",
    "        # label_pooled = jax.device_put(label_pooled, device=jax.devices(\"gpu\")[0])\n",
    "        # label_seq = jax.device_put(label_seq, device=jax.devices(\"gpu\")[0])\n",
    "        return label_pooled, label_seq\n",
    "\n",
    "    class augmenter(pygrain.RandomMapTransform):\n",
    "        def random_map(self, element: Dict[str, Any], rng: np.random.Generator) ->  Dict[str, jnp.array]:\n",
    "            image = element['image']\n",
    "            image = preprocess(image)\n",
    "            image = augments(rng.integers(0, 2**32, [2], dtype=np.uint32), image) \n",
    "            labelidx = element['label']\n",
    "            label, label_seq = labelizer(labelidx)\n",
    "            # image, label = move2gpu(image, label)\n",
    "            return {'image':image, 'label':label, 'label_seq':label_seq}\n",
    "\n",
    "    sampler = pygrain.IndexSampler(\n",
    "        num_records=len(data_source),\n",
    "        shuffle=True,\n",
    "        seed=0,\n",
    "        num_epochs=None,\n",
    "        shard_options=pygrain.ShardByJaxProcess(),\n",
    "    )\n",
    "\n",
    "    transformations = [augmenter(), pygrain.Batch(batch_size, drop_remainder=True)]\n",
    "\n",
    "    loader = pygrain.DataLoader(\n",
    "        data_source=data_source,\n",
    "        sampler=sampler,\n",
    "        operations=transformations,\n",
    "        worker_count=4,\n",
    "        read_options=pygrain.ReadOptions(8, 500),\n",
    "        worker_buffer_size=5\n",
    "        )\n",
    "    return {\n",
    "        \"loader\": loader,\n",
    "        \"null_labels\": null_labels,\n",
    "        \"null_labels_full\": null_labels_full,\n",
    "        \"embed_labels\": embed_labels,\n",
    "        \"embed_labels_full\": embed_labels_full,\n",
    "        \"length\": len(data_source),  \n",
    "        \"batch_size\": batch_size,\n",
    "        \"image_size\": image_scale\n",
    "    }\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}