Merge branch 'main' into abidlabs/main

.github/workflows/check_size.yml

@@ -14,4 +14,4 @@ jobs:
       - name: Check large files
         uses: ActionsDesk/[email protected]
         with:
-          filesizelimit: 900000 # so we can sync to HF spaces
+          filesizelimit: 9000000 # so we can sync to HF spaces

.gitignore

@@ -1 +1,3 @@
 __pycache__
+.ipynb_checkpoints
+.streamlit

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2021 The DALL·E mini Authors
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

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

app/app.py

@@ -46,7 +46,7 @@ DALL·E mini is an AI model that generates images from any prompt you give!
 <p style='text-align: center'>
 Created by Boris Dayma et al. 2021
 <br/>
-<a href="https://github.com/borisdayma/dalle-mini">GitHub</a> | <a href="https://wandb.ai/dalle-mini/dalle-mini/reports/DALL-E-mini--Vmlldzo4NjIxODA">Project Report</a>
+<a href="https://github.com/borisdayma/dalle-mini" target="_blank">GitHub</a> | <a href="https://wandb.ai/dalle-mini/dalle-mini/reports/DALL-E-mini--Vmlldzo4NjIxODA" target="_blank">Project Report</a>
 </p>
         """, unsafe_allow_html=True)
 

app/gradio/app_gradio.py

@@ -19,7 +19,7 @@ import numpy as np
 import matplotlib.pyplot as plt
 
 
-from dalle_mini.vqgan_jax.modeling_flax_vqgan import VQModel
+from vqgan_jax.modeling_flax_vqgan import VQModel
 from dalle_mini.model import CustomFlaxBartForConditionalGeneration
 
 import gradio as gr

dalle_mini/vqgan_jax/README.md

@@ -1,5 +0,0 @@
-## 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/configuration_vqgan.py

@@ -1,40 +0,0 @@
-from typing import Tuple
-
-from transformers import PretrainedConfig
-
-
-class VQGANConfig(PretrainedConfig):
-    def __init__(
-        self,
-        ch: int = 128,
-        out_ch: int = 3,
-        in_channels: int = 3,
-        num_res_blocks: int = 2,
-        resolution: int = 256,
-        z_channels: int = 256,
-        ch_mult: Tuple = (1, 1, 2, 2, 4),
-        attn_resolutions: int = (16,),
-        n_embed: int = 1024,
-        embed_dim: int = 256,
-        dropout: float = 0.0,
-        double_z: bool = False,
-        resamp_with_conv: bool = True,
-        give_pre_end: bool = False,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.ch = ch
-        self.out_ch = out_ch
-        self.in_channels = in_channels
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.z_channels = z_channels
-        self.ch_mult = list(ch_mult)
-        self.attn_resolutions = list(attn_resolutions)
-        self.n_embed = n_embed
-        self.embed_dim = embed_dim
-        self.dropout = dropout
-        self.double_z = double_z
-        self.resamp_with_conv = resamp_with_conv
-        self.give_pre_end = give_pre_end
-        self.num_resolutions = len(ch_mult)

dalle_mini/vqgan_jax/modeling_flax_vqgan.py

@@ -1,609 +0,0 @@
-# JAX implementation of VQGAN from taming-transformers https://github.com/CompVis/taming-transformers
-
-from functools import partial
-from typing import Tuple
-import math
-
-import jax
-import jax.numpy as jnp
-import numpy as np
-import flax.linen as nn
-from flax.core.frozen_dict import FrozenDict
-
-from transformers.modeling_flax_utils import FlaxPreTrainedModel
-
-from .configuration_vqgan import VQGANConfig
-
-
-class Upsample(nn.Module):
-    in_channels: int
-    with_conv: bool
-    dtype: jnp.dtype = jnp.float32
-
-    def setup(self):
-        if self.with_conv:
-            self.conv = nn.Conv(
-                self.in_channels,
-                kernel_size=(3, 3),
-                strides=(1, 1),
-                padding=((1, 1), (1, 1)),
-                dtype=self.dtype,
-            )
-
-    def __call__(self, hidden_states):
-        batch, height, width, channels = hidden_states.shape
-        hidden_states = jax.image.resize(
-            hidden_states,
-            shape=(batch, height * 2, width * 2, channels),
-            method="nearest",
-        )
-        if self.with_conv:
-            hidden_states = self.conv(hidden_states)
-        return hidden_states
-
-
-class Downsample(nn.Module):
-    in_channels: int
-    with_conv: bool
-    dtype: jnp.dtype = jnp.float32
-
-    def setup(self):
-        if self.with_conv:
-            self.conv = nn.Conv(
-                self.in_channels,
-                kernel_size=(3, 3),
-                strides=(2, 2),
-                padding="VALID",
-                dtype=self.dtype,
-            )
-
-    def __call__(self, hidden_states):
-        if self.with_conv:
-            pad = ((0, 0), (0, 1), (0, 1), (0, 0))  # pad height and width dim
-            hidden_states = jnp.pad(hidden_states, pad_width=pad)
-            hidden_states = self.conv(hidden_states)
-        else:
-            hidden_states = nn.avg_pool(hidden_states, window_shape=(2, 2), strides=(2, 2), padding="VALID")
-        return hidden_states
-
-
-class ResnetBlock(nn.Module):
-    in_channels: int
-    out_channels: int = None
-    use_conv_shortcut: bool = False
-    temb_channels: int = 512
-    dropout_prob: float = 0.0
-    dtype: jnp.dtype = jnp.float32
-
-    def setup(self):
-        self.out_channels_ = self.in_channels if self.out_channels is None else self.out_channels
-
-        self.norm1 = nn.GroupNorm(num_groups=32, epsilon=1e-6)
-        self.conv1 = nn.Conv(
-            self.out_channels_,
-            kernel_size=(3, 3),
-            strides=(1, 1),
-            padding=((1, 1), (1, 1)),
-            dtype=self.dtype,
-        )
-
-        if self.temb_channels:
-            self.temb_proj = nn.Dense(self.out_channels_, dtype=self.dtype)
-
-        self.norm2 = nn.GroupNorm(num_groups=32, epsilon=1e-6)
-        self.dropout = nn.Dropout(self.dropout_prob)
-        self.conv2 = nn.Conv(
-            self.out_channels_,
-            kernel_size=(3, 3),
-            strides=(1, 1),
-            padding=((1, 1), (1, 1)),
-            dtype=self.dtype,
-        )
-
-        if self.in_channels != self.out_channels_:
-            if self.use_conv_shortcut:
-                self.conv_shortcut = nn.Conv(
-                    self.out_channels_,
-                    kernel_size=(3, 3),
-                    strides=(1, 1),
-                    padding=((1, 1), (1, 1)),
-                    dtype=self.dtype,
-                )
-            else:
-                self.nin_shortcut = nn.Conv(
-                    self.out_channels_,
-                    kernel_size=(1, 1),
-                    strides=(1, 1),
-                    padding="VALID",
-                    dtype=self.dtype,
-                )
-
-    def __call__(self, hidden_states, temb=None, deterministic: bool = True):
-        residual = hidden_states
-        hidden_states = self.norm1(hidden_states)
-        hidden_states = nn.swish(hidden_states)
-        hidden_states = self.conv1(hidden_states)
-
-        if temb is not None:
-            hidden_states = hidden_states + self.temb_proj(nn.swish(temb))[:, :, None, None]  # TODO: check shapes
-
-        hidden_states = self.norm2(hidden_states)
-        hidden_states = nn.swish(hidden_states)
-        hidden_states = self.dropout(hidden_states, deterministic)
-        hidden_states = self.conv2(hidden_states)
-
-        if self.in_channels != self.out_channels_:
-            if self.use_conv_shortcut:
-                residual = self.conv_shortcut(residual)
-            else:
-                residual = self.nin_shortcut(residual)
-
-        return hidden_states + residual
-
-
-class AttnBlock(nn.Module):
-    in_channels: int
-    dtype: jnp.dtype = jnp.float32
-
-    def setup(self):
-        conv = partial(
-            nn.Conv, self.in_channels, kernel_size=(1, 1), strides=(1, 1), padding="VALID", dtype=self.dtype
-        )
-
-        self.norm = nn.GroupNorm(num_groups=32, epsilon=1e-6)
-        self.q, self.k, self.v = conv(), conv(), conv()
-        self.proj_out = conv()
-
-    def __call__(self, hidden_states):
-        residual = hidden_states
-        hidden_states = self.norm(hidden_states)
-
-        query = self.q(hidden_states)
-        key = self.k(hidden_states)
-        value = self.v(hidden_states)
-
-        # compute attentions
-        batch, height, width, channels = query.shape
-        query = query.reshape((batch, height * width, channels))
-        key = key.reshape((batch, height * width, channels))
-        attn_weights = jnp.einsum("...qc,...kc->...qk", query, key)
-        attn_weights = attn_weights * (int(channels) ** -0.5)
-        attn_weights = nn.softmax(attn_weights, axis=2)
-
-        ## attend to values
-        value = value.reshape((batch, height * width, channels))
-        hidden_states = jnp.einsum("...kc,...qk->...qc", value, attn_weights)
-        hidden_states = hidden_states.reshape((batch, height, width, channels))
-
-        hidden_states = self.proj_out(hidden_states)
-        hidden_states = hidden_states + residual
-        return hidden_states
-
-
-class UpsamplingBlock(nn.Module):
-    config: VQGANConfig
-    curr_res: int
-    block_idx: int
-    dtype: jnp.dtype = jnp.float32
-
-    def setup(self):
-        if self.block_idx == self.config.num_resolutions - 1:
-            block_in = self.config.ch * self.config.ch_mult[-1]
-        else:
-            block_in = self.config.ch * self.config.ch_mult[self.block_idx + 1]
-
-        block_out = self.config.ch * self.config.ch_mult[self.block_idx]
-        self.temb_ch = 0
-
-        res_blocks = []
-        attn_blocks = []
-        for _ in range(self.config.num_res_blocks + 1):
-            res_blocks.append(
-                ResnetBlock(
-                    block_in, block_out, temb_channels=self.temb_ch, dropout_prob=self.config.dropout, dtype=self.dtype
-                )
-            )
-            block_in = block_out
-            if self.curr_res in self.config.attn_resolutions:
-                attn_blocks.append(AttnBlock(block_in, dtype=self.dtype))
-
-        self.block = res_blocks
-        self.attn = attn_blocks
-
-        self.upsample = None
-        if self.block_idx != 0:
-            self.upsample = Upsample(block_in, self.config.resamp_with_conv, dtype=self.dtype)
-
-    def __call__(self, hidden_states, temb=None, deterministic: bool = True):
-        for res_block in self.block:
-            hidden_states = res_block(hidden_states, temb, deterministic=deterministic)
-            for attn_block in self.attn:
-                hidden_states = attn_block(hidden_states)
-
-        if self.upsample is not None:
-            hidden_states = self.upsample(hidden_states)
-
-        return hidden_states
-
-
-class DownsamplingBlock(nn.Module):
-    config: VQGANConfig
-    curr_res: int
-    block_idx: int
-    dtype: jnp.dtype = jnp.float32
-
-    def setup(self):
-        in_ch_mult = (1,) + tuple(self.config.ch_mult)
-        block_in = self.config.ch * in_ch_mult[self.block_idx]
-        block_out = self.config.ch * self.config.ch_mult[self.block_idx]
-        self.temb_ch = 0
-
-        res_blocks = []
-        attn_blocks = []
-        for _ in range(self.config.num_res_blocks):
-            res_blocks.append(
-                ResnetBlock(
-                    block_in, block_out, temb_channels=self.temb_ch, dropout_prob=self.config.dropout, dtype=self.dtype
-                )
-            )
-            block_in = block_out
-            if self.curr_res in self.config.attn_resolutions:
-                attn_blocks.append(AttnBlock(block_in, dtype=self.dtype))
-
-        self.block = res_blocks
-        self.attn = attn_blocks
-
-        self.downsample = None
-        if self.block_idx != self.config.num_resolutions - 1:
-            self.downsample = Downsample(block_in, self.config.resamp_with_conv, dtype=self.dtype)
-
-    def __call__(self, hidden_states, temb=None, deterministic: bool = True):
-        for res_block in self.block:
-            hidden_states = res_block(hidden_states, temb, deterministic=deterministic)
-            for attn_block in self.attn:
-                hidden_states = attn_block(hidden_states)
-
-        if self.downsample is not None:
-            hidden_states = self.downsample(hidden_states)
-
-        return hidden_states
-
-
-class MidBlock(nn.Module):
-    in_channels: int
-    temb_channels: int
-    dropout: float
-    dtype: jnp.dtype = jnp.float32
-
-    def setup(self):
-        self.block_1 = ResnetBlock(
-            self.in_channels,
-            self.in_channels,
-            temb_channels=self.temb_channels,
-            dropout_prob=self.dropout,
-            dtype=self.dtype,
-        )
-        self.attn_1 = AttnBlock(self.in_channels, dtype=self.dtype)
-        self.block_2 = ResnetBlock(
-            self.in_channels,
-            self.in_channels,
-            temb_channels=self.temb_channels,
-            dropout_prob=self.dropout,
-            dtype=self.dtype,
-        )
-
-    def __call__(self, hidden_states, temb=None, deterministic: bool = True):
-        hidden_states = self.block_1(hidden_states, temb, deterministic=deterministic)
-        hidden_states = self.attn_1(hidden_states)
-        hidden_states = self.block_2(hidden_states, temb, deterministic=deterministic)
-        return hidden_states
-
-
-class Encoder(nn.Module):
-    config: VQGANConfig
-    dtype: jnp.dtype = jnp.float32
-
-    def setup(self):
-        self.temb_ch = 0
-
-        # downsampling
-        self.conv_in = nn.Conv(
-            self.config.ch,
-            kernel_size=(3, 3),
-            strides=(1, 1),
-            padding=((1, 1), (1, 1)),
-            dtype=self.dtype,
-        )
-
-        curr_res = self.config.resolution
-        downsample_blocks = []
-        for i_level in range(self.config.num_resolutions):
-            downsample_blocks.append(DownsamplingBlock(self.config, curr_res, block_idx=i_level, dtype=self.dtype))
-
-            if i_level != self.config.num_resolutions - 1:
-                curr_res = curr_res // 2
-        self.down = downsample_blocks
-
-        # middle
-        mid_channels = self.config.ch * self.config.ch_mult[-1]
-        self.mid = MidBlock(mid_channels, self.temb_ch, self.config.dropout, dtype=self.dtype)
-
-        # end
-        self.norm_out = nn.GroupNorm(num_groups=32, epsilon=1e-6)
-        self.conv_out = nn.Conv(
-            2 * self.config.z_channels if self.config.double_z else self.config.z_channels,
-            kernel_size=(3, 3),
-            strides=(1, 1),
-            padding=((1, 1), (1, 1)),
-            dtype=self.dtype,
-        )
-
-    def __call__(self, pixel_values, deterministic: bool = True):
-        # timestep embedding
-        temb = None
-
-        # downsampling
-        hidden_states = self.conv_in(pixel_values)
-        for block in self.down:
-            hidden_states = block(hidden_states, temb, deterministic=deterministic)
-
-        # middle
-        hidden_states = self.mid(hidden_states, temb, deterministic=deterministic)
-
-        # end
-        hidden_states = self.norm_out(hidden_states)
-        hidden_states = nn.swish(hidden_states)
-        hidden_states = self.conv_out(hidden_states)
-
-        return hidden_states
-
-
-class Decoder(nn.Module):
-    config: VQGANConfig
-    dtype: jnp.dtype = jnp.float32
-
-    def setup(self):
-        self.temb_ch = 0
-
-        # compute in_ch_mult, block_in and curr_res at lowest res
-        block_in = self.config.ch * self.config.ch_mult[self.config.num_resolutions - 1]
-        curr_res = self.config.resolution // 2 ** (self.config.num_resolutions - 1)
-        self.z_shape = (1, self.config.z_channels, curr_res, curr_res)
-        print("Working with z of shape {} = {} dimensions.".format(self.z_shape, np.prod(self.z_shape)))
-
-        # z to block_in
-        self.conv_in = nn.Conv(
-            block_in,
-            kernel_size=(3, 3),
-            strides=(1, 1),
-            padding=((1, 1), (1, 1)),
-            dtype=self.dtype,
-        )
-
-        # middle
-        self.mid = MidBlock(block_in, self.temb_ch, self.config.dropout, dtype=self.dtype)
-
-        # upsampling
-        upsample_blocks = []
-        for i_level in reversed(range(self.config.num_resolutions)):
-            upsample_blocks.append(UpsamplingBlock(self.config, curr_res, block_idx=i_level, dtype=self.dtype))
-            if i_level != 0:
-                curr_res = curr_res * 2
-        self.up = list(reversed(upsample_blocks))  # reverse to get consistent order
-
-        # end
-        self.norm_out = nn.GroupNorm(num_groups=32, epsilon=1e-6)
-        self.conv_out = nn.Conv(
-            self.config.out_ch,
-            kernel_size=(3, 3),
-            strides=(1, 1),
-            padding=((1, 1), (1, 1)),
-            dtype=self.dtype,
-        )
-
-    def __call__(self, hidden_states, deterministic: bool = True):
-        # timestep embedding
-        temb = None
-
-        # z to block_in
-        hidden_states = self.conv_in(hidden_states)
-
-        # middle
-        hidden_states = self.mid(hidden_states, temb, deterministic=deterministic)
-
-        # upsampling
-        for block in reversed(self.up):
-            hidden_states = block(hidden_states, temb, deterministic=deterministic)
-
-        # end
-        if self.config.give_pre_end:
-            return hidden_states
-
-        hidden_states = self.norm_out(hidden_states)
-        hidden_states = nn.swish(hidden_states)
-        hidden_states = self.conv_out(hidden_states)
-
-        return hidden_states
-
-
-class VectorQuantizer(nn.Module):
-    """
-    see https://github.com/MishaLaskin/vqvae/blob/d761a999e2267766400dc646d82d3ac3657771d4/models/quantizer.py
-    ____________________________________________
-    Discretization bottleneck part of the VQ-VAE.
-    Inputs:
-    - n_e : number of embeddings
-    - e_dim : dimension of embedding
-    - beta : commitment cost used in loss term, beta * ||z_e(x)-sg[e]||^2
-    _____________________________________________
-    """
-
-    config: VQGANConfig
-    dtype: jnp.dtype = jnp.float32
-
-    def setup(self):
-        self.embedding = nn.Embed(self.config.n_embed, self.config.embed_dim, dtype=self.dtype)  # TODO: init
-
-    def __call__(self, hidden_states):
-        """
-        Inputs the output of the encoder network z and maps it to a discrete
-        one-hot vector that is the index of the closest embedding vector e_j
-        z (continuous) -> z_q (discrete)
-        z.shape = (batch, channel, height, width)
-        quantization pipeline:
-            1. get encoder input (B,C,H,W)
-            2. flatten input to (B*H*W,C)
-        """
-        #  flatten
-        hidden_states_flattended = hidden_states.reshape((-1, self.config.embed_dim))
-
-        # dummy op to init the weights, so we can access them below
-        self.embedding(jnp.ones((1, 1), dtype="i4"))
-
-        # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
-        emb_weights = self.variables["params"]["embedding"]["embedding"]
-        distance = (
-            jnp.sum(hidden_states_flattended ** 2, axis=1, keepdims=True)
-            + jnp.sum(emb_weights ** 2, axis=1)
-            - 2 * jnp.dot(hidden_states_flattended, emb_weights.T)
-        )
-
-        # get quantized latent vectors
-        min_encoding_indices = jnp.argmin(distance, axis=1)
-        z_q = self.embedding(min_encoding_indices).reshape(hidden_states.shape)
-
-        # reshape to (batch, num_tokens)
-        min_encoding_indices = min_encoding_indices.reshape(hidden_states.shape[0], -1)
-
-        # compute the codebook_loss (q_loss) outside the model
-        # here we return the embeddings and indices
-        return z_q, min_encoding_indices
-
-    def get_codebook_entry(self, indices, shape=None):
-        # indices are expected to be of shape (batch, num_tokens)
-        # get quantized latent vectors
-        batch, num_tokens = indices.shape
-        z_q = self.embedding(indices)
-        z_q = z_q.reshape(batch, int(math.sqrt(num_tokens)), int(math.sqrt(num_tokens)), -1)
-        return z_q
-
-
-class VQModule(nn.Module):
-    config: VQGANConfig
-    dtype: jnp.dtype = jnp.float32
-
-    def setup(self):
-        self.encoder = Encoder(self.config, dtype=self.dtype)
-        self.decoder = Decoder(self.config, dtype=self.dtype)
-        self.quantize = VectorQuantizer(self.config, dtype=self.dtype)
-        self.quant_conv = nn.Conv(
-            self.config.embed_dim,
-            kernel_size=(1, 1),
-            strides=(1, 1),
-            padding="VALID",
-            dtype=self.dtype,
-        )
-        self.post_quant_conv = nn.Conv(
-            self.config.z_channels,
-            kernel_size=(1, 1),
-            strides=(1, 1),
-            padding="VALID",
-            dtype=self.dtype,
-        )
-
-    def encode(self, pixel_values, deterministic: bool = True):
-        hidden_states = self.encoder(pixel_values, deterministic=deterministic)
-        hidden_states = self.quant_conv(hidden_states)
-        quant_states, indices = self.quantize(hidden_states)
-        return quant_states, indices
-
-    def decode(self, hidden_states, deterministic: bool = True):
-        hidden_states = self.post_quant_conv(hidden_states)
-        hidden_states = self.decoder(hidden_states, deterministic=deterministic)
-        return hidden_states
-
-    def decode_code(self, code_b):
-        hidden_states = self.quantize.get_codebook_entry(code_b)
-        hidden_states = self.decode(hidden_states)
-        return hidden_states
-
-    def __call__(self, pixel_values, deterministic: bool = True):
-        quant_states, indices = self.encode(pixel_values, deterministic)
-        hidden_states = self.decode(quant_states, deterministic)
-        return hidden_states, indices
-
-
-class VQGANPreTrainedModel(FlaxPreTrainedModel):
-    """
-    An abstract class to handle weights initialization and a simple interface
-    for downloading and loading pretrained models.
-    """
-
-    config_class = VQGANConfig
-    base_model_prefix = "model"
-    module_class: nn.Module = None
-
-    def __init__(
-        self,
-        config: VQGANConfig,
-        input_shape: Tuple = (1, 256, 256, 3),
-        seed: int = 0,
-        dtype: jnp.dtype = jnp.float32,
-        **kwargs,
-    ):
-        module = self.module_class(config=config, dtype=dtype, **kwargs)
-        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)
-
-    def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple) -> FrozenDict:
-        # init input tensors
-        pixel_values = jnp.zeros(input_shape, dtype=jnp.float32)
-        params_rng, dropout_rng = jax.random.split(rng)
-        rngs = {"params": params_rng, "dropout": dropout_rng}
-
-        return self.module.init(rngs, pixel_values)["params"]
-
-    def encode(self, pixel_values, params: dict = None, dropout_rng: jax.random.PRNGKey = None, train: bool = False):
-        # Handle any PRNG if needed
-        rngs = {"dropout": dropout_rng} if dropout_rng is not None else {}
-
-        return self.module.apply(
-            {"params": params or self.params}, jnp.array(pixel_values), not train, rngs=rngs, method=self.module.encode
-        )
-
-    def decode(self, hidden_states, params: dict = None, dropout_rng: jax.random.PRNGKey = None, train: bool = False):
-        # Handle any PRNG if needed
-        rngs = {"dropout": dropout_rng} if dropout_rng is not None else {}
-
-        return self.module.apply(
-            {"params": params or self.params},
-            jnp.array(hidden_states),
-            not train,
-            rngs=rngs,
-            method=self.module.decode,
-        )
-
-    def decode_code(self, indices, params: dict = None):
-        return self.module.apply(
-            {"params": params or self.params}, jnp.array(indices, dtype="i4"), method=self.module.decode_code
-        )
-
-    def __call__(
-        self,
-        pixel_values,
-        params: dict = None,
-        dropout_rng: jax.random.PRNGKey = None,
-        train: bool = False,
-    ):
-        # Handle any PRNG if needed
-        rngs = {"dropout": dropout_rng} if dropout_rng is not None else {}
-
-        return self.module.apply(
-            {"params": params or self.params},
-            jnp.array(pixel_values),
-            not train,
-            rngs=rngs,
-        )
-
-
-class VQModel(VQGANPreTrainedModel):
-    module_class = VQModule

dev/{notebooks/encoding → encoding}/vqgan-jax-encoding-with-captions.ipynb

@@ -50,14 +50,6 @@
     "## VQGAN-JAX model"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "id": "bb408f6c",
-   "metadata": {},
-   "source": [
-    "`dalle_mini` is a local package that contains the VQGAN-JAX model and other utilities."
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 2,
@@ -65,7 +57,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from dalle_mini.vqgan_jax.modeling_flax_vqgan import VQModel"
+    "from vqgan_jax.modeling_flax_vqgan import VQModel"
    ]
   },
   {

dev/{notebooks/encoding → encoding}/vqgan-jax-encoding-yfcc100m.ipynb

@@ -52,14 +52,6 @@
     "## VQGAN-JAX model"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "id": "bb408f6c",
-   "metadata": {},
-   "source": [
-    "`dalle_mini` is a local package that contains the VQGAN-JAX model and other utilities."
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 93,
@@ -67,7 +59,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from dalle_mini.vqgan_jax.modeling_flax_vqgan import VQModel"
+    "from vqgan_jax.modeling_flax_vqgan import VQModel"
    ]
   },
   {
@@ -1111,9 +1103,13 @@
   }
  ],
  "metadata": {
+  "interpreter": {
+   "hash": "db471c52d602b4f5f40ecaf278e88ccfef85c29d0a1a07185b0d51fc7acf4e26"
+  },
   "kernelspec": {
-   "name": "python3",
-   "display_name": "Python 3.9.0 64-bit ('Python39')"
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -1125,12 +1121,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.0"
-  },
-  "interpreter": {
-   "hash": "db471c52d602b4f5f40ecaf278e88ccfef85c29d0a1a07185b0d51fc7acf4e26"
+   "version": "3.8.10"
   }
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

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

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

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

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

dev/notebooks/README.md

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

dev/notebooks/demo/CustomBARTv4b_model-generate.ipynb

@@ -1,394 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "ewer-Q-0w2xA"
-   },
-   "source": [
-    "# Installation"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "NpsF9ipLLl2s",
-    "outputId": "10bf54aa-b89d-4e42-9777-bc97b00a5f32"
-   },
-   "outputs": [],
-   "source": [
-    "!pip install git+https://github.com/huggingface/transformers/\n",
-    "!pip install git+https://github.com/google/flax"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "M1wVkrpjU6zO"
-   },
-   "outputs": [],
-   "source": [
-    "%load_ext autoreload\n",
-    "%autoreload 2"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "t47CH1H_IOT8"
-   },
-   "source": [
-    "# Custom BART Model"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "9jQnM6S2vCpn"
-   },
-   "outputs": [],
-   "source": [
-    "# TODO: set those args in a config file\n",
-    "OUTPUT_VOCAB_SIZE = 16384 + 1  # encoded image token space + 1 for bos\n",
-    "OUTPUT_LENGTH = 256 + 1  # number of encoded tokens + 1 for bos\n",
-    "BOS_TOKEN_ID = 16384\n",
-    "BASE_MODEL = 'facebook/bart-large'"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "_eEaJVxAKpV5"
-   },
-   "outputs": [],
-   "source": [
-    "import jax\n",
-    "import flax.linen as nn\n",
-    "\n",
-    "from transformers.models.bart.modeling_flax_bart import *\n",
-    "from transformers import BartTokenizer, FlaxBartForConditionalGeneration\n",
-    "\n",
-    "class CustomFlaxBartModule(FlaxBartModule):\n",
-    "    def setup(self):\n",
-    "        # we keep shared to easily load pre-trained weights\n",
-    "        self.shared = nn.Embed(\n",
-    "            self.config.vocab_size,\n",
-    "            self.config.d_model,\n",
-    "            embedding_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),\n",
-    "            dtype=self.dtype,\n",
-    "        )\n",
-    "        # a separate embedding is used for the decoder\n",
-    "        self.decoder_embed = nn.Embed(\n",
-    "            OUTPUT_VOCAB_SIZE,\n",
-    "            self.config.d_model,\n",
-    "            embedding_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),\n",
-    "            dtype=self.dtype,\n",
-    "        )\n",
-    "        self.encoder = FlaxBartEncoder(self.config, dtype=self.dtype, embed_tokens=self.shared)\n",
-    "\n",
-    "        # the decoder has a different config\n",
-    "        decoder_config = BartConfig(self.config.to_dict())\n",
-    "        decoder_config.max_position_embeddings = OUTPUT_LENGTH\n",
-    "        decoder_config.vocab_size = OUTPUT_VOCAB_SIZE\n",
-    "        self.decoder = FlaxBartDecoder(decoder_config, dtype=self.dtype, embed_tokens=self.decoder_embed)\n",
-    "\n",
-    "class CustomFlaxBartForConditionalGenerationModule(FlaxBartForConditionalGenerationModule):\n",
-    "    def setup(self):\n",
-    "        self.model = CustomFlaxBartModule(config=self.config, dtype=self.dtype)\n",
-    "        self.lm_head = nn.Dense(\n",
-    "            OUTPUT_VOCAB_SIZE,\n",
-    "            use_bias=False,\n",
-    "            dtype=self.dtype,\n",
-    "            kernel_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),\n",
-    "        )\n",
-    "        self.final_logits_bias = self.param(\"final_logits_bias\", self.bias_init, (1, OUTPUT_VOCAB_SIZE))\n",
-    "\n",
-    "class CustomFlaxBartForConditionalGeneration(FlaxBartForConditionalGeneration):\n",
-    "    module_class = CustomFlaxBartForConditionalGenerationModule"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "S7CP9Td9m2ge",
-    "outputId": "5638ef68-9c40-46f7-90ba-a4d05b61360d"
-   },
-   "outputs": [],
-   "source": [
-    "# load pre-trained model for encoder weights\n",
-    "base_model = FlaxBartForConditionalGeneration.from_pretrained(BASE_MODEL)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "6lmynR-poceH"
-   },
-   "outputs": [],
-   "source": [
-    "# set up our new model config\n",
-    "config = BartConfig.from_pretrained(BASE_MODEL)\n",
-    "config.tie_word_embeddings = False\n",
-    "config.decoder_start_token_id = BOS_TOKEN_ID\n",
-    "config.bos_token_id = BOS_TOKEN_ID  # should not be used\n",
-    "config.pos_token_id = BOS_TOKEN_ID  # should not be used\n",
-    "#config.eos_token_id = None  # prevents generation from stopping until we reach max_length"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "_6-XKK40oEfP"
-   },
-   "outputs": [],
-   "source": [
-    "# create our model and initialize it randomly\n",
-    "model = CustomFlaxBartForConditionalGeneration(config)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "-r_hZestr-NR"
-   },
-   "outputs": [],
-   "source": [
-    "# use pretrained weights\n",
-    "model.params['model']['encoder'] = base_model.params['model']['encoder']\n",
-    "model.params['model']['shared'] = base_model.params['model']['shared']"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "5NEX8f62sVjx"
-   },
-   "outputs": [],
-   "source": [
-    "# no need for base_model anymore\n",
-    "del base_model"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "Jz032w73nHEf",
-    "outputId": "994d8e85-bff7-480b-8b69-f69dedc15c49"
-   },
-   "outputs": [],
-   "source": [
-    "# we verify that the shape has not been modified\n",
-    "model.params['final_logits_bias'].shape"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "zLl24Ez5t7x1"
-   },
-   "source": [
-    "## Inference"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "XLLA2NK3uDQr"
-   },
-   "outputs": [],
-   "source": [
-    "tokenizer = BartTokenizer.from_pretrained(BASE_MODEL)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "Ntow53I_t81D",
-    "outputId": "59289cdd-1429-4720-cc87-88810c4b99ac"
-   },
-   "outputs": [],
-   "source": [
-    "text = \"My friends are cool but they eat too many carbs.\"\n",
-    "inputs = tokenizer(text, max_length=1024, return_tensors='jax')\n",
-    "encoder_outputs = model.encode(**inputs)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "vcRNJnJ_uJOJ",
-    "outputId": "025afd54-7908-4a9c-fb59-e40bd3458711"
-   },
-   "outputs": [],
-   "source": [
-    "decoder_start_token_id = model.config.decoder_start_token_id\n",
-    "decoder_start_token_id"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "6QWmEwL_uMld"
-   },
-   "outputs": [],
-   "source": [
-    "decoder_input_ids = jnp.ones((inputs.input_ids.shape[0], 1), dtype=\"i4\") * decoder_start_token_id\n",
-    "outputs = model.decode(decoder_input_ids, encoder_outputs)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "c_ys3yWBothF",
-    "outputId": "40d4d584-e0a8-44cb-bbea-0ffa38d50a53"
-   },
-   "outputs": [],
-   "source": [
-    "outputs"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "O6s0wtB_uTC_",
-    "outputId": "bc0e9e80-e346-4e99-d28e-3f658eda1f66"
-   },
-   "outputs": [],
-   "source": [
-    "outputs.logits.shape"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "ELzemGP3uBzy",
-    "outputId": "dc12f98a-1ccf-450d-ba2a-9c29d7d14885"
-   },
-   "outputs": [],
-   "source": [
-    "outputs.logits.argmax(axis=-1)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "fQjikkGEunpx",
-    "outputId": "3dba0209-ad4e-4069-be38-6c599c677ef1"
-   },
-   "outputs": [],
-   "source": [
-    "model.config.bos_token_id, model.config.eos_token_id, model.config.pad_token_id"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "P32mJJSbrU1F"
-   },
-   "outputs": [],
-   "source": [
-    "input_ids_test = tokenizer.encode('I enjoy walking with my cute dog', return_tensors='jax')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "C7cHbIHruELT"
-   },
-   "outputs": [],
-   "source": [
-    "greedy_output = model.generate(input_ids_test, max_length=50)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "jYugh9cOuwc9",
-    "outputId": "19c3a2ee-e7bc-4f1f-9c86-06bd7337b537"
-   },
-   "outputs": [],
-   "source": [
-    "greedy_output[0]"
-   ]
-  }
- ],
- "metadata": {
-  "accelerator": "TPU",
-  "colab": {
-   "collapsed_sections": [],
-   "machine_shape": "hm",
-   "name": "CustomBARTv4b-model-generate.ipynb",
-   "provenance": []
-  },
-  "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.8.5"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 4
-}

dev/notebooks/demo/demo_notebook.ipynb

@@ -1,387 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "ewer-Q-0w2xA"
-   },
-   "source": [
-    "# Installation"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "NpsF9ipLLl2s",
-    "outputId": "10bf54aa-b89d-4e42-9777-bc97b00a5f32"
-   },
-   "outputs": [],
-   "source": [
-    "#!pip install git+https://github.com/huggingface/transformers/\n",
-    "#!pip install git+https://github.com/google/flax"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "M1wVkrpjU6zO"
-   },
-   "outputs": [],
-   "source": [
-    "%load_ext autoreload\n",
-    "%autoreload 2"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "%cd ../../vqgan-jax"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "t47CH1H_IOT8"
-   },
-   "source": [
-    "# Custom BART Model"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "9jQnM6S2vCpn"
-   },
-   "outputs": [],
-   "source": [
-    "# TODO: set those args in a config file\n",
-    "OUTPUT_VOCAB_SIZE = 16384 + 1  # encoded image token space + 1 for bos\n",
-    "OUTPUT_LENGTH = 256 + 1  # number of encoded tokens + 1 for bos\n",
-    "BOS_TOKEN_ID = 16384\n",
-    "BASE_MODEL = 'facebook/bart-large'"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "_eEaJVxAKpV5"
-   },
-   "outputs": [],
-   "source": [
-    "import jax\n",
-    "import flax.linen as nn\n",
-    "\n",
-    "from transformers.models.bart.modeling_flax_bart import *\n",
-    "from transformers import BartTokenizer, FlaxBartForConditionalGeneration\n",
-    "\n",
-    "class CustomFlaxBartModule(FlaxBartModule):\n",
-    "    def setup(self):\n",
-    "        # we keep shared to easily load pre-trained weights\n",
-    "        self.shared = nn.Embed(\n",
-    "            self.config.vocab_size,\n",
-    "            self.config.d_model,\n",
-    "            embedding_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),\n",
-    "            dtype=self.dtype,\n",
-    "        )\n",
-    "        # a separate embedding is used for the decoder\n",
-    "        self.decoder_embed = nn.Embed(\n",
-    "            OUTPUT_VOCAB_SIZE,\n",
-    "            self.config.d_model,\n",
-    "            embedding_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),\n",
-    "            dtype=self.dtype,\n",
-    "        )\n",
-    "        self.encoder = FlaxBartEncoder(self.config, dtype=self.dtype, embed_tokens=self.shared)\n",
-    "\n",
-    "        # the decoder has a different config\n",
-    "        decoder_config = BartConfig(self.config.to_dict())\n",
-    "        decoder_config.max_position_embeddings = OUTPUT_LENGTH\n",
-    "        decoder_config.vocab_size = OUTPUT_VOCAB_SIZE\n",
-    "        self.decoder = FlaxBartDecoder(decoder_config, dtype=self.dtype, embed_tokens=self.decoder_embed)\n",
-    "\n",
-    "class CustomFlaxBartForConditionalGenerationModule(FlaxBartForConditionalGenerationModule):\n",
-    "    def setup(self):\n",
-    "        self.model = CustomFlaxBartModule(config=self.config, dtype=self.dtype)\n",
-    "        self.lm_head = nn.Dense(\n",
-    "            OUTPUT_VOCAB_SIZE,\n",
-    "            use_bias=False,\n",
-    "            dtype=self.dtype,\n",
-    "            kernel_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),\n",
-    "        )\n",
-    "        self.final_logits_bias = self.param(\"final_logits_bias\", self.bias_init, (1, OUTPUT_VOCAB_SIZE))\n",
-    "\n",
-    "class CustomFlaxBartForConditionalGeneration(FlaxBartForConditionalGeneration):\n",
-    "    module_class = CustomFlaxBartForConditionalGenerationModule"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "scrolled": true
-   },
-   "outputs": [],
-   "source": [
-    "import wandb\n",
-    "run = wandb.init()\n",
-    "artifact = run.use_artifact('wandb/hf-flax-dalle-mini/model-1ef8yxby:latest', type='bart_model')\n",
-    "artifact_dir = artifact.download()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "_6-XKK40oEfP",
-    "scrolled": true
-   },
-   "outputs": [],
-   "source": [
-    "# create our model and initialize it randomly\n",
-    "model = CustomFlaxBartForConditionalGeneration.from_pretrained(artifact_dir)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model.config.forced_bos_token_id = None"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "Jz032w73nHEf",
-    "outputId": "994d8e85-bff7-480b-8b69-f69dedc15c49"
-   },
-   "outputs": [],
-   "source": [
-    "# we verify that the shape has not been modified\n",
-    "model.params['final_logits_bias'].shape"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "zLl24Ez5t7x1"
-   },
-   "source": [
-    "## Inference"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "XLLA2NK3uDQr"
-   },
-   "outputs": [],
-   "source": [
-    "tokenizer = BartTokenizer.from_pretrained(BASE_MODEL)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "input_text = ['I enjoy walking with my cute dog']*8"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "P32mJJSbrU1F"
-   },
-   "outputs": [],
-   "source": [
-    "input_ids_test = tokenizer(input_text, return_tensors='jax')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "input_ids_test"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "C7cHbIHruELT"
-   },
-   "outputs": [],
-   "source": [
-    "greedy_output = model.generate(input_ids_test['input_ids'], max_length=257)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "greedy_output[0].shape"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "jYugh9cOuwc9",
-    "outputId": "19c3a2ee-e7bc-4f1f-9c86-06bd7337b537"
-   },
-   "outputs": [],
-   "source": [
-    "greedy_output[0]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "greedy_output[0][0]"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# VGAN Jax"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import io\n",
-    "\n",
-    "import requests\n",
-    "from PIL import Image\n",
-    "import numpy as np\n",
-    "\n",
-    "import torch\n",
-    "import torchvision.transforms as T\n",
-    "import torchvision.transforms.functional as TF\n",
-    "from torchvision.transforms import InterpolationMode"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from modeling_flax_vqgan import VQModel"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def custom_to_pil(x):\n",
-    "    x = np.clip(x, 0., 1.)\n",
-    "    x = (255*x).astype(np.uint8)\n",
-    "    x = Image.fromarray(x)\n",
-    "    if not x.mode == \"RGB\":\n",
-    "        x = x.convert(\"RGB\")\n",
-    "    return x"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "Jz032w73nHEf",
-    "outputId": "994d8e85-bff7-480b-8b69-f69dedc15c49",
-    "scrolled": true
-   },
-   "outputs": [],
-   "source": [
-    "model = VQModel.from_pretrained(\"flax-community/vqgan_f16_16384\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def get_images(indices, model):\n",
-    "    indices =  indices[:, 1:]\n",
-    "    print(indices.shape)\n",
-    "    img = model.decode_code(indices)\n",
-    "    return img"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "custom_to_pil(np.asarray(get_images(jnp.expand_dims(greedy_output[0][0],0), model)[0]))"
-   ]
-  }
- ],
- "metadata": {
-  "accelerator": "TPU",
-  "colab": {
-   "collapsed_sections": [],
-   "machine_shape": "hm",
-   "name": "CustomBARTv4b-model-generate.ipynb",
-   "provenance": []
-  },
-  "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.8.5"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 4
-}

dev/notebooks/demo/model-sweep.py

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

dev/notebooks/demo/tpu-demo.ipynb

@@ -1,455 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "id": "f6d33374",
-   "metadata": {},
-   "source": [
-    "# Test notebook with CLIP scoring"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "6eb74941-bb4d-4d7e-97f1-d5a3a07672bf",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# !pip install flax transformers\n",
-    "# !git clone https://github.com/patil-suraj/vqgan-jax.git"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "41db7534-f589-4b63-9165-9c9799e1b06e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import random\n",
-    "\n",
-    "import jax\n",
-    "import flax.linen as nn\n",
-    "from flax.training.common_utils import shard\n",
-    "from flax.jax_utils import replicate, unreplicate\n",
-    "\n",
-    "from transformers.models.bart.modeling_flax_bart import *\n",
-    "from transformers import BartTokenizer, FlaxBartForConditionalGeneration\n",
-    "\n",
-    "import io\n",
-    "\n",
-    "import requests\n",
-    "from PIL import Image\n",
-    "import numpy as np\n",
-    "import matplotlib.pyplot as plt\n",
-    "\n",
-    "import torch\n",
-    "import torchvision.transforms as T\n",
-    "import torchvision.transforms.functional as TF\n",
-    "from torchvision.transforms import InterpolationMode\n",
-    "\n",
-    "jax.devices()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "d408065c",
-   "metadata": {},
-   "source": [
-    "`dalle_mini` is a local package that contains the VQGAN-JAX model by Suraj, and other utilities. You can also `cd` to the directory that contains your checkout of [`vqgan-jax`](https://github.com/patil-suraj/vqgan-jax.git)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "09295910",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from dalle_mini.vqgan_jax.modeling_flax_vqgan import VQModel\n",
-    "#%cd /content/vqgan-jax"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "b6a3462a-9004-4121-b365-3ae3aaf94dd2",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# TODO: set those args in a config file\n",
-    "OUTPUT_VOCAB_SIZE = 16384 + 1  # encoded image token space + 1 for bos\n",
-    "OUTPUT_LENGTH = 256 + 1  # number of encoded tokens + 1 for bos\n",
-    "BOS_TOKEN_ID = 16384\n",
-    "BASE_MODEL = 'facebook/bart-large-cnn'"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "bbef1afb-0b36-44a5-83f7-643d7e2c0e30",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "class CustomFlaxBartModule(FlaxBartModule):\n",
-    "    def setup(self):\n",
-    "        # we keep shared to easily load pre-trained weights\n",
-    "        self.shared = nn.Embed(\n",
-    "            self.config.vocab_size,\n",
-    "            self.config.d_model,\n",
-    "            embedding_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),\n",
-    "            dtype=self.dtype,\n",
-    "        )\n",
-    "        # a separate embedding is used for the decoder\n",
-    "        self.decoder_embed = nn.Embed(\n",
-    "            OUTPUT_VOCAB_SIZE,\n",
-    "            self.config.d_model,\n",
-    "            embedding_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),\n",
-    "            dtype=self.dtype,\n",
-    "        )\n",
-    "        self.encoder = FlaxBartEncoder(self.config, dtype=self.dtype, embed_tokens=self.shared)\n",
-    "\n",
-    "        # the decoder has a different config\n",
-    "        decoder_config = BartConfig(self.config.to_dict())\n",
-    "        decoder_config.max_position_embeddings = OUTPUT_LENGTH\n",
-    "        decoder_config.vocab_size = OUTPUT_VOCAB_SIZE\n",
-    "        self.decoder = FlaxBartDecoder(decoder_config, dtype=self.dtype, embed_tokens=self.decoder_embed)\n",
-    "\n",
-    "class CustomFlaxBartForConditionalGenerationModule(FlaxBartForConditionalGenerationModule):\n",
-    "    def setup(self):\n",
-    "        self.model = CustomFlaxBartModule(config=self.config, dtype=self.dtype)\n",
-    "        self.lm_head = nn.Dense(\n",
-    "            OUTPUT_VOCAB_SIZE,\n",
-    "            use_bias=False,\n",
-    "            dtype=self.dtype,\n",
-    "            kernel_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),\n",
-    "        )\n",
-    "        self.final_logits_bias = self.param(\"final_logits_bias\", self.bias_init, (1, OUTPUT_VOCAB_SIZE))\n",
-    "\n",
-    "class CustomFlaxBartForConditionalGeneration(FlaxBartForConditionalGeneration):\n",
-    "    module_class = CustomFlaxBartForConditionalGenerationModule"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "879320b7-eaa0-4dc9-bbf2-c81efc53301d",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import wandb\n",
-    "run = wandb.init()\n",
-    "artifact = run.use_artifact('wandb/hf-flax-dalle-mini/model-3h3x3565:latest', type='bart_model')\n",
-    "artifact_dir = artifact.download()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "e8bcff33-e95b-4c01-b162-ee857a55c3e6",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# create our model\n",
-    "tokenizer = BartTokenizer.from_pretrained(BASE_MODEL)\n",
-    "model = CustomFlaxBartForConditionalGeneration.from_pretrained(artifact_dir)\n",
-    "model.config.force_bos_token_to_be_generated = False\n",
-    "model.config.forced_bos_token_id = None\n",
-    "model.config.forced_eos_token_id = None\n",
-    "\n",
-    "# we verify that the shape has not been modified\n",
-    "model.params['final_logits_bias'].shape"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "8d5e0f14-2502-470e-9553-daee6748601f",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "vqgan = VQModel.from_pretrained(\"flax-community/vqgan_f16_16384\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "6cca395a-93c2-49bc-a3be-98287e4403d4",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def custom_to_pil(x):\n",
-    "    x = np.clip(x, 0., 1.)\n",
-    "    x = (255*x).astype(np.uint8)\n",
-    "    x = Image.fromarray(x)\n",
-    "    if not x.mode == \"RGB\":\n",
-    "        x = x.convert(\"RGB\")\n",
-    "    return x\n",
-    "\n",
-    "def generate(input, rng, params):\n",
-    "  return model.generate(\n",
-    "      **input,\n",
-    "      max_length=257,\n",
-    "      num_beams=1,\n",
-    "      do_sample=True,\n",
-    "      prng_key=rng,\n",
-    "      eos_token_id=50000,\n",
-    "      pad_token_id=50000,\n",
-    "      params=params\n",
-    "  )\n",
-    "\n",
-    "def get_images(indices, params):\n",
-    "    return vqgan.decode_code(indices, params=params)\n",
-    "\n",
-    "\n",
-    "def plot_images(images):\n",
-    "    fig = plt.figure(figsize=(40, 20))\n",
-    "    columns = 4\n",
-    "    rows = 2\n",
-    "    plt.subplots_adjust(hspace=0, wspace=0)\n",
-    "\n",
-    "    for i in range(1, columns*rows +1):\n",
-    "        fig.add_subplot(rows, columns, i)\n",
-    "        plt.imshow(images[i-1])\n",
-    "    plt.gca().axes.get_yaxis().set_visible(False)\n",
-    "    plt.show()\n",
-    "    \n",
-    "def stack_reconstructions(images):\n",
-    "    w, h = images[0].size[0], images[0].size[1]\n",
-    "    img = Image.new(\"RGB\", (len(images)*w, h))\n",
-    "    for i, img_ in enumerate(images):\n",
-    "        img.paste(img_, (i*w,0))\n",
-    "    return img"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "b1bec3d2-ef17-4feb-aa0d-b51ed2fdcd3e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "p_generate = jax.pmap(generate, \"batch\")\n",
-    "p_get_images = jax.pmap(get_images, \"batch\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "a539823a-a775-4d92-96a5-dc8b1eef69c5",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "bart_params = replicate(model.params)\n",
-    "vqgan_params = replicate(vqgan.params)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "e8b268d8-6992-422a-8373-95651474ae70",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "prompts = [\n",
-    "    \"man in blue jacket walking on pathway in between trees during daytime\",\n",
-    "    'white snow covered mountain under blue sky during daytime',\n",
-    "    'white snow covered mountain under blue sky during night',\n",
-    "    \"orange tabby cat on persons hand\",\n",
-    "    \"aerial view of beach during daytime\",\n",
-    "    \"chess pieces on chess board\",\n",
-    "    \"laptop on brown wooden table\",\n",
-    "    \"white bus on road near high rise buildings\",\n",
-    "]\n",
-    "\n",
-    "\n",
-    "prompt = [prompts[1]] * jax.device_count()\n",
-    "inputs = tokenizer(prompt, return_tensors='jax', padding=\"max_length\", truncation=True, max_length=128).data\n",
-    "inputs = shard(inputs)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "68638cfa-9a4d-4e6a-8630-91aefb627bbd",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "%%time\n",
-    "for i in range(8):\n",
-    "    key = random.randint(0, 1e7)\n",
-    "    rng = jax.random.PRNGKey(key)\n",
-    "    rngs = jax.random.split(rng, jax.local_device_count())\n",
-    "    indices = p_generate(inputs, rngs, bart_params).sequences\n",
-    "    indices = indices[:, :, 1:]\n",
-    "\n",
-    "    images = p_get_images(indices, vqgan_params)\n",
-    "    images = np.squeeze(np.asarray(images), 1)\n",
-    "    imges = [custom_to_pil(image) for image in images]\n",
-    "\n",
-    "    plt.figure(figsize=(40, 20))\n",
-    "    plt.imshow(stack_reconstructions(imges))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "b6e1060f",
-   "metadata": {},
-   "source": [
-    "## CLIP Scoring"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "c68724bc",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from transformers import CLIPProcessor, FlaxCLIPModel"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "17158e5b",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "clip = FlaxCLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
-    "processor = CLIPProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "f1b37b6d",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def hallucinate(prompt, num_images=64):\n",
-    "    prompt = [prompt] * jax.device_count()\n",
-    "    inputs = tokenizer(prompt, return_tensors='jax', padding=\"max_length\", truncation=True, max_length=128).data\n",
-    "    inputs = shard(inputs)\n",
-    "\n",
-    "    all_images = []\n",
-    "    for i in range(num_images // jax.device_count()):\n",
-    "        key = random.randint(0, 1e7)\n",
-    "        rng = jax.random.PRNGKey(key)\n",
-    "        rngs = jax.random.split(rng, jax.local_device_count())\n",
-    "        indices = p_generate(inputs, rngs, bart_params).sequences\n",
-    "        indices = indices[:, :, 1:]\n",
-    "\n",
-    "        images = p_get_images(indices, vqgan_params)\n",
-    "        images = np.squeeze(np.asarray(images), 1)\n",
-    "        for image in images:\n",
-    "            all_images.append(custom_to_pil(image))\n",
-    "    return all_images"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "831c715f",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def clip_top_k(prompt, images, k=8):\n",
-    "    inputs = processor(text=prompt, images=images, return_tensors=\"np\", padding=True)\n",
-    "    outputs = clip(**inputs)\n",
-    "    logits = outputs.logits_per_text\n",
-    "    scores = np.array(logits[0]).argsort()[-k:][::-1]\n",
-    "    return [images[score] for score in scores]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "00605e13",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "prompt = \"white snow covered mountain under blue sky during daytime\"\n",
-    "images = hallucinate(prompt)\n",
-    "selected = clip_top_k(prompt, images, k=8)\n",
-    "stack_reconstructions(selected)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "cc745da2",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "prompt = \"aerial view of beach at night\"\n",
-    "images = hallucinate(prompt)\n",
-    "selected = clip_top_k(prompt, images, k=8)\n",
-    "stack_reconstructions(selected)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "c9cc0b1d",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "prompt = \"an armchair in the shape of an avocado\"\n",
-    "images = hallucinate(prompt)\n",
-    "selected = clip_top_k(prompt, images, k=8)\n",
-    "stack_reconstructions(selected)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "574e9433",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "prompt = \"young woman riding her bike into a forest\"\n",
-    "images = hallucinate(prompt)\n",
-    "selected = clip_top_k(prompt, images, k=8)\n",
-    "stack_reconstructions(selected)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "4762c91e",
-   "metadata": {},
-   "source": [
-    "`Forest` seems to dominate. Interesting cubist interpretation in the fourth image."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "af30608a",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.8.5"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}

dev/notebooks/model/data-pipeline.ipynb

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

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

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

dev/{seq2seq/requirements.txt → requirements.txt}

@@ -10,6 +10,7 @@ jupyter
 wandb
 nltk
 optax
+git+https://github.com/patil-suraj/vqgan-jax.git@610d842dd33c739325a944102ed33acc07692dd5
 
 # Inference
 ftfy