export_unet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "97423763",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from imagen_pytorch import Unet, Imagen, ImagenTrainer\n",
    "from imagen_pytorch.data import Dataset\n",
    "import random\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import torchvision.transforms as T\n",
    "from imagen_pytorch.t5 import t5_encode_text\n",
    "import math\n",
    "import types\n",
    "from functools import partial\n",
    "\n",
    "np.random.seed(0)\n",
    "random.seed(0)\n",
    "torch.manual_seed(0)\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "# device = torch.device('cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f7b43094",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading saved imagen at version 1.17.1, but current package version is 1.18.1\n",
      "checkpoint loaded from giddy-capybara.ckpt\n"
     ]
    }
   ],
   "source": [
    "unet = Unet(\n",
    "    dim = 128, # the \"Z\" layer dimension, i.e. the number of filters the outputs to the first layer\n",
    "    cond_dim = 128,\n",
    "    dim_mults = (1, 2, 4), # the channel dimensions inside the model (multiplied by dim)\n",
    "    num_resnet_blocks = 3,\n",
    "    layer_attns = (False, True, True),\n",
    "    layer_cross_attns = (False, True, True)\n",
    ")\n",
    "\n",
    "imagen = Imagen(\n",
    "    unets = unet,\n",
    "    image_sizes = 32,\n",
    "    timesteps = 250,\n",
    "    cond_drop_prob = 0.1,\n",
    "    dynamic_thresholding=False\n",
    ")\n",
    "\n",
    "\n",
    "trainer = ImagenTrainer(imagen)\n",
    "trainer.load(\"giddy-capybara.ckpt\")\n",
    "pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "6f920eb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Replacing all instances of \"exmp1\" to avoid:\n",
    "# UnsupportedOperatorError: Exporting the operator ::expm1 to ONNX opset version 13 is not supported\n",
    "\n",
    "def faux_expm1(x):\n",
    "    return torch.exp(x) * (1 - torch.exp(-x))\n",
    "\n",
    "def beta_linear(t):\n",
    "    return -torch.log(faux_expm1(1e-4 + 10 * (t ** 2))) # <-------------- Replacing expm1\n",
    "\n",
    "def right_pad_dims_to(x, t):\n",
    "    padding_dims = x.ndim - t.ndim\n",
    "    if padding_dims <= 0:\n",
    "        return t\n",
    "    return t.view(*t.shape, *((1,) * padding_dims))\n",
    "\n",
    "def exists(val):\n",
    "    return val is not None\n",
    "\n",
    "def default(val, d):\n",
    "    if exists(val):\n",
    "        return val\n",
    "    return d() if callable(d) else d\n",
    "\n",
    "def log_snr_to_alpha_sigma(log_snr):\n",
    "    return torch.sqrt(torch.sigmoid(log_snr)), torch.sqrt(torch.sigmoid(-log_snr))\n",
    "\n",
    "def log(t, eps: float = 1e-12):\n",
    "    return torch.log(t.clamp(min = eps))\n",
    "\n",
    "def q_posterior(self, x_start, x_t, t, *, t_next = None):\n",
    "    t_next = default(t_next, lambda: (t - 1. / self.num_timesteps).clamp(min = 0.))\n",
    "\n",
    "    \"\"\" https://openreview.net/attachment?id=2LdBqxc1Yv&name=supplementary_material \"\"\"\n",
    "    log_snr = self.log_snr(t)\n",
    "    log_snr_next = self.log_snr(t_next)\n",
    "    log_snr, log_snr_next = map(partial(right_pad_dims_to, x_t), (log_snr, log_snr_next))\n",
    "\n",
    "    alpha, sigma = log_snr_to_alpha_sigma(log_snr)\n",
    "    alpha_next, sigma_next = log_snr_to_alpha_sigma(log_snr_next)\n",
    "\n",
    "    # c - as defined near eq 33\n",
    "    c = -faux_expm1(log_snr - log_snr_next) # <-------------- Replacing expm1\n",
    "    posterior_mean = alpha_next * (x_t * (1 - c) / alpha + c * x_start)\n",
    "\n",
    "    # following (eq. 33)\n",
    "    posterior_variance = (sigma_next ** 2) * c\n",
    "    posterior_log_variance_clipped = log(posterior_variance, eps = 1e-20)\n",
    "    return posterior_mean, posterior_variance, posterior_log_variance_clipped\n",
    "\n",
    "\n",
    "scheduler = imagen.noise_schedulers[0]\n",
    "scheduler.log_snr = beta_linear\n",
    "scheduler.q_posterior = types.MethodType(q_posterior, scheduler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "c520101c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# out = trainer.sample([\"a photo of a truck\"])\n",
    "# print(out)\n",
    "# img = T.ToPILImage()(out[0])\n",
    "# img.save(\"truck-4.png\")\n",
    "# img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "022bd40e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# out = imagen.sample(text_embeds=enc, text_masks=mask)\n",
    "# print(out)\n",
    "# img = T.ToPILImage()(out[0])\n",
    "# img.save(\"truck-4.png\")\n",
    "# img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "c737c97f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ImagenOnnx(torch.nn.Module):\n",
    "    def __init__(self, unet, imagen):\n",
    "        super().__init__()\n",
    "        \n",
    "        if imagen.dynamic_thresholding == True or any(imagen.dynamic_thresholding):\n",
    "            raise ValueError(\"no dynamic thresholding allowed, got:\", imagen.dynamic_thresholding)\n",
    "\n",
    "        self.unet = unet\n",
    "        self.imagen = imagen    \n",
    "            \n",
    "\n",
    "    def sample(self, text_embeds, text_mask, cond_scale,  device, use_tqdm=True):\n",
    "        batch_size = text_embeds.shape[0]\n",
    "        noise_scheduler = self.imagen.noise_schedulers[0]\n",
    "        img = torch.randn((batch_size, self.imagen.sample_channels[0], self.imagen.image_sizes[0], self.imagen.image_sizes[0]), device = device)        \n",
    "        timesteps = noise_scheduler.get_sampling_timesteps(batch_size, device = device)\n",
    "            \n",
    "        \n",
    "        for times, times_next in tqdm(timesteps, desc='sampling loop time step', total=len(timesteps), disable=not use_tqdm):       \n",
    "            self_cond = x_start if self.unet.self_cond else None\n",
    "            \n",
    "            img, x_start = self.imagen.p_sample(\n",
    "                self.unet,\n",
    "                img,\n",
    "                times,\n",
    "                t_next = times_next,\n",
    "                text_embeds = text_embeds,\n",
    "                text_mask = text_mask,\n",
    "                cond_scale = cond_scale,\n",
    "                noise_scheduler = noise_scheduler,\n",
    "                pred_objective = self.imagen.pred_objectives[0],\n",
    "                dynamic_threshold = False\n",
    "            )\n",
    "\n",
    "        img.clamp_(-1., 1.)\n",
    "\n",
    "        unnormalize_img = self.imagen.unnormalize_img(img)\n",
    "\n",
    "        return unnormalize_img\n",
    "    \n",
    "    def forward(self, img, text_embeds, text_mask, times, times_next, cond_scale):\n",
    "        return self.imagen.p_sample(\n",
    "            self.unet,\n",
    "            img,\n",
    "            times,\n",
    "            t_next = times_next,\n",
    "            text_embeds = text_embeds,\n",
    "            text_mask = text_mask,\n",
    "            cond_scale = cond_scale,\n",
    "            noise_scheduler=imagen.noise_schedulers[0],\n",
    "            pred_objective = imagen.pred_objectives[0],\n",
    "            dynamic_threshold = False\n",
    "        )\n",
    "    \n",
    "u = ImagenOnnx(unet, imagen)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "752abd54",
   "metadata": {},
   "outputs": [],
   "source": [
    "# enc, mask = t5_encode_text([\"a photo of a truck\"], return_attn_mask = True)\n",
    "# out = u.sample(enc, mask, 1., torch.device('cuda'), use_tqdm=False)\n",
    "# print(out)\n",
    "# img = T.ToPILImage()(out[0])\n",
    "# img.save(\"truck-12.png\")\n",
    "# img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "94c57b15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# out = u.forward(\n",
    "#     torch.rand(1, 3, 32, 32).cuda(),  \n",
    "#     torch.rand(1, 27, 768).cuda(), \n",
    "#     torch.ones(1, 27, dtype=bool).cuda(), \n",
    "#     torch.Tensor([0.9]).cuda(), \n",
    "#     torch.Tensor([0.896]).cuda(), \n",
    "#     torch.Tensor([1.]).cuda()\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "a61782e8",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# you will get a bunch of ONNX warnings, it's ok I checked them and they're all not an issue\n",
    "torch.onnx.export(\n",
    "    u, \n",
    "    (\n",
    "        torch.rand(1, 3, 32, 32).cuda(),  \n",
    "        torch.rand(1, 27, 768).cuda(), \n",
    "        torch.ones(1, 27, dtype=bool).cuda(), \n",
    "        torch.Tensor([0.5]).cuda(), \n",
    "        torch.Tensor([0.56]).cuda(), \n",
    "        torch.Tensor([1.1]).cuda() # cond scale of exactly 1 causes the ONNX graph to skip some vital things, so use 1.1\n",
    "    ), \n",
    "    \"toymodel/public/unet-32.onnx\", \n",
    "    input_names=['image', 'text_embeds', 'text_mask', 'timestep', 'time_next', 'cond_scale'], \n",
    "    output_names=['prediction', 'x_start'], # you can basically ignore x_start\n",
    "    dynamic_axes={\n",
    "        'image': {0: 'batch_size'},\n",
    "        'text_embeds': {0: 'batch_size', 1: 'n_tokens'},\n",
    "        'text_mask': {0: 'batch_size', 1: 'n_tokens'},\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a26c059c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8 (tags/v3.10.8:aaaf517, Oct 11 2022, 16:50:30) [MSC v.1933 64 bit (AMD64)]"
  },
  "vscode": {
   "interpreter": {
    "hash": "369f2c481f4da34e4445cda3fffd2e751bd1c4d706f27375911949ba6bb62e1c"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}