
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  (		__magic_name__							):

 """simple docstring"""
 SCREAMING_SNAKE_CASE_				:					torch.FloatTensor




class       _lowerCAmelCase  (		nn.Module							):

 """simple docstring"""

 def __init__(						self			:			Any					, SCREAMING_SNAKE_CASE__			:			List[Any]=3					, SCREAMING_SNAKE_CASE__			:			int=3					, SCREAMING_SNAKE_CASE__			:			int=("DownEncoderBlock2D",)					, SCREAMING_SNAKE_CASE__			:			int=(64,)					, SCREAMING_SNAKE_CASE__			:			Optional[Any]=2					, SCREAMING_SNAKE_CASE__			:			List[Any]=32					, SCREAMING_SNAKE_CASE__			:			Tuple="silu"					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]=True					, ):






      """simple docstring"""



      super().__init__()
      UpperCamelCase						    =					layers_per_block

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

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

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

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

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

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

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

      UpperCamelCase						    =					False






 def      __lowerCAmelCase  (						self			:			Any					, SCREAMING_SNAKE_CASE__			:			Tuple   ):






      """simple docstring"""



      UpperCamelCase						    =					x
      UpperCamelCase						    =					self.conv_in(SCREAMING_SNAKE_CASE__   )

      if self.training and self.gradient_checkpointing:

           def create_custom_forward(SCREAMING_SNAKE_CASE__			:			List[str]   ):
                def custom_forward(*SCREAMING_SNAKE_CASE__			:			int   ):
                     return module(*SCREAMING_SNAKE_CASE__   )

                return custom_forward

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

      else:
           # down
           for down_block in self.down_blocks:
                UpperCamelCase						    =					down_block(SCREAMING_SNAKE_CASE__   )

           # middle
           UpperCamelCase						    =					self.mid_block(SCREAMING_SNAKE_CASE__   )

      # post-process
      UpperCamelCase						    =					self.conv_norm_out(SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					self.conv_act(SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					self.conv_out(SCREAMING_SNAKE_CASE__   )

      return sample




class       _lowerCAmelCase  (		nn.Module							):

 """simple docstring"""

 def __init__(						self			:			Optional[Any]					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]=3					, SCREAMING_SNAKE_CASE__			:			str=3					, SCREAMING_SNAKE_CASE__			:			Tuple=("UpDecoderBlock2D",)					, SCREAMING_SNAKE_CASE__			:			Optional[int]=(64,)					, SCREAMING_SNAKE_CASE__			:			str=2					, SCREAMING_SNAKE_CASE__			:			Tuple=32					, SCREAMING_SNAKE_CASE__			:			List[Any]="silu"					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]="group"					, ):






      """simple docstring"""



      super().__init__()
      UpperCamelCase						    =					layers_per_block

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

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

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

      # mid
      UpperCamelCase						    =					UNetMidBlockaD(
          in_channels=block_out_channels[-1]					, resnet_eps=1e-6					, resnet_act_fn=SCREAMING_SNAKE_CASE__					, 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=SCREAMING_SNAKE_CASE__					, temb_channels=SCREAMING_SNAKE_CASE__					, )

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

           UpperCamelCase						    =					i == len(SCREAMING_SNAKE_CASE__   ) - 1

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

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

      UpperCamelCase						    =					False






 def      __lowerCAmelCase  (						self			:			Tuple					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]=None   ):






      """simple docstring"""



      UpperCamelCase						    =					z
      UpperCamelCase						    =					self.conv_in(SCREAMING_SNAKE_CASE__   )

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

           def create_custom_forward(SCREAMING_SNAKE_CASE__			:			Optional[int]   ):
                def custom_forward(*SCREAMING_SNAKE_CASE__			:			List[str]   ):
                     return module(*SCREAMING_SNAKE_CASE__   )

                return custom_forward

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

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

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

           # up
           for up_block in self.up_blocks:
                UpperCamelCase						    =					up_block(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )

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

      return sample




class       _lowerCAmelCase  (		nn.Module							):

 """simple docstring"""

 def __init__(						self			:			Dict					, SCREAMING_SNAKE_CASE__			:			Optional[Any]					, SCREAMING_SNAKE_CASE__			:			int					, SCREAMING_SNAKE_CASE__			:			Any					, SCREAMING_SNAKE_CASE__			:			Optional[Any]=None					, SCREAMING_SNAKE_CASE__			:			str="random"					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]=False					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]=True   ):






      """simple docstring"""



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

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

      UpperCamelCase						    =					remap
      if self.remap is not None:
           self.register_buffer('used'					, torch.tensor(np.load(self.remap   )   )   )
           UpperCamelCase						    =					self.used.shape[0]
           UpperCamelCase						    =					unknown_index  # "random" or "extra" or integer
           if self.unknown_index == "extra":
                UpperCamelCase						    =					self.re_embed
                UpperCamelCase						    =					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:
           UpperCamelCase						    =					n_e

      UpperCamelCase						    =					sane_index_shape

 def      __lowerCAmelCase  (						self			:			List[Any]					, SCREAMING_SNAKE_CASE__			:			Optional[int]   ):






      """simple docstring"""



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

 def      __lowerCAmelCase  (						self			:			List[str]					, SCREAMING_SNAKE_CASE__			:			str   ):






      """simple docstring"""



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

 def      __lowerCAmelCase  (						self			:			Any					, SCREAMING_SNAKE_CASE__			:			Optional[int]   ):






      """simple docstring"""



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

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

      UpperCamelCase						    =					self.embedding(SCREAMING_SNAKE_CASE__   ).view(z.shape   )
      UpperCamelCase						    =					None
      UpperCamelCase						    =					None

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

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

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

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

      if self.sane_index_shape:
           UpperCamelCase						    =					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      __lowerCAmelCase  (						self			:			Dict					, SCREAMING_SNAKE_CASE__			:			Any					, SCREAMING_SNAKE_CASE__			:			List[Any]   ):






      """simple docstring"""



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

      # get quantized latent vectors
      UpperCamelCase						    =					self.embedding(SCREAMING_SNAKE_CASE__   )

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

      return z_q




class       _lowerCAmelCase  (		__magic_name__							):

 """simple docstring"""

 def __init__(						self			:			Optional[Any]					, SCREAMING_SNAKE_CASE__			:			Any					, SCREAMING_SNAKE_CASE__			:			Optional[int]=False   ):






      """simple docstring"""



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

 def      __lowerCAmelCase  (						self			:			Any					, SCREAMING_SNAKE_CASE__			:			Optional[torch.Generator] = None   ):






      """simple docstring"""



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

 def      __lowerCAmelCase  (						self			:			str					, SCREAMING_SNAKE_CASE__			:			List[str]=None   ):






      """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      __lowerCAmelCase  (						self			:			int					, SCREAMING_SNAKE_CASE__			:			int					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]=[1, 2, 3]   ):






      """simple docstring"""



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






 def      __lowerCAmelCase  (						self			:			int   ):






      """simple docstring"""



      return self.mean




def  __lowerCamelCase					( _lowercase			)       ->      str:
     return "".join(chr(ord(_lowercase			) - 32			) if 'a' <= char <= 'z' else char for char in word			)


if __name__ == "__main__":
 from doctest import testmod

 testmod()


