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 15# JAX implementation of VQGAN from taming-transformers https://github.com/CompVis/taming-transformers16 17import math18from functools import partial19from typing import Tuple20 21import flax22import flax.linen as nn23import jax24import jax.numpy as jnp25from flax.core.frozen_dict import FrozenDict26 27from ..configuration_utils import ConfigMixin, flax_register_to_config28from ..utils import BaseOutput29from .modeling_flax_utils import FlaxModelMixin30 31 32@flax.struct.dataclass33class FlaxDecoderOutput(BaseOutput):34    """35    Output of decoding method.36 37    Args:38        sample (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)`):39            Decoded output sample of the model. Output of the last layer of the model.40        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):41            Parameters `dtype`42    """43 44    sample: jnp.ndarray45 46 47@flax.struct.dataclass48class FlaxAutoencoderKLOutput(BaseOutput):49    """50    Output of AutoencoderKL encoding method.51 52    Args:53        latent_dist (`FlaxDiagonalGaussianDistribution`):54            Encoded outputs of `Encoder` represented as the mean and logvar of `FlaxDiagonalGaussianDistribution`.55            `FlaxDiagonalGaussianDistribution` allows for sampling latents from the distribution.56    """57 58    latent_dist: "FlaxDiagonalGaussianDistribution"59 60 61class FlaxUpsample2D(nn.Module):62    """63    Flax implementation of 2D Upsample layer64 65    Args:66        in_channels (`int`):67            Input channels68        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):69            Parameters `dtype`70    """71 72    in_channels: int73    dtype: jnp.dtype = jnp.float3274 75    def setup(self):76        self.conv = nn.Conv(77            self.in_channels,78            kernel_size=(3, 3),79            strides=(1, 1),80            padding=((1, 1), (1, 1)),81            dtype=self.dtype,82        )83 84    def __call__(self, hidden_states):85        batch, height, width, channels = hidden_states.shape86        hidden_states = jax.image.resize(87            hidden_states,88            shape=(batch, height * 2, width * 2, channels),89            method="nearest",90        )91        hidden_states = self.conv(hidden_states)92        return hidden_states93 94 95class FlaxDownsample2D(nn.Module):96    """97    Flax implementation of 2D Downsample layer98 99    Args:100        in_channels (`int`):101            Input channels102        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):103            Parameters `dtype`104    """105 106    in_channels: int107    dtype: jnp.dtype = jnp.float32108 109    def setup(self):110        self.conv = nn.Conv(111            self.in_channels,112            kernel_size=(3, 3),113            strides=(2, 2),114            padding="VALID",115            dtype=self.dtype,116        )117 118    def __call__(self, hidden_states):119        pad = ((0, 0), (0, 1), (0, 1), (0, 0))  # pad height and width dim120        hidden_states = jnp.pad(hidden_states, pad_width=pad)121        hidden_states = self.conv(hidden_states)122        return hidden_states123 124 125class FlaxResnetBlock2D(nn.Module):126    """127    Flax implementation of 2D Resnet Block.128 129    Args:130        in_channels (`int`):131            Input channels132        out_channels (`int`):133            Output channels134        dropout (:obj:`float`, *optional*, defaults to 0.0):135            Dropout rate136        groups (:obj:`int`, *optional*, defaults to `32`):137            The number of groups to use for group norm.138        use_nin_shortcut (:obj:`bool`, *optional*, defaults to `None`):139            Whether to use `nin_shortcut`. This activates a new layer inside ResNet block140        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):141            Parameters `dtype`142    """143 144    in_channels: int145    out_channels: int = None146    dropout: float = 0.0147    groups: int = 32148    use_nin_shortcut: bool = None149    dtype: jnp.dtype = jnp.float32150 151    def setup(self):152        out_channels = self.in_channels if self.out_channels is None else self.out_channels153 154        self.norm1 = nn.GroupNorm(num_groups=self.groups, epsilon=1e-6)155        self.conv1 = nn.Conv(156            out_channels,157            kernel_size=(3, 3),158            strides=(1, 1),159            padding=((1, 1), (1, 1)),160            dtype=self.dtype,161        )162 163        self.norm2 = nn.GroupNorm(num_groups=self.groups, epsilon=1e-6)164        self.dropout_layer = nn.Dropout(self.dropout)165        self.conv2 = nn.Conv(166            out_channels,167            kernel_size=(3, 3),168            strides=(1, 1),169            padding=((1, 1), (1, 1)),170            dtype=self.dtype,171        )172 173        use_nin_shortcut = self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut174 175        self.conv_shortcut = None176        if use_nin_shortcut:177            self.conv_shortcut = nn.Conv(178                out_channels,179                kernel_size=(1, 1),180                strides=(1, 1),181                padding="VALID",182                dtype=self.dtype,183            )184 185    def __call__(self, hidden_states, deterministic=True):186        residual = hidden_states187        hidden_states = self.norm1(hidden_states)188        hidden_states = nn.swish(hidden_states)189        hidden_states = self.conv1(hidden_states)190 191        hidden_states = self.norm2(hidden_states)192        hidden_states = nn.swish(hidden_states)193        hidden_states = self.dropout_layer(hidden_states, deterministic)194        hidden_states = self.conv2(hidden_states)195 196        if self.conv_shortcut is not None:197            residual = self.conv_shortcut(residual)198 199        return hidden_states + residual200 201 202class FlaxAttentionBlock(nn.Module):203    r"""204    Flax Convolutional based multi-head attention block for diffusion-based VAE.205 206    Parameters:207        channels (:obj:`int`):208            Input channels209        num_head_channels (:obj:`int`, *optional*, defaults to `None`):210            Number of attention heads211        num_groups (:obj:`int`, *optional*, defaults to `32`):212            The number of groups to use for group norm213        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):214            Parameters `dtype`215 216    """217    channels: int218    num_head_channels: int = None219    num_groups: int = 32220    dtype: jnp.dtype = jnp.float32221 222    def setup(self):223        self.num_heads = self.channels // self.num_head_channels if self.num_head_channels is not None else 1224 225        dense = partial(nn.Dense, self.channels, dtype=self.dtype)226 227        self.group_norm = nn.GroupNorm(num_groups=self.num_groups, epsilon=1e-6)228        self.query, self.key, self.value = dense(), dense(), dense()229        self.proj_attn = dense()230 231    def transpose_for_scores(self, projection):232        new_projection_shape = projection.shape[:-1] + (self.num_heads, -1)233        # move heads to 2nd position (B, T, H * D) -> (B, T, H, D)234        new_projection = projection.reshape(new_projection_shape)235        # (B, T, H, D) -> (B, H, T, D)236        new_projection = jnp.transpose(new_projection, (0, 2, 1, 3))237        return new_projection238 239    def __call__(self, hidden_states):240        residual = hidden_states241        batch, height, width, channels = hidden_states.shape242 243        hidden_states = self.group_norm(hidden_states)244 245        hidden_states = hidden_states.reshape((batch, height * width, channels))246 247        query = self.query(hidden_states)248        key = self.key(hidden_states)249        value = self.value(hidden_states)250 251        # transpose252        query = self.transpose_for_scores(query)253        key = self.transpose_for_scores(key)254        value = self.transpose_for_scores(value)255 256        # compute attentions257        scale = 1 / math.sqrt(math.sqrt(self.channels / self.num_heads))258        attn_weights = jnp.einsum("...qc,...kc->...qk", query * scale, key * scale)259        attn_weights = nn.softmax(attn_weights, axis=-1)260 261        # attend to values262        hidden_states = jnp.einsum("...kc,...qk->...qc", value, attn_weights)263 264        hidden_states = jnp.transpose(hidden_states, (0, 2, 1, 3))265        new_hidden_states_shape = hidden_states.shape[:-2] + (self.channels,)266        hidden_states = hidden_states.reshape(new_hidden_states_shape)267 268        hidden_states = self.proj_attn(hidden_states)269        hidden_states = hidden_states.reshape((batch, height, width, channels))270        hidden_states = hidden_states + residual271        return hidden_states272 273 274class FlaxDownEncoderBlock2D(nn.Module):275    r"""276    Flax Resnet blocks-based Encoder block for diffusion-based VAE.277 278    Parameters:279        in_channels (:obj:`int`):280            Input channels281        out_channels (:obj:`int`):282            Output channels283        dropout (:obj:`float`, *optional*, defaults to 0.0):284            Dropout rate285        num_layers (:obj:`int`, *optional*, defaults to 1):286            Number of Resnet layer block287        resnet_groups (:obj:`int`, *optional*, defaults to `32`):288            The number of groups to use for the Resnet block group norm289        add_downsample (:obj:`bool`, *optional*, defaults to `True`):290            Whether to add downsample layer291        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):292            Parameters `dtype`293    """294    in_channels: int295    out_channels: int296    dropout: float = 0.0297    num_layers: int = 1298    resnet_groups: int = 32299    add_downsample: bool = True300    dtype: jnp.dtype = jnp.float32301 302    def setup(self):303        resnets = []304        for i in range(self.num_layers):305            in_channels = self.in_channels if i == 0 else self.out_channels306 307            res_block = FlaxResnetBlock2D(308                in_channels=in_channels,309                out_channels=self.out_channels,310                dropout=self.dropout,311                groups=self.resnet_groups,312                dtype=self.dtype,313            )314            resnets.append(res_block)315        self.resnets = resnets316 317        if self.add_downsample:318            self.downsamplers_0 = FlaxDownsample2D(self.out_channels, dtype=self.dtype)319 320    def __call__(self, hidden_states, deterministic=True):321        for resnet in self.resnets:322            hidden_states = resnet(hidden_states, deterministic=deterministic)323 324        if self.add_downsample:325            hidden_states = self.downsamplers_0(hidden_states)326 327        return hidden_states328 329 330class FlaxUpDecoderBlock2D(nn.Module):331    r"""332    Flax Resnet blocks-based Decoder block for diffusion-based VAE.333 334    Parameters:335        in_channels (:obj:`int`):336            Input channels337        out_channels (:obj:`int`):338            Output channels339        dropout (:obj:`float`, *optional*, defaults to 0.0):340            Dropout rate341        num_layers (:obj:`int`, *optional*, defaults to 1):342            Number of Resnet layer block343        resnet_groups (:obj:`int`, *optional*, defaults to `32`):344            The number of groups to use for the Resnet block group norm345        add_upsample (:obj:`bool`, *optional*, defaults to `True`):346            Whether to add upsample layer347        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):348            Parameters `dtype`349    """350    in_channels: int351    out_channels: int352    dropout: float = 0.0353    num_layers: int = 1354    resnet_groups: int = 32355    add_upsample: bool = True356    dtype: jnp.dtype = jnp.float32357 358    def setup(self):359        resnets = []360        for i in range(self.num_layers):361            in_channels = self.in_channels if i == 0 else self.out_channels362            res_block = FlaxResnetBlock2D(363                in_channels=in_channels,364                out_channels=self.out_channels,365                dropout=self.dropout,366                groups=self.resnet_groups,367                dtype=self.dtype,368            )369            resnets.append(res_block)370 371        self.resnets = resnets372 373        if self.add_upsample:374            self.upsamplers_0 = FlaxUpsample2D(self.out_channels, dtype=self.dtype)375 376    def __call__(self, hidden_states, deterministic=True):377        for resnet in self.resnets:378            hidden_states = resnet(hidden_states, deterministic=deterministic)379 380        if self.add_upsample:381            hidden_states = self.upsamplers_0(hidden_states)382 383        return hidden_states384 385 386class FlaxUNetMidBlock2D(nn.Module):387    r"""388    Flax Unet Mid-Block module.389 390    Parameters:391        in_channels (:obj:`int`):392            Input channels393        dropout (:obj:`float`, *optional*, defaults to 0.0):394            Dropout rate395        num_layers (:obj:`int`, *optional*, defaults to 1):396            Number of Resnet layer block397        resnet_groups (:obj:`int`, *optional*, defaults to `32`):398            The number of groups to use for the Resnet and Attention block group norm399        attn_num_head_channels (:obj:`int`, *optional*, defaults to `1`):400            Number of attention heads for each attention block401        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):402            Parameters `dtype`403    """404    in_channels: int405    dropout: float = 0.0406    num_layers: int = 1407    resnet_groups: int = 32408    attn_num_head_channels: int = 1409    dtype: jnp.dtype = jnp.float32410 411    def setup(self):412        resnet_groups = self.resnet_groups if self.resnet_groups is not None else min(self.in_channels // 4, 32)413 414        # there is always at least one resnet415        resnets = [416            FlaxResnetBlock2D(417                in_channels=self.in_channels,418                out_channels=self.in_channels,419                dropout=self.dropout,420                groups=resnet_groups,421                dtype=self.dtype,422            )423        ]424 425        attentions = []426 427        for _ in range(self.num_layers):428            attn_block = FlaxAttentionBlock(429                channels=self.in_channels,430                num_head_channels=self.attn_num_head_channels,431                num_groups=resnet_groups,432                dtype=self.dtype,433            )434            attentions.append(attn_block)435 436            res_block = FlaxResnetBlock2D(437                in_channels=self.in_channels,438                out_channels=self.in_channels,439                dropout=self.dropout,440                groups=resnet_groups,441                dtype=self.dtype,442            )443            resnets.append(res_block)444 445        self.resnets = resnets446        self.attentions = attentions447 448    def __call__(self, hidden_states, deterministic=True):449        hidden_states = self.resnets[0](hidden_states, deterministic=deterministic)450        for attn, resnet in zip(self.attentions, self.resnets[1:]):451            hidden_states = attn(hidden_states)452            hidden_states = resnet(hidden_states, deterministic=deterministic)453 454        return hidden_states455 456 457class FlaxEncoder(nn.Module):458    r"""459    Flax Implementation of VAE Encoder.460 461    This model is a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module)462    subclass. Use it as a regular Flax linen Module and refer to the Flax documentation for all matter related to463    general usage and behavior.464 465    Finally, this model supports inherent JAX features such as:466    - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)467    - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)468    - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)469    - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)470 471    Parameters:472        in_channels (:obj:`int`, *optional*, defaults to 3):473            Input channels474        out_channels (:obj:`int`, *optional*, defaults to 3):475            Output channels476        down_block_types (:obj:`Tuple[str]`, *optional*, defaults to `(DownEncoderBlock2D)`):477            DownEncoder block type478        block_out_channels (:obj:`Tuple[str]`, *optional*, defaults to `(64,)`):479            Tuple containing the number of output channels for each block480        layers_per_block (:obj:`int`, *optional*, defaults to `2`):481            Number of Resnet layer for each block482        norm_num_groups (:obj:`int`, *optional*, defaults to `32`):483            norm num group484        act_fn (:obj:`str`, *optional*, defaults to `silu`):485            Activation function486        double_z (:obj:`bool`, *optional*, defaults to `False`):487            Whether to double the last output channels488        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):489            Parameters `dtype`490    """491    in_channels: int = 3492    out_channels: int = 3493    down_block_types: Tuple[str] = ("DownEncoderBlock2D",)494    block_out_channels: Tuple[int] = (64,)495    layers_per_block: int = 2496    norm_num_groups: int = 32497    act_fn: str = "silu"498    double_z: bool = False499    dtype: jnp.dtype = jnp.float32500 501    def setup(self):502        block_out_channels = self.block_out_channels503        # in504        self.conv_in = nn.Conv(505            block_out_channels[0],506            kernel_size=(3, 3),507            strides=(1, 1),508            padding=((1, 1), (1, 1)),509            dtype=self.dtype,510        )511 512        # downsampling513        down_blocks = []514        output_channel = block_out_channels[0]515        for i, _ in enumerate(self.down_block_types):516            input_channel = output_channel517            output_channel = block_out_channels[i]518            is_final_block = i == len(block_out_channels) - 1519 520            down_block = FlaxDownEncoderBlock2D(521                in_channels=input_channel,522                out_channels=output_channel,523                num_layers=self.layers_per_block,524                resnet_groups=self.norm_num_groups,525                add_downsample=not is_final_block,526                dtype=self.dtype,527            )528            down_blocks.append(down_block)529        self.down_blocks = down_blocks530 531        # middle532        self.mid_block = FlaxUNetMidBlock2D(533            in_channels=block_out_channels[-1],534            resnet_groups=self.norm_num_groups,535            attn_num_head_channels=None,536            dtype=self.dtype,537        )538 539        # end540        conv_out_channels = 2 * self.out_channels if self.double_z else self.out_channels541        self.conv_norm_out = nn.GroupNorm(num_groups=self.norm_num_groups, epsilon=1e-6)542        self.conv_out = nn.Conv(543            conv_out_channels,544            kernel_size=(3, 3),545            strides=(1, 1),546            padding=((1, 1), (1, 1)),547            dtype=self.dtype,548        )549 550    def __call__(self, sample, deterministic: bool = True):551        # in552        sample = self.conv_in(sample)553 554        # downsampling555        for block in self.down_blocks:556            sample = block(sample, deterministic=deterministic)557 558        # middle559        sample = self.mid_block(sample, deterministic=deterministic)560 561        # end562        sample = self.conv_norm_out(sample)563        sample = nn.swish(sample)564        sample = self.conv_out(sample)565 566        return sample567 568 569class FlaxDecoder(nn.Module):570    r"""571    Flax Implementation of VAE Decoder.572 573    This model is a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module)574    subclass. Use it as a regular Flax linen Module and refer to the Flax documentation for all matter related to575    general usage and behavior.576 577    Finally, this model supports inherent JAX features such as:578    - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)579    - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)580    - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)581    - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)582 583    Parameters:584        in_channels (:obj:`int`, *optional*, defaults to 3):585            Input channels586        out_channels (:obj:`int`, *optional*, defaults to 3):587            Output channels588        up_block_types (:obj:`Tuple[str]`, *optional*, defaults to `(UpDecoderBlock2D)`):589            UpDecoder block type590        block_out_channels (:obj:`Tuple[str]`, *optional*, defaults to `(64,)`):591            Tuple containing the number of output channels for each block592        layers_per_block (:obj:`int`, *optional*, defaults to `2`):593            Number of Resnet layer for each block594        norm_num_groups (:obj:`int`, *optional*, defaults to `32`):595            norm num group596        act_fn (:obj:`str`, *optional*, defaults to `silu`):597            Activation function598        double_z (:obj:`bool`, *optional*, defaults to `False`):599            Whether to double the last output channels600        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):601            parameters `dtype`602    """603    in_channels: int = 3604    out_channels: int = 3605    up_block_types: Tuple[str] = ("UpDecoderBlock2D",)606    block_out_channels: int = (64,)607    layers_per_block: int = 2608    norm_num_groups: int = 32609    act_fn: str = "silu"610    dtype: jnp.dtype = jnp.float32611 612    def setup(self):613        block_out_channels = self.block_out_channels614 615        # z to block_in616        self.conv_in = nn.Conv(617            block_out_channels[-1],618            kernel_size=(3, 3),619            strides=(1, 1),620            padding=((1, 1), (1, 1)),621            dtype=self.dtype,622        )623 624        # middle625        self.mid_block = FlaxUNetMidBlock2D(626            in_channels=block_out_channels[-1],627            resnet_groups=self.norm_num_groups,628            attn_num_head_channels=None,629            dtype=self.dtype,630        )631 632        # upsampling633        reversed_block_out_channels = list(reversed(block_out_channels))634        output_channel = reversed_block_out_channels[0]635        up_blocks = []636        for i, _ in enumerate(self.up_block_types):637            prev_output_channel = output_channel638            output_channel = reversed_block_out_channels[i]639 640            is_final_block = i == len(block_out_channels) - 1641 642            up_block = FlaxUpDecoderBlock2D(643                in_channels=prev_output_channel,644                out_channels=output_channel,645                num_layers=self.layers_per_block + 1,646                resnet_groups=self.norm_num_groups,647                add_upsample=not is_final_block,648                dtype=self.dtype,649            )650            up_blocks.append(up_block)651            prev_output_channel = output_channel652 653        self.up_blocks = up_blocks654 655        # end656        self.conv_norm_out = nn.GroupNorm(num_groups=self.norm_num_groups, epsilon=1e-6)657        self.conv_out = nn.Conv(658            self.out_channels,659            kernel_size=(3, 3),660            strides=(1, 1),661            padding=((1, 1), (1, 1)),662            dtype=self.dtype,663        )664 665    def __call__(self, sample, deterministic: bool = True):666        # z to block_in667        sample = self.conv_in(sample)668 669        # middle670        sample = self.mid_block(sample, deterministic=deterministic)671 672        # upsampling673        for block in self.up_blocks:674            sample = block(sample, deterministic=deterministic)675 676        sample = self.conv_norm_out(sample)677        sample = nn.swish(sample)678        sample = self.conv_out(sample)679 680        return sample681 682 683class FlaxDiagonalGaussianDistribution(object):684    def __init__(self, parameters, deterministic=False):685        # Last axis to account for channels-last686        self.mean, self.logvar = jnp.split(parameters, 2, axis=-1)687        self.logvar = jnp.clip(self.logvar, -30.0, 20.0)688        self.deterministic = deterministic689        self.std = jnp.exp(0.5 * self.logvar)690        self.var = jnp.exp(self.logvar)691        if self.deterministic:692            self.var = self.std = jnp.zeros_like(self.mean)693 694    def sample(self, key):695        return self.mean + self.std * jax.random.normal(key, self.mean.shape)696 697    def kl(self, other=None):698        if self.deterministic:699            return jnp.array([0.0])700 701        if other is None:702            return 0.5 * jnp.sum(self.mean**2 + self.var - 1.0 - self.logvar, axis=[1, 2, 3])703 704        return 0.5 * jnp.sum(705            jnp.square(self.mean - other.mean) / other.var + self.var / other.var - 1.0 - self.logvar + other.logvar,706            axis=[1, 2, 3],707        )708 709    def nll(self, sample, axis=[1, 2, 3]):710        if self.deterministic:711            return jnp.array([0.0])712 713        logtwopi = jnp.log(2.0 * jnp.pi)714        return 0.5 * jnp.sum(logtwopi + self.logvar + jnp.square(sample - self.mean) / self.var, axis=axis)715 716    def mode(self):717        return self.mean718 719 720@flax_register_to_config721class FlaxAutoencoderKL(nn.Module, FlaxModelMixin, ConfigMixin):722    r"""723    Flax Implementation of Variational Autoencoder (VAE) model with KL loss from the paper Auto-Encoding Variational724    Bayes by Diederik P. Kingma and Max Welling.725 726    This model is a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module)727    subclass. Use it as a regular Flax linen Module and refer to the Flax documentation for all matter related to728    general usage and behavior.729 730    Finally, this model supports inherent JAX features such as:731    - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)732    - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)733    - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)734    - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)735 736    Parameters:737        in_channels (:obj:`int`, *optional*, defaults to 3):738            Input channels739        out_channels (:obj:`int`, *optional*, defaults to 3):740            Output channels741        down_block_types (:obj:`Tuple[str]`, *optional*, defaults to `(DownEncoderBlock2D)`):742            DownEncoder block type743        up_block_types (:obj:`Tuple[str]`, *optional*, defaults to `(UpDecoderBlock2D)`):744            UpDecoder block type745        block_out_channels (:obj:`Tuple[str]`, *optional*, defaults to `(64,)`):746            Tuple containing the number of output channels for each block747        layers_per_block (:obj:`int`, *optional*, defaults to `2`):748            Number of Resnet layer for each block749        act_fn (:obj:`str`, *optional*, defaults to `silu`):750            Activation function751        latent_channels (:obj:`int`, *optional*, defaults to `4`):752            Latent space channels753        norm_num_groups (:obj:`int`, *optional*, defaults to `32`):754            Norm num group755        sample_size (:obj:`int`, *optional*, defaults to 32):756            Sample input size757        scaling_factor (`float`, *optional*, defaults to 0.18215):758            The component-wise standard deviation of the trained latent space computed using the first batch of the759            training set. This is used to scale the latent space to have unit variance when training the diffusion760            model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the761            diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 1762            / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image763            Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) paper.764        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):765            parameters `dtype`766    """767    in_channels: int = 3768    out_channels: int = 3769    down_block_types: Tuple[str] = ("DownEncoderBlock2D",)770    up_block_types: Tuple[str] = ("UpDecoderBlock2D",)771    block_out_channels: Tuple[int] = (64,)772    layers_per_block: int = 1773    act_fn: str = "silu"774    latent_channels: int = 4775    norm_num_groups: int = 32776    sample_size: int = 32777    scaling_factor: float = 0.18215778    dtype: jnp.dtype = jnp.float32779 780    def setup(self):781        self.encoder = FlaxEncoder(782            in_channels=self.config.in_channels,783            out_channels=self.config.latent_channels,784            down_block_types=self.config.down_block_types,785            block_out_channels=self.config.block_out_channels,786            layers_per_block=self.config.layers_per_block,787            act_fn=self.config.act_fn,788            norm_num_groups=self.config.norm_num_groups,789            double_z=True,790            dtype=self.dtype,791        )792        self.decoder = FlaxDecoder(793            in_channels=self.config.latent_channels,794            out_channels=self.config.out_channels,795            up_block_types=self.config.up_block_types,796            block_out_channels=self.config.block_out_channels,797            layers_per_block=self.config.layers_per_block,798            norm_num_groups=self.config.norm_num_groups,799            act_fn=self.config.act_fn,800            dtype=self.dtype,801        )802        self.quant_conv = nn.Conv(803            2 * self.config.latent_channels,804            kernel_size=(1, 1),805            strides=(1, 1),806            padding="VALID",807            dtype=self.dtype,808        )809        self.post_quant_conv = nn.Conv(810            self.config.latent_channels,811            kernel_size=(1, 1),812            strides=(1, 1),813            padding="VALID",814            dtype=self.dtype,815        )816 817    def init_weights(self, rng: jax.random.KeyArray) -> FrozenDict:818        # init input tensors819        sample_shape = (1, self.in_channels, self.sample_size, self.sample_size)820        sample = jnp.zeros(sample_shape, dtype=jnp.float32)821 822        params_rng, dropout_rng, gaussian_rng = jax.random.split(rng, 3)823        rngs = {"params": params_rng, "dropout": dropout_rng, "gaussian": gaussian_rng}824 825        return self.init(rngs, sample)["params"]826 827    def encode(self, sample, deterministic: bool = True, return_dict: bool = True):828        sample = jnp.transpose(sample, (0, 2, 3, 1))829 830        hidden_states = self.encoder(sample, deterministic=deterministic)831        moments = self.quant_conv(hidden_states)832        posterior = FlaxDiagonalGaussianDistribution(moments)833 834        if not return_dict:835            return (posterior,)836 837        return FlaxAutoencoderKLOutput(latent_dist=posterior)838 839    def decode(self, latents, deterministic: bool = True, return_dict: bool = True):840        if latents.shape[-1] != self.config.latent_channels:841            latents = jnp.transpose(latents, (0, 2, 3, 1))842 843        hidden_states = self.post_quant_conv(latents)844        hidden_states = self.decoder(hidden_states, deterministic=deterministic)845 846        hidden_states = jnp.transpose(hidden_states, (0, 3, 1, 2))847 848        if not return_dict:849            return (hidden_states,)850 851        return FlaxDecoderOutput(sample=hidden_states)852 853    def __call__(self, sample, sample_posterior=False, deterministic: bool = True, return_dict: bool = True):854        posterior = self.encode(sample, deterministic=deterministic, return_dict=return_dict)855        if sample_posterior:856            rng = self.make_rng("gaussian")857            hidden_states = posterior.latent_dist.sample(rng)858        else:859            hidden_states = posterior.latent_dist.mode()860 861        sample = self.decode(hidden_states, return_dict=return_dict).sample862 863        if not return_dict:864            return (sample,)865 866        return FlaxDecoderOutput(sample=sample)867