remove .ipynb

demo/.ipynb_checkpoints/tpu-demo-checkpoint.ipynb

@@ -1,391 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "6eb74941-bb4d-4d7e-97f1-d5a3a07672bf",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# !pip install flax transformers\n",
-    "# !git clone https://github.com/patil-suraj/vqgan-jax.git"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 305,
-   "id": "41db7534-f589-4b63-9165-9c9799e1b06e",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "/home/surajpatil/vqgan-jax\n"
-     ]
-    },
-    {
-     "data": {
-      "text/plain": [
-       "[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),\n",
-       " TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1),\n",
-       " TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0),\n",
-       " TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1),\n",
-       " TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0),\n",
-       " TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1),\n",
-       " TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0),\n",
-       " TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]"
-      ]
-     },
-     "execution_count": 305,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "%cd ~/vqgan-jax\n",
-    "\n",
-    "import random\n",
-    "\n",
-    "\n",
-    "import jax\n",
-    "import flax.linen as nn\n",
-    "from flax.training.common_utils import shard\n",
-    "from flax.jax_utils import replicate, unreplicate\n",
-    "\n",
-    "from transformers.models.bart.modeling_flax_bart import *\n",
-    "from transformers import BartTokenizer, FlaxBartForConditionalGeneration\n",
-    "\n",
-    "import io\n",
-    "\n",
-    "import requests\n",
-    "from PIL import Image\n",
-    "import numpy as np\n",
-    "import matplotlib.pyplot as plt\n",
-    "\n",
-    "import torch\n",
-    "import torchvision.transforms as T\n",
-    "import torchvision.transforms.functional as TF\n",
-    "from torchvision.transforms import InterpolationMode\n",
-    "\n",
-    "\n",
-    "from modeling_flax_vqgan import VQModel\n",
-    "\n",
-    "jax.devices()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "id": "b6a3462a-9004-4121-b365-3ae3aaf94dd2",
-   "metadata": {},
-   "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": 3,
-   "id": "bbef1afb-0b36-44a5-83f7-643d7e2c0e30",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "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": null,
-   "id": "879320b7-eaa0-4dc9-bbf2-c81efc53301d",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import wandb\n",
-    "run = wandb.init()\n",
-    "artifact = run.use_artifact('wandb/hf-flax-dalle-mini/model-3h3x3565:v7', type='bart_model')\n",
-    "artifact_dir = artifact.download()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 164,
-   "id": "e8bcff33-e95b-4c01-b162-ee857a55c3e6",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/surajpatil/transformers/src/transformers/models/bart/configuration_bart.py:177: 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"
-     ]
-    },
-    {
-     "data": {
-      "text/plain": [
-       "(1, 16385)"
-      ]
-     },
-     "execution_count": 164,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "# create our model and initialize it randomly\n",
-    "tokenizer = BartTokenizer.from_pretrained(BASE_MODEL)\n",
-    "model = CustomFlaxBartForConditionalGeneration.from_pretrained(artifact_dir)\n",
-    "model.config.force_bos_token_to_be_generated = False\n",
-    "model.config.forced_bos_token_id = None\n",
-    "model.config.forced_eos_token_id = None\n",
-    "\n",
-    "# we verify that the shape has not been modified\n",
-    "model.params['final_logits_bias'].shape"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "id": "8d5e0f14-2502-470e-9553-daee6748601f",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9b979a72ab9e449387a89bf9b3012af5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=433.0, style=ProgressStyle(description_…"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "01730e0e9d02428ca9dad680f9fdda42",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=304307206.0, style=ProgressStyle(descri…"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\n",
-      "Working with z of shape (1, 256, 16, 16) = 65536 dimensions.\n"
-     ]
-    }
-   ],
-   "source": [
-    "vqgan = VQModel.from_pretrained(\"flax-community/vqgan_f16_16384\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 295,
-   "id": "6cca395a-93c2-49bc-a3be-98287e4403d4",
-   "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\n",
-    "\n",
-    "def generate(input, rng, params):\n",
-    "  return model.generate(\n",
-    "      **input,\n",
-    "      max_length=257,\n",
-    "      num_beams=1,\n",
-    "      do_sample=True,\n",
-    "      prng_key=rng,\n",
-    "      eos_token_id=50000,\n",
-    "      pad_token_id=50000,\n",
-    "      params=params\n",
-    "  )\n",
-    "\n",
-    "def get_images(indices, params):\n",
-    "    return vqgan.decode_code(indices, params=params)\n",
-    "\n",
-    "\n",
-    "def plot_images(images):\n",
-    "    fig = plt.figure(figsize=(40, 20))\n",
-    "    columns = 4\n",
-    "    rows = 2\n",
-    "    plt.subplots_adjust(hspace=0, wspace=0)\n",
-    "\n",
-    "    for i in range(1, columns*rows +1):\n",
-    "        fig.add_subplot(rows, columns, i)\n",
-    "        plt.imshow(images[i-1])\n",
-    "    plt.gca().axes.get_yaxis().set_visible(False)\n",
-    "    plt.show()\n",
-    "    \n",
-    "def stack_reconstructions(images):\n",
-    "    w, h = images[0].size[0], images[0].size[1]\n",
-    "    img = Image.new(\"RGB\", (len(images)*w, h))\n",
-    "    for i, img_ in enumerate(images):\n",
-    "        img.paste(img_, (i*w,0))\n",
-    "    return img"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 166,
-   "id": "b1bec3d2-ef17-4feb-aa0d-b51ed2fdcd3e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "p_generate = jax.pmap(generate, \"batch\")\n",
-    "p_get_images = jax.pmap(get_images, \"batch\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "a539823a-a775-4d92-96a5-dc8b1eef69c5",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "bart_params = replicate(model.params)\n",
-    "vqgan_params = replicate(vqgan.params)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 328,
-   "id": "e8b268d8-6992-422a-8373-95651474ae70",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "prompts = [\n",
-    "    \"man in blue jacket walking on pathway in between trees during daytime\",\n",
-    "    'white snow covered mountain under blue sky during daytime',\n",
-    "    'white snow covered mountain under blue sky during night',\n",
-    "    \"orange tabby cat on persons hand\",\n",
-    "    \"aerial view of beach during daytime\",\n",
-    "    \"chess pieces on chess board\",\n",
-    "    \"laptop on brown wooden table\",\n",
-    "    \"white bus on road near high rise buildings\",\n",
-    "]\n",
-    "\n",
-    "\n",
-    "prompt = [prompts[-1]] * 8\n",
-    "inputs = tokenizer(prompt, return_tensors='jax', padding=\"max_length\", truncation=True, max_length=128).data\n",
-    "inputs = shard(inputs)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "68638cfa-9a4d-4e6a-8630-91aefb627bbd",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "%%time\n",
-    "for i in range(8):\n",
-    "    key = random.randint(0, 1e7)\n",
-    "    rng = jax.random.PRNGKey(key)\n",
-    "    rngs = jax.random.split(rng, jax.local_device_count())\n",
-    "    indices = p_generate(inputs, rngs, bart_params).sequences\n",
-    "    indices = indices[:, :, 1:]\n",
-    "\n",
-    "    images = p_get_images(indices, vqgan_params)\n",
-    "    images = np.squeeze(np.asarray(images), 1)\n",
-    "    imges = [custom_to_pil(image) for image in images]\n",
-    "\n",
-    "    plt.figure(figsize=(40, 20))\n",
-    "    plt.imshow(stack_reconstructions(imges))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "681af54e-da10-4b8e-80d0-ebcbdf23f376",
-   "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.8.10"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}