Merge pull request #63 from borisdayma/chore-mv

README.md

@@ -22,10 +22,6 @@ You can create your own pictures with [the demo](https://huggingface.co/spaces/f
 
 Refer to [our report](https://wandb.ai/dalle-mini/dalle-mini/reports/DALL-E-mini--Vmlldzo4NjIxODA).
 
-## Where does the logo come from?
-
-The "armchair in the shape of an avocado" was used by OpenAI when releasing DALL·E to illustrate the model's capabilities. Having successful predictions on this prompt represents a big milestone to us.
-
 ## Development
 
 This section is for the adventurous people wanting to look into the code.
@@ -58,13 +54,19 @@ Use [patil-suraj/vqgan-jax](https://github.com/patil-suraj/vqgan-jax).
 
 ### Training of Seq2Seq
 
-Refer to `dev/seq2seq` folder.
+Refer to [`dev/seq2seq`](dev/seq2seq) folder.
 
 You can also adjust the [sweep configuration file](https://docs.wandb.ai/guides/sweeps) if you need to perform a hyperparameter search.
 
-### Inference
+### Inference Pipeline
+
+To generate sample predictions and understand the inference pipeline step by step, refer to [`dev/inference/inference_pipeline.ipynb`](dev/inference/inference_pipeline.ipynb).
 
-Refer to `dev/notebooks/demo`.
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/borisdayma/dalle-mini/blob/main/dev/inference/inference_pipeline.ipynb)
+
+## Where does the logo come from?
+
+The "armchair in the shape of an avocado" was used by OpenAI when releasing DALL·E to illustrate the model's capabilities. Having successful predictions on this prompt represents a big milestone to us.
 
 ## Authors
 

dev/inference/wandb-examples-from-backend.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+from PIL import Image, ImageDraw, ImageFont
+import wandb
+import os
+
+from dalle_mini.backend import ServiceError, get_images_from_backend
+from dalle_mini.helpers import captioned_strip
+
+os.environ["WANDB_SILENT"] = "true"
+os.environ["WANDB_CONSOLE"] = "off"
+
+def log_to_wandb(prompts):
+    try:
+        backend_url = os.environ["BACKEND_SERVER"]
+        for _ in range(1):        
+            for prompt in prompts:
+                print(f"Getting selections for: {prompt}")
+                # make a separate run per prompt
+                with wandb.init(
+                    entity='wandb',
+                    project='hf-flax-dalle-mini',
+                    job_type='predictions',# tags=['openai'],
+                    config={'prompt': prompt}
+                ):
+                    imgs = []
+                    selected = get_images_from_backend(prompt, backend_url)
+                    strip = captioned_strip(selected, prompt)
+                    imgs.append(wandb.Image(strip))
+                    wandb.log({"images": imgs})
+    except ServiceError as error:
+        print(f"Service unavailable, status: {error.status_code}")
+    except KeyError:
+        print("Error: BACKEND_SERVER unset")
+
+prompts = [
+    # "white snow covered mountain under blue sky during daytime",
+    # "aerial view of beach during daytime",
+    # "aerial view of beach at night",
+    # "a farmhouse surrounded by beautiful flowers",
+    # "an armchair in the shape of an avocado",
+    # "young woman riding her bike trough a forest",
+    # "a unicorn is passing by a rainbow in a field of flowers",
+    # "illustration of a baby shark swimming around corals",
+    # "painting of an oniric forest glade surrounded by tall trees",
+    # "sunset over green mountains",
+    # "a forest glade surrounded by tall trees in a sunny Spring morning",
+    # "fishing village under the moonlight in a serene sunset",
+    # "cartoon of a carrot with big eyes",
+    # "still life in the style of Kandinsky",
+    # "still life in the style of Picasso",
+    # "a graphite sketch of a gothic cathedral",
+    # "a graphite sketch of Elon Musk",
+    # "a watercolor pond with green leaves and yellow flowers",
+    # "a logo of a cute avocado armchair singing karaoke on stage in front of a crowd of strawberry shaped lamps",
+    # "happy celebration in a small village in Africa",
+    # "a logo of an armchair in the shape of an avocado"
+    # "Pele and Maradona in a hypothetical match",
+    # "Mohammed Ali and Mike Tyson in a hypothetical match",
+    # "a storefront that has the word 'openai' written on it",
+    # "a pentagonal green clock",
+    # "a collection of glasses is sitting on a table",
+    # "a small red block sitting on a large green block",
+    # "an extreme close-up view of a capybara sitting in a field",
+    # "a cross-section view of a walnut",
+    # "a professional high-quality emoji of a lovestruck cup of boba",
+    # "a photo of san francisco's golden gate bridge",
+    # "an illustration of a baby daikon radish in a tutu walking a dog",
+    # "a picture of the Eiffel tower on the Moon",
+    # "a colorful stairway to heaven",
+    "this is a detailed high-resolution scan of a human brain"
+]
+
+for _ in range(1):
+    log_to_wandb(prompts)

dev/{predictions → inference}/wandb-examples.py

@@ -4,16 +4,14 @@
 import random
 
 import jax
-import flax.linen as nn
 from flax.training.common_utils import shard
 from flax.jax_utils import replicate, unreplicate
 
 from transformers.models.bart.modeling_flax_bart import *
 from transformers import BartTokenizer, FlaxBartForConditionalGeneration
 
-import io
+import os
 
-import requests
 from PIL import Image
 import numpy as np
 import matplotlib.pyplot as plt
@@ -23,58 +21,24 @@ import torchvision.transforms as T
 import torchvision.transforms.functional as TF
 from torchvision.transforms import InterpolationMode
 
+from dalle_mini.model import CustomFlaxBartForConditionalGeneration
 from vqgan_jax.modeling_flax_vqgan import VQModel
 
-# TODO: set those args in a config file
-OUTPUT_VOCAB_SIZE = 16384 + 1  # encoded image token space + 1 for bos
-OUTPUT_LENGTH = 256 + 1  # number of encoded tokens + 1 for bos
-BOS_TOKEN_ID = 16384
-BASE_MODEL = 'facebook/bart-large-cnn'
-
-class CustomFlaxBartModule(FlaxBartModule):
-    def setup(self):
-        # we keep shared to easily load pre-trained weights
-        self.shared = nn.Embed(
-            self.config.vocab_size,
-            self.config.d_model,
-            embedding_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),
-            dtype=self.dtype,
-        )
-        # a separate embedding is used for the decoder
-        self.decoder_embed = nn.Embed(
-            OUTPUT_VOCAB_SIZE,
-            self.config.d_model,
-            embedding_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),
-            dtype=self.dtype,
-        )
-        self.encoder = FlaxBartEncoder(self.config, dtype=self.dtype, embed_tokens=self.shared)
-
-        # the decoder has a different config
-        decoder_config = BartConfig(self.config.to_dict())
-        decoder_config.max_position_embeddings = OUTPUT_LENGTH
-        decoder_config.vocab_size = OUTPUT_VOCAB_SIZE
-        self.decoder = FlaxBartDecoder(decoder_config, dtype=self.dtype, embed_tokens=self.decoder_embed)
-
-class CustomFlaxBartForConditionalGenerationModule(FlaxBartForConditionalGenerationModule):
-    def setup(self):
-        self.model = CustomFlaxBartModule(config=self.config, dtype=self.dtype)
-        self.lm_head = nn.Dense(
-            OUTPUT_VOCAB_SIZE,
-            use_bias=False,
-            dtype=self.dtype,
-            kernel_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),
-        )
-        self.final_logits_bias = self.param("final_logits_bias", self.bias_init, (1, OUTPUT_VOCAB_SIZE))
-
-class CustomFlaxBartForConditionalGeneration(FlaxBartForConditionalGeneration):
-    module_class = CustomFlaxBartForConditionalGenerationModule
-
+# ## CLIP Scoring
+from transformers import CLIPProcessor, FlaxCLIPModel
 
 import wandb
 import os
+
+from dalle_mini.helpers import captioned_strip
+
+
 os.environ["WANDB_SILENT"] = "true"
 os.environ["WANDB_CONSOLE"] = "off"
 
+# TODO: used for legacy support
+BASE_MODEL = 'facebook/bart-large-cnn'
+
 # set id to None so our latest images don't get overwritten
 id = None
 run = wandb.init(id=id,
@@ -87,8 +51,10 @@ artifact = run.use_artifact('wandb/hf-flax-dalle-mini/model-4oh3u7ca:latest', ty
 artifact_dir = artifact.download()
 
 # create our model
-tokenizer = BartTokenizer.from_pretrained(BASE_MODEL)
 model = CustomFlaxBartForConditionalGeneration.from_pretrained(artifact_dir)
+
+# TODO: legacy support (earlier models)
+tokenizer = BartTokenizer.from_pretrained(BASE_MODEL)
 model.config.force_bos_token_to_be_generated = False
 model.config.forced_bos_token_id = None
 model.config.forced_eos_token_id = None
@@ -143,9 +109,6 @@ p_get_images = jax.pmap(get_images, "batch")
 bart_params = replicate(model.params)
 vqgan_params = replicate(vqgan.params)
 
-# ## CLIP Scoring
-from transformers import CLIPProcessor, FlaxCLIPModel
-
 clip = FlaxCLIPModel.from_pretrained("openai/clip-vit-base-patch32")
 processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
@@ -170,16 +133,12 @@ def hallucinate(prompt, num_images=64):
 
 def clip_top_k(prompt, images, k=8):
     inputs = processor(text=prompt, images=images, return_tensors="np", padding=True)
+    # FIXME: image should be resized and normalized prior to being processed by CLIP
     outputs = clip(**inputs)
     logits = outputs.logits_per_text
     scores = np.array(logits[0]).argsort()[-k:][::-1]
     return [images[score] for score in scores]
 
-
-# ## Log to wandb
-
-from dalle_mini.helpers import captioned_strip
-
 def log_to_wandb(prompts):
     strips = []
     for prompt in prompts:

dev/notebooks/README.md

@@ -1,5 +0,0 @@
-# Notebooks
-
-These notebooks were used during development.
-
-TODO: This section requires some refactor and clean up.

dev/notebooks/demo/CustomBARTv4b_model-generate.ipynb

@@ -1,394 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "ewer-Q-0w2xA"
-   },
-   "source": [
-    "# Installation"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "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": null,
-   "metadata": {
-    "id": "M1wVkrpjU6zO"
-   },
-   "outputs": [],
-   "source": [
-    "%load_ext autoreload\n",
-    "%autoreload 2"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "t47CH1H_IOT8"
-   },
-   "source": [
-    "# Custom BART Model"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "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": null,
-   "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": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "S7CP9Td9m2ge",
-    "outputId": "5638ef68-9c40-46f7-90ba-a4d05b61360d"
-   },
-   "outputs": [],
-   "source": [
-    "# load pre-trained model for encoder weights\n",
-    "base_model = FlaxBartForConditionalGeneration.from_pretrained(BASE_MODEL)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "6lmynR-poceH"
-   },
-   "outputs": [],
-   "source": [
-    "# set up our new model config\n",
-    "config = BartConfig.from_pretrained(BASE_MODEL)\n",
-    "config.tie_word_embeddings = False\n",
-    "config.decoder_start_token_id = BOS_TOKEN_ID\n",
-    "config.bos_token_id = BOS_TOKEN_ID  # should not be used\n",
-    "config.pos_token_id = BOS_TOKEN_ID  # should not be used\n",
-    "#config.eos_token_id = None  # prevents generation from stopping until we reach max_length"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "_6-XKK40oEfP"
-   },
-   "outputs": [],
-   "source": [
-    "# create our model and initialize it randomly\n",
-    "model = CustomFlaxBartForConditionalGeneration(config)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "-r_hZestr-NR"
-   },
-   "outputs": [],
-   "source": [
-    "# use pretrained weights\n",
-    "model.params['model']['encoder'] = base_model.params['model']['encoder']\n",
-    "model.params['model']['shared'] = base_model.params['model']['shared']"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "5NEX8f62sVjx"
-   },
-   "outputs": [],
-   "source": [
-    "# no need for base_model anymore\n",
-    "del base_model"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "Jz032w73nHEf",
-    "outputId": "994d8e85-bff7-480b-8b69-f69dedc15c49"
-   },
-   "outputs": [],
-   "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": null,
-   "metadata": {
-    "id": "XLLA2NK3uDQr"
-   },
-   "outputs": [],
-   "source": [
-    "tokenizer = BartTokenizer.from_pretrained(BASE_MODEL)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "Ntow53I_t81D",
-    "outputId": "59289cdd-1429-4720-cc87-88810c4b99ac"
-   },
-   "outputs": [],
-   "source": [
-    "text = \"My friends are cool but they eat too many carbs.\"\n",
-    "inputs = tokenizer(text, max_length=1024, return_tensors='jax')\n",
-    "encoder_outputs = model.encode(**inputs)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "vcRNJnJ_uJOJ",
-    "outputId": "025afd54-7908-4a9c-fb59-e40bd3458711"
-   },
-   "outputs": [],
-   "source": [
-    "decoder_start_token_id = model.config.decoder_start_token_id\n",
-    "decoder_start_token_id"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "6QWmEwL_uMld"
-   },
-   "outputs": [],
-   "source": [
-    "decoder_input_ids = jnp.ones((inputs.input_ids.shape[0], 1), dtype=\"i4\") * decoder_start_token_id\n",
-    "outputs = model.decode(decoder_input_ids, encoder_outputs)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "c_ys3yWBothF",
-    "outputId": "40d4d584-e0a8-44cb-bbea-0ffa38d50a53"
-   },
-   "outputs": [],
-   "source": [
-    "outputs"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "O6s0wtB_uTC_",
-    "outputId": "bc0e9e80-e346-4e99-d28e-3f658eda1f66"
-   },
-   "outputs": [],
-   "source": [
-    "outputs.logits.shape"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "ELzemGP3uBzy",
-    "outputId": "dc12f98a-1ccf-450d-ba2a-9c29d7d14885"
-   },
-   "outputs": [],
-   "source": [
-    "outputs.logits.argmax(axis=-1)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "fQjikkGEunpx",
-    "outputId": "3dba0209-ad4e-4069-be38-6c599c677ef1"
-   },
-   "outputs": [],
-   "source": [
-    "model.config.bos_token_id, model.config.eos_token_id, model.config.pad_token_id"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "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": null,
-   "metadata": {
-    "id": "C7cHbIHruELT"
-   },
-   "outputs": [],
-   "source": [
-    "greedy_output = model.generate(input_ids_test, max_length=50)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "jYugh9cOuwc9",
-    "outputId": "19c3a2ee-e7bc-4f1f-9c86-06bd7337b537"
-   },
-   "outputs": [],
-   "source": [
-    "greedy_output[0]"
-   ]
-  }
- ],
- "metadata": {
-  "accelerator": "TPU",
-  "colab": {
-   "collapsed_sections": [],
-   "machine_shape": "hm",
-   "name": "CustomBARTv4b-model-generate.ipynb",
-   "provenance": []
-  },
-  "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": 4
-}

dev/notebooks/demo/demo_notebook.ipynb

@@ -1,387 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "ewer-Q-0w2xA"
-   },
-   "source": [
-    "# Installation"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "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": null,
-   "metadata": {
-    "id": "M1wVkrpjU6zO"
-   },
-   "outputs": [],
-   "source": [
-    "%load_ext autoreload\n",
-    "%autoreload 2"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "%cd ../../vqgan-jax"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "t47CH1H_IOT8"
-   },
-   "source": [
-    "# Custom BART Model"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "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": null,
-   "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": null,
-   "metadata": {
-    "scrolled": true
-   },
-   "outputs": [],
-   "source": [
-    "import wandb\n",
-    "run = wandb.init()\n",
-    "artifact = run.use_artifact('wandb/hf-flax-dalle-mini/model-1ef8yxby:latest', type='bart_model')\n",
-    "artifact_dir = artifact.download()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "_6-XKK40oEfP",
-    "scrolled": true
-   },
-   "outputs": [],
-   "source": [
-    "# create our model and initialize it randomly\n",
-    "model = CustomFlaxBartForConditionalGeneration.from_pretrained(artifact_dir)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model.config.forced_bos_token_id = None"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "Jz032w73nHEf",
-    "outputId": "994d8e85-bff7-480b-8b69-f69dedc15c49"
-   },
-   "outputs": [],
-   "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": null,
-   "metadata": {
-    "id": "XLLA2NK3uDQr"
-   },
-   "outputs": [],
-   "source": [
-    "tokenizer = BartTokenizer.from_pretrained(BASE_MODEL)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "input_text = ['I enjoy walking with my cute dog']*8"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "P32mJJSbrU1F"
-   },
-   "outputs": [],
-   "source": [
-    "input_ids_test = tokenizer(input_text, return_tensors='jax')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "input_ids_test"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "C7cHbIHruELT"
-   },
-   "outputs": [],
-   "source": [
-    "greedy_output = model.generate(input_ids_test['input_ids'], max_length=257)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "greedy_output[0].shape"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "jYugh9cOuwc9",
-    "outputId": "19c3a2ee-e7bc-4f1f-9c86-06bd7337b537"
-   },
-   "outputs": [],
-   "source": [
-    "greedy_output[0]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "greedy_output[0][0]"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# VGAN Jax"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "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": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from modeling_flax_vqgan import VQModel"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "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": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "Jz032w73nHEf",
-    "outputId": "994d8e85-bff7-480b-8b69-f69dedc15c49",
-    "scrolled": true
-   },
-   "outputs": [],
-   "source": [
-    "model = VQModel.from_pretrained(\"flax-community/vqgan_f16_16384\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def get_images(indices, model):\n",
-    "    indices =  indices[:, 1:]\n",
-    "    print(indices.shape)\n",
-    "    img = model.decode_code(indices)\n",
-    "    return img"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "custom_to_pil(np.asarray(get_images(jnp.expand_dims(greedy_output[0][0],0), model)[0]))"
-   ]
-  }
- ],
- "metadata": {
-  "accelerator": "TPU",
-  "colab": {
-   "collapsed_sections": [],
-   "machine_shape": "hm",
-   "name": "CustomBARTv4b-model-generate.ipynb",
-   "provenance": []
-  },
-  "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": 4
-}

dev/notebooks/demo/model-sweep.py

@@ -1,216 +0,0 @@
-#!/usr/bin/env python
-# coding: utf-8
-
-import random
-
-import jax
-import flax.linen as nn
-from flax.training.common_utils import shard
-from flax.jax_utils import replicate, unreplicate
-
-from transformers.models.bart.modeling_flax_bart import *
-from transformers import BartTokenizer, FlaxBartForConditionalGeneration
-
-from PIL import Image
-import numpy as np
-import matplotlib.pyplot as plt
-
-import torchvision.transforms as T
-import torchvision.transforms.functional as TF
-from torchvision.transforms import InterpolationMode
-
-from vqgan_jax.modeling_flax_vqgan import VQModel
-
-# TODO: set those args in a config file
-OUTPUT_VOCAB_SIZE = 16384 + 1  # encoded image token space + 1 for bos
-OUTPUT_LENGTH = 256 + 1  # number of encoded tokens + 1 for bos
-BOS_TOKEN_ID = 16384
-BASE_MODEL = 'facebook/bart-large-cnn'
-WANDB_MODEL = '3iwhu4w6'
-
-class CustomFlaxBartModule(FlaxBartModule):
-    def setup(self):
-        # we keep shared to easily load pre-trained weights
-        self.shared = nn.Embed(
-            self.config.vocab_size,
-            self.config.d_model,
-            embedding_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),
-            dtype=self.dtype,
-        )
-        # a separate embedding is used for the decoder
-        self.decoder_embed = nn.Embed(
-            OUTPUT_VOCAB_SIZE,
-            self.config.d_model,
-            embedding_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),
-            dtype=self.dtype,
-        )
-        self.encoder = FlaxBartEncoder(self.config, dtype=self.dtype, embed_tokens=self.shared)
-
-        # the decoder has a different config
-        decoder_config = BartConfig(self.config.to_dict())
-        decoder_config.max_position_embeddings = OUTPUT_LENGTH
-        decoder_config.vocab_size = OUTPUT_VOCAB_SIZE
-        self.decoder = FlaxBartDecoder(decoder_config, dtype=self.dtype, embed_tokens=self.decoder_embed)
-
-class CustomFlaxBartForConditionalGenerationModule(FlaxBartForConditionalGenerationModule):
-    def setup(self):
-        self.model = CustomFlaxBartModule(config=self.config, dtype=self.dtype)
-        self.lm_head = nn.Dense(
-            OUTPUT_VOCAB_SIZE,
-            use_bias=False,
-            dtype=self.dtype,
-            kernel_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),
-        )
-        self.final_logits_bias = self.param("final_logits_bias", self.bias_init, (1, OUTPUT_VOCAB_SIZE))
-
-class CustomFlaxBartForConditionalGeneration(FlaxBartForConditionalGeneration):
-    module_class = CustomFlaxBartForConditionalGenerationModule
-
-tokenizer = BartTokenizer.from_pretrained(BASE_MODEL)
-vqgan = VQModel.from_pretrained("flax-community/vqgan_f16_16384")
-
-def custom_to_pil(x):
-    x = np.clip(x, 0., 1.)
-    x = (255*x).astype(np.uint8)
-    x = Image.fromarray(x)
-    if not x.mode == "RGB":
-        x = x.convert("RGB")
-    return x
-
-def generate(input, rng, params):
-    return model.generate(
-        **input,
-        max_length=257,
-        num_beams=1,
-        do_sample=True,
-        prng_key=rng,
-        eos_token_id=50000,
-        pad_token_id=50000,
-        params=params,
-    )
-
-def get_images(indices, params):
-    return vqgan.decode_code(indices, params=params)
-
-def plot_images(images):
-    fig = plt.figure(figsize=(40, 20))
-    columns = 4
-    rows = 2
-    plt.subplots_adjust(hspace=0, wspace=0)
-
-    for i in range(1, columns*rows +1):
-        fig.add_subplot(rows, columns, i)
-        plt.imshow(images[i-1])
-    plt.gca().axes.get_yaxis().set_visible(False)
-    plt.show()
-    
-def stack_reconstructions(images):
-    w, h = images[0].size[0], images[0].size[1]
-    img = Image.new("RGB", (len(images)*w, h))
-    for i, img_ in enumerate(images):
-        img.paste(img_, (i*w,0))
-    return img
-
-p_generate = jax.pmap(generate, "batch")
-p_get_images = jax.pmap(get_images, "batch")
-
-# ## CLIP Scoring
-from transformers import CLIPProcessor, FlaxCLIPModel
-
-clip = FlaxCLIPModel.from_pretrained("openai/clip-vit-base-patch32")
-processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
-
-def hallucinate(prompt, num_images=64):
-    prompt = [prompt] * jax.device_count()
-    inputs = tokenizer(prompt, return_tensors='jax', padding="max_length", truncation=True, max_length=128).data
-    inputs = shard(inputs)
-
-    all_images = []
-    for i in range(num_images // jax.device_count()):
-        key = random.randint(0, 1e7)
-        rng = jax.random.PRNGKey(key)
-        rngs = jax.random.split(rng, jax.local_device_count())
-        indices = p_generate(inputs, rngs, bart_params).sequences
-        indices = indices[:, :, 1:]
-
-        images = p_get_images(indices, vqgan_params)
-        images = np.squeeze(np.asarray(images), 1)
-        for image in images:
-            all_images.append(custom_to_pil(image))
-    return all_images
-
-def clip_top_k(prompt, images, k=8):
-    inputs = processor(text=prompt, images=images, return_tensors="np", padding=True)
-    outputs = clip(**inputs)
-    logits = outputs.logits_per_text
-    scores = np.array(logits[0]).argsort()[-k:][::-1]
-    return [images[score] for score in scores]
-
-from PIL import ImageDraw, ImageFont
-
-def captioned_strip(images, caption):
-    w, h = images[0].size[0], images[0].size[1]
-    img = Image.new("RGB", (len(images)*w, h + 48))
-    for i, img_ in enumerate(images):
-        img.paste(img_, (i*w, 48))
-    draw = ImageDraw.Draw(img)
-    font = ImageFont.truetype("/usr/share/fonts/truetype/liberation2/LiberationMono-Bold.ttf", 40)
-    draw.text((20, 3), caption, (255,255,255), font=font)
-    return img
-    
-def log_to_wandb(prompts):
-    strips = []
-    for prompt in prompts:
-        print(f"Generating candidates for: {prompt}")
-        images = hallucinate(prompt, num_images=32)
-        selected = clip_top_k(prompt, images, k=8)
-        strip = captioned_strip(selected, prompt)
-        strips.append(wandb.Image(strip))
-    wandb.log({"images": strips})
-
-## Artifact loop
-
-import wandb
-import os
-os.environ["WANDB_SILENT"] = "true"
-os.environ["WANDB_CONSOLE"] = "off"
-
-id = wandb.util.generate_id()
-print(f"Logging images to wandb run id: {id}")
-
-run = wandb.init(id=id,
-        entity='wandb',
-        project="hf-flax-dalle-mini",
-        job_type="predictions",
-        resume="allow"
-)
-
-artifact = run.use_artifact('wandb/hf-flax-dalle-mini/model-3iwhu4w6:v0', type='bart_model')
-producer_run = artifact.logged_by()
-logged_artifacts = producer_run.logged_artifacts()
-
-for artifact in logged_artifacts:
-    print(f"Generating predictions with version {artifact.version}")
-    artifact_dir = artifact.download()
-
-    # create our model
-    model = CustomFlaxBartForConditionalGeneration.from_pretrained(artifact_dir)
-    model.config.force_bos_token_to_be_generated = False
-    model.config.forced_bos_token_id = None
-    model.config.forced_eos_token_id = None
-
-    bart_params = replicate(model.params)
-    vqgan_params = replicate(vqgan.params)
-
-    prompts = prompts = [
-        "white snow covered mountain under blue sky during daytime",
-        "aerial view of beach during daytime",
-        "aerial view of beach at night",
-        "an armchair in the shape of an avocado",
-        "young woman riding her bike trough a forest",
-        "rice fields by the mediterranean coast",
-        "white houses on the hill of a greek coastline",
-        "illustration of a shark with a baby shark",
-    ]
-
-    log_to_wandb(prompts)

dev/notebooks/demo/tpu-demo.ipynb

@@ -1,446 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "id": "f6d33374",
-   "metadata": {},
-   "source": [
-    "# Test notebook with CLIP scoring"
-   ]
-  },
-  {
-   "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": null,
-   "id": "41db7534-f589-4b63-9165-9c9799e1b06e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import random\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",
-    "jax.devices()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "09295910",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from vqgan_jax.modeling_flax_vqgan import VQModel"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "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-cnn'"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "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:latest', type='bart_model')\n",
-    "artifact_dir = artifact.download()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "e8bcff33-e95b-4c01-b162-ee857a55c3e6",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# create our model\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": null,
-   "id": "8d5e0f14-2502-470e-9553-daee6748601f",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "vqgan = VQModel.from_pretrained(\"flax-community/vqgan_f16_16384\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "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": null,
-   "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": null,
-   "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]] * jax.device_count()\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": "markdown",
-   "id": "b6e1060f",
-   "metadata": {},
-   "source": [
-    "## CLIP Scoring"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "c68724bc",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from transformers import CLIPProcessor, FlaxCLIPModel"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "17158e5b",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "clip = FlaxCLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
-    "processor = CLIPProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "f1b37b6d",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def hallucinate(prompt, num_images=64):\n",
-    "    prompt = [prompt] * jax.device_count()\n",
-    "    inputs = tokenizer(prompt, return_tensors='jax', padding=\"max_length\", truncation=True, max_length=128).data\n",
-    "    inputs = shard(inputs)\n",
-    "\n",
-    "    all_images = []\n",
-    "    for i in range(num_images // jax.device_count()):\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",
-    "        for image in images:\n",
-    "            all_images.append(custom_to_pil(image))\n",
-    "    return all_images"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "831c715f",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def clip_top_k(prompt, images, k=8):\n",
-    "    inputs = processor(text=prompt, images=images, return_tensors=\"np\", padding=True)\n",
-    "    outputs = clip(**inputs)\n",
-    "    logits = outputs.logits_per_text\n",
-    "    scores = np.array(logits[0]).argsort()[-k:][::-1]\n",
-    "    return [images[score] for score in scores]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "00605e13",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "prompt = \"white snow covered mountain under blue sky during daytime\"\n",
-    "images = hallucinate(prompt)\n",
-    "selected = clip_top_k(prompt, images, k=8)\n",
-    "stack_reconstructions(selected)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "cc745da2",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "prompt = \"aerial view of beach at night\"\n",
-    "images = hallucinate(prompt)\n",
-    "selected = clip_top_k(prompt, images, k=8)\n",
-    "stack_reconstructions(selected)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "c9cc0b1d",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "prompt = \"an armchair in the shape of an avocado\"\n",
-    "images = hallucinate(prompt)\n",
-    "selected = clip_top_k(prompt, images, k=8)\n",
-    "stack_reconstructions(selected)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "574e9433",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "prompt = \"young woman riding her bike into a forest\"\n",
-    "images = hallucinate(prompt)\n",
-    "selected = clip_top_k(prompt, images, k=8)\n",
-    "stack_reconstructions(selected)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "4762c91e",
-   "metadata": {},
-   "source": [
-    "`Forest` seems to dominate. Interesting cubist interpretation in the fourth image."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "af30608a",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "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.10"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}

dev/notebooks/model/data-pipeline.ipynb

@@ -1,385 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "id": "bf8fb38a",
-   "metadata": {},
-   "source": [
-    "# Data Pipeline"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "id": "9b83dcb9",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from dataclasses import dataclass, field\n",
-    "from pathlib import Path\n",
-    "\n",
-    "import datasets\n",
-    "from datasets import Dataset, load_dataset\n",
-    "import numpy as np\n",
-    "\n",
-    "from transformers import BartTokenizer\n",
-    "\n",
-    "from tqdm import tqdm\n",
-    "\n",
-    "import jax\n",
-    "import jax.numpy as jnp\n",
-    "\n",
-    "from flax.training.common_utils import shard"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "a661a89e",
-   "metadata": {},
-   "source": [
-    "File containing image paths, captions and VQGAN-encoded indices."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "id": "0e84e889",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "datafile = '/data/CC12M/images-encoded-10000.tsv'   # 9999 encoded images from CC12M"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "7fdc640b",
-   "metadata": {},
-   "source": [
-    "TODO: generate train/test splits if necessary."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "id": "cc6789b4",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Using custom data configuration default-91833df78e844785\n",
-      "Reusing dataset csv (/home/pedro/.cache/huggingface/datasets/csv/default-91833df78e844785/0.0.0/e138af468cb14e747fb46a19c787ffcfa5170c821476d20d5304287ce12bbc23)\n"
-     ]
-    }
-   ],
-   "source": [
-    "dataset = load_dataset('csv', delimiter='\\t', data_files=[datafile])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "id": "f3ed4919",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "DatasetDict({\n",
-       "    train: Dataset({\n",
-       "        features: ['image_file', 'caption', 'encoding'],\n",
-       "        num_rows: 9999\n",
-       "    })\n",
-       "})"
-      ]
-     },
-     "execution_count": 4,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "dataset"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "id": "a70c7354",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Dataset({\n",
-       "    features: ['image_file', 'caption', 'encoding'],\n",
-       "    num_rows: 9999\n",
-       "})"
-      ]
-     },
-     "execution_count": 5,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "dataset = dataset[\"train\"]\n",
-    "dataset"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "a73454cf",
-   "metadata": {},
-   "source": [
-    "We don't really need the `image_file` field for training. We'll drop it during pre-processing because we won't be able to numericalize it to a `jnp.array`, which would be required in JAX."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "7c0fa992",
-   "metadata": {},
-   "source": [
-    "## Preprocessing"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "a0e36582",
-   "metadata": {},
-   "source": [
-    "The `encoding` field contains a string representation of the encoded indices. We'll convert them to numbers. We also need to tokenize the captions."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "id": "d46f6ac5",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Setting padding=\"max_length\" as we need fixed length inputs for jitted functions\n",
-    "max_length = 256   # Read from data_args.max_source_length\n",
-    "tokenizer = BartTokenizer.from_pretrained('facebook/bart-large-cnn')\n",
-    "image_bos = 16384   # Max token is 16383 in our VQGAN configuration"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "id": "4cac6643",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def preprocess_function(examples):\n",
-    "    inputs = examples[\"caption\"]\n",
-    "#     inputs = [prefix + inp for inp in inputs]   # Do we need this?\n",
-    "    model_inputs = tokenizer(\n",
-    "        inputs, max_length=max_length, padding=\"max_length\", truncation=True, return_tensors=\"np\"\n",
-    "    )\n",
-    "\n",
-    "    model_inputs[\"labels\"] = [[image_bos] + eval(indices) for indices in examples['encoding']]\n",
-    "\n",
-    "    return model_inputs"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "id": "e6a4cb91",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "num_workers = 48     # We have 96 processors in the TPU\n",
-    "column_names = dataset.column_names\n",
-    "input_dataset = dataset.map(preprocess_function,\n",
-    "                            remove_columns=column_names,\n",
-    "                            batched=True,\n",
-    "                            num_proc=48\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "id": "a9b1b467",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuffle: bool = False):\n",
-    "    \"\"\"\n",
-    "    Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices.\n",
-    "    Shuffle batches if `shuffle` is `True`.\n",
-    "    \"\"\"\n",
-    "    steps_per_epoch = len(dataset) // batch_size\n",
-    "\n",
-    "    if shuffle:\n",
-    "        batch_idx = jax.random.permutation(rng, len(dataset))\n",
-    "    else:\n",
-    "        batch_idx = jnp.arange(len(dataset))\n",
-    "\n",
-    "    batch_idx = batch_idx[: steps_per_epoch * batch_size]  # Skip incomplete batch.\n",
-    "    batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))\n",
-    "\n",
-    "    for idx in batch_idx:\n",
-    "        batch = dataset[idx]        \n",
-    "        batch = {k: jnp.array(v) for k, v in batch.items()}\n",
-    "        batch = shard(batch)\n",
-    "        yield batch"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "id": "0a628505",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "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: Host TPU Interpreter\n"
-     ]
-    }
-   ],
-   "source": [
-    "rng = jax.random.PRNGKey(23)  # Use training_args.seed\n",
-    "batch_size = 64    # Per device\n",
-    "super_batch_size = batch_size * jax.device_count()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "id": "b3a5ce7d",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "loader = data_loader(rng, input_dataset, batch_size=super_batch_size)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "id": "67aa8f9c",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "superbatch = next(iter(loader))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 13,
-   "id": "7cd99402",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "dict_keys(['attention_mask', 'input_ids', 'labels'])"
-      ]
-     },
-     "execution_count": 13,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "superbatch.keys()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 14,
-   "id": "652a4a9e",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "8"
-      ]
-     },
-     "execution_count": 14,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "len(superbatch[\"labels\"])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 15,
-   "id": "de7de4e8",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "(8, 64, 257)"
-      ]
-     },
-     "execution_count": 15,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "superbatch[\"labels\"].shape"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "6800153b",
-   "metadata": {},
-   "source": [
-    "Any image sequence should begin with `image_bos`:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 16,
-   "id": "cfe23a71",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "assert superbatch[\"labels\"][1][5][0].item() == image_bos"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "0fb899b4",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "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.10"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}

dev/predictions/wandb-examples-from-backend.py

@@ -1,52 +0,0 @@
-#!/usr/bin/env python
-# coding: utf-8
-
-from PIL import Image, ImageDraw, ImageFont
-import wandb
-import os
-
-from dalle_mini.backend import ServiceError, get_images_from_backend
-from dalle_mini.helpers import captioned_strip
-
-os.environ["WANDB_SILENT"] = "true"
-os.environ["WANDB_CONSOLE"] = "off"
-
-# set id to None so our latest images don't get overwritten
-id = None
-run = wandb.init(id=id,
-        entity='wandb',
-        project="hf-flax-dalle-mini",
-        job_type="predictions",
-        resume="allow"
-)
-
-def log_to_wandb(prompts):
-    try:
-        backend_url = os.environ["BACKEND_SERVER"]
-
-        strips = []
-        for prompt in prompts:
-            print(f"Getting selections for: {prompt}")
-            selected = get_images_from_backend(prompt, backend_url)
-            strip = captioned_strip(selected, prompt)
-            strips.append(wandb.Image(strip))
-        wandb.log({"images": strips})
-    except ServiceError as error:
-        print(f"Service unavailable, status: {error.status_code}")
-    except KeyError:
-        print("Error: BACKEND_SERVER unset")
-
-prompts = [
-    "white snow covered mountain under blue sky during daytime",
-    "aerial view of beach during daytime",
-    "aerial view of beach at night",
-    "an armchair in the shape of an avocado",
-    "a logo of an avocado armchair playing music",
-    "young woman riding her bike trough a forest",
-    "rice fields by the mediterranean coast",
-    "white houses on the hill of a greek coastline",
-    "illustration of a shark with a baby shark",
-    "painting of an oniric forest glade surrounded by tall trees",
-]
-
-log_to_wandb(prompts)