feat: improve inference demo

tools/inference/inference_pipeline.ipynb

@@ -41,10 +41,10 @@
    "outputs": [],
    "source": [
     "# Install required libraries\n",
-    "!pip install -q transformers\n",
-    "!pip install -q git+https://github.com/patil-suraj/vqgan-jax.git\n",
-    "!pip install -q git+https://github.com/borisdayma/dalle-mini.git\n",
-    "!pip install -q wandb"
+    "#!pip install -q transformers\n",
+    "#!pip install -q git+https://github.com/patil-suraj/vqgan-jax.git\n",
+    "#!pip install -q git+https://github.com/borisdayma/dalle-mini.git\n",
+    "#!pip install -q wandb"
    ]
   },
   {
@@ -70,8 +70,8 @@
     "# Model references\n",
     "\n",
     "# dalle-mini\n",
-    "DALLE_MODEL = \"dalle-mini/dalle-mini/model-3bqwu04f:latest\"  # can be wandb artifact or 🤗 Hub or local folder\n",
-    "DALLE_COMMIT_ID = None  # used only with 🤗 hub\n",
+    "DALLE_MODEL = \"dalle-mini/dalle-mini/model-mehdx7dg:latest\"  # can be wandb artifact or 🤗 Hub or local folder\n",
+    "DALLE_COMMIT_ID = None\n",
     "\n",
     "# VQGAN model\n",
     "VQGAN_REPO = \"dalle-mini/vqgan_imagenet_f16_16384\"\n",
@@ -91,13 +91,20 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "\n",
+    "# check how many devices are available\n",
+    "jax.local_device_count()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "# type used for computation - use bfloat16 on TPU's\n",
     "dtype = jnp.bfloat16 if jax.local_device_count() == 8 else jnp.float32\n",
     "\n",
-    "# TODO:\n",
-    "# - we currently have an issue with model.generate() in bfloat16\n",
-    "# - https://github.com/google/jax/pull/9089 should fix it\n",
-    "# - remove below line and test on TPU with next release of JAX\n",
+    "# TODO: fix issue with bfloat16\n",
     "dtype = jnp.float32"
    ]
   },
@@ -115,35 +122,18 @@
    "outputs": [],
    "source": [
     "# Load models & tokenizer\n",
-    "from dalle_mini.model import DalleBart\n",
+    "from dalle_mini.model import DalleBart, DalleBartTokenizer\n",
     "from vqgan_jax.modeling_flax_vqgan import VQModel\n",
-    "from transformers import AutoTokenizer, CLIPProcessor, FlaxCLIPModel\n",
+    "from transformers import CLIPProcessor, FlaxCLIPModel\n",
     "import wandb\n",
     "\n",
     "# Load dalle-mini\n",
-    "if \":\" in DALLE_MODEL:\n",
-    "    # wandb artifact\n",
-    "    artifact = wandb.Api().artifact(DALLE_MODEL)\n",
-    "    # we only download required files (no need for opt_state which is large)\n",
-    "    model_files = [\n",
-    "        \"config.json\",\n",
-    "        \"flax_model.msgpack\",\n",
-    "        \"merges.txt\",\n",
-    "        \"special_tokens_map.json\",\n",
-    "        \"tokenizer.json\",\n",
-    "        \"tokenizer_config.json\",\n",
-    "        \"vocab.json\",\n",
-    "    ]\n",
-    "    for f in model_files:\n",
-    "        artifact.get_path(f).download(\"model\")\n",
-    "    model = DalleBart.from_pretrained(\"model\", dtype=dtype, abstract_init=True)\n",
-    "    tokenizer = AutoTokenizer.from_pretrained(\"model\")\n",
-    "else:\n",
-    "    # local folder or 🤗 Hub\n",
-    "    model = DalleBart.from_pretrained(\n",
-    "        DALLE_MODEL, revision=DALLE_COMMIT_ID, dtype=dtype, abstract_init=True\n",
-    "    )\n",
-    "    tokenizer = AutoTokenizer.from_pretrained(DALLE_MODEL, revision=DALLE_COMMIT_ID)\n",
+    "model = DalleBart.from_pretrained(\n",
+    "    DALLE_MODEL, revision=DALLE_COMMIT_ID, dtype=dtype, abstract_init=True\n",
+    ")\n",
+    "tokenizer = DalleBartTokenizer.from_pretrained(\n",
+    "    DALLE_MODEL, revision=DALLE_COMMIT_ID\n",
+    ")\n",
     "\n",
     "# Load VQGAN\n",
     "vqgan = VQModel.from_pretrained(VQGAN_REPO, revision=VQGAN_COMMIT_ID)\n",
@@ -210,7 +200,8 @@
     "        prng_key=key,\n",
     "        params=params,\n",
     "        top_k=top_k,\n",
-    "        top_p=top_p\n",
+    "        top_p=top_p,\n",
+    "        max_length=257\n",
     "    )\n",
     "\n",
     "\n",
@@ -233,7 +224,7 @@
     "id": "HmVN6IBwapBA"
    },
    "source": [
-    "Keys are passed to the model on each device to generate unique inferences per device."
+    "Keys are passed to the model on each device to generate unique inference per device."
    ]
   },
   {
@@ -247,7 +238,7 @@
     "import random\n",
     "\n",
     "# create a random key\n",
-    "seed = random.randint(0, 2 ** 32 - 1)\n",
+    "seed = random.randint(0, 2**32 - 1)\n",
     "key = jax.random.PRNGKey(seed)"
    ]
   },
@@ -299,7 +290,7 @@
    },
    "outputs": [],
    "source": [
-    "prompt = \"a red T-shirt\""
+    "prompt = \"a waterfall under the sunset\""
    ]
   },
   {
@@ -316,27 +307,19 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "id": "iFVOyYboP0L-"
-   },
+   "metadata": {},
    "source": [
-    "We repeat the prompt on each device and tokenize it."
+    "We tokenize the prompt."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "id": "Rii_FJ7POw1y"
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "# repeat the prompt on each device\n",
-    "repeated_prompts = [processed_prompt] * jax.device_count()\n",
-    "\n",
-    "# tokenize\n",
     "tokenized_prompt = tokenizer(\n",
-    "    repeated_prompts,\n",
+    "    processed_prompt,\n",
     "    return_tensors=\"jax\",\n",
     "    padding=\"max_length\",\n",
     "    truncation=True,\n",
@@ -360,24 +343,18 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "id": "2wiDtG3_SH2u"
-   },
+   "metadata": {},
    "source": [
-    "Finally we distribute the tokenized prompt onto the devices."
+    "Finally we replicate it onto each device."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "id": "AImyrxHtR9TG"
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "from flax.training.common_utils import shard\n",
-    "\n",
-    "tokenized_prompt = shard(tokenized_prompt)"
+    "tokenized_prompt = replicate(tokenized_prompt)"
    ]
   },
   {
@@ -455,6 +432,8 @@
    },
    "outputs": [],
    "source": [
+    "from flax.training.common_utils import shard\n",
+    "\n",
     "# get clip scores\n",
     "clip_inputs = processor(\n",
     "    text=[prompt] * jax.device_count(),\n",