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 18from .attention_flax import FlaxTransformer2DModel19from .resnet_flax import FlaxDownsample2D, FlaxResnetBlock2D, FlaxUpsample2D20 21 22class FlaxCrossAttnDownBlock2D(nn.Module):23    r"""24    Cross Attention 2D Downsizing block - original architecture from Unet transformers:25    https://arxiv.org/abs/2103.0610426 27    Parameters:28        in_channels (:obj:`int`):29            Input channels30        out_channels (:obj:`int`):31            Output channels32        dropout (:obj:`float`, *optional*, defaults to 0.0):33            Dropout rate34        num_layers (:obj:`int`, *optional*, defaults to 1):35            Number of attention blocks layers36        attn_num_head_channels (:obj:`int`, *optional*, defaults to 1):37            Number of attention heads of each spatial transformer block38        add_downsample (:obj:`bool`, *optional*, defaults to `True`):39            Whether to add downsampling layer before each final output40        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):41            Parameters `dtype`42    """43    in_channels: int44    out_channels: int45    dropout: float = 0.046    num_layers: int = 147    attn_num_head_channels: int = 148    add_downsample: bool = True49    use_linear_projection: bool = False50    only_cross_attention: bool = False51    dtype: jnp.dtype = jnp.float3252 53    def setup(self):54        resnets = []55        attentions = []56 57        for i in range(self.num_layers):58            in_channels = self.in_channels if i == 0 else self.out_channels59 60            res_block = FlaxResnetBlock2D(61                in_channels=in_channels,62                out_channels=self.out_channels,63                dropout_prob=self.dropout,64                dtype=self.dtype,65            )66            resnets.append(res_block)67 68            attn_block = FlaxTransformer2DModel(69                in_channels=self.out_channels,70                n_heads=self.attn_num_head_channels,71                d_head=self.out_channels // self.attn_num_head_channels,72                depth=1,73                use_linear_projection=self.use_linear_projection,74                only_cross_attention=self.only_cross_attention,75                dtype=self.dtype,76            )77            attentions.append(attn_block)78 79        self.resnets = resnets80        self.attentions = attentions81 82        if self.add_downsample:83            self.downsamplers_0 = FlaxDownsample2D(self.out_channels, dtype=self.dtype)84 85    def __call__(self, hidden_states, temb, encoder_hidden_states, deterministic=True):86        output_states = ()87 88        for resnet, attn in zip(self.resnets, self.attentions):89            hidden_states = resnet(hidden_states, temb, deterministic=deterministic)90            hidden_states = attn(hidden_states, encoder_hidden_states, deterministic=deterministic)91            output_states += (hidden_states,)92 93        if self.add_downsample:94            hidden_states = self.downsamplers_0(hidden_states)95            output_states += (hidden_states,)96 97        return hidden_states, output_states98 99 100class FlaxDownBlock2D(nn.Module):101    r"""102    Flax 2D downsizing block103 104    Parameters:105        in_channels (:obj:`int`):106            Input channels107        out_channels (:obj:`int`):108            Output channels109        dropout (:obj:`float`, *optional*, defaults to 0.0):110            Dropout rate111        num_layers (:obj:`int`, *optional*, defaults to 1):112            Number of attention blocks layers113        add_downsample (:obj:`bool`, *optional*, defaults to `True`):114            Whether to add downsampling layer before each final output115        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):116            Parameters `dtype`117    """118    in_channels: int119    out_channels: int120    dropout: float = 0.0121    num_layers: int = 1122    add_downsample: bool = True123    dtype: jnp.dtype = jnp.float32124 125    def setup(self):126        resnets = []127 128        for i in range(self.num_layers):129            in_channels = self.in_channels if i == 0 else self.out_channels130 131            res_block = FlaxResnetBlock2D(132                in_channels=in_channels,133                out_channels=self.out_channels,134                dropout_prob=self.dropout,135                dtype=self.dtype,136            )137            resnets.append(res_block)138        self.resnets = resnets139 140        if self.add_downsample:141            self.downsamplers_0 = FlaxDownsample2D(self.out_channels, dtype=self.dtype)142 143    def __call__(self, hidden_states, temb, deterministic=True):144        output_states = ()145 146        for resnet in self.resnets:147            hidden_states = resnet(hidden_states, temb, deterministic=deterministic)148            output_states += (hidden_states,)149 150        if self.add_downsample:151            hidden_states = self.downsamplers_0(hidden_states)152            output_states += (hidden_states,)153 154        return hidden_states, output_states155 156 157class FlaxCrossAttnUpBlock2D(nn.Module):158    r"""159    Cross Attention 2D Upsampling block - original architecture from Unet transformers:160    https://arxiv.org/abs/2103.06104161 162    Parameters:163        in_channels (:obj:`int`):164            Input channels165        out_channels (:obj:`int`):166            Output channels167        dropout (:obj:`float`, *optional*, defaults to 0.0):168            Dropout rate169        num_layers (:obj:`int`, *optional*, defaults to 1):170            Number of attention blocks layers171        attn_num_head_channels (:obj:`int`, *optional*, defaults to 1):172            Number of attention heads of each spatial transformer block173        add_upsample (:obj:`bool`, *optional*, defaults to `True`):174            Whether to add upsampling layer before each final output175        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):176            Parameters `dtype`177    """178    in_channels: int179    out_channels: int180    prev_output_channel: int181    dropout: float = 0.0182    num_layers: int = 1183    attn_num_head_channels: int = 1184    add_upsample: bool = True185    use_linear_projection: bool = False186    only_cross_attention: bool = False187    dtype: jnp.dtype = jnp.float32188 189    def setup(self):190        resnets = []191        attentions = []192 193        for i in range(self.num_layers):194            res_skip_channels = self.in_channels if (i == self.num_layers - 1) else self.out_channels195            resnet_in_channels = self.prev_output_channel if i == 0 else self.out_channels196 197            res_block = FlaxResnetBlock2D(198                in_channels=resnet_in_channels + res_skip_channels,199                out_channels=self.out_channels,200                dropout_prob=self.dropout,201                dtype=self.dtype,202            )203            resnets.append(res_block)204 205            attn_block = FlaxTransformer2DModel(206                in_channels=self.out_channels,207                n_heads=self.attn_num_head_channels,208                d_head=self.out_channels // self.attn_num_head_channels,209                depth=1,210                use_linear_projection=self.use_linear_projection,211                only_cross_attention=self.only_cross_attention,212                dtype=self.dtype,213            )214            attentions.append(attn_block)215 216        self.resnets = resnets217        self.attentions = attentions218 219        if self.add_upsample:220            self.upsamplers_0 = FlaxUpsample2D(self.out_channels, dtype=self.dtype)221 222    def __call__(self, hidden_states, res_hidden_states_tuple, temb, encoder_hidden_states, deterministic=True):223        for resnet, attn in zip(self.resnets, self.attentions):224            # pop res hidden states225            res_hidden_states = res_hidden_states_tuple[-1]226            res_hidden_states_tuple = res_hidden_states_tuple[:-1]227            hidden_states = jnp.concatenate((hidden_states, res_hidden_states), axis=-1)228 229            hidden_states = resnet(hidden_states, temb, deterministic=deterministic)230            hidden_states = attn(hidden_states, encoder_hidden_states, deterministic=deterministic)231 232        if self.add_upsample:233            hidden_states = self.upsamplers_0(hidden_states)234 235        return hidden_states236 237 238class FlaxUpBlock2D(nn.Module):239    r"""240    Flax 2D upsampling block241 242    Parameters:243        in_channels (:obj:`int`):244            Input channels245        out_channels (:obj:`int`):246            Output channels247        prev_output_channel (:obj:`int`):248            Output channels from the previous block249        dropout (:obj:`float`, *optional*, defaults to 0.0):250            Dropout rate251        num_layers (:obj:`int`, *optional*, defaults to 1):252            Number of attention blocks layers253        add_downsample (:obj:`bool`, *optional*, defaults to `True`):254            Whether to add downsampling layer before each final output255        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):256            Parameters `dtype`257    """258    in_channels: int259    out_channels: int260    prev_output_channel: int261    dropout: float = 0.0262    num_layers: int = 1263    add_upsample: bool = True264    dtype: jnp.dtype = jnp.float32265 266    def setup(self):267        resnets = []268 269        for i in range(self.num_layers):270            res_skip_channels = self.in_channels if (i == self.num_layers - 1) else self.out_channels271            resnet_in_channels = self.prev_output_channel if i == 0 else self.out_channels272 273            res_block = FlaxResnetBlock2D(274                in_channels=resnet_in_channels + res_skip_channels,275                out_channels=self.out_channels,276                dropout_prob=self.dropout,277                dtype=self.dtype,278            )279            resnets.append(res_block)280 281        self.resnets = resnets282 283        if self.add_upsample:284            self.upsamplers_0 = FlaxUpsample2D(self.out_channels, dtype=self.dtype)285 286    def __call__(self, hidden_states, res_hidden_states_tuple, temb, deterministic=True):287        for resnet in self.resnets:288            # pop res hidden states289            res_hidden_states = res_hidden_states_tuple[-1]290            res_hidden_states_tuple = res_hidden_states_tuple[:-1]291            hidden_states = jnp.concatenate((hidden_states, res_hidden_states), axis=-1)292 293            hidden_states = resnet(hidden_states, temb, deterministic=deterministic)294 295        if self.add_upsample:296            hidden_states = self.upsamplers_0(hidden_states)297 298        return hidden_states299 300 301class FlaxUNetMidBlock2DCrossAttn(nn.Module):302    r"""303    Cross Attention 2D Mid-level block - original architecture from Unet transformers: https://arxiv.org/abs/2103.06104304 305    Parameters:306        in_channels (:obj:`int`):307            Input channels308        dropout (:obj:`float`, *optional*, defaults to 0.0):309            Dropout rate310        num_layers (:obj:`int`, *optional*, defaults to 1):311            Number of attention blocks layers312        attn_num_head_channels (:obj:`int`, *optional*, defaults to 1):313            Number of attention heads of each spatial transformer block314        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):315            Parameters `dtype`316    """317    in_channels: int318    dropout: float = 0.0319    num_layers: int = 1320    attn_num_head_channels: int = 1321    use_linear_projection: bool = False322    dtype: jnp.dtype = jnp.float32323 324    def setup(self):325        # there is always at least one resnet326        resnets = [327            FlaxResnetBlock2D(328                in_channels=self.in_channels,329                out_channels=self.in_channels,330                dropout_prob=self.dropout,331                dtype=self.dtype,332            )333        ]334 335        attentions = []336 337        for _ in range(self.num_layers):338            attn_block = FlaxTransformer2DModel(339                in_channels=self.in_channels,340                n_heads=self.attn_num_head_channels,341                d_head=self.in_channels // self.attn_num_head_channels,342                depth=1,343                use_linear_projection=self.use_linear_projection,344                dtype=self.dtype,345            )346            attentions.append(attn_block)347 348            res_block = FlaxResnetBlock2D(349                in_channels=self.in_channels,350                out_channels=self.in_channels,351                dropout_prob=self.dropout,352                dtype=self.dtype,353            )354            resnets.append(res_block)355 356        self.resnets = resnets357        self.attentions = attentions358 359    def __call__(self, hidden_states, temb, encoder_hidden_states, deterministic=True):360        hidden_states = self.resnets[0](hidden_states, temb)361        for attn, resnet in zip(self.attentions, self.resnets[1:]):362            hidden_states = attn(hidden_states, encoder_hidden_states, deterministic=deterministic)363            hidden_states = resnet(hidden_states, temb, deterministic=deterministic)364 365        return hidden_states366