feat(inference_notebook): dalle-mini is installable

dev/inference/inference_pipeline.ipynb

@@ -6,7 +6,7 @@
       "name": "DALL·E mini - Inference pipeline.ipynb",
       "provenance": [],
       "collapsed_sections": [],
-      "authorship_tag": "ABX9TyOmaisFwTAYRR7mJmMVxzdA",
+      "authorship_tag": "ABX9TyMUjEt1XMLq+6/GhSnVFsSx",
       "include_colab_link": true
     },
     "kernelspec": {
@@ -22,6 +22,7 @@
         "49304912717a4995ae45d04a59d1f50e": {
           "model_module": "@jupyter-widgets/controls",
           "model_name": "HBoxModel",
+          "model_module_version": "1.5.0",
           "state": {
             "_view_name": "HBoxView",
             "_dom_classes": [],
@@ -42,6 +43,7 @@
         "5fd9f97986024e8db560a6737ade9e2e": {
           "model_module": "@jupyter-widgets/base",
           "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
           "state": {
             "_view_name": "LayoutView",
             "grid_template_rows": null,
@@ -93,6 +95,7 @@
         "caced43e3a4c493b98fb07cb41db045c": {
           "model_module": "@jupyter-widgets/controls",
           "model_name": "FloatProgressModel",
+          "model_module_version": "1.5.0",
           "state": {
             "_view_name": "ProgressView",
             "style": "IPY_MODEL_40c54b9454d346aabd197f2bcf189467",
@@ -116,6 +119,7 @@
         "0acc161f2e9948b68b3fc4e57ef333c9": {
           "model_module": "@jupyter-widgets/controls",
           "model_name": "HTMLModel",
+          "model_module_version": "1.5.0",
           "state": {
             "_view_name": "HTMLView",
             "style": "IPY_MODEL_7e7c488f57fc4acb8d261e2db81d61f0",
@@ -136,6 +140,7 @@
         "40c54b9454d346aabd197f2bcf189467": {
           "model_module": "@jupyter-widgets/controls",
           "model_name": "ProgressStyleModel",
+          "model_module_version": "1.5.0",
           "state": {
             "_view_name": "StyleView",
             "_model_name": "ProgressStyleModel",
@@ -151,6 +156,7 @@
         "8b25334a48244a14aa9ba0176887e655": {
           "model_module": "@jupyter-widgets/base",
           "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
           "state": {
             "_view_name": "LayoutView",
             "grid_template_rows": null,
@@ -202,6 +208,7 @@
         "7e7c488f57fc4acb8d261e2db81d61f0": {
           "model_module": "@jupyter-widgets/controls",
           "model_name": "DescriptionStyleModel",
+          "model_module_version": "1.5.0",
           "state": {
             "_view_name": "StyleView",
             "_model_name": "DescriptionStyleModel",
@@ -216,6 +223,7 @@
         "72c401062a5348b1a366dffb5a403568": {
           "model_module": "@jupyter-widgets/base",
           "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
           "state": {
             "_view_name": "LayoutView",
             "grid_template_rows": null,
@@ -267,6 +275,7 @@
         "022c124dfff348f285335732781b0887": {
           "model_module": "@jupyter-widgets/controls",
           "model_name": "HBoxModel",
+          "model_module_version": "1.5.0",
           "state": {
             "_view_name": "HBoxView",
             "_dom_classes": [],
@@ -287,6 +296,7 @@
         "a44e47e9d26c4deb81a5a11a9db92a9f": {
           "model_module": "@jupyter-widgets/base",
           "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
           "state": {
             "_view_name": "LayoutView",
             "grid_template_rows": null,
@@ -338,6 +348,7 @@
         "cd9c7016caae47c1b41fb2608c78b0bf": {
           "model_module": "@jupyter-widgets/controls",
           "model_name": "FloatProgressModel",
+          "model_module_version": "1.5.0",
           "state": {
             "_view_name": "ProgressView",
             "style": "IPY_MODEL_c22f207311cf4fb69bd9328eabfd4ebb",
@@ -361,6 +372,7 @@
         "36ff1d0fea4b47e2ae35aa6bfae6a5e8": {
           "model_module": "@jupyter-widgets/controls",
           "model_name": "HTMLModel",
+          "model_module_version": "1.5.0",
           "state": {
             "_view_name": "HTMLView",
             "style": "IPY_MODEL_037563a7eadd4ac5abb7249a2914d346",
@@ -381,6 +393,7 @@
         "c22f207311cf4fb69bd9328eabfd4ebb": {
           "model_module": "@jupyter-widgets/controls",
           "model_name": "ProgressStyleModel",
+          "model_module_version": "1.5.0",
           "state": {
             "_view_name": "StyleView",
             "_model_name": "ProgressStyleModel",
@@ -396,6 +409,7 @@
         "5a38c6d83a264bedbf7efe6e97eba953": {
           "model_module": "@jupyter-widgets/base",
           "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
           "state": {
             "_view_name": "LayoutView",
             "grid_template_rows": null,
@@ -447,6 +461,7 @@
         "037563a7eadd4ac5abb7249a2914d346": {
           "model_module": "@jupyter-widgets/controls",
           "model_name": "DescriptionStyleModel",
+          "model_module_version": "1.5.0",
           "state": {
             "_view_name": "StyleView",
             "_model_name": "DescriptionStyleModel",
@@ -461,6 +476,7 @@
         "3975e7ed0b704990b1fa05909a9bb9b6": {
           "model_module": "@jupyter-widgets/base",
           "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
           "state": {
             "_view_name": "LayoutView",
             "grid_template_rows": null,
@@ -512,6 +528,7 @@
         "f9f1fdc3819a4142b85304cd3c6358a2": {
           "model_module": "@jupyter-widgets/controls",
           "model_name": "HBoxModel",
+          "model_module_version": "1.5.0",
           "state": {
             "_view_name": "HBoxView",
             "_dom_classes": [],
@@ -532,6 +549,7 @@
         "ea9ed54e7c9d4ead8b3e1ff4cb27fa61": {
           "model_module": "@jupyter-widgets/base",
           "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
           "state": {
             "_view_name": "LayoutView",
             "grid_template_rows": null,
@@ -583,6 +601,7 @@
         "29d42e94b3b34c86a117b623da68faed": {
           "model_module": "@jupyter-widgets/controls",
           "model_name": "FloatProgressModel",
+          "model_module_version": "1.5.0",
           "state": {
             "_view_name": "ProgressView",
             "style": "IPY_MODEL_8ce4d20d004a4382afa0abdd3b1f7191",
@@ -606,6 +625,7 @@
         "8b73de7dbdfe40dbbb39fb593520b984": {
           "model_module": "@jupyter-widgets/controls",
           "model_name": "HTMLModel",
+          "model_module_version": "1.5.0",
           "state": {
             "_view_name": "HTMLView",
             "style": "IPY_MODEL_717ccef4df1f477abb51814650eb47da",
@@ -626,6 +646,7 @@
         "8ce4d20d004a4382afa0abdd3b1f7191": {
           "model_module": "@jupyter-widgets/controls",
           "model_name": "ProgressStyleModel",
+          "model_module_version": "1.5.0",
           "state": {
             "_view_name": "StyleView",
             "_model_name": "ProgressStyleModel",
@@ -641,6 +662,7 @@
         "efc4812245c8459c92e6436889b4f600": {
           "model_module": "@jupyter-widgets/base",
           "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
           "state": {
             "_view_name": "LayoutView",
             "grid_template_rows": null,
@@ -692,6 +714,7 @@
         "717ccef4df1f477abb51814650eb47da": {
           "model_module": "@jupyter-widgets/controls",
           "model_name": "DescriptionStyleModel",
+          "model_module_version": "1.5.0",
           "state": {
             "_view_name": "StyleView",
             "_model_name": "DescriptionStyleModel",
@@ -706,6 +729,7 @@
         "7dba58f0391c485a86e34e8039ec6189": {
           "model_module": "@jupyter-widgets/base",
           "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
           "state": {
             "_view_name": "LayoutView",
             "grid_template_rows": null,
@@ -804,8 +828,7 @@
       "source": [
         "!pip install -q transformers flax\n",
         "!pip install -q git+https://github.com/patil-suraj/vqgan-jax.git  # VQGAN model in JAX\n",
-        "!git clone https://github.com/borisdayma/dalle-mini  # Model files\n",
-        "%cd dalle-mini/"
+        "!pip install -q git+https://github.com/borisdayma/dalle-mini.git  # Model files"
       ],
       "execution_count": null,
       "outputs": []
@@ -833,7 +856,7 @@
         "import random\n",
         "from tqdm.notebook import tqdm, trange"
       ],
-      "execution_count": 2,
+      "execution_count": null,
       "outputs": []
     },
     {
@@ -846,7 +869,7 @@
         "DALLE_REPO = 'flax-community/dalle-mini'\n",
         "DALLE_COMMIT_ID = '4d34126d0df8bc4a692ae933e3b902a1fa8b6114'"
       ],
-      "execution_count": 3,
+      "execution_count": null,
       "outputs": []
     },
     {
@@ -871,7 +894,7 @@
         "# set a prompt\n",
         "prompt = 'picture of a waterfall under the sunset'"
       ],
-      "execution_count": 5,
+      "execution_count": null,
       "outputs": []
     },
     {
@@ -888,7 +911,7 @@
         "tokenized_prompt = tokenizer(prompt, return_tensors='jax', padding='max_length', truncation=True, max_length=128)\n",
         "tokenized_prompt"
       ],
-      "execution_count": 6,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "execute_result",
@@ -956,7 +979,7 @@
         "subkeys = jax.random.split(key, num=n_predictions)\n",
         "subkeys"
       ],
-      "execution_count": 7,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "execute_result",
@@ -1004,7 +1027,7 @@
         "encoded_images = [model.generate(**tokenized_prompt, do_sample=True, num_beams=1, prng_key=subkey) for subkey in tqdm(subkeys)]\n",
         "encoded_images[0]"
       ],
-      "execution_count": 8,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "display_data",
@@ -1099,7 +1122,7 @@
         "encoded_images = [img.sequences[..., 1:] for img in encoded_images]\n",
         "encoded_images[0]"
       ],
-      "execution_count": 9,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "execute_result",
@@ -1167,7 +1190,7 @@
       "source": [
         "encoded_images[0].shape"
       ],
-      "execution_count": 10,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "execute_result",
@@ -1204,7 +1227,7 @@
         "import numpy as np\n",
         "from PIL import Image"
       ],
-      "execution_count": 11,
+      "execution_count": null,
       "outputs": []
     },
     {
@@ -1217,7 +1240,7 @@
         "VQGAN_REPO = 'flax-community/vqgan_f16_16384'\n",
         "VQGAN_COMMIT_ID = '90cc46addd2dd8f5be21586a9a23e1b95aa506a9'"
       ],
-      "execution_count": 12,
+      "execution_count": null,
       "outputs": []
     },
     {
@@ -1233,7 +1256,7 @@
         "# set up VQGAN\n",
         "vqgan = VQModel.from_pretrained(VQGAN_REPO, revision=VQGAN_COMMIT_ID)"
       ],
-      "execution_count": 13,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "stream",
@@ -1269,7 +1292,7 @@
         "decoded_images = [vqgan.decode_code(encoded_image) for encoded_image in tqdm(encoded_images)]\n",
         "decoded_images[0]"
       ],
-      "execution_count": 14,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "display_data",
@@ -1373,7 +1396,7 @@
         "# normalize images\n",
         "clipped_images = [img.squeeze().clip(0., 1.) for img in decoded_images]"
       ],
-      "execution_count": 15,
+      "execution_count": null,
       "outputs": []
     },
     {
@@ -1385,7 +1408,7 @@
         "# convert to image\n",
         "images = [Image.fromarray(np.asarray(img * 255, dtype=np.uint8)) for img in clipped_images]"
       ],
-      "execution_count": 16,
+      "execution_count": null,
       "outputs": []
     },
     {
@@ -1402,7 +1425,7 @@
         "# display an image\n",
         "images[0]"
       ],
-      "execution_count": 17,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "execute_result",
@@ -1438,7 +1461,7 @@
       "source": [
         "from transformers import CLIPProcessor, FlaxCLIPModel"
       ],
-      "execution_count": 18,
+      "execution_count": null,
       "outputs": []
     },
     {
@@ -1474,7 +1497,7 @@
         "logits = clip(**inputs).logits_per_image\n",
         "scores = jax.nn.softmax(logits, axis=0).squeeze()  # normalize and sum all scores to 1"
       ],
-      "execution_count": 20,
+      "execution_count": null,
       "outputs": []
     },
     {
@@ -1495,7 +1518,7 @@
         "    display(images[idx])\n",
         "    print()"
       ],
-      "execution_count": 21,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "stream",
@@ -1690,7 +1713,7 @@
         "from flax.training.common_utils import shard\n",
         "from flax.jax_utils import replicate"
       ],
-      "execution_count": 22,
+      "execution_count": null,
       "outputs": []
     },
     {
@@ -1706,7 +1729,7 @@
         "# check we can access TPU's or GPU's\n",
         "jax.devices()"
       ],
-      "execution_count": 23,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "execute_result",
@@ -1744,7 +1767,7 @@
         "# one set of inputs per device\n",
         "prompt = ['picture of a waterfall under the sunset'] * jax.device_count()"
       ],
-      "execution_count": 25,
+      "execution_count": null,
       "outputs": []
     },
     {
@@ -1757,7 +1780,7 @@
         "tokenized_prompt = tokenizer(prompt, return_tensors='jax', padding='max_length', truncation=True, max_length=128).data\n",
         "tokenized_prompt = shard(tokenized_prompt)"
       ],
-      "execution_count": 26,
+      "execution_count": null,
       "outputs": []
     },
     {
@@ -1793,7 +1816,7 @@
         "def p_decode(indices, params):\n",
         "    return vqgan.decode_code(indices, params=params)"
       ],
-      "execution_count": 27,
+      "execution_count": null,
       "outputs": []
     },
     {
@@ -1834,7 +1857,7 @@
         "    for img in decoded_images:\n",
         "        images.append(Image.fromarray(np.asarray(img * 255, dtype=np.uint8)))"
       ],
-      "execution_count": 28,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "display_data",
@@ -1877,7 +1900,7 @@
         "    display(img)\n",
         "    print()"
       ],
-      "execution_count": 29,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "display_data",