





'''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 					__snake_case  (					__SCREAMING_SNAKE_CASE							):






   '''simple docstring'''





   lowerCamelCase__    							=      42



class 					__snake_case  (					nn.Module							):






   '''simple docstring'''





   def __init__(							self	,		__SCREAMING_SNAKE_CASE=3	,		__SCREAMING_SNAKE_CASE=3	,		__SCREAMING_SNAKE_CASE=("DownEncoderBlock2D",)	,		__SCREAMING_SNAKE_CASE=(6_4,)	,		__SCREAMING_SNAKE_CASE=2	,		__SCREAMING_SNAKE_CASE=3_2	,		__SCREAMING_SNAKE_CASE="silu"	,		__SCREAMING_SNAKE_CASE=True	,		):
    super().__init__()
    snake_case__     :					str    =  layers_per_block

    snake_case__     :					int    =  torch.nn.Convad(
        __SCREAMING_SNAKE_CASE	,		block_out_channels[0]	,		kernel_size=3	,		stride=1	,		padding=1	,		)

    snake_case__     :					List[Any]    =  None
    snake_case__     :					List[Any]    =  nn.ModuleList([]       )

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

     snake_case__     :					str    =  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
    snake_case__     :					Optional[Any]    =  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
    snake_case__     :					Tuple    =  nn.GroupNorm(num_channels=block_out_channels[-1]	,		num_groups=__SCREAMING_SNAKE_CASE	,		eps=1e-6       )
    snake_case__     :					Tuple    =  nn.SiLU()

    snake_case__     :					str    =  2 * out_channels if double_z else out_channels
    snake_case__     :					int    =  nn.Convad(block_out_channels[-1]	,		__SCREAMING_SNAKE_CASE	,		3	,		padding=1       )

    snake_case__     :					Union[str, Any]    =  False


   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					Optional[Any]    =  x
    snake_case__     :					int    =  self.conv_in(__SCREAMING_SNAKE_CASE       )

    if self.training and self.gradient_checkpointing:

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

      return custom_forward

     # down
     if is_torch_version(""">="""	,		"""1.11.0"""       ):
      for down_block in self.down_blocks:
       snake_case__     :					List[Any]    =  torch.utils.checkpoint.checkpoint(
           create_custom_forward(__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE	,		use_reentrant=__SCREAMING_SNAKE_CASE       )
      # middle
      snake_case__     :					List[Any]    =  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:
       snake_case__     :					Dict    =  torch.utils.checkpoint.checkpoint(create_custom_forward(__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )
      # middle
      snake_case__     :					str    =  torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block       )	,		__SCREAMING_SNAKE_CASE       )

    else:
     # down
     for down_block in self.down_blocks:
      snake_case__     :					List[str]    =  down_block(__SCREAMING_SNAKE_CASE       )

     # middle
     snake_case__     :					str    =  self.mid_block(__SCREAMING_SNAKE_CASE       )

    # post-process
    snake_case__     :					Any    =  self.conv_norm_out(__SCREAMING_SNAKE_CASE       )
    snake_case__     :					List[str]    =  self.conv_act(__SCREAMING_SNAKE_CASE       )
    snake_case__     :					str    =  self.conv_out(__SCREAMING_SNAKE_CASE       )

    return sample



class 					__snake_case  (					nn.Module							):






   '''simple docstring'''





   def __init__(							self	,		__SCREAMING_SNAKE_CASE=3	,		__SCREAMING_SNAKE_CASE=3	,		__SCREAMING_SNAKE_CASE=("UpDecoderBlock2D",)	,		__SCREAMING_SNAKE_CASE=(6_4,)	,		__SCREAMING_SNAKE_CASE=2	,		__SCREAMING_SNAKE_CASE=3_2	,		__SCREAMING_SNAKE_CASE="silu"	,		__SCREAMING_SNAKE_CASE="group"	,		):
    super().__init__()
    snake_case__     :					Any    =  layers_per_block

    snake_case__     :					Optional[Any]    =  nn.Convad(
        __SCREAMING_SNAKE_CASE	,		block_out_channels[-1]	,		kernel_size=3	,		stride=1	,		padding=1	,		)

    snake_case__     :					Union[str, Any]    =  None
    snake_case__     :					Dict    =  nn.ModuleList([]       )

    snake_case__     :					Optional[int]    =  in_channels if norm_type == """spatial""" else None

    # mid
    snake_case__     :					Tuple    =  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
    snake_case__     :					List[Any]    =  list(reversed(__SCREAMING_SNAKE_CASE       )       )
    snake_case__     :					Optional[Any]    =  reversed_block_out_channels[0]
    for i, up_block_type in enumerate(__SCREAMING_SNAKE_CASE       ):
     snake_case__     :					List[Any]    =  output_channel
     snake_case__     :					Optional[Any]    =  reversed_block_out_channels[i]

     snake_case__     :					List[str]    =  i == len(__SCREAMING_SNAKE_CASE       ) - 1

     snake_case__     :					int    =  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       )
     snake_case__     :					int    =  output_channel

    # out
    if norm_type == "spatial":
     snake_case__     :					List[Any]    =  SpatialNorm(block_out_channels[0]	,		__SCREAMING_SNAKE_CASE       )
    else:
     snake_case__     :					Any    =  nn.GroupNorm(num_channels=block_out_channels[0]	,		num_groups=__SCREAMING_SNAKE_CASE	,		eps=1e-6       )
    snake_case__     :					Tuple    =  nn.SiLU()
    snake_case__     :					Union[str, Any]    =  nn.Convad(block_out_channels[0]	,		__SCREAMING_SNAKE_CASE	,		3	,		padding=1       )

    snake_case__     :					int    =  False


   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE=None       ):
    snake_case__     :					Union[str, Any]    =  z
    snake_case__     :					Any    =  self.conv_in(__SCREAMING_SNAKE_CASE       )

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

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

      return custom_forward

     if is_torch_version(""">="""	,		"""1.11.0"""       ):
      # middle
      snake_case__     :					int    =  torch.utils.checkpoint.checkpoint(
          create_custom_forward(self.mid_block       )	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		use_reentrant=__SCREAMING_SNAKE_CASE       )
      snake_case__     :					int    =  sample.to(__SCREAMING_SNAKE_CASE       )

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

      # up
      for up_block in self.up_blocks:
       snake_case__     :					str    =  torch.utils.checkpoint.checkpoint(create_custom_forward(__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )
    else:
     # middle
     snake_case__     :					List[Any]    =  self.mid_block(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )
     snake_case__     :					List[Any]    =  sample.to(__SCREAMING_SNAKE_CASE       )

     # up
     for up_block in self.up_blocks:
      snake_case__     :					Dict    =  up_block(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )

        # post-process
    if latent_embeds is None:
     snake_case__     :					Optional[Any]    =  self.conv_norm_out(__SCREAMING_SNAKE_CASE       )
    else:
     snake_case__     :					str    =  self.conv_norm_out(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )
    snake_case__     :					Any    =  self.conv_act(__SCREAMING_SNAKE_CASE       )
    snake_case__     :					Optional[Any]    =  self.conv_out(__SCREAMING_SNAKE_CASE       )

    return sample



class 					__snake_case  (					nn.Module							):






   '''simple docstring'''





   def __init__(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE=None	,		__SCREAMING_SNAKE_CASE="random"	,		__SCREAMING_SNAKE_CASE=False	,		__SCREAMING_SNAKE_CASE=True       ):
    super().__init__()
    snake_case__     :					int    =  n_e
    snake_case__     :					Optional[int]    =  vq_embed_dim
    snake_case__     :					int    =  beta
    snake_case__     :					Optional[int]    =  legacy

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

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

    snake_case__     :					str    =  sane_index_shape

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

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

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE       ):
    # reshape z -> (batch, height, width, channel) and flatten
    snake_case__     :					Any    =  z.permute(0	,		2	,		3	,		1       ).contiguous()
    snake_case__     :					Optional[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
    snake_case__     :					Dict    =  torch.argmin(torch.cdist(__SCREAMING_SNAKE_CASE	,		self.embedding.weight       )	,		dim=1       )

    snake_case__     :					Union[str, Any]    =  self.embedding(__SCREAMING_SNAKE_CASE       ).view(z.shape       )
    snake_case__     :					List[str]    =  None
    snake_case__     :					Union[str, Any]    =  None

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

    # preserve gradients
    snake_case__     :					Any    =  z + (z_q - z).detach()

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

    if self.remap is not None:
     snake_case__     :					List[Any]    =  min_encoding_indices.reshape(z.shape[0]	,		-1       )  # add batch axis
     snake_case__     :					str    =  self.remap_to_used(__SCREAMING_SNAKE_CASE       )
     snake_case__     :					str    =  min_encoding_indices.reshape(-1	,		1       )  # flatten

    if self.sane_index_shape:
     snake_case__     :					Tuple    =  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 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ):
    # shape specifying (batch, height, width, channel)
    if self.remap is not None:
     snake_case__     :					List[Any]    =  indices.reshape(shape[0]	,		-1       )  # add batch axis
     snake_case__     :					Optional[int]    =  self.unmap_to_all(__SCREAMING_SNAKE_CASE       )
     snake_case__     :					Optional[Any]    =  indices.reshape(-1       )  # flatten again

    # get quantized latent vectors
    snake_case__     :					int    =  self.embedding(__SCREAMING_SNAKE_CASE       )

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

    return z_q



class 					__snake_case  (					__SCREAMING_SNAKE_CASE							):






   '''simple docstring'''





   def __init__(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE=False       ):
    snake_case__     :					Tuple    =  parameters
    snake_case__  ,							snake_case__     :					Any    =  torch.chunk(__SCREAMING_SNAKE_CASE	,		2	,		dim=1       )
    snake_case__     :					Union[str, Any]    =  torch.clamp(self.logvar	,		-30.0	,		20.0       )
    snake_case__     :					Optional[int]    =  deterministic
    snake_case__     :					Optional[int]    =  torch.exp(0.5 * self.logvar       )
    snake_case__     :					Any    =  torch.exp(self.logvar       )
    if self.deterministic:
     snake_case__     :					List[str]    =  torch.zeros_like(
         self.mean	,		device=self.parameters.device	,		dtype=self.parameters.dtype       )

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE = None       ):
    # make sure sample is on the same device as the parameters and has same dtype
    snake_case__     :					Dict    =  randn_tensor(
        self.mean.shape	,		generator=__SCREAMING_SNAKE_CASE	,		device=self.parameters.device	,		dtype=self.parameters.dtype       )
    snake_case__     :					Optional[int]    =  self.mean + self.std * sample
    return x

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE=None       ):
    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 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE=[1, 2, 3]       ):
    if self.deterministic:
     return torch.Tensor([0.0]       )
    snake_case__     :					Any    =  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 		__UpperCamelCase						(							self       ):
    return self.mean









"""simple docstring"""







import importlib
import inspect
import os
import re


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_config_docstrings.py
A		=						'src/transformers'


# This is to make sure the transformers module imported is the one in the repo.
A		=						importlib.util.spec_from_file_location(
    'transformers',
    os.path.join(PATH_TO_TRANSFORMERS, '__init__.py'),
    submodule_search_locations=[PATH_TO_TRANSFORMERS],
)
A		=						spec.loader.load_module()

A		=						transformers.models.auto.configuration_auto.CONFIG_MAPPING

# Regex pattern used to find the checkpoint mentioned in the docstring of `config_class`.
# For example, `[bert-base-uncased](https://huggingface.co/bert-base-uncased)`
A		=						re.compile('\[(.+?)\]\((https://huggingface\.co/.+?)\)')


A		=						{
    'CLIPConfigMixin',
    'DecisionTransformerConfigMixin',
    'EncoderDecoderConfigMixin',
    'RagConfigMixin',
    'SpeechEncoderDecoderConfigMixin',
    'VisionEncoderDecoderConfigMixin',
    'VisionTextDualEncoderConfigMixin',
}

def  __SCREAMING_SNAKE_CASE		(			):
						"""simple docstring"""






						snake_case   :			Union[str, Any]						    = []

						for config_class in list(CONFIG_MAPPING.values() ):
												snake_case   :			List[str]						    = False

												# source code of `config_class`
												snake_case   :			Optional[int]						    = inspect.getsource(lowerCamelCase_ )
												snake_case   :			Optional[int]						    = _re_checkpoint.findall(lowerCamelCase_ )

												for checkpoint in checkpoints:
																		# Each `checkpoint` is a tuple of a checkpoint name and a checkpoint link.
																		# For example, `('bert-base-uncased', 'https://huggingface.co/bert-base-uncased')`
																		snake_case    ,			snake_case   :			Tuple						    = checkpoint

																		# verify the checkpoint name corresponds to the checkpoint link
																		snake_case   :			Any						    = f'''https://huggingface.co/{ckpt_name}'''
																		if ckpt_link == ckpt_link_from_name:
																								snake_case   :			List[str]						    = True
																								break

												snake_case   :			Any						    = config_class.__name__
												if not checkpoint_found and name not in CONFIG_CLASSES_TO_IGNORE_FOR_DOCSTRING_CHECKPOINT_CHECK:
																		configs_without_checkpoint.append(lowerCamelCase_ )

						if len(lowerCamelCase_ ) > 0:
												snake_case   :			List[str]						    = "\n".join(sorted(lowerCamelCase_ ) )
												raise ValueError(f'''The following configurations don\'t contain any valid checkpoint:\n{message}''' )


if __name__ == "__main__":
			check_config_docstrings_have_checkpoints()



