



from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..utils import maybe_allow_in_graph
from .activations import get_activation
from .attention_processor import Attention
from .embeddings import CombinedTimestepLabelEmbeddings

@maybe_allow_in_graph
class   UpperCAmelCase_    (  nn.Module			):

						'''simple docstring'''




						def __init__(       self ,				_A ,				_A ,				_A ,				_A=0.0 ,				_A = None ,				_A = "geglu" ,				_A = None ,				_A = False ,				_A = False ,				_A = False ,				_A = False ,				_A = True ,				_A = "layer_norm" ,				_A = False ,				):







													'''simple docstring'''

													super().__init__()
													__SCREAMING_SNAKE_CASE										=   only_cross_attention

													__SCREAMING_SNAKE_CASE										=   (num_embeds_ada_norm is not None) and norm_type == '''ada_norm_zero'''
													__SCREAMING_SNAKE_CASE										=   (num_embeds_ada_norm is not None) and norm_type == '''ada_norm'''

													if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
																				raise ValueError(
																				    f"""`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"""
																				    f""" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."""						)

													# Define 3 blocks. Each block has its own normalization layer.
													# 1. Self-Attn
													if self.use_ada_layer_norm:
																				__SCREAMING_SNAKE_CASE										=   AdaLayerNorm(A__ ,				A__						)
													elif self.use_ada_layer_norm_zero:
																				__SCREAMING_SNAKE_CASE										=   AdaLayerNormZero(A__ ,				A__						)
													else:
																				__SCREAMING_SNAKE_CASE										=   nn.LayerNorm(A__ ,				elementwise_affine=A__						)
													__SCREAMING_SNAKE_CASE										=   Attention(
													    query_dim=A__ ,				heads=A__ ,				dim_head=A__ ,				dropout=A__ ,				bias=A__ ,				cross_attention_dim=cross_attention_dim if only_cross_attention else None ,				upcast_attention=A__ ,				)

													# 2. Cross-Attn
													if cross_attention_dim is not None or double_self_attention:
																				# We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
																				# I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
																				# the second cross attention block.
																				__SCREAMING_SNAKE_CASE										=   (
																				    AdaLayerNorm(A__ ,				A__						)
																				    if self.use_ada_layer_norm
																				    else nn.LayerNorm(A__ ,				elementwise_affine=A__						)
																				)
																				__SCREAMING_SNAKE_CASE										=   Attention(
																				    query_dim=A__ ,				cross_attention_dim=cross_attention_dim if not double_self_attention else None ,				heads=A__ ,				dim_head=A__ ,				dropout=A__ ,				bias=A__ ,				upcast_attention=A__ ,				)  # is self-attn if encoder_hidden_states is none
													else:
																				__SCREAMING_SNAKE_CASE										=   None
																				__SCREAMING_SNAKE_CASE										=   None

													# 3. Feed-forward
													__SCREAMING_SNAKE_CASE										=   nn.LayerNorm(A__ ,				elementwise_affine=A__						)
													__SCREAMING_SNAKE_CASE										=   FeedForward(A__ ,				dropout=A__ ,				activation_fn=A__ ,				final_dropout=A__						)

													# let chunk size default to None
													__SCREAMING_SNAKE_CASE										=   None
													__SCREAMING_SNAKE_CASE										=   0






						def 			_A      (       self ,				_A ,				_A						):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE										=   chunk_size
													__SCREAMING_SNAKE_CASE										=   dim






						def 			_A      (       self ,				_A ,				_A = None ,				_A = None ,				_A = None ,				_A = None ,				_A = None ,				_A = None ,				):







													'''simple docstring'''

													if self.use_ada_layer_norm:
																				__SCREAMING_SNAKE_CASE										=   self.norma(A__ ,				A__						)
													elif self.use_ada_layer_norm_zero:
																				__SCREAMING_SNAKE_CASE										=   self.norma(
																				    A__ ,				A__ ,				A__ ,				hidden_dtype=hidden_states.dtype						)
													else:
																				__SCREAMING_SNAKE_CASE										=   self.norma(A__						)

													__SCREAMING_SNAKE_CASE										=   cross_attention_kwargs if cross_attention_kwargs is not None else {}

													__SCREAMING_SNAKE_CASE										=   self.attna(
													    A__ ,				encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None ,				attention_mask=A__ ,				**A__ ,				)
													if self.use_ada_layer_norm_zero:
																				__SCREAMING_SNAKE_CASE										=   gate_msa.unsqueeze(1						) * attn_output
													__SCREAMING_SNAKE_CASE										=   attn_output + hidden_states

													# 2. Cross-Attention
													if self.attna is not None:
																				__SCREAMING_SNAKE_CASE										=   (
																				    self.norma(A__ ,				A__						) if self.use_ada_layer_norm else self.norma(A__						)
																				)

																				__SCREAMING_SNAKE_CASE										=   self.attna(
																				    A__ ,				encoder_hidden_states=A__ ,				attention_mask=A__ ,				**A__ ,				)
																				__SCREAMING_SNAKE_CASE										=   attn_output + hidden_states

													# 3. Feed-forward
													__SCREAMING_SNAKE_CASE										=   self.norma(A__						)

													if self.use_ada_layer_norm_zero:
																				__SCREAMING_SNAKE_CASE										=   norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]

													if self._chunk_size is not None:
																				# "feed_forward_chunk_size" can be used to save memory
																				if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
																											raise ValueError(
																											    f"""`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."""						)

																				__SCREAMING_SNAKE_CASE										=   norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
																				__SCREAMING_SNAKE_CASE										=   torch.cat(
																				    [self.ff(A__						) for hid_slice in norm_hidden_states.chunk(A__ ,				dim=self._chunk_dim						)] ,				dim=self._chunk_dim ,				)
													else:
																				__SCREAMING_SNAKE_CASE										=   self.ff(A__						)

													if self.use_ada_layer_norm_zero:
																				__SCREAMING_SNAKE_CASE										=   gate_mlp.unsqueeze(1						) * ff_output

													__SCREAMING_SNAKE_CASE										=   ff_output + hidden_states

													return hidden_states







class   UpperCAmelCase_    (  nn.Module			):

						'''simple docstring'''




						def __init__(       self ,				_A ,				_A = None ,				_A = 4 ,				_A = 0.0 ,				_A = "geglu" ,				_A = False ,				):







													'''simple docstring'''

													super().__init__()
													__SCREAMING_SNAKE_CASE										=   int(dim * mult						)
													__SCREAMING_SNAKE_CASE										=   dim_out if dim_out is not None else dim

													if activation_fn == "gelu":
																				__SCREAMING_SNAKE_CASE										=   GELU(A__ ,				A__						)
													if activation_fn == "gelu-approximate":
																				__SCREAMING_SNAKE_CASE										=   GELU(A__ ,				A__ ,				approximate='tanh'						)
													elif activation_fn == "geglu":
																				__SCREAMING_SNAKE_CASE										=   GEGLU(A__ ,				A__						)
													elif activation_fn == "geglu-approximate":
																				__SCREAMING_SNAKE_CASE										=   ApproximateGELU(A__ ,				A__						)

													__SCREAMING_SNAKE_CASE										=   nn.ModuleList([]						)
													# project in
													self.net.append(A__						)
													# project dropout
													self.net.append(nn.Dropout(A__						)						)
													# project out
													self.net.append(nn.Linear(A__ ,				A__						)						)
													# FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
													if final_dropout:
																				self.net.append(nn.Dropout(A__						)						)






						def 			_A      (       self ,				_A						):







													'''simple docstring'''

													for module in self.net:
																				__SCREAMING_SNAKE_CASE										=   module(A__						)
													return hidden_states







class   UpperCAmelCase_    (  nn.Module			):

						'''simple docstring'''




						def __init__(       self ,				_A ,				_A ,				_A = "none"						):







													'''simple docstring'''

													super().__init__()
													__SCREAMING_SNAKE_CASE										=   nn.Linear(A__ ,				A__						)
													__SCREAMING_SNAKE_CASE										=   approximate






						def 			_A      (       self ,				_A						):







													'''simple docstring'''

													if gate.device.type != "mps":
																				return F.gelu(A__ ,				approximate=self.approximate						)
													# mps: gelu is not implemented for float16
													return F.gelu(gate.to(dtype=torch.floataa						) ,				approximate=self.approximate						).to(dtype=gate.dtype						)






						def 			_A      (       self ,				_A						):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE										=   self.proj(A__						)
													__SCREAMING_SNAKE_CASE										=   self.gelu(A__						)
													return hidden_states







class   UpperCAmelCase_    (  nn.Module			):

						'''simple docstring'''




						def __init__(       self ,				_A ,				_A						):







													'''simple docstring'''

													super().__init__()
													__SCREAMING_SNAKE_CASE										=   nn.Linear(A__ ,				dim_out * 2						)






						def 			_A      (       self ,				_A						):







													'''simple docstring'''

													if gate.device.type != "mps":
																				return F.gelu(A__						)
													# mps: gelu is not implemented for float16
													return F.gelu(gate.to(dtype=torch.floataa						)						).to(dtype=gate.dtype						)






						def 			_A      (       self ,				_A						):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE										=   self.proj(A__						).chunk(2 ,				dim=-1						)
													return hidden_states * self.gelu(A__						)







class   UpperCAmelCase_    (  nn.Module			):

						'''simple docstring'''




						def __init__(       self ,				_A ,				_A						):







													'''simple docstring'''

													super().__init__()
													__SCREAMING_SNAKE_CASE										=   nn.Linear(A__ ,				A__						)






						def 			_A      (       self ,				_A						):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE										=   self.proj(A__						)
													return x * torch.sigmoid(1.7_0_2 * x						)







class   UpperCAmelCase_    (  nn.Module			):

						'''simple docstring'''




						def __init__(       self ,				_A ,				_A						):







													'''simple docstring'''

													super().__init__()
													__SCREAMING_SNAKE_CASE										=   nn.Embedding(A__ ,				A__						)
													__SCREAMING_SNAKE_CASE										=   nn.SiLU()
													__SCREAMING_SNAKE_CASE										=   nn.Linear(A__ ,				embedding_dim * 2						)
													__SCREAMING_SNAKE_CASE										=   nn.LayerNorm(A__ ,				elementwise_affine=A__						)






						def 			_A      (       self ,				_A ,				_A						):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE										=   self.linear(self.silu(self.emb(A__						)						)						)
													__SCREAMING_SNAKE_CASE										=   torch.chunk(A__ ,				2						)
													__SCREAMING_SNAKE_CASE										=   self.norm(A__						) * (1 + scale) + shift
													return x







class   UpperCAmelCase_    (  nn.Module			):

						'''simple docstring'''




						def __init__(       self ,				_A ,				_A						):







													'''simple docstring'''

													super().__init__()

													__SCREAMING_SNAKE_CASE										=   CombinedTimestepLabelEmbeddings(A__ ,				A__						)

													__SCREAMING_SNAKE_CASE										=   nn.SiLU()
													__SCREAMING_SNAKE_CASE										=   nn.Linear(A__ ,				6 * embedding_dim ,				bias=A__						)
													__SCREAMING_SNAKE_CASE										=   nn.LayerNorm(A__ ,				elementwise_affine=A__ ,				eps=1e-6						)






						def 			_A      (       self ,				_A ,				_A ,				_A ,				_A=None						):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE										=   self.linear(self.silu(self.emb(A__ ,				A__ ,				hidden_dtype=A__						)						)						)
													__SCREAMING_SNAKE_CASE										=   emb.chunk(6 ,				dim=1						)
													__SCREAMING_SNAKE_CASE										=   self.norm(A__						) * (1 + scale_msa[:, None]) + shift_msa[:, None]
													return x, gate_msa, shift_mlp, scale_mlp, gate_mlp







class   UpperCAmelCase_    (  nn.Module			):

						'''simple docstring'''




						def __init__(       self ,				_A ,				_A ,				_A ,				_A = None ,				_A = 1e-5						):







													'''simple docstring'''

													super().__init__()
													__SCREAMING_SNAKE_CASE										=   num_groups
													__SCREAMING_SNAKE_CASE										=   eps

													if act_fn is None:
																				__SCREAMING_SNAKE_CASE										=   None
													else:
																				__SCREAMING_SNAKE_CASE										=   get_activation(A__						)

													__SCREAMING_SNAKE_CASE										=   nn.Linear(A__ ,				out_dim * 2						)






						def 			_A      (       self ,				_A ,				_A						):







													'''simple docstring'''

													if self.act:
																				__SCREAMING_SNAKE_CASE										=   self.act(A__						)
													__SCREAMING_SNAKE_CASE										=   self.linear(A__						)
													__SCREAMING_SNAKE_CASE										=   emb[:, :, None, None]
													__SCREAMING_SNAKE_CASE										=   emb.chunk(2 ,				dim=1						)

													__SCREAMING_SNAKE_CASE										=   F.group_norm(A__ ,				self.num_groups ,				eps=self.eps						)
													__SCREAMING_SNAKE_CASE										=   x * (1 + scale) + shift
													return x














import functools
import operator

from ...configuration_utils import PretrainedConfig
from ...utils import logging


_UpperCamelCase			:	List[Any]								=logging.get_logger(__name__)

_UpperCamelCase			:	Dict								={
    'asapp/sew-tiny-100k': 'https://huggingface.co/asapp/sew-tiny-100k/resolve/main/config.json',
    # See all SEW models at https://huggingface.co/models?filter=sew
}






class 						UpperCAmelCase__ (							__snake_case							):
		__snake_case      :       Tuple     =   "sew"







		def __init__(					self       ,A__=32       ,A__=768       ,A__=12       ,A__=12       ,A__=3072       ,A__=2       ,A__="gelu"       ,A__=0.1       ,A__=0.1       ,A__=0.1       ,A__=0.0       ,A__=0.1       ,A__=0.1       ,A__=0.02       ,A__=1E-5       ,A__="group"       ,A__="gelu"       ,A__=(64, 128, 128, 128, 128, 256, 256, 256, 256, 512, 512, 512, 512)       ,A__=(5, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1)       ,A__=(10, 3, 1, 3, 1, 3, 1, 3, 1, 2, 1, 2, 1)       ,A__=False       ,A__=128       ,A__=16       ,A__=True       ,A__=0.05       ,A__=10       ,A__=2       ,A__=0.0       ,A__=10       ,A__=0       ,A__="mean"       ,A__=False       ,A__=False       ,A__=256       ,A__=0       ,A__=1       ,A__=2       ,**A__       ,):
						super().__init__(**A__       ,pad_token_id=A__       ,bos_token_id=A__       ,eos_token_id=A__		)
						_A	:				str   =      hidden_size
						_A	:				Optional[Any]   =      feat_extract_norm
						_A	:				Any   =      feat_extract_activation
						_A	:				Optional[int]   =      list(A__		)
						_A	:				str   =      list(A__		)
						_A	:				Optional[Any]   =      list(A__		)
						_A	:				List[Any]   =      conv_bias
						_A	:				Optional[int]   =      num_conv_pos_embeddings
						_A	:				Any   =      num_conv_pos_embedding_groups
						_A	:				Union[str, Any]   =      len(self.conv_dim		)
						_A	:				List[Any]   =      num_hidden_layers
						_A	:				str   =      intermediate_size
						_A	:				List[Any]   =      squeeze_factor
						_A	:				Tuple   =      hidden_act
						_A	:				Union[str, Any]   =      num_attention_heads
						_A	:				Optional[int]   =      hidden_dropout
						_A	:				Optional[int]   =      attention_dropout
						_A	:				Optional[int]   =      activation_dropout
						_A	:				Tuple   =      feat_proj_dropout
						_A	:				Optional[Any]   =      final_dropout
						_A	:				str   =      layerdrop
						_A	:				int   =      layer_norm_eps
						_A	:				List[str]   =      initializer_range
						_A	:				List[Any]   =      vocab_size

						if (
						    (len(self.conv_stride		) != self.num_feat_extract_layers)
						    or (len(self.conv_kernel		) != self.num_feat_extract_layers)
						    or (len(self.conv_dim		) != self.num_feat_extract_layers)
						):
										raise ValueError(
										    '''Configuration for convolutional layers is incorrect.'''
										    '''It is required that `len(config.conv_dim)` == `len(config.conv_stride)` == `len(config.conv_kernel)`,'''
										    f"""but is `len(config.conv_dim) = {len(self.conv_dim		)}`, `len(config.conv_stride)"""
										    f"""= {len(self.conv_stride		)}`, `len(config.conv_kernel) = {len(self.conv_kernel		)}`."""		)

						# fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
						_A	:				Union[str, Any]   =      apply_spec_augment
						_A	:				str   =      mask_time_prob
						_A	:				Any   =      mask_time_length
						_A	:				int   =      mask_time_min_masks
						_A	:				str   =      mask_feature_prob
						_A	:				Optional[Any]   =      mask_feature_length
						_A	:				Dict   =      mask_feature_min_masks

						# ctc loss
						_A	:				Optional[int]   =      ctc_loss_reduction
						_A	:				List[str]   =      ctc_zero_infinity

						# sequence classification
						_A	:				List[Any]   =      use_weighted_layer_sum
						_A	:				Optional[int]   =      classifier_proj_size







		@property
		def  A__      (					self		):
						return functools.reduce(operator.mul       ,self.conv_stride       ,1		)



