dalle2_laion_alpha.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y8vri4p7Alx1"
      },
      "source": [
        "### Acknowledgements\n",
        "\n",
        "First, a huge thank you to [LucidRains](https://github.com/lucidrains) for his amazing work creating dalle2 model repo.\n",
        "\n",
        "Thank you to [Nousr](https://twitter.com/nousr_) and [Aidan](https://github.com/Veldrovive) for their work on the training scripts for the prior and decoder. As well, thank you to [Romain](https://github.com/rom1504) for his work on enabling multi-node training and for his part in creating the dataset.\n",
        "\n",
        "This was trained thanks to the generous donation of compute time from [LAION](https://laion.ai/) and its sponsors, especially [StabilityAI](https://stability.ai/) for the servers used in this project.\n",
        "This model was trained using the aesthetic subset of the [LAION2B dataset](https://laion.ai/blog/laion-5b/), details of the run can be found on [wandb](https://wandb.ai/nousr_laion/dalle2_train_decoder/reports/Decoder-Training--VmlldzoyMjEyMjcw)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qHpA7GtYAlx3"
      },
      "source": [
        "#### This is a WIP notebook.\n",
        "It is missing the upsamplers so it can only produce 64x64 images and has only be trained for 0.5% of what OpenAI did for their DALLE2. It can still produce some impressive results, but is not nearly at the level you might see on the news yet.\n",
        "\n",
        "While the upsamplers have yet to be trained, we have substituted [SwinIR](https://github.com/JingyunLiang/SwinIR) to produce 256x256 images. This tends to flatten details, but still generally improves the quality of the output."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LZmRFKPkAlx3"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sjnUmrV2Alx4"
      },
      "source": [
        "### Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "lPklyP0pg0mF"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "!pip install -q ipywidgets\n",
        "import ipywidgets as widgets\n",
        "import sys\n",
        "import os\n",
        "import json\n",
        "\n",
        "decoder_versions = [{\n",
        "    \"name\": \"Original\",\n",
        "    \"dalle2_install_path\": \"git+https://github.com/Veldrovive/DALLE2-pytorch@f4b687798d367fc434d8127ab31141f0fea0db26\",\n",
        "    \"decoder_path\": \"https://huggingface.co/Veldrovive/DA-VINC-E/resolve/main/text_conditioned_epoch_34.pth\",\n",
        "    \"config_path\": \"https://huggingface.co/Veldrovive/DA-VINC-E/raw/main/text_conditioned_config.json\"\n",
        "},{\n",
        "    \"name\": \"New 1B (Aesthetic)\",\n",
        "    \"dalle2_install_path\": \"dalle2_pytorch==0.15.4\",\n",
        "    \"decoder_path\": \"https://huggingface.co/laion/DALLE2-PyTorch/resolve/main/decoder/small_32gpus/latest.pth\",\n",
        "    \"config_path\": \"https://huggingface.co/laion/DALLE2-PyTorch/raw/main/decoder/small_32gpus/decoder_config.json\"\n",
        "},{\n",
        "    \"name\": \"New 1.5B (Aesthetic)\",\n",
        "    \"dalle2_install_path\": \"dalle2_pytorch==0.15.4\",\n",
        "    \"decoder_path\": \"https://huggingface.co/laion/DALLE2-PyTorch/resolve/main/decoder/1.5B/latest.pth\",\n",
        "    \"config_path\": \"https://huggingface.co/laion/DALLE2-PyTorch/raw/main/decoder/1.5B/decoder_config.json\"\n",
        "},{\n",
        "    \"name\": \"New 1.5B (Laion2B)\",\n",
        "    \"dalle2_install_path\": \"dalle2_pytorch==0.15.4\",\n",
        "    \"decoder_path\": \"https://huggingface.co/laion/DALLE2-PyTorch/resolve/main/decoder/1.5B_laion2B/latest.pth\",\n",
        "    \"config_path\": \"https://huggingface.co/laion/DALLE2-PyTorch/raw/main/decoder/1.5B_laion2B/decoder_config.json\"\n",
        "},{\n",
        "    \"name\": \"Upsampler\",\n",
        "    \"dalle2_install_path\": \"git+https://github.com/Veldrovive/DALLE2-pytorch@b2549a4d17244dab09e7a9496a9cb6330b7d3070\",\n",
        "    \"decoder\": [\n",
        "        {\n",
        "            \"unets\": [0],\n",
        "            \"model_path\": \"https://huggingface.co/laion/DALLE2-PyTorch/resolve/main/decoder/1.5B_laion2B/latest.pth\",\n",
        "            \"config_path\": \"https://huggingface.co/laion/DALLE2-PyTorch/raw/main/decoder/1.5B_laion2B/decoder_config.json\"\n",
        "        },\n",
        "        {\n",
        "            \"unets\": [1],\n",
        "            \"model_path\": \"https://huggingface.co/Veldrovive/upsamplers/resolve/main/working/latest.pth\",\n",
        "            \"config_path\": \"https://huggingface.co/Veldrovive/upsamplers/raw/main/working/decoder_config.json\"\n",
        "        }\n",
        "    ],\n",
        "    \"prior\": {\n",
        "        \"model_path\": \"https://huggingface.co/zenglishuci/conditioned-prior/resolve/main/vit-l-14/prior_aes_finetune.pth\",\n",
        "        \"config_path\": \"\"\n",
        "    }\n",
        "}]\n",
        "\n",
        "decoder_options = [version[\"name\"] for version in decoder_versions]\n",
        "\n",
        "def load_state():\n",
        "    state_path = \"script_state.json\"\n",
        "    try:\n",
        "        assert os.path.exists(state_path)\n",
        "        with open(state_path, \"r\") as f:\n",
        "            state = json.load(f)\n",
        "        # Make sure the save config is up to date. You might think this is a stupid system but...\n",
        "        decoder = state[\"decoder\"]\n",
        "        if decoder is not None:\n",
        "          current_decoder_name = decoder[\"name\"]\n",
        "          try:\n",
        "              current_decoder_index = decoder_options.index(current_decoder_name)\n",
        "              state[\"decoder\"] = decoder_versions[current_decoder_index]\n",
        "          except ValueError:\n",
        "              print(\"The decoder you were using no longer exists. Please pick a new option.\")\n",
        "              state[\"decoder\"] = None\n",
        "\n",
        "        # Check if models are where they say they are\n",
        "        for filekey in [\"decoder\", \"decoder_config\", \"prior\", \"prior_config\"]:\n",
        "            path = state[\"model_paths\"][filekey]\n",
        "            if path is not None and not os.path.exists(path):\n",
        "                print(f\"{filekey} not found in expected place. Removing decoder config.\")\n",
        "                state[\"decoder\"] = None\n",
        "                state[\"model_paths\"] = {\n",
        "                    \"decoder\": None,\n",
        "                    \"decoder_config\": None,\n",
        "                    \"prior\": None,\n",
        "                    \"prior_config\": None\n",
        "                }\n",
        "                save_state()\n",
        "    except Exception as e:\n",
        "        state = {\n",
        "            \"text_input\": '',\n",
        "            \"text_repeat\": 3,\n",
        "            \"prior_conditioning\": 1.0,\n",
        "            \"img_repeat\": 1,\n",
        "            \"decoder_conditioning\": 3.5,\n",
        "            \"include_prompt_checkbox\": True,\n",
        "            \"upsample_checkbox\": True,\n",
        "            \"decoder\": None,\n",
        "            \"model_paths\": {\n",
        "                \"decoder\": None,\n",
        "                \"decoder_config\": None,\n",
        "                \"prior\": None,\n",
        "                \"prior_config\": None\n",
        "            }\n",
        "        }\n",
        "    return state\n",
        "\n",
        "current_state = load_state()\n",
        "\n",
        "def save_state():\n",
        "    global current_state\n",
        "    state_path = \"script_state.json\"\n",
        "    with open(state_path, \"w\") as f:\n",
        "        json.dump(current_state, f)\n",
        "\n",
        "def choice_equal(new_choice):\n",
        "    global current_state\n",
        "    if current_state[\"decoder\"] is None:\n",
        "        return False\n",
        "    return current_state[\"decoder\"][\"decoder_path\"] == new_choice[\"decoder_path\"]\n",
        "\n",
        "def dalle2_imported():\n",
        "    return \"dalle2_pytorch\" in sys.modules\n",
        "\n",
        "chosen_decoder = current_state[\"decoder\"] if current_state[\"decoder\"] is not None else decoder_versions[-1]\n",
        "\n",
        "decoder_version_dropdown = widgets.Dropdown(\n",
        "    options=decoder_options,\n",
        "    value=chosen_decoder[\"name\"],\n",
        "    description='Decoder:',\n",
        "    disabled=False,\n",
        ")\n",
        "\n",
        "start_setup_button = widgets.Button(\n",
        "    description=\"Setup\"\n",
        ")\n",
        "\n",
        "redownload_button = widgets.Button (\n",
        "    description=\"Force Update Models\"\n",
        ")\n",
        "\n",
        "main_layout = widgets.VBox([decoder_version_dropdown, start_setup_button, redownload_button])\n",
        "\n",
        "def setup(decoder_version_name):\n",
        "    global current_state\n",
        "    global chosen_decoder\n",
        "    new_choice = decoder_versions[decoder_options.index(decoder_version_name)]\n",
        "    current_choice = current_state[\"decoder\"]\n",
        "\n",
        "    new_choice_equal = choice_equal(new_choice)\n",
        "    already_imported = dalle2_imported()\n",
        "\n",
        "    requires_restart = not new_choice_equal and already_imported  # The wrong dalle2_pytorch version is already imported\n",
        "    requires_download = not new_choice_equal  # The wrong decoder version is downloaded\n",
        "\n",
        "    print(f\"You are using the model {new_choice['name']} which will be downloaded from {new_choice['decoder_path']}\\n\")\n",
        "    if requires_restart:\n",
        "        print(\"You environment already has dalle2 imported and collab requires a restart for you to be able to import the new version.\")\n",
        "        print(\"Restart your runtime to proceed.\")\n",
        "    elif requires_download:\n",
        "        print(\"The models are not downloaded. They will be when you proceed.\")\n",
        "    else:\n",
        "        print(\"You are ready to run inference. If you suspect your models are out of date, force update them.\")\n",
        "\n",
        "    chosen_decoder = new_choice\n",
        "\n",
        "\n",
        "out = widgets.interactive_output(setup, { 'decoder_version_name': decoder_version_dropdown })\n",
        "display(main_layout, out)\n",
        "\n",
        "def download_models(current_choice):\n",
        "    model_dir = \"./models\"\n",
        "    os.makedirs(model_dir, exist_ok=True)\n",
        "\n",
        "    # Download decoder\n",
        "    print(\"Downloading decoder and decoder config\")\n",
        "    decoder_url = current_choice[\"decoder_path\"]\n",
        "    decoder_config_url = current_choice[\"config_path\"]\n",
        "\n",
        "    decoder_path = os.path.join(model_dir, \"decoder.pth\")\n",
        "    decoder_config_path = os.path.join(model_dir, \"decoder_config.json\")\n",
        "\n",
        "    !curl -L {decoder_url} > {decoder_path}\n",
        "    !curl -L {decoder_config_url} > {decoder_config_path}\n",
        "\n",
        "    # Download prior\n",
        "    print(\"Downloading prior and prior config\")\n",
        "    prior_url = \"https://huggingface.co/zenglishuci/conditioned-prior/resolve/main/vit-l-14/prior_aes_finetune.pth\"\n",
        "\n",
        "    prior_path = os.path.join(model_dir, \"prior.pth\")\n",
        "\n",
        "    !curl -L {prior_url} > {prior_path}\n",
        "    return decoder_path, decoder_config_path, prior_path, None\n",
        "\n",
        "def install_dependencies(state):\n",
        "    print(\"Installing dependencies\")\n",
        "    dalle2_install_path = state[\"decoder\"][\"dalle2_install_path\"]\n",
        "    print(f\"Installing dalle2 version {dalle2_install_path}\")\n",
        "    !pip install -q {dalle2_install_path}\n",
        "\n",
        "    !pip install -q matplotlib\n",
        "\n",
        "    print(\"Do not worry if you get the error `fatal: destination path 'SwinIR' already exists and is not an empty directory.`\")\n",
        "    print(\"That just means SwinIR is already installed and I'm too lazy to do the check myself\")\n",
        "    !git clone https://github.com/JingyunLiang/SwinIR.git\n",
        "    !pip install -q timm\n",
        "    !pip install -q opencv-python\n",
        "\n",
        "def start_setup(b):\n",
        "    global current_state\n",
        "    global chosen_decoder\n",
        "    current_choice = current_state[\"decoder\"]\n",
        "\n",
        "    new_choice_equal = choice_equal(chosen_decoder)\n",
        "    already_imported = dalle2_imported()\n",
        "\n",
        "    requires_restart = not new_choice_equal and already_imported  # The wrong dalle2_pytorch version is already imported\n",
        "    requires_download = not new_choice_equal  # The wrong decoder version is downloaded\n",
        "\n",
        "    if requires_restart:\n",
        "        raise Exception(\"Please restart your runtime before trying to set up the environment\")\n",
        "\n",
        "    if requires_download:\n",
        "        try:\n",
        "            decoder_path, decoder_config_path, prior_path, prior_config_path = download_models(chosen_decoder)\n",
        "        except Exception as e:\n",
        "            print(\"Model download was interrupted. Manually delete all models or environment may be corrupted.\")\n",
        "            current_state[\"decoder\"] = None\n",
        "            save_state()\n",
        "            raise e\n",
        "        current_state[\"decoder\"] = chosen_decoder\n",
        "        current_state[\"model_paths\"] = {\n",
        "                \"decoder\": decoder_path,\n",
        "                \"decoder_config\": decoder_config_path,\n",
        "                \"prior\": prior_path,\n",
        "                \"prior_config\": prior_config_path\n",
        "            }\n",
        "        save_state()\n",
        "\n",
        "    install_dependencies(current_state)\n",
        "\n",
        "start_setup_button.on_click(start_setup)\n",
        "\n",
        "def force_download(b):\n",
        "    global current_state\n",
        "    current_choice = current_state[\"decoder\"]\n",
        "    updated_choice = decoder_versions[decoder_options.index(current_choice[\"name\"])]\n",
        "    chosen_decoder = updated_choice\n",
        "    try:\n",
        "        decoder_path, decoder_config_path, prior_path, prior_config_path = download_models(chosen_decoder)\n",
        "    except Exception as e:\n",
        "        print(\"Model download was interrupted. Force update models or environment may be corrupted.\")\n",
        "        current_state[\"decoder\"] = None\n",
        "        save_state()\n",
        "        raise e\n",
        "    current_state[\"decoder\"] = chosen_decoder\n",
        "    current_state[\"model_paths\"] = {\n",
        "            \"decoder\": decoder_path,\n",
        "            \"decoder_config\": decoder_config_path,\n",
        "            \"prior\": prior_path,\n",
        "            \"prior_config\": prior_config_path\n",
        "        }\n",
        "    save_state()\n",
        "    install_dependencies(current_state)\n",
        "\n",
        "redownload_button.on_click(force_download)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ZwesxZ3uAlx4"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "# I would suggest running on a remote machine https://research.google.com/colaboratory/local-runtimes.html\n",
        "import shutil\n",
        "import torch\n",
        "import os\n",
        "import importlib\n",
        "\n",
        "\n",
        "from dalle2_pytorch import DALLE2, DiffusionPriorNetwork, DiffusionPrior, OpenAIClipAdapter, train_configs\n",
        "from dalle2_pytorch.tokenizer import tokenizer\n",
        "\n",
        "import numpy as np\n",
        "\n",
        "\n",
        "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6D9RjtGoAlx5"
      },
      "source": [
        "### Load the decoder"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "WjdOURcWAlx5"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "\n",
        "def conditioned_on_text(config):\n",
        "    try:\n",
        "        return config.decoder.unets[0].cond_on_text_encodings\n",
        "    except AttributeError:\n",
        "        pass\n",
        "\n",
        "    try:\n",
        "        return config.decoder.condition_on_text_encodings\n",
        "    except AttributeError:\n",
        "        pass\n",
        "\n",
        "    return False\n",
        "\n",
        "decoder_text_conditioned = False\n",
        "clip_config = None\n",
        "def load_decoder(decoder_state_dict_path, config_file_path):\n",
        "  config = train_configs.TrainDecoderConfig.from_json_path(config_file_path)\n",
        "  global decoder_text_conditioned\n",
        "  decoder_text_conditioned = conditioned_on_text(config)\n",
        "  global clip_config\n",
        "  clip_config = config.decoder.clip\n",
        "  config.decoder.clip = None\n",
        "  print(\"Decoder conditioned on text\", decoder_text_conditioned)\n",
        "  decoder = config.decoder.create().to(device)\n",
        "  decoder_state_dict = torch.load(decoder_state_dict_path, map_location='cpu')\n",
        "  decoder.load_state_dict(decoder_state_dict, strict=False)\n",
        "  del decoder_state_dict\n",
        "  decoder.eval()\n",
        "  return decoder\n",
        "decoder = load_decoder(current_state[\"model_paths\"][\"decoder\"], current_state[\"model_paths\"][\"decoder_config\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ab__xeOeAlx5"
      },
      "source": [
        "### Load the prior"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ZIGSaY2rAlx5"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "\n",
        "def load_prior(model_path):\n",
        "  prior_network = DiffusionPriorNetwork(\n",
        "    dim=768,\n",
        "    depth=24,\n",
        "    dim_head=64,\n",
        "    heads=32,\n",
        "    normformer=True,\n",
        "    attn_dropout=5e-2,\n",
        "    ff_dropout=5e-2,\n",
        "    num_time_embeds=1,\n",
        "    num_image_embeds=1,\n",
        "    num_text_embeds=1,\n",
        "    num_timesteps=1000,\n",
        "    ff_mult=4\n",
        "  )\n",
        "\n",
        "  diffusion_prior = DiffusionPrior(\n",
        "    net=prior_network,\n",
        "    clip=OpenAIClipAdapter(\"ViT-L/14\"),\n",
        "    image_embed_dim=768,\n",
        "    timesteps=1000,\n",
        "    cond_drop_prob=0.1,\n",
        "    loss_type=\"l2\",\n",
        "    condition_on_text_encodings=True,\n",
        "  ).to(device)\n",
        "\n",
        "  state_dict = torch.load(model_path, map_location='cpu')\n",
        "  if 'ema_model' in state_dict:\n",
        "    print('Loading EMA Model')\n",
        "    diffusion_prior.load_state_dict(state_dict['ema_model'], strict=True)\n",
        "  else:\n",
        "    print('Loading Standard Model')\n",
        "    diffusion_prior.load_state_dict(state_dict['model'], strict=False)\n",
        "  del state_dict\n",
        "  return diffusion_prior\n",
        "diffusion_prior = load_prior(current_state[\"model_paths\"][\"prior\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KxnLvVQxAlx5"
      },
      "source": [
        "### Load Clip"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ikdfRY9yAlx6"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "clip = None\n",
        "if clip_config is not None:\n",
        "  clip = clip_config.create()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uhSMszwrAlx6"
      },
      "source": [
        "### Add helpers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "sLOGTt6MAlx6"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "from PIL import Image, ImageDraw, ImageFont\n",
        "\n",
        "class ImageText(object):\n",
        "    def __init__(self, filename_or_size, mode='RGBA', background=(0, 0, 0, 0), encoding='utf8'):\n",
        "        if isinstance(filename_or_size, str):\n",
        "            self.filename = filename_or_size\n",
        "            self.image = Image.open(self.filename)\n",
        "            self.size = self.image.size\n",
        "        elif isinstance(filename_or_size, (list, tuple)):\n",
        "            self.size = filename_or_size\n",
        "            self.image = Image.new(mode, self.size, color=background)\n",
        "            self.filename = None\n",
        "        self.draw = ImageDraw.Draw(self.image)\n",
        "        self.encoding = encoding\n",
        "\n",
        "    def save(self, filename=None):\n",
        "        self.image.save(filename or self.filename)\n",
        "\n",
        "    def get_font_size(self, text, font, max_width=None, max_height=None):\n",
        "        if max_width is None and max_height is None:\n",
        "            raise ValueError('You need to pass max_width or max_height')\n",
        "        font_size = 1\n",
        "        text_size = self.get_text_size(font, font_size, text)\n",
        "        if (max_width is not None and text_size[0] > max_width) or \\\n",
        "           (max_height is not None and text_size[1] > max_height):\n",
        "            raise ValueError(\"Text can't be filled in only (%dpx, %dpx)\" % \\\n",
        "                    text_size)\n",
        "        while True:\n",
        "            if (max_width is not None and text_size[0] >= max_width) or \\\n",
        "               (max_height is not None and text_size[1] >= max_height):\n",
        "                return font_size - 1\n",
        "            font_size += 1\n",
        "            text_size = self.get_text_size(font, font_size, text)\n",
        "\n",
        "    def write_text(self, xy, text, font_filename, font_size=11,\n",
        "                   color=(0, 0, 0), max_width=None, max_height=None):\n",
        "        x, y = xy\n",
        "        # if isinstance(text, str):\n",
        "        #     text = text.decode(self.encoding)\n",
        "        if font_size == 'fill' and \\\n",
        "           (max_width is not None or max_height is not None):\n",
        "            font_size = self.get_font_size(text, font_filename, max_width,\n",
        "                                           max_height)\n",
        "        text_size = self.get_text_size(font_filename, font_size, text)\n",
        "        font = ImageFont.truetype(font_filename, font_size)\n",
        "        # font = ImageFont.load_default()\n",
        "        if x == 'center':\n",
        "            x = (self.size[0] - text_size[0]) / 2\n",
        "        if y == 'center':\n",
        "            y = (self.size[1] - text_size[1]) / 2\n",
        "        self.draw.text((x, y), text, font=font, fill=color)\n",
        "        return text_size\n",
        "\n",
        "    def get_text_size(self, font_filename, font_size, text):\n",
        "        font = ImageFont.truetype(font_filename, font_size)\n",
        "        return font.getsize(text)\n",
        "\n",
        "    def write_text_box(self, xy, text, box_width, font_filename,\n",
        "                       font_size=11, color=(0, 0, 0), place='left',\n",
        "                       justify_last_line=False):\n",
        "        x, y = xy\n",
        "        lines = []\n",
        "        line = []\n",
        "        words = text.split()\n",
        "        for word in words:\n",
        "            new_line = ' '.join(line + [word])\n",
        "            size = self.get_text_size(font_filename, font_size, new_line)\n",
        "            text_height = size[1]\n",
        "            if size[0] <= box_width:\n",
        "                line.append(word)\n",
        "            else:\n",
        "                lines.append(line)\n",
        "                line = [word]\n",
        "        if line:\n",
        "            lines.append(line)\n",
        "        lines = [' '.join(line) for line in lines if line]\n",
        "        height = y\n",
        "        for index, line in enumerate(lines):\n",
        "            if place == 'left':\n",
        "                self.write_text((x, height), line, font_filename, font_size,\n",
        "                                color)\n",
        "            elif place == 'right':\n",
        "                total_size = self.get_text_size(font_filename, font_size, line)\n",
        "                x_left = x + box_width - total_size[0]\n",
        "                self.write_text((x_left, height), line, font_filename,\n",
        "                                font_size, color)\n",
        "            elif place == 'center':\n",
        "                total_size = self.get_text_size(font_filename, font_size, line)\n",
        "                x_left = int(x + ((box_width - total_size[0]) / 2))\n",
        "                self.write_text((x_left, height), line, font_filename,\n",
        "                                font_size, color)\n",
        "            elif place == 'justify':\n",
        "                words = line.split()\n",
        "                if (index == len(lines) - 1 and not justify_last_line) or \\\n",
        "                   len(words) == 1:\n",
        "                    self.write_text((x, height), line, font_filename, font_size,\n",
        "                                    color)\n",
        "                    continue\n",
        "                line_without_spaces = ''.join(words)\n",
        "                total_size = self.get_text_size(font_filename, font_size,\n",
        "                                                line_without_spaces)\n",
        "                space_width = (box_width - total_size[0]) / (len(words) - 1.0)\n",
        "                start_x = x\n",
        "                for word in words[:-1]:\n",
        "                    self.write_text((start_x, height), word, font_filename,\n",
        "                                    font_size, color)\n",
        "                    word_size = self.get_text_size(font_filename, font_size,\n",
        "                                                    word)\n",
        "                    start_x += word_size[0] + space_width\n",
        "                last_word_size = self.get_text_size(font_filename, font_size,\n",
        "                                                    words[-1])\n",
        "                last_word_x = x + box_width - last_word_size[0]\n",
        "                self.write_text((last_word_x, height), words[-1], font_filename,\n",
        "                                font_size, color)\n",
        "            height += text_height\n",
        "        return (box_width, height - y)\n",
        "\n",
        "def download_font():\n",
        "  if not os.path.exists(\"./Arial.ttf\"):\n",
        "    !wget https://github.com/matomo-org/travis-scripts/raw/master/fonts/Arial.ttf\n",
        "  return \"./Arial.ttf\"\n",
        "\n",
        "\n",
        "def map_images(np_images, prior_repeat, decoder_repeat, prompts, upscale=4):\n",
        "  # Match the images to their prompts\n",
        "  # Format [{ prompt: STRING, images: [\n",
        "  #  { prior_index: INT, decoder_index: INT, img: NP_ARR[64, 64, 3] }\n",
        "  # ] }]\n",
        "  image_map = {}\n",
        "  curr_index = 0\n",
        "  for prompt in prompts:\n",
        "    for prior_index in range(prior_repeat):\n",
        "      for decoder_index in range(decoder_repeat):\n",
        "        img = np_images[curr_index]\n",
        "        if prompt not in image_map:\n",
        "          image_map[prompt] = []\n",
        "        if isinstance(img, np.ndarray):\n",
        "          image = Image.fromarray(np.uint8(img * 255))\n",
        "          image = image.resize([dim * upscale for dim in image.size])\n",
        "        else:\n",
        "          image = img\n",
        "        image_map[prompt].append({\n",
        "            \"prior_index\": prior_index,\n",
        "            \"decoder_index\": decoder_index,\n",
        "            \"img\": image\n",
        "        })\n",
        "        curr_index += 1\n",
        "  return image_map\n",
        "\n",
        "def format_image_grid(img_array):\n",
        "  example_image = img_array[0][\"img\"]\n",
        "  max_prior_index = max((inst[\"prior_index\"] for inst in img_array))\n",
        "  cols = max_prior_index + 1\n",
        "  max_decoder_index = max((inst[\"decoder_index\"] for inst in img_array))\n",
        "  rows = max_decoder_index + 1\n",
        "\n",
        "  w, h = example_image.size\n",
        "  grid = Image.new('RGB', size=(cols*w, rows*h))\n",
        "  grid_w, grid_h = grid.size\n",
        "\n",
        "  for img_data in img_array:\n",
        "    x_pos = img_data[\"prior_index\"] * w\n",
        "    y_pos = img_data[\"decoder_index\"] * h\n",
        "    grid.paste(img_data[\"img\"], box=(x_pos, y_pos))\n",
        "  return grid\n",
        "\n",
        "def format_prompt_image(prompt, grid_img, font_size=20, horizontal_padding=10, vertical_padding=10):\n",
        "  grid_w, grid_h = grid_img.size\n",
        "  prompt_img = ImageText((grid_w, 2000), background=(255, 255, 255, 255))\n",
        "  # font = ImageFont.load(\"arial.pil\")\n",
        "  # font = ImageFont.load_default()\n",
        "  font_path = download_font()\n",
        "  text_w, text_h = prompt_img.write_text_box((horizontal_padding, vertical_padding), prompt, box_width=grid_w - horizontal_padding, font_filename=font_path, font_size=font_size, color=(0, 0, 0), place='center')\n",
        "  text_img = prompt_img.image\n",
        "  text_img = text_img.crop((0, 0, grid_w, text_h + 2*vertical_padding))\n",
        "  full_img = Image.new('RGB', (grid_w, text_img.size[1] + grid_h))\n",
        "  full_img.paste(text_img, (0, 0))\n",
        "  full_img.paste(grid_img, (0, text_img.size[1]))\n",
        "  return full_img\n",
        "\n",
        "def save_images(output_dir, np_images):\n",
        "  os.makedirs(output_dir, exist_ok=True)\n",
        "  for i, np_img in enumerate(np_images):\n",
        "    image = Image.fromarray(np.uint8(np_img * 255))\n",
        "    output_path = os.path.join(output_dir, f'{i}.png')\n",
        "    image.save(output_path)\n",
        "\n",
        "def upscale_dir(input_dir):\n",
        "  !python SwinIR/main_test_swinir.py --task real_sr --model_path experiments/pretrained_models/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pth --folder_lq {input_dir} --scale 4 --large_model\n",
        "  results_dir = \"./results/swinir_real_sr_x4_large\"\n",
        "  output_files = sorted([file for file in os.listdir(results_dir) if '.png' in file], key=lambda e: int(e.split('_')[0]))\n",
        "  upscale_dims = (256, 256, 3)\n",
        "  images = [None] * len(output_files)\n",
        "  for i, filename in enumerate(output_files):\n",
        "    pil_img = Image.open(os.path.join(results_dir, filename))\n",
        "    images[i] = pil_img\n",
        "#   shutil.rmtree(results_dir)\n",
        "  !rm {results_dir}\n",
        "  return images"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a70fTZfEAlx6"
      },
      "source": [
        "### Start Inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "SV1GRZxpAlx6",
        "scrolled": false
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "from IPython.display import display, clear_output\n",
        "from ipywidgets import interact\n",
        "import torchvision.transforms as T\n",
        "try:\n",
        "  from google.colab import files\n",
        "  can_download = True\n",
        "except ImportError:\n",
        "  can_download = False\n",
        "\n",
        "text_input = widgets.Textarea(\n",
        "    placeholder='Prompts separated by new lines...',\n",
        "    value=str(current_state[\"text_input\"]),\n",
        "    disabled=False,\n",
        "    continuous_update=True,\n",
        "    layout={'width': 'auto'},\n",
        "    rows=10\n",
        ")\n",
        "\n",
        "file_input = widgets.FileUpload(\n",
        "#     description=\"Image Variation Upload\",\n",
        "    multiple=False\n",
        ")\n",
        "file_input.observe(lambda: print(\"Uploaded\"), names=\"_test\")\n",
        "file_reset = widgets.Button(\n",
        "    description=\"Clear Uploads\"\n",
        ")\n",
        "file_box = widgets.HBox([file_input, file_reset])\n",
        "\n",
        "def reset_file_input(b):\n",
        "    global file_input\n",
        "    new_file_input = widgets.FileUpload(\n",
        "        multiple=False\n",
        "    )\n",
        "    file_input = new_file_input\n",
        "    file_box.children = (new_file_input, file_box.children[1])\n",
        "    render_layout()\n",
        "file_reset.on_click(reset_file_input)\n",
        "\n",
        "textbox_box = widgets.VBox([text_input, file_box], layout={'border': '2px solid grey'})\n",
        "\n",
        "prior_label = widgets.HTML(value=\"<b>Prior Options:</b> Set how many sample to take from the prior and what conditioning scale to use.\")\n",
        "text_repeat = widgets.IntSlider(\n",
        "    value=int(current_state[\"text_repeat\"]),\n",
        "    min=1,\n",
        "    max=10,\n",
        "    step=1,\n",
        "    description='Text Repeat',\n",
        "    disabled=False,\n",
        "    continuous_update=True,\n",
        "    orientation='horizontal',\n",
        "    readout=True,\n",
        "    readout_format='d',\n",
        "    style={'description_width': 'initial'},\n",
        "    layout=widgets.Layout(width=\"100%\")\n",
        ")\n",
        "\n",
        "prior_conditioning = widgets.FloatSlider(\n",
        "    value=float(current_state[\"prior_conditioning\"]),\n",
        "    min=0.0,\n",
        "    max=10.0,\n",
        "    step=0.1,\n",
        "    description='Prior Cond Scale',\n",
        "    disabled=False,\n",
        "    continuous_update=True,\n",
        "    orientation='horizontal',\n",
        "    readout=True,\n",
        "    readout_format='.1f',\n",
        "    style={'description_width': 'initial'},\n",
        "    layout=widgets.Layout(width=\"100%\")\n",
        ")\n",
        "prior_options_box = widgets.VBox([prior_label, text_repeat, prior_conditioning], layout=widgets.Layout(border=\"2px solid grey\", padding=\"5px 10px\", flex=\"1 0 auto\"))\n",
        "\n",
        "decoder_label = widgets.HTML(value=\"<b>Decoder Options:</b> Set how many sample to take from the decoder and what conditioning scale to use.\")\n",
        "img_repeat = widgets.IntSlider(\n",
        "    value=int(current_state[\"img_repeat\"]),\n",
        "    min=1,\n",
        "    max=10,\n",
        "    step=1,\n",
        "    description='Img Repeat',\n",
        "    disabled=False,\n",
        "    continuous_update=True,\n",
        "    orientation='horizontal',\n",
        "    readout=True,\n",
        "    readout_format='d',\n",
        "    style={'description_width': 'initial'},\n",
        "    layout=widgets.Layout(width=\"100%\")\n",
        ")\n",
        "\n",
        "decoder_conditioning = widgets.FloatSlider(\n",
        "    value=float(current_state[\"decoder_conditioning\"]),\n",
        "    min=0.0,\n",
        "    max=10.0,\n",
        "    step=0.1,\n",
        "    description='Decoder Cond Scale',\n",
        "    disabled=False,\n",
        "    continuous_update=True,\n",
        "    orientation='horizontal',\n",
        "    readout=True,\n",
        "    readout_format='.1f',\n",
        "    style={'description_width': 'initial'},\n",
        "    layout=widgets.Layout(width=\"100%\")\n",
        ")\n",
        "decoder_options_box = widgets.VBox([decoder_label, img_repeat, decoder_conditioning], layout=widgets.Layout(border=\"2px solid grey\", padding=\"5px 10px\", flex=\"1 0 auto\"))\n",
        "main_options_box = widgets.HBox([prior_options_box, decoder_options_box], layout=widgets.Layout(width=\"100%\"))\n",
        "\n",
        "include_prompt_checkbox = widgets.Checkbox(\n",
        "    value=bool(current_state[\"include_prompt_checkbox\"]),\n",
        "    description='Show prompt in output image',\n",
        "    disabled=False,\n",
        "    indent=False\n",
        ")\n",
        "\n",
        "upsample_checkbox = widgets.Checkbox(\n",
        "    value=bool(current_state['upsample_checkbox']),\n",
        "    description = 'Upsample to 256x256 with SwinIR',\n",
        "    disabled=False,\n",
        "    indent=False\n",
        ")\n",
        "\n",
        "meta_options_box = widgets.VBox([include_prompt_checkbox, upsample_checkbox])\n",
        "\n",
        "button = widgets.Button(\n",
        "    description=\"Start\"\n",
        ")\n",
        "final_options_box = widgets.HBox([meta_options_box, button], layout=widgets.Layout(justify_content=\"space-between\", border=\"2px solid grey\", padding=\"10px 30px\"))\n",
        "\n",
        "main_layout = widgets.VBox([textbox_box, main_options_box, final_options_box])\n",
        "\n",
        "def get_prompts():\n",
        "    import json\n",
        "    from itertools import zip_longest\n",
        "    import io\n",
        "    text = text_input.value\n",
        "    try:\n",
        "        prompts_array = json.loads(text)\n",
        "        assert isinstance(prompts_array, list)\n",
        "        text_prompts = prompts_array\n",
        "    except Exception as e:\n",
        "        pass\n",
        "\n",
        "    try:\n",
        "        text_prompts = list(filter(lambda v: len(v) > 0, text.split(\"\\n\")))\n",
        "    except Exception as e:\n",
        "        print(\"Failed to read as text with newlines\", e)\n",
        "        return []\n",
        "\n",
        "    files = file_input.value\n",
        "    file = None\n",
        "    if len(files) > 0:\n",
        "        file_name, file_info = list(files.items())[0]\n",
        "        image_pil = Image.open(io.BytesIO(file_info['content'])).convert('RGB')\n",
        "        transforms = T.Compose([\n",
        "            T.CenterCrop(min(image_pil.size)),\n",
        "            T.Resize(clip.image_size)\n",
        "        ])\n",
        "        image_pil = transforms(image_pil)\n",
        "        file = (file_name, image_pil)\n",
        "\n",
        "    return (text_prompts, file)\n",
        "\n",
        "def f(text_input, text_repeat, prior_conditioning, img_repeat, decoder_conditioning, include_prompt_checkbox, upsample_checkbox, image):\n",
        "  text_prompts, image_prompt = get_prompts()\n",
        "  total_images = len(text_prompts) * text_repeat * img_repeat\n",
        "\n",
        "  global current_state\n",
        "  current_state = {\n",
        "      **current_state,\n",
        "      \"text_input\": text_input,\n",
        "      \"text_repeat\": text_repeat,\n",
        "      \"prior_conditioning\": prior_conditioning,\n",
        "      \"img_repeat\": img_repeat,\n",
        "      \"decoder_conditioning\": decoder_conditioning,\n",
        "      \"include_prompt_checkbox\": include_prompt_checkbox,\n",
        "      \"upsample_checkbox\": upsample_checkbox,\n",
        "  }\n",
        "\n",
        "  def get_prompt_text(index):\n",
        "    text_prompt = text_prompts[index]\n",
        "    return f\"Prompt {index}: \\\"{text_prompt}\\\"\"\n",
        "\n",
        "  if image_prompt is not None:\n",
        "    print(\"Taking variation of image\")\n",
        "    display(image_prompt[1])\n",
        "  output_strings = []\n",
        "  output_strings.append(f\"Using model: {current_state['decoder']['name']}\")\n",
        "  output_strings.append(f\"Total output images: {total_images}\")\n",
        "  output_strings.append(\"\")\n",
        "  output_strings.extend([get_prompt_text(index) for index in range(len(text_prompts))])\n",
        "  output_strings.append(\"\")\n",
        "  output_strings.append(\"Including prompt text in output image\" if include_prompt_checkbox else \"Not including prompt text in output image\")\n",
        "  output_strings.append(f\"Prior Conditioning Scale: {prior_conditioning}\")\n",
        "  output_strings.append(f\"Decoder Conditioning Scale: {decoder_conditioning}\")\n",
        "  print('\\n'.join(output_strings))\n",
        "  save_state()\n",
        "\n",
        "def render_layout():\n",
        "    clear_output()\n",
        "    out = widgets.interactive_output(f, {'text_input': text_input, 'text_repeat': text_repeat, 'prior_conditioning': prior_conditioning, 'img_repeat': img_repeat, 'decoder_conditioning': decoder_conditioning, 'include_prompt_checkbox': include_prompt_checkbox, 'upsample_checkbox': upsample_checkbox, 'image': file_input })\n",
        "    display(main_layout, out)\n",
        "\n",
        "def get_image_embeddings(prompt_tokens, prompt_image, text_rep: int, prior_cond_scale: float):\n",
        "    if prompt_image is None:\n",
        "        print(\"Computing embedings using prior\")\n",
        "        with torch.no_grad():\n",
        "            image_embed = diffusion_prior.sample(prompt_tokens, cond_scale = prior_cond_scale).cpu().numpy()\n",
        "    else:\n",
        "        print(\"Computing embeddings from example image\")\n",
        "        image_tensor = T.ToTensor()(prompt_image[1]).unsqueeze_(0).to(device)\n",
        "        unbatched_image_embed, _ = clip.embed_image(image_tensor)\n",
        "        image_embed = torch.zeros(len(prompt_tokens), unbatched_image_embed.shape[-1])\n",
        "        for i in range(len(prompt_tokens)):\n",
        "            image_embed[i] = unbatched_image_embed\n",
        "        image_embed = image_embed.cpu().numpy()\n",
        "    return image_embed\n",
        "\n",
        "def on_start(_, recall_embeddings=False, recall_images=False):\n",
        "  if os.path.exists(\"./output\"):\n",
        "    shutil.rmtree(\"./output\")\n",
        "  render_layout()\n",
        "  prompts, prompt_image = get_prompts()\n",
        "  prior_cond_scale = prior_conditioning.value\n",
        "  decoder_cond_scale = decoder_conditioning.value\n",
        "  text_rep = text_repeat.value\n",
        "  img_rep = img_repeat.value\n",
        "  include_prompt = include_prompt_checkbox.value\n",
        "  upsample = upsample_checkbox.value\n",
        "\n",
        "  prior_text_input = []\n",
        "  for prompt in prompts:\n",
        "    for _ in range(text_rep):\n",
        "      prior_text_input.append(prompt)\n",
        "\n",
        "  tokens = tokenizer.tokenize(prior_text_input).to(device)\n",
        "  if recall_embeddings:\n",
        "    print(\"Loading embeddings\")\n",
        "    image_embed = np.load('img_emb_prior.npy')\n",
        "  else:\n",
        "    image_embed = get_image_embeddings(tokens, prompt_image, text_rep, prior_cond_scale)\n",
        "    np.save('img_emb_prior.npy', image_embed)\n",
        "\n",
        "  embeddings = np.repeat(image_embed, img_rep, axis=0)\n",
        "  embeddings = torch.from_numpy(embeddings).float().to(device)\n",
        "  if recall_images:\n",
        "    print(\"Loading images\")\n",
        "    images = np.load('images_decoder.npy')\n",
        "  else:\n",
        "    print(\"Running decoder\")\n",
        "    with torch.no_grad():\n",
        "      if decoder_text_conditioned:\n",
        "        print(\"Generating clip embeddings\")\n",
        "        _, text_encoding, text_mask = clip.embed_text(tokens)\n",
        "        images = decoder.sample(embeddings, text_encodings = text_encoding, text_mask = text_mask, cond_scale = decoder_cond_scale)\n",
        "      else:\n",
        "        print(\"Not generating clip embeddings\")\n",
        "        images = decoder.sample(embeddings, text = None, cond_scale = decoder_cond_scale)\n",
        "    images = images.cpu().permute(0, 2, 3, 1).numpy()\n",
        "    np.save('images_decoder.npy', images)\n",
        "\n",
        "  if upsample:\n",
        "    save_images('output/images', images)\n",
        "    images = upscale_dir('output/images')\n",
        "\n",
        "  img_map = map_images(images, text_rep, img_rep, prompts, upscale=4)\n",
        "  for index, (prompt, imgs) in enumerate(img_map.items()):\n",
        "    img = format_image_grid(imgs)\n",
        "    if include_prompt:\n",
        "      img = format_prompt_image(prompt, img, font_size=20, horizontal_padding=5, vertical_padding=5)\n",
        "    display(img)\n",
        "    if can_download:\n",
        "      download_button = widgets.Button(\n",
        "        description=\"Download\"\n",
        "      )\n",
        "      display(download_button)\n",
        "      download_button.on_click(lambda b: files.download(f\"./output/example_{index}.png\"))\n",
        "    os.makedirs(\"./output\", exist_ok=True)\n",
        "    img.save(f\"./output/example_{index}.png\")\n",
        "\n",
        "button.on_click(on_start)\n",
        "render_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1worYBzsAlx7"
      },
      "source": [
        "### Rerank Test"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "pG8ErLZ6Alx7"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "\"\"\"\n",
        "Here we take the prompt, generate n number of embeddings and rerank them by cosine similarity to the text embedding,\n",
        "then take a linspace of N and sample the decoder with those embeddings to see the variation in the performance of the prior\n",
        "\"\"\"\n",
        "def rerank_and_sample(image_embeddings, text_embedding, samples=None, strategy=\"top\"):\n",
        "  if samples is None:\n",
        "    samples = len(image_embeddings)\n",
        "\n",
        "  def similarity(image_embedding, text_embedding):\n",
        "    image_embedding = image_embedding / np.linalg.norm(image_embedding)\n",
        "    text_embedding = text_embedding / np.linalg.norm(text_embedding)\n",
        "    return np.inner(image_embedding, text_embedding)\n",
        "\n",
        "  reranked = sorted(list(image_embeddings), key=lambda img_emb: similarity(img_emb, text_embedding))\n",
        "  if strategy == \"top\":\n",
        "    sampled_embeddings = np.array(reranked[-samples:])\n",
        "  elif strategy == \"even\":\n",
        "    sample_indices = np.linspace(0, len(reranked) - 1, num=samples, dtype=int)\n",
        "    sampled_embeddings = np.array([reranked[i] for i in sample_indices])\n",
        "  rankings = [similarity(emb, text_embedding) for emb in sampled_embeddings]\n",
        "  print(rankings, rankings[0], rankings[-1])\n",
        "  return sampled_embeddings\n",
        "\n",
        "def rerank_test(num_samples, num_trials):\n",
        "  prompts = get_prompts()\n",
        "  prior_cond_scale = prior_conditioning.value\n",
        "  decoder_cond_scale = decoder_conditioning.value\n",
        "  text_rep = text_repeat.value\n",
        "  img_rep = img_repeat.value\n",
        "  include_prompt = include_prompt_checkbox.value\n",
        "\n",
        "  prompt_tokens = tokenizer.tokenize(prompts).to(device)\n",
        "  prompt_embeddings, prompt_encoding, prompt_mask = clip.embed_text(prompt_tokens)\n",
        "  prompt_embeddings = prompt_embeddings.cpu().numpy()\n",
        "\n",
        "  prior_prompts = []\n",
        "  for prompt in prompts:\n",
        "    for _ in range(num_trials):\n",
        "      prior_prompts.append(prompt)\n",
        "\n",
        "  sample_prompts = []\n",
        "  for prompt in prompts:\n",
        "    for _ in range(num_samples):\n",
        "      sample_prompts.append(prompt)\n",
        "\n",
        "  tokens = tokenizer.tokenize(prior_prompts).to(device)\n",
        "  with torch.no_grad():\n",
        "      image_embed = diffusion_prior.sample(tokens, cond_scale = prior_cond_scale).cpu().numpy()\n",
        "  print(f\"Generated {len(image_embed)} image embeddings\")\n",
        "\n",
        "  image_embed = np.split(image_embed, len(prompts))\n",
        "  reranked_embeddings = []\n",
        "  for i, embedding_set in enumerate(image_embed):\n",
        "    reranked_embeddings.append(rerank_and_sample(embedding_set, prompt_embeddings[i], samples=num_samples))\n",
        "\n",
        "  sampled_embedding_array = np.concatenate(reranked_embeddings)\n",
        "  print(f\"After reranking there are {len(sampled_embedding_array)} image embeddings\")\n",
        "  sampled_embedding_tensor = torch.from_numpy(sampled_embedding_array).to(device)\n",
        "\n",
        "  sample_tokens = tokenizer.tokenize(sample_prompts).to(device)\n",
        "  with torch.no_grad():\n",
        "    if decoder_text_conditioned:\n",
        "      print(\"Generating clip embeddings\")\n",
        "      _, text_encoding, text_mask = clip.embed_text(sample_tokens)\n",
        "      images = decoder.sample(sampled_embedding_tensor, text_encodings = text_encoding, text_mask = text_mask, cond_scale = decoder_cond_scale)\n",
        "    else:\n",
        "      print(\"Not generating clip embeddings\")\n",
        "      images = decoder.sample(sampled_embedding_tensor, text = None, cond_scale = decoder_cond_scale)\n",
        "  np_images = images.cpu().permute(0, 2, 3, 1)\n",
        "  img_map = map_images(np_images, num_samples, 1, prompts, upscale=4)\n",
        "  for index, (prompt, imgs) in enumerate(img_map.items()):\n",
        "    img = format_image_grid(imgs)\n",
        "    if include_prompt:\n",
        "      img = format_prompt_image(prompt, img, font_size=20, horizontal_padding=5, vertical_padding=5)\n",
        "    display(img)\n",
        "    if can_download:\n",
        "      download_button = widgets.Button(\n",
        "        description=\"Download\"\n",
        "      )\n",
        "      display(download_button)\n",
        "      download_button.on_click(lambda b: files.download(f\"./output/example_{index}.png\"))\n",
        "    os.makedirs(\"./reranked_output\", exist_ok=True)\n",
        "    img.save(f\"./reranked_output/example_{index}.png\")\n",
        "\n",
        "rerank_test(5, 50)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "DALLE2-Inference-Demo.ipynb",
      "provenance": []
    },
    "gpuClass": "standard",
    "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.5"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}