Merge pull request #36 from borisdayma/chore-cleanup

Chore: repo cleanup
Former-commit-id: 8f1d9a59610dff1f33c1d2b51ed75184776f1a26

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

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 +0,0 @@
-# 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
-
-# Inference
-ftfy

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