



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     (	lowerCAmelCase__   ):
 lowerCAmelCase__       					=  42



class  __lowerCAmelCase     (	nn.Module   ):




 def __init__(						self	,				__UpperCAmelCase=3	,				__UpperCAmelCase=3	,				__UpperCAmelCase=("DownEncoderBlock2D",)	,				__UpperCAmelCase=(64,)	,				__UpperCAmelCase=2	,				__UpperCAmelCase=32	,				__UpperCAmelCase="silu"	,				__UpperCAmelCase=True	,				):



    '''simple docstring'''
    super().__init__()
    __lowerCamelCase													=    layers_per_block

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

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

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

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

    # mid
    __lowerCamelCase													=    UNetMidBlockaD(
        in_channels=block_out_channels[-1]	,				resnet_eps=1E-6	,				resnet_act_fn=__UpperCAmelCase	,				output_scale_factor=1	,				resnet_time_scale_shift='''default'''	,				attention_head_dim=block_out_channels[-1]	,				resnet_groups=__UpperCAmelCase	,				temb_channels=__UpperCAmelCase	,				)

    # out
    __lowerCamelCase													=    nn.GroupNorm(num_channels=block_out_channels[-1]	,				num_groups=__UpperCAmelCase	,				eps=1E-6				)
    __lowerCamelCase													=    nn.SiLU()

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

    __lowerCamelCase													=    False




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    x
    __lowerCamelCase													=    self.conv_in(__UpperCAmelCase				)

    if self.training and self.gradient_checkpointing:

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

          return custom_forward

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

    else:
       # down
       for down_block in self.down_blocks:
          __lowerCamelCase													=    down_block(__UpperCAmelCase				)

       # middle
       __lowerCamelCase													=    self.mid_block(__UpperCAmelCase				)

    # post-process
    __lowerCamelCase													=    self.conv_norm_out(__UpperCAmelCase				)
    __lowerCamelCase													=    self.conv_act(__UpperCAmelCase				)
    __lowerCamelCase													=    self.conv_out(__UpperCAmelCase				)

    return sample



class  __lowerCAmelCase     (	nn.Module   ):




 def __init__(						self	,				__UpperCAmelCase=3	,				__UpperCAmelCase=3	,				__UpperCAmelCase=("UpDecoderBlock2D",)	,				__UpperCAmelCase=(64,)	,				__UpperCAmelCase=2	,				__UpperCAmelCase=32	,				__UpperCAmelCase="silu"	,				__UpperCAmelCase="group"	,				):



    '''simple docstring'''
    super().__init__()
    __lowerCamelCase													=    layers_per_block

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

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

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

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

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

       __lowerCamelCase													=    i == len(__UpperCAmelCase				) - 1

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

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

    __lowerCamelCase													=    False




 def        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase=None				):



    '''simple docstring'''
    __lowerCamelCase													=    z
    __lowerCamelCase													=    self.conv_in(__UpperCAmelCase				)

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

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

          return custom_forward

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

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

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

       # up
       for up_block in self.up_blocks:
          __lowerCamelCase													=    up_block(__UpperCAmelCase	,				__UpperCAmelCase				)

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

    return sample



class  __lowerCAmelCase     (	nn.Module   ):




 def __init__(						self	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase=None	,				__UpperCAmelCase="random"	,				__UpperCAmelCase=False	,				__UpperCAmelCase=True				):



    '''simple docstring'''
    super().__init__()
    __lowerCamelCase													=    n_e
    __lowerCamelCase													=    vq_embed_dim
    __lowerCamelCase													=    beta
    __lowerCamelCase													=    legacy

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

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

    __lowerCamelCase													=    sane_index_shape




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



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




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



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




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    # reshape z -> (batch, height, width, channel) and flatten
    __lowerCamelCase													=    z.permute(0	,				2	,				3	,				1				).contiguous()
    __lowerCamelCase													=    z.view(-1	,				self.vq_embed_dim				)

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

    __lowerCamelCase													=    self.embedding(__UpperCAmelCase				).view(z.shape				)
    __lowerCamelCase													=    None
    __lowerCamelCase													=    None

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

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

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

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

    if self.sane_index_shape:
       __lowerCamelCase													=    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	,				__UpperCAmelCase	,				__UpperCAmelCase				):



    '''simple docstring'''
    # shape specifying (batch, height, width, channel)
    if self.remap is not None:
       __lowerCamelCase													=    indices.reshape(shape[0]	,				-1				)  # add batch axis
       __lowerCamelCase													=    self.unmap_to_all(__UpperCAmelCase				)
       __lowerCamelCase													=    indices.reshape(-1				)  # flatten again

    # get quantized latent vectors
    __lowerCamelCase													=    self.embedding(__UpperCAmelCase				)

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

    return z_q



class  __lowerCAmelCase     (	lowerCAmelCase__   ):




 def __init__(						self	,				__UpperCAmelCase	,				__UpperCAmelCase=False				):



    '''simple docstring'''
    __lowerCamelCase													=    parameters
    __lowerCamelCase							,__lowerCamelCase													=    torch.chunk(__UpperCAmelCase	,				2	,				dim=1				)
    __lowerCamelCase													=    torch.clamp(self.logvar	,				-30.0	,				20.0				)
    __lowerCamelCase													=    deterministic
    __lowerCamelCase													=    torch.exp(0.5 * self.logvar				)
    __lowerCamelCase													=    torch.exp(self.logvar				)
    if self.deterministic:
       __lowerCamelCase													=    __lowerCamelCase													=    torch.zeros_like(
           self.mean	,				device=self.parameters.device	,				dtype=self.parameters.dtype				)




 def        lowerCamelCase      (						self	,				__UpperCAmelCase = None				):



    '''simple docstring'''
    # make sure sample is on the same device as the parameters and has same dtype
    __lowerCamelCase													=    randn_tensor(
        self.mean.shape	,				generator=__UpperCAmelCase	,				device=self.parameters.device	,				dtype=self.parameters.dtype				)
    __lowerCamelCase													=    self.mean + self.std * sample
    return x




 def        lowerCamelCase      (						self	,				__UpperCAmelCase=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	,				__UpperCAmelCase	,				__UpperCAmelCase=[1, 2, 3]				):



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




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    return self.mean











import inspect
import re

from transformers.utils import direct_transformers_import


# 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_			=					direct_transformers_import(PATH_TO_TRANSFORMERS)

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(R"""\[(.+?)\]\((https://huggingface\.co/.+?)\)""")


a_			=					{
    """DecisionTransformerConfig""",
    """EncoderDecoderConfig""",
    """MusicgenConfig""",
    """RagConfig""",
    """SpeechEncoderDecoderConfig""",
    """TimmBackboneConfig""",
    """VisionEncoderDecoderConfig""",
    """VisionTextDualEncoderConfig""",
    """LlamaConfig""",
}





def 		a__			(		_UpperCamelCase			: Union[str, Any]							):
   __lowerCamelCase													=    None

   # source code of `config_class`
   __lowerCamelCase													=    inspect.getsource(_UpperCamelCase							)
   __lowerCamelCase													=    _re_checkpoint.findall(_UpperCamelCase							)

   # Each `checkpoint` is a tuple of a checkpoint name and a checkpoint link.
   # For example, `('bert-base-uncased', 'https://huggingface.co/bert-base-uncased')`
   for ckpt_name, ckpt_link in checkpoints:
      # allow the link to end with `/`
      if ckpt_link.endswith('''/'''							):
         __lowerCamelCase													=    ckpt_link[:-1]

      # verify the checkpoint name corresponds to the checkpoint link
      __lowerCamelCase													=    F"""https://huggingface.co/{ckpt_name}"""
      if ckpt_link == ckpt_link_from_name:
         __lowerCamelCase													=    ckpt_name
         break

   return checkpoint





def 		a__			(		):
   __lowerCamelCase													=    []

   for config_class in list(CONFIG_MAPPING.values()							):
      # Skip deprecated models
      if "models.deprecated" in config_class.__module__:
         continue
      __lowerCamelCase													=    get_checkpoint_from_config_class(_UpperCamelCase							)

      __lowerCamelCase													=    config_class.__name__
      if checkpoint is None and name not in CONFIG_CLASSES_TO_IGNORE_FOR_DOCSTRING_CHECKPOINT_CHECK:
         configs_without_checkpoint.append(_UpperCamelCase							)

   if len(_UpperCamelCase							) > 0:
      __lowerCamelCase													=    '''\n'''.join(sorted(_UpperCamelCase							)							)
      raise ValueError(F"""The following configurations don't contain any valid checkpoint:\n{message}"""							)


if __name__ == "__main__":
  check_config_docstrings_have_checkpoints()






