





from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .modeling_utils import ModelMixin



@dataclass
class   _SCREAMING_SNAKE_CASE		(      __SCREAMING_SNAKE_CASE       ):







 '''simple docstring'''
 lowercase_						=     42



class   _SCREAMING_SNAKE_CASE		(      __SCREAMING_SNAKE_CASE       ,				__SCREAMING_SNAKE_CASE       ):







 '''simple docstring'''

 @register_to_config
 def __init__(self							:  Dict			,					UpperCAmelCase_							:  int = 16			,					UpperCAmelCase_							:  int = 88			,					UpperCAmelCase_							:  Optional[int] = None			,					UpperCAmelCase_							:  Optional[int] = None			,					UpperCAmelCase_							:  int = 1			,					UpperCAmelCase_							:  float = 0.0			,					UpperCAmelCase_							:  int = 32			,					UpperCAmelCase_							:  Optional[int] = None			,					UpperCAmelCase_							:  bool = False			,					UpperCAmelCase_							:  Optional[int] = None			,					UpperCAmelCase_							:  str = "geglu"			,					UpperCAmelCase_							:  bool = True			,					UpperCAmelCase_							:  bool = True			,					)			->Optional[Any]:



   '''simple docstring'''



   super().__init__()
   lowerCamelCase__: int         =num_attention_heads
   lowerCamelCase__: Union[str, Any]         =attention_head_dim
   lowerCamelCase__: Optional[int]         =num_attention_heads * attention_head_dim

   lowerCamelCase__: Tuple         =in_channels

   lowerCamelCase__: Optional[Any]         =torch.nn.GroupNorm(num_groups=UpperCAmelCase_			,					num_channels=UpperCAmelCase_			,					eps=1E-6			,					affine=UpperCAmelCase_)
   lowerCamelCase__: List[Any]         =nn.Linear(UpperCAmelCase_			,					UpperCAmelCase_)

   # 3. Define transformers blocks
   lowerCamelCase__: Dict         =nn.ModuleList(
       [
           BasicTransformerBlock(
               UpperCAmelCase_			,					UpperCAmelCase_			,					UpperCAmelCase_			,					dropout=UpperCAmelCase_			,					cross_attention_dim=UpperCAmelCase_			,					activation_fn=UpperCAmelCase_			,					attention_bias=UpperCAmelCase_			,					double_self_attention=UpperCAmelCase_			,					norm_elementwise_affine=UpperCAmelCase_			,					)
           for d in range(UpperCAmelCase_)
       ])

   lowerCamelCase__: Tuple         =nn.Linear(UpperCAmelCase_			,					UpperCAmelCase_)







 def 				SCREAMING_SNAKE_CASE_   (self							:  int			,					UpperCAmelCase_							:  Tuple			,					UpperCAmelCase_							:  Union[str, Any]=None			,					UpperCAmelCase_							:  Optional[Any]=None			,					UpperCAmelCase_							:  Dict=None			,					UpperCAmelCase_							:  Tuple=1			,					UpperCAmelCase_							:  Any=None			,					UpperCAmelCase_							:  bool = True			,					)			->Dict:



   '''simple docstring'''



   lowerCamelCase__    ,   lowerCamelCase__    ,   lowerCamelCase__    ,   lowerCamelCase__: Dict         =hidden_states.shape
   lowerCamelCase__: Any         =batch_frames // num_frames

   lowerCamelCase__: Optional[Any]         =hidden_states

   lowerCamelCase__: List[str]         =hidden_states[None, :].reshape(UpperCAmelCase_			,					UpperCAmelCase_			,					UpperCAmelCase_			,					UpperCAmelCase_			,					UpperCAmelCase_)
   lowerCamelCase__: List[str]         =hidden_states.permute(0			,					2			,					1			,					3			,					4)

   lowerCamelCase__: List[Any]         =self.norm(UpperCAmelCase_)
   lowerCamelCase__: List[str]         =hidden_states.permute(0			,					3			,					4			,					2			,					1).reshape(batch_size * height * width			,					UpperCAmelCase_			,					UpperCAmelCase_)

   lowerCamelCase__: List[Any]         =self.proj_in(UpperCAmelCase_)

   # 2. Blocks
   for block in self.transformer_blocks:
     lowerCamelCase__: int         =block(
         UpperCAmelCase_			,					encoder_hidden_states=UpperCAmelCase_			,					timestep=UpperCAmelCase_			,					cross_attention_kwargs=UpperCAmelCase_			,					class_labels=UpperCAmelCase_			,					)

   # 3. Output
   lowerCamelCase__: Any         =self.proj_out(UpperCAmelCase_)
   lowerCamelCase__: Union[str, Any]         =(
       hidden_states[None, None, :]
       .reshape(UpperCAmelCase_			,					UpperCAmelCase_			,					UpperCAmelCase_			,					UpperCAmelCase_			,					UpperCAmelCase_)
       .permute(0			,					3			,					4			,					1			,					2)
       .contiguous()
   )
   lowerCamelCase__: List[Any]         =hidden_states.reshape(UpperCAmelCase_			,					UpperCAmelCase_			,					UpperCAmelCase_			,					UpperCAmelCase_)

   lowerCamelCase__: int         =hidden_states + residual

   if not return_dict:
     return (output,)

   return TransformerTemporalModelOutput(sample=UpperCAmelCase_)




from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_speech_available, is_torch_available


_lowerCamelCase      :					int					    =      {
    '''configuration_audio_spectrogram_transformer''': [
        '''AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''ASTConfig''',
    ]
}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _lowerCamelCase      :					Tuple					    =      [
       '''AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST''',
       '''ASTForAudioClassification''',
       '''ASTModel''',
       '''ASTPreTrainedModel''',
   ]

try:
   if not is_speech_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _lowerCamelCase      :					Optional[Any]					    =      ['''ASTFeatureExtractor''']

if TYPE_CHECKING:
   from .configuration_audio_spectrogram_transformer import (
       AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
       ASTConfig,
   )

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_audio_spectrogram_transformer import (
          AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
          ASTForAudioClassification,
          ASTModel,
          ASTPreTrainedModel,
      )

   try:
      if not is_speech_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .feature_extraction_audio_spectrogram_transformer import ASTFeatureExtractor


else:
   import sys

   _lowerCamelCase      :					Optional[int]					    =      _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)


