Simplified_DALLE_mini_Inference_pipeline_jax.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Copy of DALL·E mini - Inference pipeline.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sadnow/ESRGAN-UltraFast-CLIP-Guided-Diffusion-Colab/blob/main/Simplified_DALLE_mini_Inference_pipeline_jax.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "118UKH5bWCGa"
      },
      "source": [
        "# DALL·E mini - Inference pipeline\n",
        "\n",
        "*Generate images from a text prompt*\n",
        "\n",
        "<img src=\"https://github.com/borisdayma/dalle-mini/blob/main/img/logo.png?raw=true\" width=\"200\">\n",
        "\n",
        "This notebook illustrates [DALL·E mini](https://github.com/borisdayma/dalle-mini) inference pipeline.\n",
        "\n",
        "Just want to play? Use [the demo](https://huggingface.co/spaces/flax-community/dalle-mini) (temporarily in beta on Huging Face Spaces but soon to be open to all).\n",
        "\n",
        "For more understanding of the model, refer to [the report](https://wandb.ai/dalle-mini/dalle-mini/reports/DALL-E-mini--Vmlldzo4NjIxODA)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1Q1aBuyzlxnn"
      },
      "source": [
        "#@markdown ##Check GPU\n",
        "#@markdown ***This Notebook supports \"Run All. Once everything is set to your liking, you can do `Ctrl+F9`!***\n",
        "import torch\n",
        "# Check the GPU status\n",
        "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
        "print('Using device:', device)\n",
        "\n",
        "enable_error_checking = False#@param {type:\"boolean\"}\n",
        "if enable_error_checking:\n",
        "  !nvidia-smi\n",
        "else:\n",
        "  !nvidia-smi\n",
        "  !nvidia-smi -i 0 -e 0\n",
        "#@markdown *Special thanks to sportracer48 and his Discord channel. If you want to make AI animations, he has the meats in closed beta.* https://www.patreon.com/sportsracer48"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uzjAM2GBYpZX"
      },
      "source": [
        "#@markdown ##Install dependencies\n",
        "!pip install -q transformers flax\n",
        "!pip install -q git+https://github.com/patil-suraj/vqgan-jax.git  # VQGAN model in JAX\n",
        "!pip install -q git+https://github.com/borisdayma/dalle-mini.git  # Model files"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "phQ9bhjRkgAZ"
      },
      "source": [
        "## Generate encoded images\n",
        "\n",
        "We generate prediction samples from a text prompt using `flax-community/dalle-mini` model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yyT4tk5EbsdO"
      },
      "source": [
        "\n",
        "from dalle_mini.model import CustomFlaxBartForConditionalGeneration\n",
        "from transformers import BartTokenizer\n",
        "import jax\n",
        "import random\n",
        "from tqdm.notebook import tqdm, trange\n",
        "#####\n",
        "# make sure we use compatible versions\n",
        "DALLE_REPO = 'flax-community/dalle-mini'\n",
        "DALLE_COMMIT_ID = '4d34126d0df8bc4a692ae933e3b902a1fa8b6114'\n",
        "#####\n",
        "# set up tokenizer and model\n",
        "tokenizer = BartTokenizer.from_pretrained(DALLE_REPO, revision=DALLE_COMMIT_ID)\n",
        "model = CustomFlaxBartForConditionalGeneration.from_pretrained(DALLE_REPO, revision=DALLE_COMMIT_ID)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IBFSuYbSgIf9"
      },
      "source": [
        "# set a prompt\n",
        "_prompt = 'Sad cow' #@param {type:\"string\"}\n",
        "prompt = _prompt\n",
        "############################################################\n",
        "# tokenize the prompt\n",
        "tokenized_prompt = tokenizer(prompt, return_tensors='jax', padding='max_length', truncation=True, max_length=128)\n",
        "tokenized_prompt\n",
        "############################################################\n",
        "_n_predictions =  8#@param{type:\"raw\"}\n",
        "n_predictions = _n_predictions\n",
        "\n",
        "# create random keys\n",
        "seed = random.randint(0, 2**32-1)\n",
        "key = jax.random.PRNGKey(seed)\n",
        "subkeys = jax.random.split(key, num=n_predictions)\n",
        "subkeys\n",
        "#####################################################\n",
        "# generate sample predictions\n",
        "encoded_images = [model.generate(**tokenized_prompt, do_sample=True, num_beams=1, prng_key=subkey) for subkey in tqdm(subkeys)]\n",
        "encoded_images[0]\n",
        "####################################################\n",
        "# remove first token (BOS)\n",
        "encoded_images = [img.sequences[..., 1:] for img in encoded_images]\n",
        "encoded_images[0]\n",
        "#####################################################\n",
        "encoded_images[0].shape\n",
        "#####################################################\n",
        "####################################################\n",
        "#Decode images\n",
        "from vqgan_jax.modeling_flax_vqgan import VQModel\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "#######################################################\n",
        "# make sure we use compatible versions\n",
        "VQGAN_REPO = 'flax-community/vqgan_f16_16384'\n",
        "VQGAN_COMMIT_ID = '90cc46addd2dd8f5be21586a9a23e1b95aa506a9'\n",
        "#######################################################\n",
        "# set up VQGAN\n",
        "vqgan = VQModel.from_pretrained(VQGAN_REPO, revision=VQGAN_COMMIT_ID)\n",
        "########################################################\n",
        "# decode images\n",
        "decoded_images = [vqgan.decode_code(encoded_image) for encoded_image in tqdm(encoded_images)]\n",
        "decoded_images[0]\n",
        "##########################################################\n",
        "# normalize images\n",
        "clipped_images = [img.squeeze().clip(0., 1.) for img in decoded_images]\n",
        "############################################################\n",
        "# convert to image\n",
        "images = [Image.fromarray(np.asarray(img * 255, dtype=np.uint8)) for img in clipped_images]\n",
        "##############################################################\n",
        "# display an image\n",
        "images[0]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_Y5dqFj7prMQ"
      },
      "source": [
        "Notes:\n",
        "\n",
        "* `0`: BOS, special token representing the beginning of a sequence\n",
        "* `2`: EOS, special token representing the end of a sequence\n",
        "* `1`: special token representing the padding of a sequence when requesting a specific length"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K3m2lxNOrQWG"
      },
      "source": [
        "The first token (`16384`) is a special token representing the start of a sequence in the decoder (not part of the image codebook)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "037WuNLXsCoq"
      },
      "source": [
        "The generated images are now represented by 256 tokens."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HDiVxh1wfrJ1"
      },
      "source": [
        "## Decode images\n",
        "\n",
        "The generated images need to be decoded with `flax-community/vqgan_f16_16384`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4biKhmc8_DmX"
      },
      "source": [
        "## Rank images with CLIP\n",
        "\n",
        "We use `openai/clip-vit-base-patch32` to rank generated images against the prompt."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v6TbUsfN_IYp"
      },
      "source": [
        "#_prompt = 'Bridge to the end of the world' #@param {type:\"string\"}\n",
        "#prompt = _prompt\n",
        "###########################################################\n",
        "from transformers import CLIPProcessor, FlaxCLIPModel\n",
        "##########################################################\n",
        "# set up model and processor\n",
        "clip = FlaxCLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
        "processor = CLIPProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
        "#############################################################\n",
        "# evaluate scores\n",
        "inputs = processor(text=prompt, images=images, return_tensors='np')\n",
        "logits = clip(**inputs).logits_per_image\n",
        "scores = jax.nn.softmax(logits, axis=0).squeeze()  # normalize and sum all scores to 1\n",
        "###################################################\n",
        "# rank images by score\n",
        "print(f'Prompt: {prompt}\\n')\n",
        "for idx in scores.argsort()[::-1]:\n",
        "    print(f'Score: {scores[idx]}')\n",
        "    display(images[idx])\n",
        "    print()\n",
        "###################################################"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ezs7ucnfMAHy"
      },
      "source": [
        "The CLIP processor tokenizes text and pre-processes images (resize to 224x224 and normalize) as required per the CLIP model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lYyOyZGH9dpb"
      },
      "source": [
        "## Leverage JAX for faster inference\n",
        "\n",
        "[JAX](https://github.com/google/jax) uses XLA to compile code to GPU/TPU, leading to faster inference.\n",
        "\n",
        "Even with only 1 GPU, we can benefit from impressive speedups, even more after first inference (where the compilation happens)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "k1hec1qcX7dE"
      },
      "source": [
        "from functools import partial\n",
        "from flax.training.common_utils import shard\n",
        "from flax.jax_utils import replicate\n",
        "import gc\n",
        "#gc.collect()\n",
        "#torch.cuda.empty_cache()\n",
        "#XLA_PYTHON_CLIENT_MEM_FRACTION=.80\n",
        "#####\n",
        "# check we can access TPU's or GPU's\n",
        "jax.devices()\n",
        "#####\n",
        "# replicate parameters on all devices\n",
        "dalle_params = replicate(model.params)\n",
        "vqgan_params = replicate(vqgan.params)\n",
        "#####\n",
        "# one set of inputs per device\n",
        "prompt = [_prompt] * jax.device_count()\n",
        "#####\n",
        "# tokenize prompts and shard them across available devices\n",
        "tokenized_prompt = tokenizer(prompt, return_tensors='jax', padding='max_length', truncation=True, max_length=128).data\n",
        "tokenized_prompt = shard(tokenized_prompt)\n",
        "#####\n",
        "# parallelize and compile functions\n",
        "\n",
        "# function to generate encoded images\n",
        "@partial(jax.pmap, axis_name=\"batch\")\n",
        "def p_generate(tokenized_prompt, key, params):\n",
        "    return model.generate(\n",
        "        **tokenized_prompt,\n",
        "        do_sample=True,\n",
        "        num_beams=1,\n",
        "        prng_key=key,\n",
        "        params=params\n",
        "    )\n",
        "\n",
        "# function to decode images\n",
        "@partial(jax.pmap, axis_name=\"batch\")\n",
        "def p_decode(indices, params):\n",
        "    return vqgan.decode_code(indices, params=params)\n",
        "#####\n",
        "# generate images with compiled functions\n",
        "_jax_predictions =  1#@param{type:\"raw\"}\n",
        "n_predictions = _jax_predictions\n",
        "images = []\n",
        "\n",
        "for i in trange(n_predictions // jax.device_count()):\n",
        "    key, *subkeys = jax.random.split(key, jax.device_count() + 1)\n",
        "    subkeys = jax.numpy.stack(subkeys)\n",
        "\n",
        "    encoded_images = p_generate(tokenized_prompt, subkeys, dalle_params)\n",
        "    encoded_images = encoded_images.sequences[..., 1:]\n",
        "\n",
        "    decoded_images = p_decode(encoded_images, vqgan_params)\n",
        "    decoded_images = decoded_images.clip(0., 1.).reshape((-1, 256, 256, 3))\n",
        "\n",
        "    for img in decoded_images:\n",
        "        images.append(Image.fromarray(np.asarray(img * 255, dtype=np.uint8)))\n",
        "#####\n",
        "for img in images:\n",
        "    display(img)\n",
        "    print()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "leO2RvYMLj5h"
      },
      "source": [
        "We use `pmap` to compile the functions with XLA and perform operations in parallel on multiple devices."
      ]
    }
  ]
}