declare-lab/tango2

1# Copyright 2023 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import flax.linen as nn16import jax.numpy as jnp17 18 19class FlaxAttention(nn.Module):20    r"""21    A Flax multi-head attention module as described in: https://arxiv.org/abs/1706.0376222 23    Parameters:24        query_dim (:obj:`int`):25            Input hidden states dimension26        heads (:obj:`int`, *optional*, defaults to 8):27            Number of heads28        dim_head (:obj:`int`, *optional*, defaults to 64):29            Hidden states dimension inside each head30        dropout (:obj:`float`, *optional*, defaults to 0.0):31            Dropout rate32        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):33            Parameters `dtype`34 35    """36    query_dim: int37    heads: int = 838    dim_head: int = 6439    dropout: float = 0.040    dtype: jnp.dtype = jnp.float3241 42    def setup(self):43        inner_dim = self.dim_head * self.heads44        self.scale = self.dim_head**-0.545 46        # Weights were exported with old names {to_q, to_k, to_v, to_out}47        self.query = nn.Dense(inner_dim, use_bias=False, dtype=self.dtype, name="to_q")48        self.key = nn.Dense(inner_dim, use_bias=False, dtype=self.dtype, name="to_k")49        self.value = nn.Dense(inner_dim, use_bias=False, dtype=self.dtype, name="to_v")50 51        self.proj_attn = nn.Dense(self.query_dim, dtype=self.dtype, name="to_out_0")52 53    def reshape_heads_to_batch_dim(self, tensor):54        batch_size, seq_len, dim = tensor.shape55        head_size = self.heads56        tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size)57        tensor = jnp.transpose(tensor, (0, 2, 1, 3))58        tensor = tensor.reshape(batch_size * head_size, seq_len, dim // head_size)59        return tensor60 61    def reshape_batch_dim_to_heads(self, tensor):62        batch_size, seq_len, dim = tensor.shape63        head_size = self.heads64        tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)65        tensor = jnp.transpose(tensor, (0, 2, 1, 3))66        tensor = tensor.reshape(batch_size // head_size, seq_len, dim * head_size)67        return tensor68 69    def __call__(self, hidden_states, context=None, deterministic=True):70        context = hidden_states if context is None else context71 72        query_proj = self.query(hidden_states)73        key_proj = self.key(context)74        value_proj = self.value(context)75 76        query_states = self.reshape_heads_to_batch_dim(query_proj)77        key_states = self.reshape_heads_to_batch_dim(key_proj)78        value_states = self.reshape_heads_to_batch_dim(value_proj)79 80        # compute attentions81        attention_scores = jnp.einsum("b i d, b j d->b i j", query_states, key_states)82        attention_scores = attention_scores * self.scale83        attention_probs = nn.softmax(attention_scores, axis=2)84 85        # attend to values86        hidden_states = jnp.einsum("b i j, b j d -> b i d", attention_probs, value_states)87        hidden_states = self.reshape_batch_dim_to_heads(hidden_states)88        hidden_states = self.proj_attn(hidden_states)89        return hidden_states90 91 92class FlaxBasicTransformerBlock(nn.Module):93    r"""94    A Flax transformer block layer with `GLU` (Gated Linear Unit) activation function as described in:95    https://arxiv.org/abs/1706.0376296 97 98    Parameters:99        dim (:obj:`int`):100            Inner hidden states dimension101        n_heads (:obj:`int`):102            Number of heads103        d_head (:obj:`int`):104            Hidden states dimension inside each head105        dropout (:obj:`float`, *optional*, defaults to 0.0):106            Dropout rate107        only_cross_attention (`bool`, defaults to `False`):108            Whether to only apply cross attention.109        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):110            Parameters `dtype`111    """112    dim: int113    n_heads: int114    d_head: int115    dropout: float = 0.0116    only_cross_attention: bool = False117    dtype: jnp.dtype = jnp.float32118 119    def setup(self):120        # self attention (or cross_attention if only_cross_attention is True)121        self.attn1 = FlaxAttention(self.dim, self.n_heads, self.d_head, self.dropout, dtype=self.dtype)122        # cross attention123        self.attn2 = FlaxAttention(self.dim, self.n_heads, self.d_head, self.dropout, dtype=self.dtype)124        self.ff = FlaxFeedForward(dim=self.dim, dropout=self.dropout, dtype=self.dtype)125        self.norm1 = nn.LayerNorm(epsilon=1e-5, dtype=self.dtype)126        self.norm2 = nn.LayerNorm(epsilon=1e-5, dtype=self.dtype)127        self.norm3 = nn.LayerNorm(epsilon=1e-5, dtype=self.dtype)128 129    def __call__(self, hidden_states, context, deterministic=True):130        # self attention131        residual = hidden_states132        if self.only_cross_attention:133            hidden_states = self.attn1(self.norm1(hidden_states), context, deterministic=deterministic)134        else:135            hidden_states = self.attn1(self.norm1(hidden_states), deterministic=deterministic)136        hidden_states = hidden_states + residual137 138        # cross attention139        residual = hidden_states140        hidden_states = self.attn2(self.norm2(hidden_states), context, deterministic=deterministic)141        hidden_states = hidden_states + residual142 143        # feed forward144        residual = hidden_states145        hidden_states = self.ff(self.norm3(hidden_states), deterministic=deterministic)146        hidden_states = hidden_states + residual147 148        return hidden_states149 150 151class FlaxTransformer2DModel(nn.Module):152    r"""153    A Spatial Transformer layer with Gated Linear Unit (GLU) activation function as described in:154    https://arxiv.org/pdf/1506.02025.pdf155 156 157    Parameters:158        in_channels (:obj:`int`):159            Input number of channels160        n_heads (:obj:`int`):161            Number of heads162        d_head (:obj:`int`):163            Hidden states dimension inside each head164        depth (:obj:`int`, *optional*, defaults to 1):165            Number of transformers block166        dropout (:obj:`float`, *optional*, defaults to 0.0):167            Dropout rate168        use_linear_projection (`bool`, defaults to `False`): tbd169        only_cross_attention (`bool`, defaults to `False`): tbd170        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):171            Parameters `dtype`172    """173    in_channels: int174    n_heads: int175    d_head: int176    depth: int = 1177    dropout: float = 0.0178    use_linear_projection: bool = False179    only_cross_attention: bool = False180    dtype: jnp.dtype = jnp.float32181 182    def setup(self):183        self.norm = nn.GroupNorm(num_groups=32, epsilon=1e-5)184 185        inner_dim = self.n_heads * self.d_head186        if self.use_linear_projection:187            self.proj_in = nn.Dense(inner_dim, dtype=self.dtype)188        else:189            self.proj_in = nn.Conv(190                inner_dim,191                kernel_size=(1, 1),192                strides=(1, 1),193                padding="VALID",194                dtype=self.dtype,195            )196 197        self.transformer_blocks = [198            FlaxBasicTransformerBlock(199                inner_dim,200                self.n_heads,201                self.d_head,202                dropout=self.dropout,203                only_cross_attention=self.only_cross_attention,204                dtype=self.dtype,205            )206            for _ in range(self.depth)207        ]208 209        if self.use_linear_projection:210            self.proj_out = nn.Dense(inner_dim, dtype=self.dtype)211        else:212            self.proj_out = nn.Conv(213                inner_dim,214                kernel_size=(1, 1),215                strides=(1, 1),216                padding="VALID",217                dtype=self.dtype,218            )219 220    def __call__(self, hidden_states, context, deterministic=True):221        batch, height, width, channels = hidden_states.shape222        residual = hidden_states223        hidden_states = self.norm(hidden_states)224        if self.use_linear_projection:225            hidden_states = hidden_states.reshape(batch, height * width, channels)226            hidden_states = self.proj_in(hidden_states)227        else:228            hidden_states = self.proj_in(hidden_states)229            hidden_states = hidden_states.reshape(batch, height * width, channels)230 231        for transformer_block in self.transformer_blocks:232            hidden_states = transformer_block(hidden_states, context, deterministic=deterministic)233 234        if self.use_linear_projection:235            hidden_states = self.proj_out(hidden_states)236            hidden_states = hidden_states.reshape(batch, height, width, channels)237        else:238            hidden_states = hidden_states.reshape(batch, height, width, channels)239            hidden_states = self.proj_out(hidden_states)240 241        hidden_states = hidden_states + residual242        return hidden_states243 244 245class FlaxFeedForward(nn.Module):246    r"""247    Flax module that encapsulates two Linear layers separated by a non-linearity. It is the counterpart of PyTorch's248    [`FeedForward`] class, with the following simplifications:249    - The activation function is currently hardcoded to a gated linear unit from:250    https://arxiv.org/abs/2002.05202251    - `dim_out` is equal to `dim`.252    - The number of hidden dimensions is hardcoded to `dim * 4` in [`FlaxGELU`].253 254    Parameters:255        dim (:obj:`int`):256            Inner hidden states dimension257        dropout (:obj:`float`, *optional*, defaults to 0.0):258            Dropout rate259        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):260            Parameters `dtype`261    """262    dim: int263    dropout: float = 0.0264    dtype: jnp.dtype = jnp.float32265 266    def setup(self):267        # The second linear layer needs to be called268        # net_2 for now to match the index of the Sequential layer269        self.net_0 = FlaxGEGLU(self.dim, self.dropout, self.dtype)270        self.net_2 = nn.Dense(self.dim, dtype=self.dtype)271 272    def __call__(self, hidden_states, deterministic=True):273        hidden_states = self.net_0(hidden_states)274        hidden_states = self.net_2(hidden_states)275        return hidden_states276 277 278class FlaxGEGLU(nn.Module):279    r"""280    Flax implementation of a Linear layer followed by the variant of the gated linear unit activation function from281    https://arxiv.org/abs/2002.05202.282 283    Parameters:284        dim (:obj:`int`):285            Input hidden states dimension286        dropout (:obj:`float`, *optional*, defaults to 0.0):287            Dropout rate288        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):289            Parameters `dtype`290    """291    dim: int292    dropout: float = 0.0293    dtype: jnp.dtype = jnp.float32294 295    def setup(self):296        inner_dim = self.dim * 4297        self.proj = nn.Dense(inner_dim * 2, dtype=self.dtype)298 299    def __call__(self, hidden_states, deterministic=True):300        hidden_states = self.proj(hidden_states)301        hidden_linear, hidden_gelu = jnp.split(hidden_states, 2, axis=2)302        return hidden_linear * nn.gelu(hidden_gelu)303