





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 							__SCREAMING_SNAKE_CASE       (  _a ,		_a							):






							@register_to_config
							def __init__(    self	,			__lowerCAmelCase = 128	,			__lowerCAmelCase = 256	,			__lowerCAmelCase = 2000.0	,			__lowerCAmelCase = 768	,			__lowerCAmelCase = 12	,			__lowerCAmelCase = 12	,			__lowerCAmelCase = 64	,			__lowerCAmelCase = 2048	,			__lowerCAmelCase = 0.1	,			):
														super().__init__()

														UpperCamelCase__					  =      nn.Sequential(
														    nn.Linear(__lowerCAmelCase	,			d_model * 4	,			bias=__lowerCAmelCase	)	,			nn.SiLU()	,			nn.Linear(d_model * 4	,			d_model * 4	,			bias=__lowerCAmelCase	)	,			nn.SiLU()	,			)

														UpperCamelCase__					  =      nn.Embedding(__lowerCAmelCase	,			__lowerCAmelCase	)
														UpperCamelCase__					  =      False

														UpperCamelCase__					  =      nn.Linear(__lowerCAmelCase	,			__lowerCAmelCase	,			bias=__lowerCAmelCase	)

														UpperCamelCase__					  =      nn.Dropout(p=__lowerCAmelCase	)

														UpperCamelCase__					  =      nn.ModuleList()
														for lyr_num in range(__lowerCAmelCase	):
																					# FiLM conditional T5 decoder
																					UpperCamelCase__					  =      DecoderLayer(d_model=__lowerCAmelCase	,			d_kv=__lowerCAmelCase	,			num_heads=__lowerCAmelCase	,			d_ff=__lowerCAmelCase	,			dropout_rate=__lowerCAmelCase	)
																					self.decoders.append(__lowerCAmelCase	)

														UpperCamelCase__					  =      TaLayerNorm(__lowerCAmelCase	)

														UpperCamelCase__					  =      nn.Dropout(p=__lowerCAmelCase	)
														UpperCamelCase__					  =      nn.Linear(__lowerCAmelCase	,			__lowerCAmelCase	,			bias=__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase	):
														UpperCamelCase__					  =      torch.mul(query_input.unsqueeze(-1	)	,			key_input.unsqueeze(-2	)	)
														return mask.unsqueeze(-3	)







							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	):
														UpperCamelCase__  ,  UpperCamelCase__  ,  UpperCamelCase__					  =      decoder_input_tokens.shape
														assert decoder_noise_time.shape == (batch,)

														# decoder_noise_time is in [0, 1), so rescale to expected timing range.
														UpperCamelCase__					  =      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	)

														UpperCamelCase__					  =      self.conditioning_emb(__lowerCAmelCase	).unsqueeze(1	)

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

														UpperCamelCase__					  =      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.
														UpperCamelCase__					  =      torch.broadcast_to(
														    torch.arange(__lowerCAmelCase	,			device=decoder_input_tokens.device	)	,			(batch, seq_length)	,			)

														UpperCamelCase__					  =      self.position_encoding(__lowerCAmelCase	)

														UpperCamelCase__					  =      self.continuous_inputs_projection(__lowerCAmelCase	)
														inputs += position_encodings
														UpperCamelCase__					  =      self.dropout(__lowerCAmelCase	)

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

														# Translate encoding masks to encoder-decoder masks.
														UpperCamelCase__					  =      [(x, self.encoder_decoder_mask(__lowerCAmelCase	,			__lowerCAmelCase	)) for x, y in encodings_and_masks]

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

														for lyr in self.decoders:
																					UpperCamelCase__					  =      lyr(
																					    __lowerCAmelCase	,			conditioning_emb=__lowerCAmelCase	,			encoder_hidden_states=__lowerCAmelCase	,			encoder_attention_mask=__lowerCAmelCase	,			)[0]

														UpperCamelCase__					  =      self.decoder_norm(__lowerCAmelCase	)
														UpperCamelCase__					  =      self.post_dropout(__lowerCAmelCase	)

														UpperCamelCase__					  =      self.spec_out(__lowerCAmelCase	)
														return spec_out






class 							__SCREAMING_SNAKE_CASE       (  nn.Module							):






							def __init__(    self	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase=1E-6	):
														super().__init__()
														UpperCamelCase__					  =      nn.ModuleList()

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

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

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







							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase=None	,			__lowerCAmelCase=None	,			__lowerCAmelCase=None	,			__lowerCAmelCase=None	,			__lowerCAmelCase=None	,			):
														UpperCamelCase__					  =      self.layer[0](
														    __lowerCAmelCase	,			conditioning_emb=__lowerCAmelCase	,			attention_mask=__lowerCAmelCase	,			)

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

																					UpperCamelCase__					  =      self.layer[1](
																					    __lowerCAmelCase	,			key_value_states=__lowerCAmelCase	,			attention_mask=__lowerCAmelCase	,			)

														# Apply Film Conditional Feed Forward layer
														UpperCamelCase__					  =      self.layer[-1](__lowerCAmelCase	,			__lowerCAmelCase	)

														return (hidden_states,)






class 							__SCREAMING_SNAKE_CASE       (  nn.Module							):






							def __init__(    self	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	):
														super().__init__()
														UpperCamelCase__					  =      TaLayerNorm(__lowerCAmelCase	)
														UpperCamelCase__					  =      TaFiLMLayer(in_features=d_model * 4	,			out_features=__lowerCAmelCase	)
														UpperCamelCase__					  =      Attention(query_dim=__lowerCAmelCase	,			heads=__lowerCAmelCase	,			dim_head=__lowerCAmelCase	,			out_bias=__lowerCAmelCase	,			scale_qk=__lowerCAmelCase	)
														UpperCamelCase__					  =      nn.Dropout(__lowerCAmelCase	)







							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase=None	,			__lowerCAmelCase=None	,			):
														# pre_self_attention_layer_norm
														UpperCamelCase__					  =      self.layer_norm(__lowerCAmelCase	)

														if conditioning_emb is not None:
																					UpperCamelCase__					  =      self.FiLMLayer(__lowerCAmelCase	,			__lowerCAmelCase	)

														# Self-attention block
														UpperCamelCase__					  =      self.attention(__lowerCAmelCase	)

														UpperCamelCase__					  =      hidden_states + self.dropout(__lowerCAmelCase	)

														return hidden_states






class 							__SCREAMING_SNAKE_CASE       (  nn.Module							):






							def __init__(    self	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	):
														super().__init__()
														UpperCamelCase__					  =      Attention(query_dim=__lowerCAmelCase	,			heads=__lowerCAmelCase	,			dim_head=__lowerCAmelCase	,			out_bias=__lowerCAmelCase	,			scale_qk=__lowerCAmelCase	)
														UpperCamelCase__					  =      TaLayerNorm(__lowerCAmelCase	,			eps=__lowerCAmelCase	)
														UpperCamelCase__					  =      nn.Dropout(__lowerCAmelCase	)







							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase=None	,			__lowerCAmelCase=None	,			):
														UpperCamelCase__					  =      self.layer_norm(__lowerCAmelCase	)
														UpperCamelCase__					  =      self.attention(
														    __lowerCAmelCase	,			encoder_hidden_states=__lowerCAmelCase	,			attention_mask=attention_mask.squeeze(1	)	,			)
														UpperCamelCase__					  =      hidden_states + self.dropout(__lowerCAmelCase	)
														return layer_output






class 							__SCREAMING_SNAKE_CASE       (  nn.Module							):






							def __init__(    self	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	):
														super().__init__()
														UpperCamelCase__					  =      TaDenseGatedActDense(d_model=__lowerCAmelCase	,			d_ff=__lowerCAmelCase	,			dropout_rate=__lowerCAmelCase	)
														UpperCamelCase__					  =      TaFiLMLayer(in_features=d_model * 4	,			out_features=__lowerCAmelCase	)
														UpperCamelCase__					  =      TaLayerNorm(__lowerCAmelCase	,			eps=__lowerCAmelCase	)
														UpperCamelCase__					  =      nn.Dropout(__lowerCAmelCase	)







							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase=None	):
														UpperCamelCase__					  =      self.layer_norm(__lowerCAmelCase	)
														if conditioning_emb is not None:
																					UpperCamelCase__					  =      self.film(__lowerCAmelCase	,			__lowerCAmelCase	)

														UpperCamelCase__					  =      self.DenseReluDense(__lowerCAmelCase	)
														UpperCamelCase__					  =      hidden_states + self.dropout(__lowerCAmelCase	)
														return hidden_states






class 							__SCREAMING_SNAKE_CASE       (  nn.Module							):






							def __init__(    self	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	):
														super().__init__()
														UpperCamelCase__					  =      nn.Linear(__lowerCAmelCase	,			__lowerCAmelCase	,			bias=__lowerCAmelCase	)
														UpperCamelCase__					  =      nn.Linear(__lowerCAmelCase	,			__lowerCAmelCase	,			bias=__lowerCAmelCase	)
														UpperCamelCase__					  =      nn.Linear(__lowerCAmelCase	,			__lowerCAmelCase	,			bias=__lowerCAmelCase	)
														UpperCamelCase__					  =      nn.Dropout(__lowerCAmelCase	)
														UpperCamelCase__					  =      NewGELUActivation()







							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	):
														UpperCamelCase__					  =      self.act(self.wi_a(__lowerCAmelCase	)	)
														UpperCamelCase__					  =      self.wi_a(__lowerCAmelCase	)
														UpperCamelCase__					  =      hidden_gelu * hidden_linear
														UpperCamelCase__					  =      self.dropout(__lowerCAmelCase	)

														UpperCamelCase__					  =      self.wo(__lowerCAmelCase	)
														return hidden_states






class 							__SCREAMING_SNAKE_CASE       (  nn.Module							):






							def __init__(    self	,			__lowerCAmelCase	,			__lowerCAmelCase=1E-6	):
														super().__init__()
														UpperCamelCase__					  =      nn.Parameter(torch.ones(__lowerCAmelCase	)	)
														UpperCamelCase__					  =      eps







							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	):
														# T5 uses a layer_norm which only scales and doesn't shift, which is also known as Root Mean
														# Square Layer Normalization https://arxiv.org/abs/1910.07467 thus variance is calculated
														# w/o mean and there is no bias. Additionally we want to make sure that the accumulation for
														# half-precision inputs is done in fp32

														UpperCamelCase__					  =      hidden_states.to(torch.floataa	).pow(2	).mean(-1	,			keepdim=__lowerCAmelCase	)
														UpperCamelCase__					  =      hidden_states * torch.rsqrt(variance + self.variance_epsilon	)

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

														return self.weight * hidden_states






class 							__SCREAMING_SNAKE_CASE       (  nn.Module							):






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	):
														return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi	) * (input + 0.04_4715 * torch.pow(__lowerCAmelCase	,			3.0	))	))






class 							__SCREAMING_SNAKE_CASE       (  nn.Module							):






							def __init__(    self	,			__lowerCAmelCase	,			__lowerCAmelCase	):
														super().__init__()
														UpperCamelCase__					  =      nn.Linear(__lowerCAmelCase	,			out_features * 2	,			bias=__lowerCAmelCase	)







							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase	):
														UpperCamelCase__					  =      self.scale_bias(__lowerCAmelCase	)
														UpperCamelCase__  ,  UpperCamelCase__					  =      torch.chunk(__lowerCAmelCase	,			2	,			-1	)
														UpperCamelCase__					  =      x * (1 + scale) + shift
														return x











import copy

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


UpperCamelCase__           =			logging.get_logger(__name__)

class 							__SCREAMING_SNAKE_CASE       (  _a							):
							snake_case  :   Optional[Any]      =							"""encoder-decoder"""
							snake_case  :   Optional[int]      =							True






							def __init__(    self	,			**__lowerCAmelCase	):
														super().__init__(**__lowerCAmelCase	)
														assert (
														    "encoder" in kwargs and "decoder" in kwargs
														), "Config has to be initialized with encoder and decoder config"
														UpperCamelCase__					  =      kwargs.pop("""encoder"""	)
														UpperCamelCase__					  =      encoder_config.pop("""model_type"""	)
														UpperCamelCase__					  =      kwargs.pop("""decoder"""	)
														UpperCamelCase__					  =      decoder_config.pop("""model_type"""	)

														from ..auto.configuration_auto import AutoConfig

														UpperCamelCase__					  =      AutoConfig.for_model(__lowerCAmelCase	,			**__lowerCAmelCase	)
														UpperCamelCase__					  =      AutoConfig.for_model(__lowerCAmelCase	,			**__lowerCAmelCase	)
														UpperCamelCase__					  =      True






							@classmethod
							def 					_lowerCamelCase		(    cls	,			__lowerCAmelCase	,			__lowerCAmelCase	,			**__lowerCAmelCase	):
														logger.info("""Set `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config"""	)
														UpperCamelCase__					  =      True
														UpperCamelCase__					  =      True

														return cls(encoder=encoder_config.to_dict()	,			decoder=decoder_config.to_dict()	,			**__lowerCAmelCase	)







							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      copy.deepcopy(self.__dict__	)
														UpperCamelCase__					  =      self.encoder.to_dict()
														UpperCamelCase__					  =      self.decoder.to_dict()
														UpperCamelCase__					  =      self.__class__.model_type
														return output




