
"""simple docstring"""







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 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




 UpperCamelCase_    :  torch.FloatTensor



class 		UpperCAmelCase_						(nn.Module    ):
 """simple docstring"""




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



      """simple docstring"""

      super().__init__()
      UpperCAmelCase_    :     Dict            =	layers_per_block

      UpperCAmelCase_    :     Union[str, Any]            =	torch.nn.Convad(
          a_  , block_out_channels[0]  , kernel_size=3  , stride=1  , padding=1  , )

      UpperCAmelCase_    :     Optional[int]            =	None
      UpperCAmelCase_    :     Tuple            =	nn.ModuleList([]			)

      # down
      UpperCAmelCase_    :     str            =	block_out_channels[0]
      for i, down_block_type in enumerate(a_			):
           UpperCAmelCase_    :     Union[str, Any]            =	output_channel
           UpperCAmelCase_    :     int            =	block_out_channels[i]
           UpperCAmelCase_    :     Optional[Any]            =	i == len(a_			) - 1

           UpperCAmelCase_    :     Dict            =	get_down_block(
               a_  , num_layers=self.layers_per_block  , in_channels=a_  , out_channels=a_  , add_downsample=not is_final_block  , resnet_eps=1E-6  , downsample_padding=0  , resnet_act_fn=a_  , resnet_groups=a_  , attention_head_dim=a_  , temb_channels=a_  , )
           self.down_blocks.append(a_			)

      # mid
      UpperCAmelCase_    :     Optional[Any]            =	UNetMidBlockaD(
          in_channels=block_out_channels[-1]  , resnet_eps=1E-6  , resnet_act_fn=a_  , output_scale_factor=1  , resnet_time_scale_shift="""default"""  , attention_head_dim=block_out_channels[-1]  , resnet_groups=a_  , temb_channels=a_  , )

      # out
      UpperCAmelCase_    :     List[str]            =	nn.GroupNorm(num_channels=block_out_channels[-1]  , num_groups=a_  , eps=1E-6			)
      UpperCAmelCase_    :     int            =	nn.SiLU()

      UpperCAmelCase_    :     str            =	2 * out_channels if double_z else out_channels
      UpperCAmelCase_    :     Tuple            =	nn.Convad(block_out_channels[-1]  , a_  , 3  , padding=1			)

      UpperCAmelCase_    :     Dict            =	False






 def 	a   (			self							:			Tuple  , a_							:			Dict			)->       List[Any]:



      """simple docstring"""

      UpperCAmelCase_    :     Union[str, Any]            =	x
      UpperCAmelCase_    :     Tuple            =	self.conv_in(a_			)

      if self.training and self.gradient_checkpointing:

           def create_custom_forward(a_							:			Tuple			):
                def custom_forward(*a_							:			List[Any]			):
                     return module(*a_			)

                return custom_forward

           # down
           if is_torch_version(""">="""  , """1.11.0"""			):
                for down_block in self.down_blocks:
                     UpperCAmelCase_    :     Tuple            =	torch.utils.checkpoint.checkpoint(
                         create_custom_forward(a_			)  , a_  , use_reentrant=a_			)
                # middle
                UpperCAmelCase_    :     List[Any]            =	torch.utils.checkpoint.checkpoint(
                    create_custom_forward(self.mid_block			)  , a_  , use_reentrant=a_			)
           else:
                for down_block in self.down_blocks:
                     UpperCAmelCase_    :     int            =	torch.utils.checkpoint.checkpoint(create_custom_forward(a_			)  , a_			)
                # middle
                UpperCAmelCase_    :     Union[str, Any]            =	torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block			)  , a_			)

      else:
           # down
           for down_block in self.down_blocks:
                UpperCAmelCase_    :     Dict            =	down_block(a_			)

           # middle
           UpperCAmelCase_    :     List[Any]            =	self.mid_block(a_			)

      # post-process
      UpperCAmelCase_    :     List[str]            =	self.conv_norm_out(a_			)
      UpperCAmelCase_    :     Union[str, Any]            =	self.conv_act(a_			)
      UpperCAmelCase_    :     Union[str, Any]            =	self.conv_out(a_			)

      return sample



class 		UpperCAmelCase_						(nn.Module    ):
 """simple docstring"""




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



      """simple docstring"""

      super().__init__()
      UpperCAmelCase_    :     List[Any]            =	layers_per_block

      UpperCAmelCase_    :     Any            =	nn.Convad(
          a_  , block_out_channels[-1]  , kernel_size=3  , stride=1  , padding=1  , )

      UpperCAmelCase_    :     Any            =	None
      UpperCAmelCase_    :     Tuple            =	nn.ModuleList([]			)

      UpperCAmelCase_    :     List[str]            =	in_channels if norm_type == """spatial""" else None

      # mid
      UpperCAmelCase_    :     List[str]            =	UNetMidBlockaD(
          in_channels=block_out_channels[-1]  , resnet_eps=1E-6  , resnet_act_fn=a_  , 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=a_  , temb_channels=a_  , )

      # up
      UpperCAmelCase_    :     Tuple            =	list(reversed(a_			)			)
      UpperCAmelCase_    :     Union[str, Any]            =	reversed_block_out_channels[0]
      for i, up_block_type in enumerate(a_			):
           UpperCAmelCase_    :     Optional[Any]            =	output_channel
           UpperCAmelCase_    :     str            =	reversed_block_out_channels[i]

           UpperCAmelCase_    :     str            =	i == len(a_			) - 1

           UpperCAmelCase_    :     int            =	get_up_block(
               a_  , num_layers=self.layers_per_block + 1  , in_channels=a_  , out_channels=a_  , prev_output_channel=a_  , add_upsample=not is_final_block  , resnet_eps=1E-6  , resnet_act_fn=a_  , resnet_groups=a_  , attention_head_dim=a_  , temb_channels=a_  , resnet_time_scale_shift=a_  , )
           self.up_blocks.append(a_			)
           UpperCAmelCase_    :     Union[str, Any]            =	output_channel

      # out
      if norm_type == "spatial":
           UpperCAmelCase_    :     Optional[int]            =	SpatialNorm(block_out_channels[0]  , a_			)
      else:
           UpperCAmelCase_    :     Union[str, Any]            =	nn.GroupNorm(num_channels=block_out_channels[0]  , num_groups=a_  , eps=1E-6			)
      UpperCAmelCase_    :     int            =	nn.SiLU()
      UpperCAmelCase_    :     Dict            =	nn.Convad(block_out_channels[0]  , a_  , 3  , padding=1			)

      UpperCAmelCase_    :     int            =	False






 def 	a   (			self							:			List[str]  , a_							:			Optional[int]  , a_							:			Any=None			)->       Tuple:



      """simple docstring"""

      UpperCAmelCase_    :     Any            =	z
      UpperCAmelCase_    :     Optional[Any]            =	self.conv_in(a_			)

      UpperCAmelCase_    :     Optional[Any]            =	next(iter(self.up_blocks.parameters()			)			).dtype
      if self.training and self.gradient_checkpointing:

           def create_custom_forward(a_							:			str			):
                def custom_forward(*a_							:			str			):
                     return module(*a_			)

                return custom_forward

           if is_torch_version(""">="""  , """1.11.0"""			):
                # middle
                UpperCAmelCase_    :     Dict            =	torch.utils.checkpoint.checkpoint(
                    create_custom_forward(self.mid_block			)  , a_  , a_  , use_reentrant=a_			)
                UpperCAmelCase_    :     List[Any]            =	sample.to(a_			)

                # up
                for up_block in self.up_blocks:
                     UpperCAmelCase_    :     Union[str, Any]            =	torch.utils.checkpoint.checkpoint(
                         create_custom_forward(a_			)  , a_  , a_  , use_reentrant=a_			)
           else:
                # middle
                UpperCAmelCase_    :     Any            =	torch.utils.checkpoint.checkpoint(
                    create_custom_forward(self.mid_block			)  , a_  , a_			)
                UpperCAmelCase_    :     int            =	sample.to(a_			)

                # up
                for up_block in self.up_blocks:
                     UpperCAmelCase_    :     Optional[int]            =	torch.utils.checkpoint.checkpoint(create_custom_forward(a_			)  , a_  , a_			)
      else:
           # middle
           UpperCAmelCase_    :     List[Any]            =	self.mid_block(a_  , a_			)
           UpperCAmelCase_    :     Dict            =	sample.to(a_			)

           # up
           for up_block in self.up_blocks:
                UpperCAmelCase_    :     Optional[Any]            =	up_block(a_  , a_			)

        # post-process
      if latent_embeds is None:
           UpperCAmelCase_    :     List[Any]            =	self.conv_norm_out(a_			)
      else:
           UpperCAmelCase_    :     Dict            =	self.conv_norm_out(a_  , a_			)
      UpperCAmelCase_    :     List[Any]            =	self.conv_act(a_			)
      UpperCAmelCase_    :     Any            =	self.conv_out(a_			)

      return sample



class 		UpperCAmelCase_						(nn.Module    ):
 """simple docstring"""




 def __init__(			self							:			List[Any]  , a_							:			Optional[Any]  , a_							:			Optional[Any]  , a_							:			str  , a_							:			Dict=None  , a_							:			List[str]="random"  , a_							:			List[Any]=False  , a_							:			List[str]=True			)->       List[str]:



      """simple docstring"""

      super().__init__()
      UpperCAmelCase_    :     Optional[int]            =	n_e
      UpperCAmelCase_    :     List[Any]            =	vq_embed_dim
      UpperCAmelCase_    :     Tuple            =	beta
      UpperCAmelCase_    :     Union[str, Any]            =	legacy

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

      UpperCAmelCase_    :     List[str]            =	sane_index_shape




 def 	a   (			self							:			Optional[Any]  , a_							:			str			)->       Any:



      """simple docstring"""

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




 def 	a   (			self							:			Tuple  , a_							:			List[str]			)->       Optional[int]:



      """simple docstring"""

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




 def 	a   (			self							:			str  , a_							:			Tuple			)->       Optional[Any]:



      """simple docstring"""

      UpperCAmelCase_    :     List[Any]            =	z.permute(0  , 2  , 3  , 1			).contiguous()
      UpperCAmelCase_    :     Any            =	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_    :     Optional[Any]            =	torch.argmin(torch.cdist(a_  , self.embedding.weight			)  , dim=1			)

      UpperCAmelCase_    :     int            =	self.embedding(a_			).view(z.shape			)
      UpperCAmelCase_    :     Optional[Any]            =	None
      UpperCAmelCase_    :     str            =	None

      # compute loss for embedding
      if not self.legacy:
           UpperCAmelCase_    :     Optional[int]            =	self.beta * torch.mean((z_q.detach() - z) ** 2			) + torch.mean((z_q - z.detach()) ** 2			)
      else:
           UpperCAmelCase_    :     Tuple            =	torch.mean((z_q.detach() - z) ** 2			) + self.beta * torch.mean((z_q - z.detach()) ** 2			)

      # preserve gradients
      UpperCAmelCase_    :     Tuple            =	z + (z_q - z).detach()

      # reshape back to match original input shape
      UpperCAmelCase_    :     Union[str, Any]            =	z_q.permute(0  , 3  , 1  , 2			).contiguous()

      if self.remap is not None:
           UpperCAmelCase_    :     Union[str, Any]            =	min_encoding_indices.reshape(z.shape[0]  , -1			)  # add batch axis
           UpperCAmelCase_    :     Dict            =	self.remap_to_used(a_			)
           UpperCAmelCase_    :     Optional[int]            =	min_encoding_indices.reshape(-1  , 1			)  # flatten

      if self.sane_index_shape:
           UpperCAmelCase_    :     str            =	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							:			List[Any]  , a_							:			List[str]  , a_							:			List[str]			)->       int:



      """simple docstring"""

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

      # get quantized latent vectors
      UpperCAmelCase_    :     List[Any]            =	self.embedding(a_			)

      if shape is not None:
           UpperCAmelCase_    :     List[str]            =	z_q.view(a_			)
           # reshape back to match original input shape
           UpperCAmelCase_    :     List[str]            =	z_q.permute(0  , 3  , 1  , 2			).contiguous()

      return z_q



class 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




 def __init__(			self							:			Optional[int]  , a_							:			List[str]  , a_							:			str=False			)->       List[Any]:



      """simple docstring"""

      UpperCAmelCase_    :     str            =	parameters
      UpperCAmelCase_       ,UpperCAmelCase_    :     Union[str, Any]            =	torch.chunk(a_  , 2  , dim=1			)
      UpperCAmelCase_    :     int            =	torch.clamp(self.logvar  , -30.0  , 20.0			)
      UpperCAmelCase_    :     Tuple            =	deterministic
      UpperCAmelCase_    :     List[Any]            =	torch.exp(0.5 * self.logvar			)
      UpperCAmelCase_    :     Dict            =	torch.exp(self.logvar			)
      if self.deterministic:
           UpperCAmelCase_    :     Tuple            =	torch.zeros_like(
               self.mean  , device=self.parameters.device  , dtype=self.parameters.dtype			)




 def 	a   (			self							:			List[str]  , a_							:			Optional[torch.Generator] = None			)->       torch.FloatTensor:



      """simple docstring"""

      UpperCAmelCase_    :     int            =	randn_tensor(
          self.mean.shape  , generator=a_  , device=self.parameters.device  , dtype=self.parameters.dtype			)
      UpperCAmelCase_    :     Optional[int]            =	self.mean + self.std * sample
      return x




 def 	a   (			self							:			Dict  , a_							:			Optional[int]=None			)->       List[Any]:



      """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							:			List[str]  , a_							:			Optional[int]  , a_							:			Union[str, Any]=[1, 2, 3]			)->       Tuple:



      """simple docstring"""

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






 def 	a   (			self							:			Tuple			)->       Dict:



      """simple docstring"""

      return self.mean



"""simple docstring"""







from typing import Optional
from urllib.parse import quote

import huggingface_hub as hfh
from packaging import version




def       A_      (       lowercase    ,			lowercase    ,			lowercase = None   )  ->				str:


     """simple docstring"""







     if version.parse(hfh.__version__   ).release < version.parse("""0.11.0"""   ).release:
          # old versions of hfh don't url-encode the file path
          UpperCAmelCase_    :     Tuple            =	quote(lowercase   )
     return hfh.hf_hub_url(lowercase    ,			lowercase    ,			repo_type="""dataset"""    ,			revision=lowercase   )

