readme: ## VQGAN-f16-16384
  
  ### Model Description
  
  This is a Flax/JAX implementation of VQGAN, which learns a codebook of context-rich visual parts by leveraging both the use of convolutional methods and transformers. It was introduced in [Taming Transformers for High-Resolution Image Synthesis](https://compvis.github.io/taming-transformers/) ([CVPR paper](https://openaccess.thecvf.com/content/CVPR2021/html/Esser_Taming_Transformers_for_High-Resolution_Image_Synthesis_CVPR_2021_paper.html)).
  
  The model allows the encoding of images as a fixed-length sequence of tokens taken from the codebook.
  
  This version of the model uses a reduction factor `f=16` and a vocabulary of `13,384` tokens.
  
  As an example of how the reduction factor works, images of size `256x256` are encoded to sequences of `256` tokens: `256/16 * 256/16`. Images of `512x512` would result in sequences of `1024` tokens.
  
  ### Datasets Used for Training
  
  * ImageNet. We didn't train this model from scratch. Instead, we started from [a checkpoint pre-trained on ImageNet](https://heibox.uni-heidelberg.de/d/a7530b09fed84f80a887/).
  * [Conceptual Captions 3M](https://ai.google.com/research/ConceptualCaptions/) (CC3M).
  * [OpenAI subset of YFCC100M](https://github.com/openai/CLIP/blob/main/data/yfcc100m.md).
  
  We fine-tuned on CC3M and YFCC100M to improve the encoding quality of people and faces, which are not very well represented in ImageNet. We used a subset of 2,268,720 images from CC3M and YFCC100M for this purpose.
  
  ### Training Process
  
  Finetuning was performed in PyTorch using [taming-transformers](https://github.com/CompVis/taming-transformers). The full training process and model preparation includes these steps:
  
  * Pre-training on ImageNet. Previously performed. We used [this checkpoint](https://heibox.uni-heidelberg.de/d/a7530b09fed84f80a887).
  * Fine-tuning, [Part 1](https://wandb.ai/wandb/hf-flax-dalle-mini/runs/2021-07-09T15-33-11_dalle_vqgan?workspace=user-borisd13).
  * Fine-tuning, [Part 2](https://wandb.ai/wandb/hf-flax-dalle-mini/runs/2021-07-09T21-42-07_dalle_vqgan?workspace=user-borisd13) – continuation from Part 1. The final checkpoint was uploaded to [boris/vqgan_f16_16384](https://huggingface.co/boris/vqgan_f16_16384).
  * Conversion to JAX, which is the model described in this card.
  
  ### How to Use
  
  The checkpoint can be loaded using [Suraj Patil's implementation](https://github.com/patil-suraj/vqgan-jax) of `VQModel`.
  
  * Example notebook, heavily based in work by [Suraj](https://huggingface.co/valhalla): [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/borisdayma/dalle-mini/blob/main/dev/vqgan/JAX_VQGAN_f16_16384_Reconstruction.ipynb)
  
  * Batch encoding using JAX `pmap`, complete example including data loading with PyTorch:
  
  ```python
  # VQGAN-JAX - pmap encoding HowTo
  
  import numpy as np
  
  # For data loading
  import torch
  import torchvision.transforms.functional as TF
  from torch.utils.data import Dataset, DataLoader
  from torchvision.datasets.folder import default_loader
  from torchvision.transforms import InterpolationMode
  
  # For data saving
  from pathlib import Path
  import pandas as pd
  from tqdm import tqdm
  
  import jax
  from jax import pmap
  
  from vqgan_jax.modeling_flax_vqgan import VQModel
  
  ## Params and arguments
  
  # List of paths containing images to encode
  image_list = '/sddata/dalle-mini/CC12M/10k.tsv'
  output_tsv = 'output.tsv'                # Encoded results
  batch_size = 64
  num_workers = 4    # TPU v3-8s have 96 cores, so feel free to increase this number when necessary
  
  # Load model
  model = VQModel.from_pretrained("flax-community/vqgan_f16_16384")
  
  ## Data Loading.
  
  # Simple torch Dataset to load images from paths.
  # You can use your own pipeline instead.
  class ImageDataset(Dataset):
      def __init__(self, image_list_path: str, image_size: int, max_items=None):
          """
          :param image_list_path: Path to a file containing a list of all images. We assume absolute paths for now.
          :param image_size: Image size. Source images will be resized and center-cropped.
          :max_items: Limit dataset size for debugging
          """
          self.image_list = pd.read_csv(image_list_path, sep='\t', header=None)
          if max_items is not None: self.image_list = self.image_list[:max_items]
          self.image_size = image_size
          
      def __len__(self):
          return len(self.image_list)
      
      def _get_raw_image(self, i):
          image_path = Path(self.image_list.iloc[i][0])
          return default_loader(image_path)
      
      def resize_image(self, image):
          s = min(image.size)
          r = self.image_size / s
          s = (round(r * image.size[1]), round(r * image.size[0]))
          image = TF.resize(image, s, interpolation=InterpolationMode.LANCZOS)
          image = TF.center_crop(image, output_size = 2 * [self.image_size])
          image = np.expand_dims(np.array(image), axis=0)
          return image
      
      def __getitem__(self, i):
          image = self._get_r