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

README.md

@@ -1,42 +1,76 @@
-## DALL-E Mini - Generate image from text
+---
+title: Dalle Mini
+emoji: 🎨
+colorFrom: red
+colorTo: blue
+sdk: streamlit
+app_file: app/app.py
+pinned: false
+---
 
-## Tentative Strategy of training (proposed by Luke and Suraj)
+# DALL-E Mini
 
-### Data: 
-* [Conceptual 12M](https://github.com/google-research-datasets/conceptual-12m) Dataset (already loaded and preprocessed in TPU VM by Luke).
-* [YFCC100M Subset](https://github.com/openai/CLIP/blob/main/data/yfcc100m.md)
-* [Coneptual Captions 3M](https://github.com/google-research-datasets/conceptual-captions)
+_Generate images from a text prompt_
 
-### Architecture: 
-  * Use the Taming Transformers VQ-GAN (with 16384 tokens)
-  * Use a seq2seq (language encoder --> image decoder) model with a pretrained non-autoregressive encoder (e.g. BERT) and an autoregressive decoder (like GPT). 
+TODO: add some cool example
 
-### Remaining Architecture Questions: 
-  * Whether to freeze the text encoder?
-  * Whether to finetune the VQ-GAN?
-  * Which text encoder to use (e.g. BERT, RoBERTa, etc.)?
-  * Hyperparameter choices for the decoder (e.g. positional embedding, initialization, etc.)
+## [Create my own images with the demo →](TODO)
 
-## TODO
+## How does it work?
 
-* experiment with flax/jax and setup of the TPU instance that we should get shortly
-* work on dataset loading - [see suggested datasets](https://discuss.huggingface.co/t/dall-e-mini-version/7324/4)
-* Optionally create the OpenAI YFCC100M subset (see [this post](https://discuss.huggingface.co/t/dall-e-mini-version/7324/30?u=boris))
-* work on text/image encoding
-* concatenate inputs (not sure if we need fixed length for text or use a special token separating text & image)
-* adapt training script
-* create inference function
-* integrate CLIP for better results (only if we have the time)
-* work on a demo (streamlit or colab or maybe just HF widget)
-* document (set up repo on model hub per instructions, start on README writeup…)
-* help with coordinating activities & progress
+Refer to [our report](TODO).
 
+## Development
 
-## Dependencies Installation
-You should create a new python virtual environment and install the project dependencies inside the virtual env. You need to use the `-f` (`--find-links`) option for `pip` to be able to find the appropriate `libtpu` required for the TPU hardware:
+This section is for the adventurous people wanting to look into the code.
+
+### Dependencies Installation
+
+The root folder and associated `requirements.txt` is only for the app.
+
+You will find necessary requirements in each sub-section.
+
+You should create a new python virtual environment and install the project dependencies inside the virtual env. You need to use the `-f` (`--find-links`) option for `pip` to be able to find the appropriate `libtpu` required for the TPU hardware.
+
+Adapt the installation to your own hardware and follow library installation instructions.
 
 ```
 $ pip install -r requirements.txt -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
 ```
 
 If you use `conda`, you can create the virtual env and install everything using: `conda env update -f environments.yaml`
+
+### Training of VQGAN
+
+The VQGAN was trained using [taming-transformers](https://github.com/CompVis/taming-transformers).
+
+We recommend using the latest version available.
+
+### Conversion of VQGAN to JAX
+
+Use [patil-suraj/vqgan-jax](https://github.com/patil-suraj/vqgan-jax).
+
+### Training of Seq2Seq
+
+Refer to `seq2seq` folder (some parameters may have been hardcoded for convenience when training on our TPU VM).
+
+### Inference
+
+Refer to the demo notebooks.
+TODO: add links
+
+## Authors
+
+- [Boris Dayma](https://github.com/borisdayma)
+- [Suraj Patil](https://github.com/patil-suraj)
+- [Pedro Cuenca](https://github.com/pcuenca)
+- [Khalid Saifullah](https://github.com/khalidsaifullaah)
+- [Tanishq Abraham](https://github.com/tmabraham)
+- [Phúc Lê Khắc](https://github.com/lkhphuc)
+- [Luke Melas](https://github.com/lukemelas)
+- [Ritobrata Ghosh](https://github.com/ghosh-r)
+
+## 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

app/app_gradio.py

@@ -0,0 +1,194 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# Uncomment to run on cpu
+#import os
+#os.environ["JAX_PLATFORM_NAME"] = "cpu"
+
+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 requests
+from PIL import Image
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+from dalle_mini.vqgan_jax.modeling_flax_vqgan import VQModel
+
+import gradio as gr
+
+
+# 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 = 'flax-community/dalle-mini'
+
+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
+
+# create our model
+# FIXME: Save tokenizer to hub so we can load from there
+tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
+model = CustomFlaxBartForConditionalGeneration.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
+
+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")
+
+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")
+print("Initialize FlaxCLIPModel")
+processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+print("Initialize CLIPProcessor")
+
+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]
+
+def captioned_strip(images, caption):
+    increased_h = 0 if caption is None else 48
+    w, h = images[0].size[0], images[0].size[1]
+    img = Image.new("RGB", (len(images)*w, h + increased_h))
+    for i, img_ in enumerate(images):
+        img.paste(img_, (i*w, increased_h))
+
+    if caption is not None:
+        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 run_inference(prompt, num_images=64, num_preds=8):
+    images = hallucinate(prompt, num_images=num_images)
+    images = clip_top_k(prompt, images, k=num_preds)
+    predictions_strip = captioned_strip(images, None)
+    return predictions_strip
+
+gr.Interface(run_inference, 
+    inputs=[gr.inputs.Textbox(label='Prompt')], #, gr.inputs.Slider(1,64,1,8, label='Candidates to generate'), gr.inputs.Slider(1,8,1,1, label='Best predictions to show')], 
+    outputs=gr.outputs.Image(label='Generated image'), 
+    title='DALLE-mini - HuggingFace Community Week',
+    description='This is a demo of the DALLE-mini model trained with Jax/Flax on TPU v3-8s during the HuggingFace Community Week',
+    article="<p style='text-align: center'> DALLE-mini by Boris Dayma et al. | <a href='https://github.com/borisdayma/dalle-mini'>GitHub</a></p>",
+    layout='vertical',
+    theme='huggingface',
+    examples=[['an armchair in the shape of an avocado']],
+    server_port=8999).launch(share=True)

app/requirements.txt

@@ -1,11 +0,0 @@
-# Requirements for huggingface spaces
--f https://storage.googleapis.com/jax-releases/jax_releases.html
-jax[cuda111]
-flax
-requests
--e git+https://github.com/huggingface/transformers.git@master#egg=transformers
--e git+https://github.com/huggingface/datasets.git@master#egg=datasets
-flax
-jupyter
-wandb
-ftfy

dalle_mini/vqgan_jax/README.md

@@ -0,0 +1,5 @@
+## vqgan-jax
+
+Files copied from [patil-suraj/vqgan-jax](https://github.com/patil-suraj/vqgan-jax/tree/main/vqgan_jax)
+
+Required for VQGAN Jax model.

dalle_mini/vqgan_jax/convert_pt_model_to_jax.py

@@ -1,109 +0,0 @@
-import re
-
-import jax.numpy as jnp
-from flax.traverse_util import flatten_dict, unflatten_dict
-
-import torch
-
-from modeling_flax_vqgan import VQModel
-from configuration_vqgan import VQGANConfig
-
-
-regex = r"\w+[.]\d+"
-
-
-def rename_key(key):
-    pats = re.findall(regex, key)
-    for pat in pats:
-        key = key.replace(pat, "_".join(pat.split(".")))
-    return key
-
-
-# Adapted from https://github.com/huggingface/transformers/blob/ff5cdc086be1e0c3e2bbad8e3469b34cffb55a85/src/transformers/modeling_flax_pytorch_utils.py#L61
-def convert_pytorch_state_dict_to_flax(pt_state_dict, flax_model):
-    # convert pytorch tensor to numpy
-    pt_state_dict = {k: v.numpy() for k, v in pt_state_dict.items()}
-
-    random_flax_state_dict = flatten_dict(flax_model.params)
-    flax_state_dict = {}
-
-    remove_base_model_prefix = (flax_model.base_model_prefix not in flax_model.params) and (
-        flax_model.base_model_prefix in set([k.split(".")[0] for k in pt_state_dict.keys()])
-    )
-    add_base_model_prefix = (flax_model.base_model_prefix in flax_model.params) and (
-        flax_model.base_model_prefix not in set([k.split(".")[0] for k in pt_state_dict.keys()])
-    )
-
-    # Need to change some parameters name to match Flax names so that we don't have to fork any layer
-    for pt_key, pt_tensor in pt_state_dict.items():
-        pt_tuple_key = tuple(pt_key.split("."))
-
-        has_base_model_prefix = pt_tuple_key[0] == flax_model.base_model_prefix
-        require_base_model_prefix = (flax_model.base_model_prefix,) + pt_tuple_key in random_flax_state_dict
-
-        if remove_base_model_prefix and has_base_model_prefix:
-            pt_tuple_key = pt_tuple_key[1:]
-        elif add_base_model_prefix and require_base_model_prefix:
-            pt_tuple_key = (flax_model.base_model_prefix,) + pt_tuple_key
-
-        # Correctly rename weight parameters
-        if (
-            "norm" in pt_key
-            and (pt_tuple_key[-1] == "bias")
-            and (pt_tuple_key[:-1] + ("bias",) in random_flax_state_dict)
-        ):
-            pt_tensor = pt_tensor[None, None, None, :]
-        elif (
-            "norm" in pt_key
-            and (pt_tuple_key[-1] == "bias")
-            and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
-        ):
-            pt_tuple_key = pt_tuple_key[:-1] + ("scale",)
-            pt_tensor = pt_tensor[None, None, None, :]
-        elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
-            pt_tuple_key = pt_tuple_key[:-1] + ("scale",)
-            pt_tensor = pt_tensor[None, None, None, :]
-        if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
-            pt_tuple_key = pt_tuple_key[:-1] + ("embedding",)
-        elif pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4 and pt_tuple_key not in random_flax_state_dict:
-            # conv layer
-            pt_tuple_key = pt_tuple_key[:-1] + ("kernel",)
-            pt_tensor = pt_tensor.transpose(2, 3, 1, 0)
-        elif pt_tuple_key[-1] == "weight" and pt_tuple_key not in random_flax_state_dict:
-            # linear layer
-            pt_tuple_key = pt_tuple_key[:-1] + ("kernel",)
-            pt_tensor = pt_tensor.T
-        elif pt_tuple_key[-1] == "gamma":
-            pt_tuple_key = pt_tuple_key[:-1] + ("weight",)
-        elif pt_tuple_key[-1] == "beta":
-            pt_tuple_key = pt_tuple_key[:-1] + ("bias",)
-
-        if pt_tuple_key in random_flax_state_dict:
-            if pt_tensor.shape != random_flax_state_dict[pt_tuple_key].shape:
-                raise ValueError(
-                    f"PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape "
-                    f"{random_flax_state_dict[pt_tuple_key].shape}, but is {pt_tensor.shape}."
-                )
-
-        # also add unexpected weight so that warning is thrown
-        flax_state_dict[pt_tuple_key] = jnp.asarray(pt_tensor)
-
-    return unflatten_dict(flax_state_dict)
-
-
-def convert_model(config_path, pt_state_dict_path, save_path):
-    config = VQGANConfig.from_pretrained(config_path)
-    model = VQModel(config)
-
-    state_dict = torch.load(pt_state_dict_path, map_location="cpu")["state_dict"]
-    keys = list(state_dict.keys())
-    for key in keys:
-        if key.startswith("loss"):
-            state_dict.pop(key)
-            continue
-        renamed_key = rename_key(key)
-        state_dict[renamed_key] = state_dict.pop(key)
-
-    state = convert_pytorch_state_dict_to_flax(state_dict, model)
-    model.params = unflatten_dict(state)
-    model.save_pretrained(save_path)

requirements.txt

@@ -1,12 +1,13 @@
-# 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
+# Requirements for huggingface spaces
+-f https://storage.googleapis.com/jax-releases/jax_releases.html
+jax[cuda111]
+flax
 requests
-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
 flax
 jupyter
-
-# Inference
+wandb
 ftfy
+streamlit
+gradio

seq2seq/requirements.txt

@@ -1,8 +1,15 @@
-datasets >= 1.1.3
-jax>=0.2.8
-jaxlib>=0.1.59
-flax>=0.3.4
-optax>=0.0.8
-tensorboard
-nltk
+# 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
+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
+flax
+jupyter
 wandb
+nltk
+optax
+
+# Inference
+ftfy