{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ewer-Q-0w2xA"
   },
   "source": [
    "# Installation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "NpsF9ipLLl2s",
    "outputId": "10bf54aa-b89d-4e42-9777-bc97b00a5f32"
   },
   "outputs": [],
   "source": [
    "#!pip install git+https://github.com/huggingface/transformers/\n",
    "#!pip install git+https://github.com/google/flax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "M1wVkrpjU6zO"
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/tmabraham/vqgan-jax\n"
     ]
    }
   ],
   "source": [
    "%cd ../../vqgan-jax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t47CH1H_IOT8"
   },
   "source": [
    "# Custom BART Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "9jQnM6S2vCpn"
   },
   "outputs": [],
   "source": [
    "# TODO: set those args in a config file\n",
    "OUTPUT_VOCAB_SIZE = 16384 + 1  # encoded image token space + 1 for bos\n",
    "OUTPUT_LENGTH = 256 + 1  # number of encoded tokens + 1 for bos\n",
    "BOS_TOKEN_ID = 16384\n",
    "BASE_MODEL = 'facebook/bart-large'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "_eEaJVxAKpV5"
   },
   "outputs": [],
   "source": [
    "import jax\n",
    "import flax.linen as nn\n",
    "\n",
    "from transformers.models.bart.modeling_flax_bart import *\n",
    "from transformers import BartTokenizer, FlaxBartForConditionalGeneration\n",
    "\n",
    "class CustomFlaxBartModule(FlaxBartModule):\n",
    "    def setup(self):\n",
    "        # we keep shared to easily load pre-trained weights\n",
    "        self.shared = nn.Embed(\n",
    "            self.config.vocab_size,\n",
    "            self.config.d_model,\n",
    "            embedding_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),\n",
    "            dtype=self.dtype,\n",
    "        )\n",
    "        # a separate embedding is used for the decoder\n",
    "        self.decoder_embed = nn.Embed(\n",
    "            OUTPUT_VOCAB_SIZE,\n",
    "            self.config.d_model,\n",
    "            embedding_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),\n",
    "            dtype=self.dtype,\n",
    "        )\n",
    "        self.encoder = FlaxBartEncoder(self.config, dtype=self.dtype, embed_tokens=self.shared)\n",
    "\n",
    "        # the decoder has a different config\n",
    "        decoder_config = BartConfig(self.config.to_dict())\n",
    "        decoder_config.max_position_embeddings = OUTPUT_LENGTH\n",
    "        decoder_config.vocab_size = OUTPUT_VOCAB_SIZE\n",
    "        self.decoder = FlaxBartDecoder(decoder_config, dtype=self.dtype, embed_tokens=self.decoder_embed)\n",
    "\n",
    "class CustomFlaxBartForConditionalGenerationModule(FlaxBartForConditionalGenerationModule):\n",
    "    def setup(self):\n",
    "        self.model = CustomFlaxBartModule(config=self.config, dtype=self.dtype)\n",
    "        self.lm_head = nn.Dense(\n",
    "            OUTPUT_VOCAB_SIZE,\n",
    "            use_bias=False,\n",
    "            dtype=self.dtype,\n",
    "            kernel_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),\n",
    "        )\n",
    "        self.final_logits_bias = self.param(\"final_logits_bias\", self.bias_init, (1, OUTPUT_VOCAB_SIZE))\n",
    "\n",
    "class CustomFlaxBartForConditionalGeneration(FlaxBartForConditionalGeneration):\n",
    "    module_class = CustomFlaxBartForConditionalGenerationModule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mtmabraham\u001b[0m (use `wandb login --relogin` to force relogin)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                Tracking run with wandb version 0.10.33<br/>\n",
       "                Syncing run <strong style=\"color:#cdcd00\">serene-resonance-1</strong> to <a href=\"https://wandb.ai\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
       "                Project page: <a href=\"https://wandb.ai/tmabraham/vqgan-jax\" target=\"_blank\">https://wandb.ai/tmabraham/vqgan-jax</a><br/>\n",
       "                Run page: <a href=\"https://wandb.ai/tmabraham/vqgan-jax/runs/1cm35ims\" target=\"_blank\">https://wandb.ai/tmabraham/vqgan-jax/runs/1cm35ims</a><br/>\n",
       "                Run data is saved locally in <code>/home/tmabraham/vqgan-jax/wandb/run-20210715_030616-1cm35ims</code><br/><br/>\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Downloading large artifact model-1ef8yxby:v1, 1674.97MB. 2 files... Done. 0:0:0\n"
     ]
    }
   ],
   "source": [
    "import wandb\n",
    "run = wandb.init()\n",
    "artifact = run.use_artifact('wandb/hf-flax-dalle-mini/model-1ef8yxby:v1', type='bart_model')\n",
    "artifact_dir = artifact.download()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "_6-XKK40oEfP",
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tmabraham/dalle-mini/src/transformers/src/transformers/models/bart/configuration_bart.py:180: UserWarning: Please make sure the config includes `forced_bos_token_id=16384` in future versions.The config can simply be saved and uploaded again to be fixed.\n",
      "  warnings.warn(\n",
      "INFO:absl:Starting the local TPU driver.\n",
      "INFO:absl:Unable to initialize backend 'tpu_driver': Not found: Unable to find driver in registry given worker: local://\n",
      "INFO:absl:Unable to initialize backend 'gpu': Not found: Could not find registered platform with name: \"cuda\". Available platform names are: TPU Interpreter Host\n"
     ]
    }
   ],
   "source": [
    "# create our model and initialize it randomly\n",
    "model = CustomFlaxBartForConditionalGeneration.from_pretrained(artifact_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Jz032w73nHEf",
    "outputId": "994d8e85-bff7-480b-8b69-f69dedc15c49"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 16385)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# we verify that the shape has not been modified\n",
    "model.params['final_logits_bias'].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zLl24Ez5t7x1"
   },
   "source": [
    "## Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "XLLA2NK3uDQr"
   },
   "outputs": [],
   "source": [
    "tokenizer = BartTokenizer.from_pretrained(BASE_MODEL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "P32mJJSbrU1F"
   },
   "outputs": [],
   "source": [
    "input_ids_test = tokenizer.encode('I enjoy walking with my cute dog', return_tensors='jax')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[    0,   100,  2254,  3051,    19,   127, 11962,  2335,\n",
       "                  2]], dtype=int32)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_ids_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "C7cHbIHruELT"
   },
   "outputs": [],
   "source": [
    "greedy_output = model.generate(input_ids_test, max_length=257)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "jYugh9cOuwc9",
    "outputId": "19c3a2ee-e7bc-4f1f-9c86-06bd7337b537"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[16384, 16384, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
       "              10042]], dtype=int32)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "greedy_output[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# VGAN Jax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "\n",
    "import requests\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torchvision.transforms as T\n",
    "import torchvision.transforms.functional as TF\n",
    "from torchvision.transforms import InterpolationMode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from modeling_flax_vqgan import VQModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def custom_to_pil(x):\n",
    "    x = np.clip(x, 0., 1.)\n",
    "    x = (255*x).astype(np.uint8)\n",
    "    x = Image.fromarray(x)\n",
    "    if not x.mode == \"RGB\":\n",
    "        x = x.convert(\"RGB\")\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Jz032w73nHEf",
    "outputId": "994d8e85-bff7-480b-8b69-f69dedc15c49"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Working with z of shape (1, 256, 16, 16) = 65536 dimensions.\n"
     ]
    }
   ],
   "source": [
    "model = VQModel.from_pretrained(\"valhalla/vqgan-imagenet-f16-1024\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_images(indices, model):\n",
    "    indices =  indices[:, 1:]\n",
    "    model.decode_code(indices)\n",
    "    return indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Working with z of shape (1, 256, 16, 16) = 65536 dimensions.\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAEACAIAAAD9XIvPAAAAF0lEQVR4nGP4//8/EwMDwygexaN45GEA7ucE/J1FRrMAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=1x256 at 0x7FE6389B6280>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "custom_to_pil(np.asarray(get_images(greedy_output[0], model)[0]))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "TPU",
  "colab": {
   "collapsed_sections": [],
   "machine_shape": "hm",
   "name": "CustomBARTv4b-model-generate.ipynb",
   "provenance": []
  },
  "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}