

from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block





@dataclass
class  __lowerCAmelCase     (					SCREAMING_SNAKE_CASE							):
  _a									=  42





class  __lowerCAmelCase     (					nn.Module							):




  def __init__(				self							,   lowerCAmelCase=3							,   lowerCAmelCase=3							,   lowerCAmelCase=("DownEncoderBlock2D",)							,   lowerCAmelCase=(64,)							,   lowerCAmelCase=2							,   lowerCAmelCase=32							,   lowerCAmelCase="silu"							,   lowerCAmelCase=True							,   )  -> List[str]:
        '''simple docstring'''






        super().__init__()
        _lowercase	       =layers_per_block

        _lowercase	       =torch.nn.Convad(
            lowerCAmelCase							,   block_out_channels[0]							,   kernel_size=3							,   stride=1							,   padding=1							,   )

        _lowercase	       =None
        _lowercase	       =nn.ModuleList([]					)

        # down
        _lowercase	       =block_out_channels[0]
        for i, down_block_type in enumerate(lowerCAmelCase					):
              _lowercase	       =output_channel
              _lowercase	       =block_out_channels[i]
              _lowercase	       =i == len(lowerCAmelCase					) - 1

              _lowercase	       =get_down_block(
                  lowerCAmelCase							,   num_layers=self.layers_per_block							,   in_channels=lowerCAmelCase							,   out_channels=lowerCAmelCase							,   add_downsample=not is_final_block							,   resnet_eps=1e-6							,   downsample_padding=0							,   resnet_act_fn=lowerCAmelCase							,   resnet_groups=lowerCAmelCase							,   attention_head_dim=lowerCAmelCase							,   temb_channels=lowerCAmelCase							,   )
              self.down_blocks.append(lowerCAmelCase					)

        # mid
        _lowercase	       =UNetMidBlockaD(
            in_channels=block_out_channels[-1]							,   resnet_eps=1e-6							,   resnet_act_fn=lowerCAmelCase							,   output_scale_factor=1							,   resnet_time_scale_shift='default'							,   attention_head_dim=block_out_channels[-1]							,   resnet_groups=lowerCAmelCase							,   temb_channels=lowerCAmelCase							,   )

        # out
        _lowercase	       =nn.GroupNorm(num_channels=block_out_channels[-1]							,   num_groups=lowerCAmelCase							,   eps=1e-6					)
        _lowercase	       =nn.SiLU()

        _lowercase	       =2 * out_channels if double_z else out_channels
        _lowercase	       =nn.Convad(block_out_channels[-1]							,   lowerCAmelCase							,   3							,   padding=1					)

        _lowercase	       =False




  def   A__     (				self							,   lowerCAmelCase					)  -> List[str]:
        '''simple docstring'''






        _lowercase	       =x
        _lowercase	       =self.conv_in(lowerCAmelCase					)

        if self.training and self.gradient_checkpointing:

              def create_custom_forward(lowerCAmelCase					):
                    def custom_forward(*lowerCAmelCase					):
                          return module(*lowerCAmelCase					)

                    return custom_forward

              # down
              if is_torch_version('>='							,   '1.11.0'					):
                    for down_block in self.down_blocks:
                          _lowercase	       =torch.utils.checkpoint.checkpoint(
                              create_custom_forward(lowerCAmelCase					)							,   lowerCAmelCase							,   use_reentrant=lowerCAmelCase					)
                    # middle
                    _lowercase	       =torch.utils.checkpoint.checkpoint(
                        create_custom_forward(self.mid_block					)							,   lowerCAmelCase							,   use_reentrant=lowerCAmelCase					)
              else:
                    for down_block in self.down_blocks:
                          _lowercase	       =torch.utils.checkpoint.checkpoint(create_custom_forward(lowerCAmelCase					)							,   lowerCAmelCase					)
                    # middle
                    _lowercase	       =torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block					)							,   lowerCAmelCase					)

        else:
              # down
              for down_block in self.down_blocks:
                    _lowercase	       =down_block(lowerCAmelCase					)

              # middle
              _lowercase	       =self.mid_block(lowerCAmelCase					)

        # post-process
        _lowercase	       =self.conv_norm_out(lowerCAmelCase					)
        _lowercase	       =self.conv_act(lowerCAmelCase					)
        _lowercase	       =self.conv_out(lowerCAmelCase					)

        return sample





class  __lowerCAmelCase     (					nn.Module							):




  def __init__(				self							,   lowerCAmelCase=3							,   lowerCAmelCase=3							,   lowerCAmelCase=("UpDecoderBlock2D",)							,   lowerCAmelCase=(64,)							,   lowerCAmelCase=2							,   lowerCAmelCase=32							,   lowerCAmelCase="silu"							,   lowerCAmelCase="group"							,   )  -> List[Any]:
        '''simple docstring'''






        super().__init__()
        _lowercase	       =layers_per_block

        _lowercase	       =nn.Convad(
            lowerCAmelCase							,   block_out_channels[-1]							,   kernel_size=3							,   stride=1							,   padding=1							,   )

        _lowercase	       =None
        _lowercase	       =nn.ModuleList([]					)

        _lowercase	       =in_channels if norm_type == 'spatial' else None

        # mid
        _lowercase	       =UNetMidBlockaD(
            in_channels=block_out_channels[-1]							,   resnet_eps=1e-6							,   resnet_act_fn=lowerCAmelCase							,   output_scale_factor=1							,   resnet_time_scale_shift='default' if norm_type == 'group' else norm_type							,   attention_head_dim=block_out_channels[-1]							,   resnet_groups=lowerCAmelCase							,   temb_channels=lowerCAmelCase							,   )

        # up
        _lowercase	       =list(reversed(lowerCAmelCase					)					)
        _lowercase	       =reversed_block_out_channels[0]
        for i, up_block_type in enumerate(lowerCAmelCase					):
              _lowercase	       =output_channel
              _lowercase	       =reversed_block_out_channels[i]

              _lowercase	       =i == len(lowerCAmelCase					) - 1

              _lowercase	       =get_up_block(
                  lowerCAmelCase							,   num_layers=self.layers_per_block + 1							,   in_channels=lowerCAmelCase							,   out_channels=lowerCAmelCase							,   prev_output_channel=lowerCAmelCase							,   add_upsample=not is_final_block							,   resnet_eps=1e-6							,   resnet_act_fn=lowerCAmelCase							,   resnet_groups=lowerCAmelCase							,   attention_head_dim=lowerCAmelCase							,   temb_channels=lowerCAmelCase							,   resnet_time_scale_shift=lowerCAmelCase							,   )
              self.up_blocks.append(lowerCAmelCase					)
              _lowercase	       =output_channel

        # out
        if norm_type == "spatial":
              _lowercase	       =SpatialNorm(block_out_channels[0]							,   lowerCAmelCase					)
        else:
              _lowercase	       =nn.GroupNorm(num_channels=block_out_channels[0]							,   num_groups=lowerCAmelCase							,   eps=1e-6					)
        _lowercase	       =nn.SiLU()
        _lowercase	       =nn.Convad(block_out_channels[0]							,   lowerCAmelCase							,   3							,   padding=1					)

        _lowercase	       =False




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase=None					)  -> Optional[int]:
        '''simple docstring'''






        _lowercase	       =z
        _lowercase	       =self.conv_in(lowerCAmelCase					)

        _lowercase	       =next(iter(self.up_blocks.parameters()					)					).dtype
        if self.training and self.gradient_checkpointing:

              def create_custom_forward(lowerCAmelCase					):
                    def custom_forward(*lowerCAmelCase					):
                          return module(*lowerCAmelCase					)

                    return custom_forward

              if is_torch_version('>='							,   '1.11.0'					):
                    # middle
                    _lowercase	       =torch.utils.checkpoint.checkpoint(
                        create_custom_forward(self.mid_block					)							,   lowerCAmelCase							,   lowerCAmelCase							,   use_reentrant=lowerCAmelCase					)
                    _lowercase	       =sample.to(lowerCAmelCase					)

                    # up
                    for up_block in self.up_blocks:
                          _lowercase	       =torch.utils.checkpoint.checkpoint(
                              create_custom_forward(lowerCAmelCase					)							,   lowerCAmelCase							,   lowerCAmelCase							,   use_reentrant=lowerCAmelCase					)
              else:
                    # middle
                    _lowercase	       =torch.utils.checkpoint.checkpoint(
                        create_custom_forward(self.mid_block					)							,   lowerCAmelCase							,   lowerCAmelCase					)
                    _lowercase	       =sample.to(lowerCAmelCase					)

                    # up
                    for up_block in self.up_blocks:
                          _lowercase	       =torch.utils.checkpoint.checkpoint(create_custom_forward(lowerCAmelCase					)							,   lowerCAmelCase							,   lowerCAmelCase					)
        else:
              # middle
              _lowercase	       =self.mid_block(lowerCAmelCase							,   lowerCAmelCase					)
              _lowercase	       =sample.to(lowerCAmelCase					)

              # up
              for up_block in self.up_blocks:
                    _lowercase	       =up_block(lowerCAmelCase							,   lowerCAmelCase					)

        # post-process
        if latent_embeds is None:
              _lowercase	       =self.conv_norm_out(lowerCAmelCase					)
        else:
              _lowercase	       =self.conv_norm_out(lowerCAmelCase							,   lowerCAmelCase					)
        _lowercase	       =self.conv_act(lowerCAmelCase					)
        _lowercase	       =self.conv_out(lowerCAmelCase					)

        return sample





class  __lowerCAmelCase     (					nn.Module							):




  def __init__(				self							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase=None							,   lowerCAmelCase="random"							,   lowerCAmelCase=False							,   lowerCAmelCase=True					)  -> int:
        '''simple docstring'''






        super().__init__()
        _lowercase	       =n_e
        _lowercase	       =vq_embed_dim
        _lowercase	       =beta
        _lowercase	       =legacy

        _lowercase	       =nn.Embedding(self.n_e							,   self.vq_embed_dim					)
        self.embedding.weight.data.uniform_(-1.0 / self.n_e							,   1.0 / self.n_e					)

        _lowercase	       =remap
        if self.remap is not None:
              self.register_buffer('used'							,   torch.tensor(np.load(self.remap					)					)					)
              _lowercase	       =self.used.shape[0]
              _lowercase	       =unknown_index  # "random" or "extra" or integer
              if self.unknown_index == "extra":
                    _lowercase	       =self.re_embed
                    _lowercase	       =self.re_embed + 1
              print(
                  F'''Remapping {self.n_e} indices to {self.re_embed} indices. '''
                  F'''Using {self.unknown_index} for unknown indices.'''					)
        else:
              _lowercase	       =n_e

        _lowercase	       =sane_index_shape




  def   A__     (				self							,   lowerCAmelCase					)  -> Optional[Any]:
        '''simple docstring'''






        _lowercase	       =inds.shape
        assert len(lowerCAmelCase					) > 1
        _lowercase	       =inds.reshape(ishape[0]							,   -1					)
        _lowercase	       =self.used.to(lowerCAmelCase					)
        _lowercase	       =(inds[:, :, None] == used[None, None, ...]).long()
        _lowercase	       =match.argmax(-1					)
        _lowercase	       =match.sum(2					) < 1
        if self.unknown_index == "random":
              _lowercase	       =torch.randint(0							,   self.re_embed							,   size=new[unknown].shape					).to(device=new.device					)
        else:
              _lowercase	       =self.unknown_index
        return new.reshape(lowerCAmelCase					)




  def   A__     (				self							,   lowerCAmelCase					)  -> int:
        '''simple docstring'''






        _lowercase	       =inds.shape
        assert len(lowerCAmelCase					) > 1
        _lowercase	       =inds.reshape(ishape[0]							,   -1					)
        _lowercase	       =self.used.to(lowerCAmelCase					)
        if self.re_embed > self.used.shape[0]:  # extra token
              _lowercase	       =0  # simply set to zero
        _lowercase	       =torch.gather(used[None, :][inds.shape[0] * [0], :]							,   1							,   lowerCAmelCase					)
        return back.reshape(lowerCAmelCase					)




  def   A__     (				self							,   lowerCAmelCase					)  -> str:
        '''simple docstring'''






        _lowercase	       =z.permute(0							,   2							,   3							,   1					).contiguous()
        _lowercase	       =z.view(-1							,   self.vq_embed_dim					)

        # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
        _lowercase	       =torch.argmin(torch.cdist(lowerCAmelCase							,   self.embedding.weight					)							,   dim=1					)

        _lowercase	       =self.embedding(lowerCAmelCase					).view(z.shape					)
        _lowercase	       =None
        _lowercase	       =None

        # compute loss for embedding
        if not self.legacy:
              _lowercase	       =self.beta * torch.mean((z_q.detach() - z) ** 2					) + torch.mean((z_q - z.detach()) ** 2					)
        else:
              _lowercase	       =torch.mean((z_q.detach() - z) ** 2					) + self.beta * torch.mean((z_q - z.detach()) ** 2					)

        # preserve gradients
        _lowercase	       =z + (z_q - z).detach()

        # reshape back to match original input shape
        _lowercase	       =z_q.permute(0							,   3							,   1							,   2					).contiguous()

        if self.remap is not None:
              _lowercase	       =min_encoding_indices.reshape(z.shape[0]							,   -1					)  # add batch axis
              _lowercase	       =self.remap_to_used(lowerCAmelCase					)
              _lowercase	       =min_encoding_indices.reshape(-1							,   1					)  # flatten

        if self.sane_index_shape:
              _lowercase	       =min_encoding_indices.reshape(z_q.shape[0]							,   z_q.shape[2]							,   z_q.shape[3]					)

        return z_q, loss, (perplexity, min_encodings, min_encoding_indices)




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase					)  -> Optional[Any]:
        '''simple docstring'''






        if self.remap is not None:
              _lowercase	       =indices.reshape(shape[0]							,   -1					)  # add batch axis
              _lowercase	       =self.unmap_to_all(lowerCAmelCase					)
              _lowercase	       =indices.reshape(-1					)  # flatten again

        # get quantized latent vectors
        _lowercase	       =self.embedding(lowerCAmelCase					)

        if shape is not None:
              _lowercase	       =z_q.view(lowerCAmelCase					)
              # reshape back to match original input shape
              _lowercase	       =z_q.permute(0							,   3							,   1							,   2					).contiguous()

        return z_q





class  __lowerCAmelCase     (					SCREAMING_SNAKE_CASE							):




  def __init__(				self							,   lowerCAmelCase							,   lowerCAmelCase=False					)  -> Dict:
        '''simple docstring'''






        _lowercase	       =parameters
        _lowercase       ,       _lowercase	       =torch.chunk(lowerCAmelCase							,   2							,   dim=1					)
        _lowercase	       =torch.clamp(self.logvar							,   -30.0							,   20.0					)
        _lowercase	       =deterministic
        _lowercase	       =torch.exp(0.5 * self.logvar					)
        _lowercase	       =torch.exp(self.logvar					)
        if self.deterministic:
              _lowercase	       =_lowercase	       =torch.zeros_like(
                  self.mean							,   device=self.parameters.device							,   dtype=self.parameters.dtype					)




  def   A__     (				self							,   lowerCAmelCase = None					)  -> torch.FloatTensor:
        '''simple docstring'''






        _lowercase	       =randn_tensor(
            self.mean.shape							,   generator=lowerCAmelCase							,   device=self.parameters.device							,   dtype=self.parameters.dtype					)
        _lowercase	       =self.mean + self.std * sample
        return x




  def   A__     (				self							,   lowerCAmelCase=None					)  -> Optional[int]:
        '''simple docstring'''






        if self.deterministic:
              return torch.Tensor([0.0]					)
        else:
              if other is None:
                    return 0.5 * torch.sum(torch.pow(self.mean							,   2					) + self.var - 1.0 - self.logvar							,   dim=[1, 2, 3]					)
              else:
                    return 0.5 * torch.sum(
                        torch.pow(self.mean - other.mean							,   2					) / other.var
                        + self.var / other.var
                        - 1.0
                        - self.logvar
                        + other.logvar							,   dim=[1, 2, 3]							,   )




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase=[1, 2, 3]					)  -> str:
        '''simple docstring'''






        if self.deterministic:
              return torch.Tensor([0.0]					)
        _lowercase	       =np.log(2.0 * np.pi					)
        return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean							,   2					) / self.var							,   dim=lowerCAmelCase					)




  def   A__     (				self					)  -> str:
        '''simple docstring'''






        return self.mean





import torch
import torch.nn as nn
from transformers.modeling_utils import ModuleUtilsMixin
from transformers.models.ta.modeling_ta import TaBlock, TaConfig, TaLayerNorm

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin





class  __lowerCAmelCase     (					SCREAMING_SNAKE_CASE			,			SCREAMING_SNAKE_CASE			,			SCREAMING_SNAKE_CASE							):




  @register_to_config
  def __init__(				self							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase = False							,   )  -> Tuple:
        '''simple docstring'''






        super().__init__()

        _lowercase	       =nn.Embedding(lowerCAmelCase							,   lowerCAmelCase					)

        _lowercase	       =nn.Embedding(lowerCAmelCase							,   lowerCAmelCase					)
        _lowercase	       =False

        _lowercase	       =nn.Dropout(p=lowerCAmelCase					)

        _lowercase	       =TaConfig(
            vocab_size=lowerCAmelCase							,   d_model=lowerCAmelCase							,   num_heads=lowerCAmelCase							,   d_kv=lowerCAmelCase							,   d_ff=lowerCAmelCase							,   dropout_rate=lowerCAmelCase							,   feed_forward_proj=lowerCAmelCase							,   is_decoder=lowerCAmelCase							,   is_encoder_decoder=lowerCAmelCase							,   )

        _lowercase	       =nn.ModuleList()
        for lyr_num in range(lowerCAmelCase					):
              _lowercase	       =TaBlock(lowerCAmelCase					)
              self.encoders.append(lowerCAmelCase					)

        _lowercase	       =TaLayerNorm(lowerCAmelCase					)
        _lowercase	       =nn.Dropout(p=lowerCAmelCase					)




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase					)  -> Dict:
        '''simple docstring'''






        _lowercase	       =self.token_embedder(lowerCAmelCase					)

        _lowercase	       =encoder_input_tokens.shape[1]
        _lowercase	       =torch.arange(lowerCAmelCase							,   device=encoder_input_tokens.device					)
        x += self.position_encoding(lowerCAmelCase					)

        _lowercase	       =self.dropout_pre(lowerCAmelCase					)

        # inverted the attention mask
        _lowercase	       =encoder_input_tokens.size()
        _lowercase	       =self.get_extended_attention_mask(lowerCAmelCase							,   lowerCAmelCase					)

        for lyr in self.encoders:
              _lowercase	       =lyr(lowerCAmelCase							,   lowerCAmelCase					)[0]
        _lowercase	       =self.layer_norm(lowerCAmelCase					)

        return self.dropout_post(lowerCAmelCase					), encoder_inputs_mask


