


"""simple docstring"""



import math

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from .attention_processor import Attention
from .embeddings import get_timestep_embedding
from .modeling_utils import ModelMixin
class 					__magic_name__		(			__UpperCAmelCase			,						__UpperCAmelCase     ):






       @register_to_config
       def __init__(   self  :    Union[str, Any]  ,     snake_case__  :    int = 1_2_8  ,     snake_case__  :    int = 2_5_6  ,     snake_case__  :    float = 20_00.0  ,     snake_case__  :    int = 7_6_8  ,     snake_case__  :    int = 1_2  ,     snake_case__  :    int = 1_2  ,     snake_case__  :    int = 6_4  ,     snake_case__  :    int = 2_0_4_8  ,     snake_case__  :    float = 0.1  ,     ):

              '''simple docstring'''






              super().__init__()

              lowercase		:List[str]       							=				nn.Sequential(
                  nn.Linear(snake_case__  ,     d_model * 4  ,     bias=snake_case__				)  ,     nn.SiLU()  ,     nn.Linear(d_model * 4  ,     d_model * 4  ,     bias=snake_case__				)  ,     nn.SiLU()  ,     )

              lowercase		:Dict       							=				nn.Embedding(snake_case__  ,     snake_case__				)
              lowercase		:List[str]       							=				False

              lowercase		:List[Any]       							=				nn.Linear(snake_case__  ,     snake_case__  ,     bias=snake_case__				)

              lowercase		:int       							=				nn.Dropout(p=snake_case__				)

              lowercase		:Optional[Any]       							=				nn.ModuleList()
              for lyr_num in range(snake_case__				):
                     # FiLM conditional T5 decoder
                     lowercase		:List[str]       							=				DecoderLayer(d_model=snake_case__  ,     d_kv=snake_case__  ,     num_heads=snake_case__  ,     d_ff=snake_case__  ,     dropout_rate=snake_case__				)
                     self.decoders.append(snake_case__				)

              lowercase		:Union[str, Any]       							=				TaLayerNorm(snake_case__				)

              lowercase		:Dict       							=				nn.Dropout(p=snake_case__				)
              lowercase		:Optional[int]       							=				nn.Linear(snake_case__  ,     snake_case__  ,     bias=snake_case__				)






       def  __snake_case						(   self  :    List[Any]  ,     snake_case__  :    List[Any]  ,     snake_case__  :    Any				):

              '''simple docstring'''






              lowercase		:Union[str, Any]       							=				torch.mul(query_input.unsqueeze(-1				)  ,     key_input.unsqueeze(-2				)				)
              return mask.unsqueeze(-3				)







       def  __snake_case						(   self  :    Optional[Any]  ,     snake_case__  :    List[Any]  ,     snake_case__  :    Optional[int]  ,     snake_case__  :    Optional[Any]				):

              '''simple docstring'''






              lowercase							,  lowercase							,  lowercase		:str       							=				decoder_input_tokens.shape
              assert decoder_noise_time.shape == (batch,)

              # decoder_noise_time is in [0, 1), so rescale to expected timing range.
              lowercase		:Optional[int]       							=				get_timestep_embedding(
                  decoder_noise_time * self.config.max_decoder_noise_time  ,     embedding_dim=self.config.d_model  ,     max_period=self.config.max_decoder_noise_time  ,     ).to(dtype=self.dtype				)

              lowercase		:List[str]       							=				self.conditioning_emb(snake_case__				).unsqueeze(1				)

              assert conditioning_emb.shape == (batch, 1, self.config.d_model * 4)

              lowercase		:Optional[int]       							=				decoder_input_tokens.shape[1]

              # If we want to use relative positions for audio context, we can just offset
              # this sequence by the length of encodings_and_masks.
              lowercase		:Any       							=				torch.broadcast_to(
                  torch.arange(snake_case__  ,     device=decoder_input_tokens.device				)  ,     (batch, seq_length)  ,     )

              lowercase		:List[str]       							=				self.position_encoding(snake_case__				)

              lowercase		:List[str]       							=				self.continuous_inputs_projection(snake_case__				)
              inputs += position_encodings
              lowercase		:Tuple       							=				self.dropout(snake_case__				)

              # decoder: No padding present.
              lowercase		:Optional[Any]       							=				torch.ones(
                  decoder_input_tokens.shape[:2]  ,     device=decoder_input_tokens.device  ,     dtype=inputs.dtype				)

              # Translate encoding masks to encoder-decoder masks.
              lowercase		:Any       							=				[(x, self.encoder_decoder_mask(snake_case__  ,     snake_case__				)) for x, y in encodings_and_masks]

              # cross attend style: concat encodings
              lowercase		:List[str]       							=				torch.cat([x[0] for x in encodings_and_encdec_masks]  ,     dim=1				)
              lowercase		:Dict       							=				torch.cat([x[1] for x in encodings_and_encdec_masks]  ,     dim=-1				)

              for lyr in self.decoders:
                     lowercase		:Any       							=				lyr(
                         snake_case__  ,     conditioning_emb=snake_case__  ,     encoder_hidden_states=snake_case__  ,     encoder_attention_mask=snake_case__  ,     )[0]

              lowercase		:Dict       							=				self.decoder_norm(snake_case__				)
              lowercase		:Union[str, Any]       							=				self.post_dropout(snake_case__				)

              lowercase		:Tuple       							=				self.spec_out(snake_case__				)
              return spec_out







class 					__magic_name__		(			nn.Module     ):






       def __init__(   self  :    List[Any]  ,     snake_case__  :    Optional[Any]  ,     snake_case__  :    int  ,     snake_case__  :    int  ,     snake_case__  :    Dict  ,     snake_case__  :    List[str]  ,     snake_case__  :    Union[str, Any]=1e-6				):

              '''simple docstring'''






              super().__init__()
              lowercase		:Optional[int]       							=				nn.ModuleList()

              # cond self attention: layer 0
              self.layer.append(
                  TaLayerSelfAttentionCond(d_model=snake_case__  ,     d_kv=snake_case__  ,     num_heads=snake_case__  ,     dropout_rate=snake_case__				)				)

              # cross attention: layer 1
              self.layer.append(
                  TaLayerCrossAttention(
                      d_model=snake_case__  ,     d_kv=snake_case__  ,     num_heads=snake_case__  ,     dropout_rate=snake_case__  ,     layer_norm_epsilon=snake_case__  ,     )				)

              # Film Cond MLP + dropout: last layer
              self.layer.append(
                  TaLayerFFCond(d_model=snake_case__  ,     d_ff=snake_case__  ,     dropout_rate=snake_case__  ,     layer_norm_epsilon=snake_case__				)				)







       def  __snake_case						(   self  :    Tuple  ,     snake_case__  :    str  ,     snake_case__  :    List[str]=None  ,     snake_case__  :    Optional[int]=None  ,     snake_case__  :    List[Any]=None  ,     snake_case__  :    Union[str, Any]=None  ,     snake_case__  :    Union[str, Any]=None  ,     ):

              '''simple docstring'''






              lowercase		:List[Any]       							=				self.layer[0](
                  snake_case__  ,     conditioning_emb=snake_case__  ,     attention_mask=snake_case__  ,     )

              if encoder_hidden_states is not None:
                     lowercase		:Dict       							=				torch.where(encoder_attention_mask > 0  ,     0  ,     -1e1_0				).to(
                         encoder_hidden_states.dtype				)

                     lowercase		:Optional[Any]       							=				self.layer[1](
                         snake_case__  ,     key_value_states=snake_case__  ,     attention_mask=snake_case__  ,     )

              # Apply Film Conditional Feed Forward layer
              lowercase		:List[Any]       							=				self.layer[-1](snake_case__  ,     snake_case__				)

              return (hidden_states,)







class 					__magic_name__		(			nn.Module     ):






       def __init__(   self  :    Optional[Any]  ,     snake_case__  :    List[str]  ,     snake_case__  :    Dict  ,     snake_case__  :    Union[str, Any]  ,     snake_case__  :    str				):

              '''simple docstring'''






              super().__init__()
              lowercase		:Any       							=				TaLayerNorm(snake_case__				)
              lowercase		:Optional[Any]       							=				TaFiLMLayer(in_features=d_model * 4  ,     out_features=snake_case__				)
              lowercase		:int       							=				Attention(query_dim=snake_case__  ,     heads=snake_case__  ,     dim_head=snake_case__  ,     out_bias=snake_case__  ,     scale_qk=snake_case__				)
              lowercase		:Any       							=				nn.Dropout(snake_case__				)







       def  __snake_case						(   self  :    Union[str, Any]  ,     snake_case__  :    str  ,     snake_case__  :    Tuple=None  ,     snake_case__  :    List[str]=None  ,     ):

              '''simple docstring'''






              lowercase		:Any       							=				self.layer_norm(snake_case__				)

              if conditioning_emb is not None:
                     lowercase		:Optional[int]       							=				self.FiLMLayer(snake_case__  ,     snake_case__				)

              # Self-attention block
              lowercase		:Union[str, Any]       							=				self.attention(snake_case__				)

              lowercase		:Tuple       							=				hidden_states + self.dropout(snake_case__				)

              return hidden_states







class 					__magic_name__		(			nn.Module     ):






       def __init__(   self  :    Optional[int]  ,     snake_case__  :    Any  ,     snake_case__  :    Optional[Any]  ,     snake_case__  :    Optional[Any]  ,     snake_case__  :    Optional[Any]  ,     snake_case__  :    str				):

              '''simple docstring'''






              super().__init__()
              lowercase		:Tuple       							=				Attention(query_dim=snake_case__  ,     heads=snake_case__  ,     dim_head=snake_case__  ,     out_bias=snake_case__  ,     scale_qk=snake_case__				)
              lowercase		:List[Any]       							=				TaLayerNorm(snake_case__  ,     eps=snake_case__				)
              lowercase		:List[Any]       							=				nn.Dropout(snake_case__				)







       def  __snake_case						(   self  :    int  ,     snake_case__  :    str  ,     snake_case__  :    Union[str, Any]=None  ,     snake_case__  :    Optional[int]=None  ,     ):

              '''simple docstring'''






              lowercase		:str       							=				self.layer_norm(snake_case__				)
              lowercase		:Any       							=				self.attention(
                  snake_case__  ,     encoder_hidden_states=snake_case__  ,     attention_mask=attention_mask.squeeze(1				)  ,     )
              lowercase		:Union[str, Any]       							=				hidden_states + self.dropout(snake_case__				)
              return layer_output







class 					__magic_name__		(			nn.Module     ):






       def __init__(   self  :    List[str]  ,     snake_case__  :    Optional[int]  ,     snake_case__  :    List[str]  ,     snake_case__  :    Any  ,     snake_case__  :    str				):

              '''simple docstring'''






              super().__init__()
              lowercase		:Optional[int]       							=				TaDenseGatedActDense(d_model=snake_case__  ,     d_ff=snake_case__  ,     dropout_rate=snake_case__				)
              lowercase		:List[str]       							=				TaFiLMLayer(in_features=d_model * 4  ,     out_features=snake_case__				)
              lowercase		:Optional[int]       							=				TaLayerNorm(snake_case__  ,     eps=snake_case__				)
              lowercase		:Optional[Any]       							=				nn.Dropout(snake_case__				)







       def  __snake_case						(   self  :    Tuple  ,     snake_case__  :    List[str]  ,     snake_case__  :    Any=None				):

              '''simple docstring'''






              lowercase		:Any       							=				self.layer_norm(snake_case__				)
              if conditioning_emb is not None:
                     lowercase		:Optional[Any]       							=				self.film(snake_case__  ,     snake_case__				)

              lowercase		:Optional[Any]       							=				self.DenseReluDense(snake_case__				)
              lowercase		:List[Any]       							=				hidden_states + self.dropout(snake_case__				)
              return hidden_states







class 					__magic_name__		(			nn.Module     ):






       def __init__(   self  :    int  ,     snake_case__  :    Tuple  ,     snake_case__  :    str  ,     snake_case__  :    List[str]				):

              '''simple docstring'''






              super().__init__()
              lowercase		:Any       							=				nn.Linear(snake_case__  ,     snake_case__  ,     bias=snake_case__				)
              lowercase		:Optional[int]       							=				nn.Linear(snake_case__  ,     snake_case__  ,     bias=snake_case__				)
              lowercase		:List[Any]       							=				nn.Linear(snake_case__  ,     snake_case__  ,     bias=snake_case__				)
              lowercase		:int       							=				nn.Dropout(snake_case__				)
              lowercase		:List[Any]       							=				NewGELUActivation()







       def  __snake_case						(   self  :    Any  ,     snake_case__  :    List[Any]				):

              '''simple docstring'''






              lowercase		:Optional[Any]       							=				self.act(self.wi_a(snake_case__				)				)
              lowercase		:Dict       							=				self.wi_a(snake_case__				)
              lowercase		:str       							=				hidden_gelu * hidden_linear
              lowercase		:List[Any]       							=				self.dropout(snake_case__				)

              lowercase		:Tuple       							=				self.wo(snake_case__				)
              return hidden_states







class 					__magic_name__		(			nn.Module     ):






       def __init__(   self  :    Tuple  ,     snake_case__  :    List[str]  ,     snake_case__  :    Union[str, Any]=1e-6				):

              '''simple docstring'''






              super().__init__()
              lowercase		:Optional[int]       							=				nn.Parameter(torch.ones(snake_case__				)				)
              lowercase		:List[Any]       							=				eps







       def  __snake_case						(   self  :    Optional[int]  ,     snake_case__  :    Optional[int]				):

              '''simple docstring'''






              lowercase		:str       							=				hidden_states.to(torch.floataa				).pow(2				).mean(-1  ,     keepdim=snake_case__				)
              lowercase		:Any       							=				hidden_states * torch.rsqrt(variance + self.variance_epsilon				)

              # convert into half-precision if necessary
              if self.weight.dtype in [torch.floataa, torch.bfloataa]:
                     lowercase		:int       							=				hidden_states.to(self.weight.dtype				)

              return self.weight * hidden_states







class 					__magic_name__		(			nn.Module     ):






       def  __snake_case						(   self  :    Optional[int]  ,     snake_case__  :    torch.Tensor				):

              '''simple docstring'''






              return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi				) * (input + 0.04_47_15 * torch.pow(snake_case__  ,     3.0				))				))







class 					__magic_name__		(			nn.Module     ):






       def __init__(   self  :    List[Any]  ,     snake_case__  :    Dict  ,     snake_case__  :    Optional[int]				):

              '''simple docstring'''






              super().__init__()
              lowercase		:Optional[int]       							=				nn.Linear(snake_case__  ,     out_features * 2  ,     bias=snake_case__				)







       def  __snake_case						(   self  :    int  ,     snake_case__  :    Optional[int]  ,     snake_case__  :    List[str]				):

              '''simple docstring'''






              lowercase		:str       							=				self.scale_bias(snake_case__				)
              lowercase							,  lowercase		:Union[str, Any]       							=				torch.chunk(snake_case__  ,     2  ,     -1				)
              lowercase		:Union[str, Any]       							=				x * (1 + scale) + shift
              return x





"""simple docstring"""



from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


UpperCAmelCase      		=			logging.get_logger(__name__)

UpperCAmelCase      		=			{
    '''facebook/deit-base-distilled-patch16-224''': (
        '''https://huggingface.co/facebook/deit-base-patch16-224/resolve/main/config.json'''
    ),
    # See all DeiT models at https://huggingface.co/models?filter=deit
}
class 					__magic_name__		(			__UpperCAmelCase     ):
       __A			:      Tuple        =   "deit"






       def __init__(   self  :    List[str]  ,     snake_case__  :    int=7_6_8  ,     snake_case__  :    List[Any]=1_2  ,     snake_case__  :    Dict=1_2  ,     snake_case__  :    List[Any]=3_0_7_2  ,     snake_case__  :    Optional[int]="gelu"  ,     snake_case__  :    Dict=0.0  ,     snake_case__  :    str=0.0  ,     snake_case__  :    List[str]=0.02  ,     snake_case__  :    Optional[int]=1e-1_2  ,     snake_case__  :    Any=2_2_4  ,     snake_case__  :    Optional[int]=1_6  ,     snake_case__  :    int=3  ,     snake_case__  :    str=True  ,     snake_case__  :    Optional[int]=1_6  ,     **snake_case__  :    List[Any]  ,     ):

              '''simple docstring'''






              super().__init__(**snake_case__				)

              lowercase		:List[str]       							=				hidden_size
              lowercase		:Union[str, Any]       							=				num_hidden_layers
              lowercase		:Optional[Any]       							=				num_attention_heads
              lowercase		:Optional[Any]       							=				intermediate_size
              lowercase		:int       							=				hidden_act
              lowercase		:int       							=				hidden_dropout_prob
              lowercase		:Dict       							=				attention_probs_dropout_prob
              lowercase		:Optional[int]       							=				initializer_range
              lowercase		:Tuple       							=				layer_norm_eps
              lowercase		:Optional[int]       							=				image_size
              lowercase		:Union[str, Any]       							=				patch_size
              lowercase		:Optional[Any]       							=				num_channels
              lowercase		:Tuple       							=				qkv_bias
              lowercase		:Optional[Any]       							=				encoder_stride







class 					__magic_name__		(			__UpperCAmelCase     ):
       __A			:      Optional[int]        =   version.parse("1.11"     )






       @property
       def  __snake_case						(   self  :    int				):

              '''simple docstring'''






              return OrderedDict(
                  [
                      ('''pixel_values''', {0: '''batch''', 1: '''num_channels''', 2: '''height''', 3: '''width'''}),
                  ]				)







       @property
       def  __snake_case						(   self  :    Optional[int]				):

              '''simple docstring'''






              return 1e-4

