assignment10_small.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cde55194",
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65095138",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14c07255",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "tfkl = tf.keras.layers\n",
    "import tensorflow_probability as tfp\n",
    "tfd = tfp.distributions\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt \n",
    "import tensorflow_addons as tfa\n",
    "\n",
    "from data.utils import parse_image_example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "135e5fcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def scale_data(images, labels):\n",
    "    return 2*images - 1, labels\n",
    "\n",
    "def descale_data(images):\n",
    "    return (images + 1) / 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cc57bd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 512\n",
    "\n",
    "(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()\n",
    "train_labels = train_labels.astype(np.int32)\n",
    "test_labels = test_labels.astype(np.int32)\n",
    "\n",
    "train_images = np.pad(train_images[..., None], ((0, 0), (2, 2), (2, 2), (0, 0))).astype(np.float32) / 255.\n",
    "test_images = np.pad(test_images[..., None], ((0, 0), (2, 2), (2, 2), (0, 0))).astype(np.float32) / 255.\n",
    "\n",
    "train_labels += 1\n",
    "test_labels += 1\n",
    "# note that scale_forward is already mapped onto the dataset here.\n",
    "# for sampling etc., we need to be mindful to apply scale_backward.\n",
    "# although it doesn't matter much for MNIST or single-channel images in general,\n",
    "# as colormaps are scaled automatically\n",
    "train_data = tf.data.Dataset.from_tensor_slices((train_images, train_labels)).shuffle(60000).batch(batch_size, drop_remainder=True).map(scale_data)\n",
    "test_data = tf.data.Dataset.from_tensor_slices((test_images, test_labels)).batch(32).map(scale_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f351a8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_images = np.concatenate([batch for (batch, _) in iter(test_data)], axis=0)\n",
    "\n",
    "plt.figure(figsize=(15,15))\n",
    "for ind, img in enumerate(descale_data(test_images[:64])):\n",
    "    plt.subplot(8, 8, ind+1)\n",
    "    plt.imshow(img, vmin=0, vmax=1, cmap=\"Greys\")\n",
    "    plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40fecff8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters from the paper: t=1000, betas=np.linspace(0.0001, 0.02, tmax)\n",
    "# I reduced t to 200 for faster sampling, seems to be ok\n",
    "tmax = 200\n",
    "betas = np.linspace(0.0001, 0.1, tmax).astype(np.float32)\n",
    "alphas = 1 - betas\n",
    "alphas_bar = np.cumprod(alphas)\n",
    "\n",
    "plt.plot(alphas_bar)\n",
    "plt.plot(np.sqrt(alphas_bar))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c739a2df",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# here we can look at the forward process slowly turning data to noise\n",
    "step = 20\n",
    "for t_index, alpha_bar in enumerate(alphas_bar[::step]):\n",
    "    noise_scale = np.sqrt(1 - alpha_bar)\n",
    "    noisy_imgs = np.sqrt(alpha_bar) * test_images[:64] + noise_scale*np.random.normal(size=test_images[:64].shape)\n",
    "    \n",
    "    plt.figure(figsize=(15, 15))\n",
    "    for ind, image in enumerate(descale_data(noisy_imgs)):\n",
    "        plt.subplot(8, 8, ind+1)\n",
    "        plt.imshow(image, vmin=0, vmax=1, cmap=\"Greys\")\n",
    "        plt.axis(\"off\")\n",
    "        plt.suptitle(\"t = {}: Shrinkage {}; Noise scale: {}\".format(t_index*step+1, np.sqrt(alpha_bar), noise_scale))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81a293fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Diffusion(tf.keras.Model):\n",
    "    def __init__(self, inputs, outputs, alphas, alphas_bar, betas, t_max, **kwargs):\n",
    "        super().__init__(inputs, outputs, **kwargs)\n",
    "        self.loss_tracker = tf.keras.metrics.Mean(\"loss\")\n",
    "        \n",
    "        self.alphas_tensor = tf.convert_to_tensor(alphas, dtype=tf.float32)\n",
    "        self.alphas_bar_tensor = tf.convert_to_tensor(alphas_bar, dtype=tf.float32)\n",
    "        self.betas_tensor = tf.convert_to_tensor(betas, dtype=tf.float32)\n",
    "        self.tmax = t_max\n",
    "        \n",
    "        drop_p = 0.1\n",
    "        self.drop_p = drop_p\n",
    "        self.drop_dist = tfd.Bernoulli(probs=1-drop_p)\n",
    "\n",
    "    \n",
    "    def diffusion_loss(self, image_batch, labels, training=None):\n",
    "        # this samples a batch_size tensor of t values\n",
    "        mask = self.drop_dist.sample(tf.shape(image_batch)[0])\n",
    "        labels = mask * labels\n",
    "        labels = tf.one_hot(labels, depth=11)\n",
    "        \n",
    "        sampled_ts = tf.random.uniform([tf.shape(image_batch)[0]], 0, self.tmax, dtype=tf.int32)\n",
    "        target_epsilons = tf.random.normal(tf.shape(image_batch))\n",
    "\n",
    "        batch_alphas_bar = tf.gather(self.alphas_bar_tensor, sampled_ts)[:, None, None, None]\n",
    "        noisy_batch = (tf.math.sqrt(batch_alphas_bar) * image_batch \n",
    "                       + tf.math.sqrt(1 - batch_alphas_bar) * target_epsilons)\n",
    "        # turns it into batch x 1 matrix and normalizes\n",
    "        normalized_ts = tf.cast(sampled_ts, tf.float32)[:, None] / self.tmax\n",
    "        \n",
    "        #%%%%%%%%%% TO UPDATE \n",
    "        # this function needs to receive conditioning inputs (e.g. class labels) and hand them to the model call\n",
    "        # for classifier-free guidance, the labels needs to be occasionally (randomly) dropped out\n",
    "        noise_prediction = self([noisy_batch, normalized_ts, labels], training=training)\n",
    "        loss = tf.reduce_mean(tf.reduce_sum((target_epsilons - noise_prediction)**2,\n",
    "                                            axis=[1,2,3]))\n",
    "        return loss\n",
    "\n",
    "    def train_step(self, data):\n",
    "        #%%%%%%%%% TO UPDATE\n",
    "        # your dataset should include tuples of images,labels.\n",
    "        # then you can do\n",
    "        # image_batch, label_batch = data\n",
    "        # and pass both to the loss (labels required for conditioning)\n",
    "        image_batch, label_batch = data\n",
    "        \n",
    "        # same thing in test_step!\n",
    "        with tf.GradientTape() as tape:\n",
    "            loss = self.diffusion_loss(image_batch, label_batch, training=True)\n",
    "        gradients = tape.gradient(loss, self.trainable_variables)\n",
    "        optimizer.apply_gradients(zip(gradients, self.trainable_variables))\n",
    "\n",
    "        self.loss_tracker.update_state(loss)\n",
    "        return {\"loss\": self.loss_tracker.result()}\n",
    "    \n",
    "    def test_step(self, data):\n",
    "        image_batch, label_batch = data\n",
    "        loss = self.diffusion_loss(image_batch, label_batch, training=False)\n",
    "\n",
    "        self.loss_tracker.update_state(loss)\n",
    "        return {\"loss\": self.loss_tracker.result()}\n",
    "    \n",
    "    @tf.function\n",
    "    def langevin_step(self, sample, t_step, classes, guidance, beta_version=\"one\"):\n",
    "        z = tf.random.normal(tf.shape(sample))\n",
    "        if t_step == 0:\n",
    "            z = tf.zeros_like(z)\n",
    "\n",
    "        t_normalized = tf.cast(t_step, tf.float32) * tf.ones([tf.shape(sample)[0], 1]) / self.tmax\n",
    "        #%%%%%%%%% TO UPDATE\n",
    "        # this function needs to take class inputs from langevin_sampler\n",
    "        # and put them into the model.\n",
    "        #\n",
    "        # for simple class-conditioning, that is enough.\n",
    "        #\n",
    "        # for guided diffusion, you will want to use\n",
    "        # (1+ w)*model_call(with_condition) - w*model_call(without_condition)\n",
    "        # without_condition refers to \"dropped out\" conditioning\n",
    "        class_inp = tf.one_hot(classes, depth=11)\n",
    "        no_inp = tf.zeros_like(classes)\n",
    "        no_inp = tf.one_hot(no_inp, depth=11)\n",
    "        conditional = self([sample, t_normalized, class_inp], training=False)\n",
    "        unconditional = self([sample, t_normalized, no_inp], training=False)\n",
    "        \n",
    "        model_output = (1 + guidance) * conditional - guidance * unconditional\n",
    "\n",
    "        alpha_here = self.alphas_tensor[t_step]\n",
    "        alpha_bar_here = self.alphas_bar_tensor[t_step]\n",
    "\n",
    "        if t_step == 0:\n",
    "            sigma = tf.convert_to_tensor(0., tf.float32)\n",
    "        else:\n",
    "            # these are two different choice for beta described in the paper\n",
    "            if beta_version == \"one\":\n",
    "                sigma = tf.math.sqrt(self.betas_tensor[t_step])\n",
    "            elif beta_version == \"two\":\n",
    "                sigma = tf.math.sqrt((1 - self.alphas_bar_tensor[t_step-1]) / (1 - alpha_bar_here) \n",
    "                                     * self.betas_tensor[t_step])\n",
    "\n",
    "        noise = sigma * z\n",
    "        sample = (1/tf.math.sqrt(alpha_here) \n",
    "                  * (sample - (1 - alpha_here)/tf.math.sqrt(1 - alpha_bar_here) * model_output))\n",
    "        sample = sample + noise\n",
    "        return sample\n",
    "    \n",
    "    def langevin_sampler(self, classes, guidance, beta_version=\"one\"):\n",
    "        #%%%%%%%%%%%%% TO UPDATE\n",
    "        # for conditional generation, you most likely want to change this function\n",
    "        # to accept class inputs.\n",
    "        # maybe even remove n_samples, and instead use len(class_inputs) as the number of samples.\n",
    "        # assuming you are passing a batch_size vector of classes.\n",
    "        n_samples = tf.shape(classes)[0]\n",
    "        sample = tf.random.normal((n_samples,) + self.input_shape[0][1:])\n",
    "        for t_step in tf.range(self.tmax)[::-1]:\n",
    "            sample = self.langevin_step(sample, t_step, classes, guidance, beta_version)\n",
    "\n",
    "        return sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf03a7de",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is an upsample -> conv layers, that could be used in the decoder instead of transposed convolution\n",
    "class UpsampleConv2D(tfkl.Layer):\n",
    "    def __init__(self, n_filters, filter_size, strides=1, padding=\"same\", **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.conv = tfkl.Conv2D(n_filters, filter_size, padding=padding, name=self.name + \"_conv\")\n",
    "        if strides > 1:\n",
    "            self.upsample = tfkl.UpSampling2D(size=strides, interpolation=\"bilinear\", name=self.name + \"_upsample\")\n",
    "        self.strides = strides\n",
    "    \n",
    "    def call(self, inputs):\n",
    "        if self.strides > 1:\n",
    "            upsampled = self.upsample(inputs)\n",
    "        else:\n",
    "            upsampled = inputs\n",
    "        return self.conv(upsampled)\n",
    "\n",
    "\n",
    "# Normalization -> Activation -> Convolution\n",
    "# often referred to as \"pre-activation\", popular for residual networks.\n",
    "# if you are interested: https://arxiv.org/pdf/1603.05027.pdf\n",
    "class NormActConv(tfkl.Layer):\n",
    "    def __init__(self,\n",
    "                 n_filters,\n",
    "                 mode,\n",
    "                 strides=1,\n",
    "                 activation=tf.nn.gelu,\n",
    "                 **kwargs):\n",
    "        if mode not in [\"conv\", \"transpose\", \"upconv\"]:\n",
    "            raise ValueError(\"Invalid mode; valid choices are 'conv', 'transpose', 'upconv'.\")\n",
    "        \n",
    "        super().__init__(**kwargs)\n",
    "        if mode == \"conv\":\n",
    "            layer_fn = tfkl.Conv2D\n",
    "        elif mode == \"transpose\":\n",
    "            layer_fn = tfkl.Conv2DTranspose\n",
    "        else:\n",
    "            layer_fn = UpsampleConv2D\n",
    "\n",
    "        # NOTE harcoded 3x3 filter size, could change this\n",
    "        self.conv = layer_fn(n_filters, 3, strides=strides,\n",
    "                padding=\"same\", name=self.name + \"_conv_main\")\n",
    "        self.activation = activation\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        self.normalization = tfa.layers.GroupNormalization(groups=input_shape[-1]//8,\n",
    "                                                           name=self.name + \"_normalization\")\n",
    "        # these components are for adaptive normalization\n",
    "        self.time_mult = tfkl.Dense(input_shape[-1], kernel_initializer=tf.keras.initializers.Zeros(),\n",
    "                                    name=self.name + \"_adanorm_scale\")\n",
    "        self.time_add = tfkl.Dense(input_shape[-1], kernel_initializer=tf.keras.initializers.Zeros(),\n",
    "                                    name=self.name + \"_adanorm_shift\")\n",
    "        self.class_mult = tfkl.Dense(input_shape[-1], kernel_initializer=tf.keras.initializers.Zeros(),\n",
    "                                    name=self.name + \"_adanorm_cscale\")\n",
    "        self.class_add = tfkl.Dense(input_shape[-1], kernel_initializer=tf.keras.initializers.Zeros(),\n",
    "                                    name=self.name + \"_adanorm_cshift\")\n",
    "        \n",
    "        #%%%%%%%%% TO UPDATE\n",
    "        # you can add another set of shift, scale layers for the class-conditioning input\n",
    "        \n",
    "    def call(self, inputs, time, classes):\n",
    "        \n",
    "        if self.normalization:\n",
    "            normed = self.normalization(inputs)\n",
    "        else:\n",
    "            normed = inputs\n",
    "            \n",
    "        # adaptive normalization implements conditioning on time\n",
    "        normed = (1 + self.time_mult(time)) * normed + self.time_add(time)\n",
    "        normed = (1 + self.class_mult(classes)) * normed + self.class_add(classes)\n",
    "        #%%%%%%%%%%%%% TO UPDATE\n",
    "        # the layer should also receive a class input and apply adaptive normalization.\n",
    "        # the same way it is done for time!\n",
    "        \n",
    "        if self.activation:\n",
    "            acted = self.activation(normed)\n",
    "        else:\n",
    "            acted = normed\n",
    "        \n",
    "        conved = self.conv(acted)\n",
    "        return conved\n",
    "\n",
    "\n",
    "# residual block with two convolutional layers\n",
    "class ResidualBlock(tfkl.Layer):\n",
    "    def __init__(self,\n",
    "                 n_filters,\n",
    "                 mode,\n",
    "                 strides=1,\n",
    "                 activation=tf.nn.gelu,\n",
    "                 **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.main_layer1 = NormActConv(n_filters, mode, strides,\n",
    "                                      activation=activation,\n",
    "                                      name=self.name + \"_main1\")\n",
    "        \n",
    "        self.main_layer2 = NormActConv(n_filters, mode, 1,\n",
    "                                      activation=activation,\n",
    "                                      name=self.name + \"_main2\")\n",
    "        \n",
    "        if mode == \"conv\":\n",
    "            shortcut_fn = tfkl.Conv2D\n",
    "        elif mode == \"upconv\":\n",
    "            shortcut_fn = UpsampleConv2D\n",
    "        else:\n",
    "            shortcut_fn = tfkl.Conv2DTranspose\n",
    "        self.shortcut = shortcut_fn(n_filters, 1, strides=strides,\n",
    "                                    name=self.name + \"_shortcut\")\n",
    "        self.strides = strides\n",
    "        self.n_filters = n_filters\n",
    "        \n",
    "    def call(self, inputs, time, classes):\n",
    "        #%%%%%%%%%%% TO UPDATE\n",
    "        # this layer needs to receive class conditioning input\n",
    "        # and pass it to the main layers\n",
    "        l1 = self.main_layer1(inputs, time, classes)\n",
    "        l2 = self.main_layer2(l1, time, classes)\n",
    "        \n",
    "        shortcut = self.shortcut(inputs)\n",
    "        \n",
    "        # residual connection is re-scaled by 1/sqrt(2).\n",
    "        # this helps.\n",
    "        out = (1/tf.math.sqrt(2.)) * (l2 + shortcut)\n",
    "        \n",
    "        return out\n",
    "\n",
    "\n",
    "    \n",
    "#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%\n",
    "# THESE ARE CURRENTLY NOT USED IN THE MODEL!!\n",
    "# so if you don't use them, you don't need to update them for conditioning, duh\n",
    "\n",
    "# a \"level\" is a series of residual blocks operating at the same resolution \n",
    "# and with the same number of filters\n",
    "class DownLevel(tfkl.Layer):\n",
    "    def __init__(self, n_filters, n_blocks, strides=1, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.blocks = [ResidualBlock(n_filters, \"conv\", strides=strides if ind==0 else 1,\n",
    "                                     name=self.name + \"_block{}\".format(ind))\n",
    "                       for ind in range(n_blocks)]\n",
    "        \n",
    "    def call(self, inputs, time):\n",
    "        #%%%%%%%%%% TO UPDATE\n",
    "        # this layer needs to receive class conditioning input\n",
    "        # and pass it to the residual blocks\n",
    "        for block in self.blocks:\n",
    "            inputs = block(inputs, time)\n",
    "            \n",
    "        return inputs\n",
    "    \n",
    "    \n",
    "class UpLevel(tfkl.Layer):\n",
    "    def __init__(self, n_filters, n_blocks, strides=1, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        # NOTE could change \"transpose\" to \"upconv\" to use upsampling + convolution\n",
    "        self.blocks = [ResidualBlock(n_filters, \"transpose\", strides=strides if ind==0 else 1,\n",
    "                                     name=self.name + \"_block{}\".format(ind))\n",
    "                       for ind in range(n_blocks)]\n",
    "        \n",
    "    def call(self, inputs, time):\n",
    "        #%%%%%%%%%% TO UPDATE\n",
    "        # this layer needs to receive class conditioning input\n",
    "        # and pass it to the residual blocks\n",
    "        for block in self.blocks:\n",
    "            inputs = block(inputs, time)\n",
    "            \n",
    "        return inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "131de9db",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow_addons as tfa\n",
    "\n",
    "\n",
    "#%%%%%%%%%%%%%%%%% TO UPDATE\n",
    "# this function needs to receive class conditioning input similar to time\n",
    "# Same in the function below (decoder stack)\n",
    "def residual_stack_encoder(inputs, t_input, class_input,\n",
    "                           filters, strides, blocks,\n",
    "                           mode, name):\n",
    "    # this collects outputs of all residual blocks, to be used in the decoder for skip connections\n",
    "    all_outputs = []  \n",
    "    # one initial convolution to add some channels\n",
    "    outputs = tfkl.Conv2D(filters[0], 3, padding=\"same\")(inputs)\n",
    "    for level_ind, (level_filters, level_stride, level_blocks) in enumerate(zip(filters, strides, blocks)):\n",
    "        for block_ind in range(level_blocks):\n",
    "            # %%%%%%%%%%%%%%% TO UPDATE\n",
    "            # these layers need to receive class conditiong input, just like they are receiving time\n",
    "            outputs = ResidualBlock(level_filters,\n",
    "                                    mode, \n",
    "                                    strides=level_stride if block_ind == 0 else 1,\n",
    "                                    name=\"_\".join([name, str(level_ind+1), str(block_ind+1)]))(outputs, t_input, class_input)\n",
    "            all_outputs.append(outputs)\n",
    "        \n",
    "    return outputs, all_outputs\n",
    "\n",
    "def residual_stack_decoder(inputs, t_input, class_input, all_hidden,\n",
    "                           filters, strides, blocks,\n",
    "                           mode, name):\n",
    "    outputs = inputs\n",
    "    global_ind = 0\n",
    "    for level_ind, (level_filters, level_stride, level_blocks) in enumerate(zip(filters, strides, blocks)):\n",
    "        for block_ind in range(level_blocks):\n",
    "            \n",
    "            # this part handles skip connections handling from the encoder.\n",
    "            # this is bad code! very awkward!\n",
    "            # sorry, but you shouldn't need to touch this\n",
    "            if global_ind > 0:\n",
    "                if outputs.shape[1] != all_hidden[global_ind].shape[1]:\n",
    "                    all_hidden[global_ind] = tfkl.AvgPool2D(padding=\"same\")(all_hidden[global_ind])\n",
    "                outputs = tfkl.Concatenate(axis=-1)([outputs, all_hidden[global_ind]])\n",
    "            global_ind += 1\n",
    "            \n",
    "            # back to sanity\n",
    "            outputs = ResidualBlock(level_filters,\n",
    "                                    mode, \n",
    "                                    strides=level_stride if block_ind == 0 else 1,\n",
    "                                    name=\"_\".join([name, str(level_ind+1), str(block_ind+1)]))(outputs, t_input, class_input)\n",
    "        \n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cb3c84b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def positional_encoding_v2(input_t, n_freqs):\n",
    "    # input_t: b x 1\n",
    "    # n_freqs: scalar\n",
    "    # create frequencies up to below nyquist frequency\n",
    "    frequencies = tf.convert_to_tensor(np.geomspace(0.1, tmax//2, n_freqs).astype(np.float32))\n",
    "    sines = tf.math.sin(2*np.pi*frequencies*input_t)  # b x n_freqs\n",
    "    cosines =  tf.math.cos(2*np.pi*frequencies*input_t)  # b x n_freqs\n",
    "    return tf.concat([sines, cosines], axis=-1)  # b x 2*n_freqs\n",
    "\n",
    "\n",
    "inp = tf.keras.Input(test_images.shape[1:])\n",
    "t_input = tf.keras.Input((1,))\n",
    "t_encoded = positional_encoding_v2(t_input, n_freqs=32)[:, None, None, :]\n",
    "t_encoded = tfkl.Dense(32, tf.nn.gelu)(t_encoded)\n",
    "t_encoded = tfkl.Dense(32, tf.nn.gelu)(t_encoded)\n",
    "\n",
    "\n",
    "#%%%%%%%%%% TO UPDATE\n",
    "# add a class input just like time. you could receive it as single indices,\n",
    "# or already apply one_hot outside the model, and here receive input with shape (10,) (assuming 10 classes)\n",
    "\n",
    "# positional encoding does not make sense for clases. so just use bunch of Dense layers.\n",
    "# if you receive classes as indices, your first layer should be an Embedding\n",
    "# if you receive one-hot, just use Dense\n",
    "class_input = tf.keras.Input((11,))\n",
    "class_encoded = tf.keras.layers.Dense(32, tf.nn.gelu)(class_input)\n",
    "class_encoded = tf.keras.layers.Dense(32, tf.nn.gelu)(class_encoded)\n",
    "class_encoded = tf.keras.layers.Dense(32, tf.nn.gelu)(class_encoded)[:, None, None, :]\n",
    "\n",
    "\n",
    "# you can make this smaller if it takes too long\n",
    "blocks_per_level = [2, 2, 2, 2]\n",
    "filters = [32, 48, 64, 96]\n",
    "strides = [1, 2, 2, 2]\n",
    "encoder_output, all_hidden = residual_stack_encoder(\n",
    "    inp, t_encoded, class_encoded,\n",
    "    filters,\n",
    "    strides,\n",
    "    blocks_per_level,\n",
    "    \"conv\", \"encoder\")\n",
    "\n",
    "# IF you  get some strange errors talking about upsampling:\n",
    "# try using \"transpose\" instead of \"upconv\"\n",
    "decoder_output = residual_stack_decoder(\n",
    "    encoder_output, t_encoded, class_encoded, list(reversed(all_hidden)),\n",
    "    reversed(filters),\n",
    "    strides,\n",
    "    reversed(blocks_per_level),\n",
    "    \"upconv\", \"decoder\")\n",
    "\n",
    "# final layer goes back to data shape\n",
    "decoder_final = tfkl.Conv2D(inp.shape[-1], 1)(decoder_output)\n",
    "\n",
    "score_model = Diffusion([inp, t_input, class_input], decoder_final, alphas, alphas_bar, betas, tmax)\n",
    "score_model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "000b7e9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# use fewer steps if it takes too long\n",
    "train_steps = 20000\n",
    "n_data = 60000\n",
    "n_epochs = train_steps // (n_data // batch_size)\n",
    "lr = tf.optimizers.schedules.CosineDecay(0.005, train_steps)\n",
    "\n",
    "# NOTE  if you use significantly fewer steps, consider reducing the ema_momentum\n",
    "# or turn it off completely (use_ema=False)\n",
    "optimizer = tf.optimizers.Adam(lr, use_ema=True, ema_momentum=0.99)\n",
    "\n",
    "score_model.compile(optimizer=optimizer, jit_compile=True)\n",
    "\n",
    "\n",
    "# a note on training, the loss is quite noisy and unreliable.\n",
    "# it's also hard to interpret.\n",
    "# for this setup you can expect values around 20 or so.\n",
    "# that should lead to decent samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b609ad7",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#%%%%%%%%%% TO UPDATE\n",
    "# as the sampler is called here, you will want to create & provide some class values for conditioning\n",
    "class ImageGenCallback(tf.keras.callbacks.Callback):\n",
    "    def __init__(self, frequency, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.frequency = frequency\n",
    "        \n",
    "    def on_epoch_begin(self, epoch, logs=None):\n",
    "        if not epoch % self.frequency:\n",
    "            for guidance in [-1., 0., 1., 3., 10., 30.]:\n",
    "                classes = tf.convert_to_tensor(np.repeat(np.arange(1, 11), 10).astype(np.int32))\n",
    "                generated_batch = descale_data(self.model.langevin_sampler(classes, guidance))\n",
    "\n",
    "                plt.figure(figsize=(15,15))\n",
    "                for ind, image in enumerate(generated_batch):\n",
    "                    plt.subplot(10, 10, ind+1)\n",
    "                    plt.imshow(image, cmap=\"Greys\", vmin=0, vmax=1)\n",
    "                    plt.axis(\"off\")\n",
    "                plt.suptitle(\"Random generations -- guidance {}\".format(guidance))\n",
    "                plt.show()\n",
    "\n",
    "\n",
    "score_model.fit(train_data, validation_data=test_data, epochs=n_epochs, callbacks=ImageGenCallback(20))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7529f6f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "score_model.save_weights(\"weights/weights_diffusion_mnist_conditional.hdf5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adf6bb01",
   "metadata": {},
   "outputs": [],
   "source": [
    "classes = tf.convert_to_tensor(np.repeat(np.arange(10), 10).astype(np.int32))\n",
    "generated_batch = descale_data(score_model.langevin_sampler(classes))\n",
    "plt.figure(figsize=(15,15))\n",
    "for ind, image in enumerate(generated_batch):\n",
    "    plt.subplot(10, 10, ind+1)\n",
    "    plt.imshow(image, vmin=0, vmax=1, cmap=\"Greys\")\n",
    "    plt.axis(\"off\")\n",
    "plt.suptitle(\"Random generations\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c966c4a3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}