Merge branch 'main' of https://github.com/borisdayma/dalle-mini into fix-opt_state

.github/workflows/sync_to_hub_debug.yml

@@ -0,0 +1,17 @@
+name: Deploy to debug app
+
+on:
+  # to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+jobs:
+  sync-to-hub-debug:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          fetch-depth: 0
+      - name: Push to hub
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: git push --force https://boris:[email protected]/spaces/flax-community/dalle-mini-debug +HEAD:main

CITATION.cff

@@ -0,0 +1,44 @@
+# YAML 1.2
+---
+abstract: "DALL·E mini is a JAX/Flax reimplementation of OpenAI's DALL·E that requires much smaller hardware resources. By simplifying the architecture and model memory requirements, as well as leveraging open-source code and pre-trained models, we were able to create a model that is 27 times smaller than the original DALL·E and train it on a single TPU v3-8 for only 3 days. DALL·E mini achieves impressive results, albeit of a lower quality than the original system. It can be used for exploration and further experimentation on commodity hardware."
+authors: 
+  -
+    family-names: Dayma
+    given-names: Boris
+  -
+    family-names: Patil
+    given-names: Suraj
+  -
+    family-names: Cuenca
+    given-names: Pedro
+  -
+    family-names: Saifullah
+    given-names: Khalid
+  -
+    family-names: Abraham
+    given-names: Tanishq
+  -
+    family-names: "Lê Khắc"
+    given-names: "Phúc"
+  -
+    family-names: Melas
+    given-names: Luke
+  -
+    family-names: Ghosh
+    given-names: Ritobrata
+cff-version: "1.1.0"
+date-released: 2021-07-29
+identifiers: 
+keywords: 
+  - dalle
+  - "text-to-image generation"
+  - transformer
+  - "zero-shot"
+  - JAX
+license: "Apache-2.0"
+doi: 10.5281/zenodo.5146400
+message: "If you use this project, please cite it using these metadata."
+repository-code: "https://github.com/borisdayma/dalle-mini"
+title: "DALL·E Mini"
+version: "v0.1-alpha"
+...

README.md

@@ -1,8 +1,8 @@
 ---
 title: DALL·E mini
 emoji: 🥑
-colorFrom: red
-colorTo: purple
+colorFrom: yellow
+colorTo: green
 sdk: streamlit
 app_file: app/app.py
 pinned: false
@@ -16,7 +16,7 @@ _Generate images from a text prompt_
 
 Our logo was generated with DALL·E mini using the prompt "logo of an armchair in the shape of an avocado".
 
-You can create your own pictures with [the demo](https://huggingface.co/spaces/flax-community/dalle-mini) (temporarily in beta on Huging Face Spaces but soon to be open to all).
+You can create your own pictures with [the demo](https://huggingface.co/spaces/flax-community/dalle-mini).
 
 ## How does it work?
 
@@ -26,8 +26,6 @@ Refer to [our report](https://wandb.ai/dalle-mini/dalle-mini/reports/DALL-E-mini
 
 ### Dependencies Installation
 
-The root folder and associated [`requirements.txt`](./requirements.txt) is only for the app.
-
 For development, use [`dev/requirements.txt`](dev/requirements.txt) or [`dev/environment.yaml`](dev/environment.yaml).
 
 ### Training of VQGAN
@@ -52,7 +50,16 @@ To generate sample predictions and understand the inference pipeline step by ste
 
 [![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?
+## FAQ
+
+### Where to find the latest models?
+
+Trained models are on 🤗 Model Hub:
+
+- [VQGAN-f16-16384](https://huggingface.co/flax-community/vqgan_f16_16384) for encoding/decoding images
+- [DALL·E mini](https://huggingface.co/flax-community/dalle-mini) for generating images from a text prompt
+
+### 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.
 
@@ -70,4 +77,66 @@ The "armchair in the shape of an avocado" was used by OpenAI when releasing DALL
 ## Acknowledgements
 
 - 🤗 Hugging Face for organizing [the FLAX/JAX community week](https://github.com/huggingface/transformers/tree/master/examples/research_projects/jax-projects)
-- Google Cloud team for providing access to TPU's
+- Google [TPU Research Cloud (TRC) program](https://sites.research.google/trc/) for providing computing resources
+- [Weights & Biases](https://wandb.com/) for providing the infrastructure for experiment tracking and model management
+
+## Citing DALL·E mini
+
+If you find DALL·E mini useful in your research or wish to refer, please use the following BibTeX entry.
+
+```
+@misc{Dayma_DALL·E_Mini_2021,
+author = {Dayma, Boris and Patil, Suraj and Cuenca, Pedro and Saifullah, Khalid and Abraham, Tanishq and Lê Khắc, Phúc and Melas, Luke and Ghosh, Ritobrata},
+doi = {10.5281/zenodo.5146400},
+month = {7},
+title = {DALL·E Mini},
+url = {https://github.com/borisdayma/dalle-mini},
+year = {2021}
+}
+```
+
+## References
+
+```
+@misc{ramesh2021zeroshot,
+      title={Zero-Shot Text-to-Image Generation}, 
+      author={Aditya Ramesh and Mikhail Pavlov and Gabriel Goh and Scott Gray and Chelsea Voss and Alec Radford and Mark Chen and Ilya Sutskever},
+      year={2021},
+      eprint={2102.12092},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+```
+
+```
+@misc{esser2021taming,
+      title={Taming Transformers for High-Resolution Image Synthesis}, 
+      author={Patrick Esser and Robin Rombach and Björn Ommer},
+      year={2021},
+      eprint={2012.09841},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+```
+
+```
+@misc{lewis2019bart,
+      title={BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension}, 
+      author={Mike Lewis and Yinhan Liu and Naman Goyal and Marjan Ghazvininejad and Abdelrahman Mohamed and Omer Levy and Ves Stoyanov and Luke Zettlemoyer},
+      year={2019},
+      eprint={1910.13461},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL}
+}
+```
+
+```
+@misc{radford2021learning,
+      title={Learning Transferable Visual Models From Natural Language Supervision}, 
+      author={Alec Radford and Jong Wook Kim and Chris Hallacy and Aditya Ramesh and Gabriel Goh and Sandhini Agarwal and Girish Sastry and Amanda Askell and Pamela Mishkin and Jack Clark and Gretchen Krueger and Ilya Sutskever},
+      year={2021},
+      eprint={2103.00020},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+```

app/app.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python
 # coding: utf-8
 
-import random
 from dalle_mini.backend import ServiceError, get_images_from_backend
 
 import streamlit as st
@@ -55,12 +54,31 @@ st.subheader('Generate images from text')
 
 prompt = st.text_input("What do you want to see?")
 
-#TODO: I think there's an issue where we can't run twice the same inference (not due to caching) - may need to use st.form
-
+test = st.empty()
 DEBUG = False
 if prompt != "" or (should_run_again and prompt != ""):
     container = st.empty()
-    container.markdown(f"Generating predictions for: **{prompt}**")
+    # The following mimics `streamlit.info()`.
+    # I tried to get the secondary background color using `components.streamlit.config.get_options_for_section("theme")["secondaryBackgroundColor"]`
+    # but it returns None.
+    container.markdown(f"""
+        <style> p {{ margin:0 }} div {{ margin:0 }} </style>
+        <div data-stale="false" class="element-container css-1e5imcs e1tzin5v1">
+        <div class="stAlert">
+        <div role="alert" data-baseweb="notification" class="st-ae st-af st-ag st-ah st-ai st-aj st-ak st-g3 st-am st-b8 st-ao st-ap st-aq st-ar st-as st-at st-au st-av st-aw st-ax st-ay st-az st-b9 st-b1 st-b2 st-b3 st-b4 st-b5 st-b6">
+        <div class="st-b7">
+        <div class="css-whx05o e13vu3m50">
+        <div data-testid="stMarkdownContainer" class="css-1ekf893 e16nr0p30">
+                <img src="https://raw.githubusercontent.com/borisdayma/dalle-mini/main/app/img/loading.gif" width="30"/>
+                Generating predictions for: <b>{prompt}</b>
+        </div>
+        </div>
+        </div>
+        </div>
+        </div>
+        </div>
+        <small><i>Predictions may take up to 40s under high load. Please stand by.</i></small>
+    """, unsafe_allow_html=True)
 
     try:
         backend_url = st.secrets["BACKEND_SERVER"]

app/dalle_mini

@@ -1 +0,0 @@
-../dalle_mini/

app/gradio/dalle_mini

@@ -1 +0,0 @@
-../../dalle_mini/

dalle_mini/text.py

@@ -0,0 +1,272 @@
+"""
+Utilities for processing text.
+"""
+
+import requests
+from pathlib import Path
+from unidecode import unidecode
+
+import re, math, random, html
+import ftfy
+
+WIKI_STATS_URL = "https://github.com/borisdayma/wikipedia-word-frequency/raw/feat-update/results/enwiki-20210820-words-frequency.txt"
+WIKI_STATS_LOCAL = Path(WIKI_STATS_URL).parts[-1]
+
+# based on wiki word occurence
+person_token = [("a person", 282265), ("someone", 121194), ("somebody", 12219)]
+temp_token = "xtokx"  # avoid repeating chars
+
+
+def get_wiki_file():
+    if not Path(WIKI_STATS_LOCAL).exists():
+        r = requests.get(WIKI_STATS_URL, stream=True)
+        with open(WIKI_STATS_LOCAL, "wb") as fd:
+            for chunk in r.iter_content(chunk_size=128):
+                fd.write(chunk)
+    return WIKI_STATS_LOCAL
+
+
+class HashtagProcessor:
+    # Adapted from wordninja library
+    # We use our wikipedia word count + a good heuristic to make it work
+    def __init__(self):
+        self._word_cost = (
+            l.split()[0] for l in Path(get_wiki_file()).read_text().splitlines()
+        )
+        self._word_cost = {
+            str(k): math.log(float(i + 1)) for i, k in enumerate(self._word_cost)
+        }
+        self._max_word = max(len(x) for x in self._word_cost.keys())
+        self._SPLIT_RE = re.compile("[^a-zA-Z0-9']+")
+
+    def __call__(self, s):
+        """Uses dynamic programming to infer the location of spaces in a string without spaces."""
+        l = [self._split(x) for x in self._SPLIT_RE.split(s)]
+        return " ".join([item for sublist in l for item in sublist])
+
+    def _split(self, s):
+        # Find the best match for the i first characters, assuming cost has
+        # been built for the i-1 first characters.
+        # Returns a pair (match_cost, match_length).
+        def best_match(i):
+            candidates = enumerate(reversed(cost[max(0, i - self._max_word) : i]))
+            return min(
+                (c + self._word_cost.get(s[i - k - 1 : i].lower(), 9e999), k + 1)
+                for k, c in candidates
+            )
+
+        # Build the cost array
+        cost = [0]
+        for i in range(1, len(s) + 1):
+            c, k = best_match(i)
+            cost.append(c)
+
+        # Backtrack to recover the minimal-cost string.
+        out = []
+        i = len(s)
+        while i > 0:
+            c, k = best_match(i)
+            assert c == cost[i]
+            newToken = True
+            if not s[i - k : i] == "'":  # ignore a lone apostrophe
+                if len(out) > 0:
+                    # re-attach split 's and split digits
+                    if out[-1] == "'s" or (
+                        s[i - 1].isdigit() and out[-1][0].isdigit()
+                    ):  # digit followed by digit
+                        out[-1] = (
+                            s[i - k : i] + out[-1]
+                        )  # combine current token with previous token
+                        newToken = False
+
+            if newToken:
+                out.append(s[i - k : i])
+
+            i -= k
+
+        return reversed(out)
+
+
+def replace_person_token(t):
+    "Used for CC12M"
+    t = re.sub("<person>([,\s]*(and)*[,\s]*<person>)+", " people ", t)
+    while "<person>" in t:
+        t = t.replace(
+            "<person>", f" {random.choices(*tuple(zip(*person_token)))[0]} ", 1
+        )
+    return t
+
+
+def fix_html(t):
+    "Adapted from fastai"
+    t = (
+        t.replace("#39;", "'")
+        .replace("&amp;", "&")
+        .replace("amp;", "&")
+        .replace("#146;", "'")
+        .replace("nbsp;", " ")
+        .replace("#36;", "$")
+        .replace("\\n", "\n")
+        .replace("quot;", "'")
+        .replace("<br />", "\n")
+        .replace('\\"', '"')
+        .replace("<unk>", " ")
+        .replace(" @.@ ", ".")
+        .replace(" @-@ ", "-")
+    )
+    return html.unescape(t)
+
+
+def replace_punctuation_with_commas(t):
+    return re.sub("""([()[\].,|:;?!=+~\-])""", ",", t)
+
+
+def simplify_quotes(t):
+    return re.sub("""['"`]""", ' " ', t)
+
+
+def merge_quotes(t):
+    return re.sub('(\s*"+\s*)+', ' " ', t)
+
+
+def remove_comma_numbers(t):
+    def _f(t):
+        return re.sub("(\d),(\d{3})", r"\1\2", t)
+
+    return _f(_f(t))
+
+
+def pre_process_dot_numbers(t):
+    return re.sub("(\d)\.(\d)", fr"\1{temp_token}dot{temp_token}\2", t)
+
+
+def post_process_dot_numbers(t):
+    return re.sub(f"{temp_token}dot{temp_token}", ".", t)
+
+
+def pre_process_quotes(t):
+    # allows quotes only for 's, 't, 'd, 'm, 'll, 're, 've
+    return re.sub(
+        r"'(?=([stdm]|(ll)|(re)|(ve)|(ll))\b)", fr"{temp_token}quote{temp_token}", t
+    )
+
+
+def post_process_quotes(t):
+    return re.sub(f"{temp_token}quote{temp_token}", "'", t)
+
+
+def merge_commas(t):
+    return re.sub("(\s*,+\s*)+", ", ", t)
+
+
+def add_space_after_commas(t):
+    return re.sub(",", ", ", t)
+
+
+def handle_special_chars(t):
+    "Handle special characters"
+    # replace "-" with a space when between words without space
+    t = re.sub("([a-zA-Z])-([a-zA-Z])", r"\1 \2", t)
+    # always add space around &
+    return re.sub("&", " & ", t)
+
+
+def expand_hashtags(t, hashtag_processor):
+    "Remove # and try to split words"
+    return re.sub("#(\w+)", lambda m: hashtag_processor(m.group(1)), t)
+
+
+_re_ignore_chars = """[_#\/\\%]"""
+
+
+def ignore_chars(t):
+    "Ignore useless characters"
+    return re.sub(_re_ignore_chars, " ", t)
+
+
+def remove_extra_spaces(t):
+    "Remove extra spaces (including \t and \n)"
+    return re.sub("\s+", " ", t)
+
+
+def remove_repeating_chars(t):
+    "If the same character is present 4+ times (not 3 because of roman 'VIII'), replace with single instance"
+    return re.sub(r"(\D)(\1{3,})", r"\1", t)
+
+
+def remove_urls(t):
+    return re.sub(r"http\S+", "", t)
+
+
+def remove_html_tags(t):
+    return re.sub("<[^<]+?>", "", t)
+
+
+def remove_first_last_commas(t):
+    t = t.strip()
+    t = t[:-1] if t and t[-1] == "," else t
+    t = t[1:] if t and t[0] == "," else t
+    return t.strip()
+
+
+def remove_wiki_ref(t):
+    t = re.sub(r"\A\s*\[\d+\]", "", t)
+    return re.sub(r"\[\d+\]\s*\Z", "", t)
+
+
+class TextNormalizer:
+    "Normalize text"
+
+    def __init__(self):
+        self._hashtag_processor = HashtagProcessor()
+
+    def __call__(self, t, clip=False):
+
+        # fix some characters
+        t = ftfy.fix_text(t)
+        # fix html
+        t = fix_html(t)
+        if not clip:
+            # decode and simplify text: see unidecode library
+            t = unidecode(t)
+        # lower case
+        t = t.lower()
+        # replace <PERSON> (for CC12M)
+        t = replace_person_token(t)
+        # remove wiki reference (for WIT)
+        t = remove_wiki_ref(t)
+        # remove html tags
+        t = remove_html_tags(t)
+        # remove urls
+        t = remove_urls(t)
+        # remove commas in numbers
+        t = remove_comma_numbers(t)
+        if not clip:
+            # handle dots in numbers and quotes - Part 1
+            t = pre_process_dot_numbers(t)
+            t = pre_process_quotes(t)
+            # handle special characters
+            t = handle_special_chars(t)
+            # handle hashtags
+            t = expand_hashtags(t, self._hashtag_processor)
+            # ignore useless characters
+            t = ignore_chars(t)
+            # simplify quotes
+            t = simplify_quotes(t)
+            # all punctuation becomes commas
+            t = replace_punctuation_with_commas(t)
+            # handle dots in numbers and quotes - Part 2
+            t = post_process_dot_numbers(t)
+            t = post_process_quotes(t)
+            # handle repeating characters
+            t = remove_repeating_chars(t)
+            # merge commas
+            t = merge_commas(t)
+            # merge quotes
+            t = merge_quotes(t)
+        # remove multiple spaces
+        t = remove_extra_spaces(t)
+        # remove first and last comma
+        t = remove_first_last_commas(t)
+        # always start with a space
+        return f" {t}" if not clip else t

dev/README.md

@@ -0,0 +1,122 @@
+# Development Instructions for TPU
+
+## Setup
+
+- Apply to the [TRC program](https://sites.research.google/trc/) for free TPU credits if you're elligible.
+- Follow the [Cloud TPU VM User's Guide](https://cloud.google.com/tpu/docs/users-guide-tpu-vm) to set up gcloud.
+- Verify `gcloud config list`, in particular account, project & zone.
+- Create a TPU VM per the guide and connect to it.
+
+When needing a larger disk:
+
+- Create a balanced persistent disk (SSD, so pricier than default HDD but much faster): `gcloud compute disks create DISK_NAME --size SIZE_IN_GB --type pd-balanced`
+- Attach the disk to your instance by adding `--data-disk source=REF` per ["Adding a persistent disk to a TPU VM" guide](https://cloud.google.com/tpu/docs/setup-persistent-disk), eg `gcloud alpha compute tpus tpu-vm create INSTANCE_NAME --accelerator-type=v3-8 --version=v2-alpha --data-disk source=projects/tpu-toys/zones/europe-west4-a/disks/DISK_NAME`
+- Format the partition as described in the guide.
+- Make sure to set up automatic remount of disk at restart.
+
+## Connect VS Code
+
+- Find external IP in the UI or with `gcloud alpha compute tpus tpu-vm describe INSTANCE_NAME`
+- Verify you can connect in terminal with `ssh EXTERNAL_IP -i ~/.ssh/google_compute_engine`
+- Add the same command as ssh host in VS Code.
+- Check config file
+
+  ```
+  Host INSTANCE_NAME
+    HostName EXTERNAL_IP
+    IdentityFile ~/.ssh/google_compute_engine
+  ```
+
+## Environment configuration
+
+### Use virtual environments (optional)
+
+We recommend using virtual environments (such as conda, venv or pyenv-virtualenv).
+
+If you want to use `pyenv` and `pyenv-virtualenv`:
+
+- Installation
+
+  - [Set up build environment](https://github.com/pyenv/pyenv/wiki#suggested-build-environment)
+  - Use [pyenv-installer](https://github.com/pyenv/pyenv-installer): `curl https://pyenv.run | bash`
+  - bash set-up:
+
+    ```bash
+    echo '\n'\
+        '# pyenv setup \n'\
+        'export PYENV_ROOT="$HOME/.pyenv" \n'\
+        'export PATH="$PYENV_ROOT/bin:$PATH" \n'\
+        'eval "$(pyenv init --path)" \n'\
+        'eval "$(pyenv init -)" \n'\
+        'eval "$(pyenv virtualenv-init -)"' >> ~/.bashrc
+    ```
+
+- Usage
+
+  - Install a python version: `pyenv install X.X.X`
+  - Create a virtual environment: `pyenv virtualenv 3.9.6 dalle_env`
+  - Activate: `pyenv activate dalle_env`
+
+    Note: you can auto-activate your environment at a location with `echo dalle_env >> .python-version`
+
+### Tools
+
+- Git
+
+  - `git config --global user.email "[email protected]"
+  - `git config --global user.name "First Last"
+
+- Github CLI
+
+  - See [installation instructions](https://github.com/cli/cli/blob/trunk/docs/install_linux.md)
+  - `gh auth login`
+
+- Direnv
+
+  - Install direnv: `sudo apt-get update && sudo apt-get install direnv`
+  - bash set-up:
+
+    ```bash
+    echo -e '\n'\
+        '# direnv setup \n'\
+        'eval "$(direnv hook bash)" \n' >> ~/.bashrc
+    ```
+
+### Set up repo
+
+- Clone repo: `gh repo clone borisdayma/dalle-mini`
+- If using `pyenv-virtualenv`, auto-activate env: `echo dalle_env >> .python-version`
+
+## Environment
+
+- Install the following (use it later to update our dev requirements.txt)
+
+```
+requests
+pillow
+jupyterlab
+ipywidgets
+
+-e ../datasets[streaming]
+-e ../transformers
+-e ../webdataset
+
+# JAX
+--find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
+jax[tpu]>=0.2.16
+flax
+```
+
+- `transformers-cli login`
+
+---
+
+- set `HF_HOME="/mnt/disks/persist/cache/huggingface"` in `/etc/environment` and ensure you have required permissions, then restart.
+
+## Working with datasets or models
+
+- Install [Git LFS](https://github.com/git-lfs/git-lfs/wiki/Installation)
+- Clone a dataset without large files: `GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/datasets/.../...`
+- Use a local [credential store](https://git-scm.com/book/en/v2/Git-Tools-Credential-Storage) for caching credentials
+- Track specific extentions: `git lfs track "*.ext"`
+- See files tracked with LFS with `git lfs ls-files`

dev/encoding/vqgan-jax-encoding-streaming.ipynb

@@ -0,0 +1,562 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "d0b72877",
+   "metadata": {},
+   "source": [
+    "# VQGAN JAX Encoding for 🤗 Datasets in streaming mode"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ba7b31e6",
+   "metadata": {},
+   "source": [
+    "This notebook shows how to pre-encode images to token sequences using JAX, VQGAN and 🤗 Datasets in streaming mode.\n",
+    "\n",
+    "This example uses our YFCC100M dataset, but it should be easy to adapt to any other image/caption dataset in the huggingface hub."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3b59489e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import io\n",
+    "\n",
+    "import requests\n",
+    "from PIL import Image\n",
+    "import numpy as np\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "import torch\n",
+    "import torchvision.transforms as T\n",
+    "import torchvision.transforms.functional as TF\n",
+    "from torchvision.transforms import InterpolationMode\n",
+    "import os\n",
+    "\n",
+    "import jax\n",
+    "from jax import pmap"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c7c4c1e6",
+   "metadata": {},
+   "source": [
+    "## Dataset and Parameters"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d45a289e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import datasets\n",
+    "from datasets import Dataset, load_dataset"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f26e4f18",
+   "metadata": {},
+   "source": [
+    "We'll use the `validation` set for testing. Adjust accordingly."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "28893c3e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset = load_dataset('dalle-mini/YFCC100M_OpenAI_subset', use_auth_token=True, streaming=True, split='validation')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "33861477",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "\n",
+    "yfcc100m = Path.home()/'data'/'YFCC100M_OpenAI_subset'\n",
+    "yfcc100m_output = yfcc100m/'encoded'      # Output directory for encoded files"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6e7b71c4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "batch_size = 128     # Per device\n",
+    "num_workers = 16     # Unused in streaming mode"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0793c26a",
+   "metadata": {},
+   "source": [
+    "### Data preparation"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "86415769",
+   "metadata": {},
+   "source": [
+    "* Images: we transform them so they are center-cropped and square, all of the same size so we can build batches for TPU/GPU processing.\n",
+    "* Captions: we extract a single `caption` column from the source data, by concatenating the cleaned title and description.\n",
+    "\n",
+    "These transformations are done using the Datasets `map` function. In the case of streaming datasets, transformations will run as needed instead of pre-processing the dataset at once."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0fdf1851",
+   "metadata": {},
+   "source": [
+    "This helper function is used to decode images from the bytes retrieved in `streaming` mode."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5bbca804",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from PIL import Image\n",
+    "import io\n",
+    "\n",
+    "def get_image(byte_stream):\n",
+    "    image = Image.open(io.BytesIO(byte_stream))\n",
+    "    return image.convert('RGB')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b435290b",
+   "metadata": {},
+   "source": [
+    "Image processing"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7e73dfa3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def center_crop(image, max_size=256):\n",
+    "    # Note: we allow upscaling too. We should exclude small images.    \n",
+    "    image = TF.resize(image, max_size, interpolation=InterpolationMode.LANCZOS)\n",
+    "    image = TF.center_crop(image, output_size=2 * [max_size])\n",
+    "    return image\n",
+    "\n",
+    "preprocess_image = T.Compose([\n",
+    "    get_image,\n",
+    "    center_crop,\n",
+    "    T.ToTensor(),\n",
+    "    lambda t: t.permute(1, 2, 0)   # Reorder, we need dimensions last\n",
+    "])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1e3ac8de",
+   "metadata": {},
+   "source": [
+    "Caption preparation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "aadb4d23",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import string\n",
+    "\n",
+    "def create_caption(title, description):\n",
+    "    title = title.strip()\n",
+    "    description = description.strip()\n",
+    "    if len(title) > 0 and title[-1] not in '.!?': title += '.'\n",
+    "    return f'{title} {description}'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3c4522b9",
+   "metadata": {},
+   "source": [
+    "And this is the basic transformation function to use in `map`. We don't really need the `key`, but we'll keep it for reference. Since we are returning a new dictionary (as opposed to adding entries to the input), this also removes any metadata columns we don't need."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2566ff68",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def prepare_item(item):\n",
+    "    return {\n",
+    "        'key': item['key'],\n",
+    "        'caption': create_caption(item['title_clean'], item['description_clean']),\n",
+    "        'image': preprocess_image(item['img'])\n",
+    "    }"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e519e475",
+   "metadata": {},
+   "source": [
+    "Unlike when using non-streaming datasets, the following operation completes immediately in streaming mode. In streaming mode, `num_proc` is not supported."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10d7750e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "prepared_dataset = dataset.map(prepare_item, batched=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a8595539",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "item = next(iter(prepared_dataset))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "04a6eeb4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert(list(item.keys()) == ['key', 'caption', 'image'])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "40d3115f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "item['image'].shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dd844e1c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "T.ToPILImage()(item['image'].permute(2, 0, 1))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "44d50a51",
+   "metadata": {},
+   "source": [
+    "### Torch DataLoader"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "17a4bbc6",
+   "metadata": {},
+   "source": [
+    "We'll create a PyTorch DataLoader for convenience. This allows us to easily take batches of our desired size.\n",
+    "\n",
+    "We won't be using parallel processing of the DataLoader for now, as the items will be retrieved on the fly. We could attempt to do it using these recommendations: https://pytorch.org/docs/stable/data.html#multi-process-data-loading. For performance considerations, please refer to this thread: https://discuss.huggingface.co/t/allow-streaming-of-large-datasets-with-image-audio/8062/13"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e1c08b7e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from torch.utils.data import DataLoader"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6a296677",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "torch_dataset = prepared_dataset.with_format(\"torch\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "29ab13bc",
+   "metadata": {},
+   "source": [
+    "**Note**: according to my tests, `num_workers` is not compatible with Datasets in streaming mode. Processes deadlock and there's no progress."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e2df5e13",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataloader = DataLoader(torch_dataset, batch_size=batch_size * jax.device_count())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c15e3783",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "batch = next(iter(dataloader))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "71d027fe",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "batch['image'].shape"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a354472b",
+   "metadata": {},
+   "source": [
+    "## VQGAN-JAX model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2fcf01d7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from vqgan_jax.modeling_flax_vqgan import VQModel"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9daa636d",
+   "metadata": {},
+   "source": [
+    "We'll use a VQGAN trained with Taming Transformers and converted to a JAX model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "47a8b818",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "model = VQModel.from_pretrained(\"flax-community/vqgan_f16_16384\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "62ad01c3",
+   "metadata": {},
+   "source": [
+    "## Encoding"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "20357f74",
+   "metadata": {},
+   "source": [
+    "Encoding is really simple using `shard` to automatically distribute \"superbatches\" across devices, and `pmap`. This is all it takes to create our encoding function, that will be jitted on first use."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6686b004",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from flax.training.common_utils import shard\n",
+    "from functools import partial"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "322a4619",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "@partial(jax.pmap, axis_name=\"batch\")\n",
+    "def encode(batch):\n",
+    "    # Not sure if we should `replicate` params, does not seem to have any effect\n",
+    "    _, indices = model.encode(batch)\n",
+    "    return indices"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "14375a41",
+   "metadata": {},
+   "source": [
+    "### Encoding loop"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ff6c10d4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import pandas as pd\n",
+    "\n",
+    "def encode_captioned_dataset(dataloader, output_dir, save_every=14):\n",
+    "    output_dir.mkdir(parents=True, exist_ok=True)\n",
+    "        \n",
+    "    # Saving strategy:\n",
+    "    # - Create a new file every so often to prevent excessive file seeking.\n",
+    "    # - Save each batch after processing.\n",
+    "    # - Keep the file open until we are done with it.\n",
+    "    file = None        \n",
+    "    for n, batch in enumerate(tqdm(iter(dataloader))):\n",
+    "        if (n % save_every == 0):\n",
+    "            if file is not None:\n",
+    "                file.close()\n",
+    "            split_num = n // save_every\n",
+    "            file = open(output_dir/f'split_{split_num:05x}.jsonl', 'w')\n",
+    "\n",
+    "        images = batch[\"image\"].numpy()\n",
+    "        images = shard(images.squeeze())\n",
+    "        encoded = encode(images)\n",
+    "        encoded = encoded.reshape(-1, encoded.shape[-1])\n",
+    "\n",
+    "        keys = batch[\"key\"]\n",
+    "        captions = batch[\"caption\"]\n",
+    "\n",
+    "        encoded_as_string = list(map(lambda item: np.array2string(item, separator=',', max_line_width=50000, formatter={'int':lambda x: str(x)}), encoded))\n",
+    "        batch_df = pd.DataFrame.from_dict({\"key\": keys, \"caption\": captions, \"encoding\": encoded_as_string})\n",
+    "        batch_df.to_json(file, orient='records', lines=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "09ff75a3",
+   "metadata": {},
+   "source": [
+    "Create a new file every 318 iterations. This should produce splits of ~500 MB each, when using a total batch size of 1024."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "96222bb4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "save_every = 318"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7704863d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "28it [01:17,  1.60s/it]"
+     ]
+    }
+   ],
+   "source": [
+    "encode_captioned_dataset(dataloader, yfcc100m_output, save_every=save_every)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e266a70a",
+   "metadata": {},
+   "source": [
+    "This is ~10-15 slower than local encoding from an SSD. For performance considerations, see the discussion at https://discuss.huggingface.co/t/allow-streaming-of-large-datasets-with-image-audio/8062/13."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8953dd84",
+   "metadata": {},
+   "source": [
+    "----"
+   ]
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "db471c52d602b4f5f40ecaf278e88ccfef85c29d0a1a07185b0d51fc7acf4e26"
+  },
+  "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/encoding/vqgan-jax-encoding-webdataset.ipynb

@@ -0,0 +1,461 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "d0b72877",
+   "metadata": {},
+   "source": [
+    "# VQGAN JAX Encoding for `webdataset`"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ba7b31e6",
+   "metadata": {},
+   "source": [
+    "This notebook shows how to pre-encode images to token sequences using JAX, VQGAN and a dataset in the [`webdataset` format](https://webdataset.github.io/webdataset/).\n",
+    "\n",
+    "This example uses a small subset of YFCC100M we created for testing, but it should be easy to adapt to any other image/caption dataset in the `webdataset` format."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3b59489e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "import torch\n",
+    "import torchvision.transforms as T\n",
+    "import torchvision.transforms.functional as TF\n",
+    "from torchvision.transforms import InterpolationMode\n",
+    "import math\n",
+    "\n",
+    "import webdataset as wds\n",
+    "\n",
+    "import jax\n",
+    "from jax import pmap"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c7c4c1e6",
+   "metadata": {},
+   "source": [
+    "## Dataset and Parameters"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9822850f",
+   "metadata": {},
+   "source": [
+    "The following is the list of shards we'll process. We hardcode the length of data so that we can see nice progress bars using `tqdm`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1265dbfe",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "shards = 'https://huggingface.co/datasets/dalle-mini/YFCC100M_OpenAI_subset/resolve/main/data/shard-{0000..0008}.tar'\n",
+    "length = 8320"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7e38fa14",
+   "metadata": {},
+   "source": [
+    "If we are extra cautious or our server is unreliable, we can enable retries by providing a custom `curl` retrieval command:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4c8c5960",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Enable curl retries to try to work around temporary network / server errors.\n",
+    "# This shouldn't be necessary when using reliable servers.\n",
+    "# shards = f'pipe:curl -s --retry 5 --retry-delay 5 -L {shards} || true'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "13c6631b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "\n",
+    "# Output directory for encoded files\n",
+    "encoded_output = Path.home()/'data'/'wds'/'encoded'\n",
+    "\n",
+    "batch_size = 128           # Per device\n",
+    "num_workers = 8            # For parallel processing"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3435fb85",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "bs = batch_size * jax.device_count()    # You can use a smaller size while testing\n",
+    "batches = math.ceil(length / bs)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "88598e4b",
+   "metadata": {},
+   "source": [
+    "Image processing"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "669b35df",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def center_crop(image, max_size=256):\n",
+    "    # Note: we allow upscaling too. We should exclude small images.    \n",
+    "    image = TF.resize(image, max_size, interpolation=InterpolationMode.LANCZOS)\n",
+    "    image = TF.center_crop(image, output_size=2 * [max_size])\n",
+    "    return image\n",
+    "\n",
+    "preprocess_image = T.Compose([\n",
+    "    center_crop,\n",
+    "    T.ToTensor(),\n",
+    "    lambda t: t.permute(1, 2, 0)   # Reorder, we need dimensions last\n",
+    "])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a185e90c",
+   "metadata": {},
+   "source": [
+    "Caption preparation.\n",
+    "\n",
+    "Note that we receive the contents of the `json` structure, which will be replaced by the string we return.\n",
+    "If we want to keep other fields inside `json`, we can add `caption` as a new field."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "423ee10e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def create_caption(item):\n",
+    "    title = item['title_clean'].strip()\n",
+    "    description = item['description_clean'].strip()\n",
+    "    if len(title) > 0 and title[-1] not in '.!?': title += '.'\n",
+    "    return f'{title} {description}'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8d3a95db",
+   "metadata": {},
+   "source": [
+    "When an error occurs (a download is disconnected, an image cannot be decoded, etc) the process stops with an exception. We can use one of the exception handlers provided by the `webdataset` library, such as `wds.warn_and_continue` or `wds.ignore_and_continue` to ignore the offending entry and keep iterating.\n",
+    "\n",
+    "**IMPORTANT WARNING:** Do not use error handlers to ignore exceptions until you have tested that your processing pipeline works fine. Otherwise, the process will continue trying to find a valid entry, and it will consume your whole dataset without doing any work.\n",
+    "\n",
+    "We can also create our custom exception handler as demonstrated here:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "369d9719",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# UNUSED - Log exceptions to a file\n",
+    "def ignore_and_log(exn):\n",
+    "    with open('errors.txt', 'a') as f:\n",
+    "        f.write(f'{repr(exn)}\\n')\n",
+    "    return True"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "27de1414",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Or simply use `wds.ignore_and_continue`\n",
+    "exception_handler = wds.warn_and_continue"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5149b6d5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset = wds.WebDataset(shards,\n",
+    "                         length=batches,              # Hint so `len` is implemented\n",
+    "                         shardshuffle=False,          # Keep same order for encoded files for easier bookkeeping. Set to `True` for training.\n",
+    "                         handler=exception_handler,   # Ignore read errors instead of failing.\n",
+    ")\n",
+    "\n",
+    "dataset = (dataset           \n",
+    "      .decode('pil')                     # decode image with PIL\n",
+    "#       .map_dict(jpg=preprocess_image, json=create_caption, handler=exception_handler)    # Process fields with functions defined above\n",
+    "      .map_dict(jpg=preprocess_image, json=create_caption)    # Process fields with functions defined above\n",
+    "      .to_tuple('__key__', 'jpg', 'json') # filter to keep only key (for reference), image, caption.\n",
+    "      .batched(bs))                      # better to batch in the dataset (but we could also do it in the dataloader) - this arg does not affect speed and we could remove it"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8cac98cb",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "keys, images, captions = next(iter(dataset))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cd268fbf",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "images.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c24693c0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "T.ToPILImage()(images[0].permute(2, 0, 1))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "44d50a51",
+   "metadata": {},
+   "source": [
+    "### Torch DataLoader"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e2df5e13",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dl = torch.utils.data.DataLoader(dataset, batch_size=None, num_workers=num_workers)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a354472b",
+   "metadata": {},
+   "source": [
+    "## VQGAN-JAX model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2fcf01d7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from vqgan_jax.modeling_flax_vqgan import VQModel"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9daa636d",
+   "metadata": {},
+   "source": [
+    "We'll use a VQGAN trained with Taming Transformers and converted to a JAX model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "47a8b818",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "model = VQModel.from_pretrained(\"flax-community/vqgan_f16_16384\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "62ad01c3",
+   "metadata": {},
+   "source": [
+    "## Encoding"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "20357f74",
+   "metadata": {},
+   "source": [
+    "Encoding is really simple using `shard` to automatically distribute \"superbatches\" across devices, and `pmap`. This is all it takes to create our encoding function, that will be jitted on first use."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6686b004",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from flax.training.common_utils import shard\n",
+    "from functools import partial"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "322a4619",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "@partial(jax.pmap, axis_name=\"batch\")\n",
+    "def encode(batch):\n",
+    "    # Not sure if we should `replicate` params, does not seem to have any effect\n",
+    "    _, indices = model.encode(batch)\n",
+    "    return indices"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "14375a41",
+   "metadata": {},
+   "source": [
+    "### Encoding loop"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ff6c10d4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import pandas as pd\n",
+    "\n",
+    "def encode_captioned_dataset(dataloader, output_dir, save_every=14):\n",
+    "    output_dir.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "    # Saving strategy:\n",
+    "    # - Create a new file every so often to prevent excessive file seeking.\n",
+    "    # - Save each batch after processing.\n",
+    "    # - Keep the file open until we are done with it.\n",
+    "    file = None        \n",
+    "    for n, (keys, images, captions) in enumerate(tqdm(dataloader)):\n",
+    "        if (n % save_every == 0):\n",
+    "            if file is not None:\n",
+    "                file.close()\n",
+    "            split_num = n // save_every\n",
+    "            file = open(output_dir/f'split_{split_num:05x}.jsonl', 'w')\n",
+    "\n",
+    "        images = shard(images.numpy().squeeze())\n",
+    "        encoded = encode(images)\n",
+    "        encoded = encoded.reshape(-1, encoded.shape[-1])\n",
+    "\n",
+    "        encoded_as_string = list(map(lambda item: np.array2string(item, separator=',', max_line_width=50000, formatter={'int':lambda x: str(x)}), encoded))\n",
+    "        batch_df = pd.DataFrame.from_dict({\"key\": keys, \"caption\": captions, \"encoding\": encoded_as_string})\n",
+    "        batch_df.to_json(file, orient='records', lines=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "09ff75a3",
+   "metadata": {},
+   "source": [
+    "Create a new file every 318 iterations. This should produce splits of ~500 MB each, when using a total batch size of 1024."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "96222bb4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "save_every = 318"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7704863d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "encode_captioned_dataset(dl, encoded_output, save_every=save_every)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8953dd84",
+   "metadata": {},
+   "source": [
+    "----"
+   ]
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "db471c52d602b4f5f40ecaf278e88ccfef85c29d0a1a07185b0d51fc7acf4e26"
+  },
+  "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/inference/dalle_mini

@@ -1 +0,0 @@
-../../dalle_mini

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",

dev/requirements.txt

@@ -1,10 +1,8 @@
-# Note: install with the following command:
-# pip install -r requirements.txt -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
-# Otherwise it won't find the appropriate libtpu_nightly
 requests
+-f https://storage.googleapis.com/jax-releases/libtpu_releases.html
 jax[tpu]>=0.2.16
--e git+https://github.com/huggingface/transformers.git@master#egg=transformers
--e git+https://github.com/huggingface/datasets.git@master#egg=datasets
+transformers
+datasets
 flax
 jupyter
 wandb

requirements.txt

@@ -1,2 +0,0 @@
-# Requirements for huggingface spaces
-streamlit>=0.84.2

setup.cfg

@@ -0,0 +1,16 @@
+[metadata]
+name = dalle_mini
+version = attr: dalle_mini.__version__
+description = DALL·E mini - Generate images from a text prompt
+long_description = file: README.md
+long_description_content_type = text/markdown
+url = https://github.com/borisdayma/dalle-mini
+project_urls =
+    Bug Tracker = https://github.com/borisdayma/dalle-mini/issues
+
+[options]
+packages = find:
+install_requires =
+    transformers
+    jax
+    flax

setup.py

@@ -0,0 +1,4 @@
+from setuptools import setup
+
+if __name__ == "__main__":
+    setup()