






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


	'''simple docstring'''





	@register_to_config
	def __init__(     self      ,      _A = 1_2_8      ,      _A = 2_5_6      ,      _A = 2000.0      ,      _A = 7_6_8      ,      _A = 1_2      ,      _A = 1_2      ,      _A = 6_4      ,      _A = 2_0_4_8      ,      _A = 0.1      ,      ):



		'''simple docstring'''

		super().__init__()

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

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

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

		_SCREAMING_SNAKE_CASE						=nn.Dropout(p=__lowerCAmelCase					)

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

		_SCREAMING_SNAKE_CASE						=TaLayerNorm(__lowerCAmelCase					)

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


	def 				UpperCamelCase_  (     self      ,      _A      ,      _A					):



		'''simple docstring'''

		_SCREAMING_SNAKE_CASE						=torch.mul(query_input.unsqueeze(-1					)      ,      key_input.unsqueeze(-2					)					)
		return mask.unsqueeze(-3					)







	def 				UpperCamelCase_  (     self      ,      _A      ,      _A      ,      _A					):



		'''simple docstring'''

		_SCREAMING_SNAKE_CASE		,     _SCREAMING_SNAKE_CASE		,     _SCREAMING_SNAKE_CASE						=decoder_input_tokens.shape
		assert decoder_noise_time.shape == (batch,)

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

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

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

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

		_SCREAMING_SNAKE_CASE						=self.position_encoding(__lowerCAmelCase					)

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

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

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

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

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

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

		_SCREAMING_SNAKE_CASE						=self.spec_out(__lowerCAmelCase					)
		return spec_out




class 	__UpperCAmelCase		(				nn.Module	):


	'''simple docstring'''





	def __init__(     self      ,      _A      ,      _A      ,      _A      ,      _A      ,      _A      ,      _A=1E-6					):



		'''simple docstring'''

		super().__init__()
		_SCREAMING_SNAKE_CASE						=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 				UpperCamelCase_  (     self      ,      _A      ,      _A=None      ,      _A=None      ,      _A=None      ,      _A=None      ,      _A=None      ,      ):



		'''simple docstring'''

		_SCREAMING_SNAKE_CASE						=self.layer[0](
		    __lowerCAmelCase      ,      conditioning_emb=__lowerCAmelCase      ,      attention_mask=__lowerCAmelCase      ,      )

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

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

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

		return (hidden_states,)




class 	__UpperCAmelCase		(				nn.Module	):


	'''simple docstring'''





	def __init__(     self      ,      _A      ,      _A      ,      _A      ,      _A					):



		'''simple docstring'''

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







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



		'''simple docstring'''

		_SCREAMING_SNAKE_CASE						=self.layer_norm(__lowerCAmelCase					)

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

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

		_SCREAMING_SNAKE_CASE						=hidden_states + self.dropout(__lowerCAmelCase					)

		return hidden_states




class 	__UpperCAmelCase		(				nn.Module	):


	'''simple docstring'''





	def __init__(     self      ,      _A      ,      _A      ,      _A      ,      _A      ,      _A					):



		'''simple docstring'''

		super().__init__()
		_SCREAMING_SNAKE_CASE						=Attention(query_dim=__lowerCAmelCase      ,      heads=__lowerCAmelCase      ,      dim_head=__lowerCAmelCase      ,      out_bias=__lowerCAmelCase      ,      scale_qk=__lowerCAmelCase					)
		_SCREAMING_SNAKE_CASE						=TaLayerNorm(__lowerCAmelCase      ,      eps=__lowerCAmelCase					)
		_SCREAMING_SNAKE_CASE						=nn.Dropout(__lowerCAmelCase					)







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



		'''simple docstring'''

		_SCREAMING_SNAKE_CASE						=self.layer_norm(__lowerCAmelCase					)
		_SCREAMING_SNAKE_CASE						=self.attention(
		    __lowerCAmelCase      ,      encoder_hidden_states=__lowerCAmelCase      ,      attention_mask=attention_mask.squeeze(1					)      ,      )
		_SCREAMING_SNAKE_CASE						=hidden_states + self.dropout(__lowerCAmelCase					)
		return layer_output




class 	__UpperCAmelCase		(				nn.Module	):


	'''simple docstring'''





	def __init__(     self      ,      _A      ,      _A      ,      _A      ,      _A					):



		'''simple docstring'''

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







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



		'''simple docstring'''

		_SCREAMING_SNAKE_CASE						=self.layer_norm(__lowerCAmelCase					)
		if conditioning_emb is not None:
			_SCREAMING_SNAKE_CASE						=self.film(__lowerCAmelCase      ,      __lowerCAmelCase					)

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




class 	__UpperCAmelCase		(				nn.Module	):


	'''simple docstring'''





	def __init__(     self      ,      _A      ,      _A      ,      _A					):



		'''simple docstring'''

		super().__init__()
		_SCREAMING_SNAKE_CASE						=nn.Linear(__lowerCAmelCase      ,      __lowerCAmelCase      ,      bias=__lowerCAmelCase					)
		_SCREAMING_SNAKE_CASE						=nn.Linear(__lowerCAmelCase      ,      __lowerCAmelCase      ,      bias=__lowerCAmelCase					)
		_SCREAMING_SNAKE_CASE						=nn.Linear(__lowerCAmelCase      ,      __lowerCAmelCase      ,      bias=__lowerCAmelCase					)
		_SCREAMING_SNAKE_CASE						=nn.Dropout(__lowerCAmelCase					)
		_SCREAMING_SNAKE_CASE						=NewGELUActivation()







	def 				UpperCamelCase_  (     self      ,      _A					):



		'''simple docstring'''

		_SCREAMING_SNAKE_CASE						=self.act(self.wi_a(__lowerCAmelCase					)					)
		_SCREAMING_SNAKE_CASE						=self.wi_a(__lowerCAmelCase					)
		_SCREAMING_SNAKE_CASE						=hidden_gelu * hidden_linear
		_SCREAMING_SNAKE_CASE						=self.dropout(__lowerCAmelCase					)

		_SCREAMING_SNAKE_CASE						=self.wo(__lowerCAmelCase					)
		return hidden_states




class 	__UpperCAmelCase		(				nn.Module	):


	'''simple docstring'''





	def __init__(     self      ,      _A      ,      _A=1E-6					):



		'''simple docstring'''

		super().__init__()
		_SCREAMING_SNAKE_CASE						=nn.Parameter(torch.ones(__lowerCAmelCase					)					)
		_SCREAMING_SNAKE_CASE						=eps







	def 				UpperCamelCase_  (     self      ,      _A					):



		'''simple docstring'''

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

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

		return self.weight * hidden_states




class 	__UpperCAmelCase		(				nn.Module	):


	'''simple docstring'''





	def 				UpperCamelCase_  (     self      ,      _A					):



		'''simple docstring'''

		return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi					) * (input + 0.044715 * torch.pow(__lowerCAmelCase      ,      3.0					))					))




class 	__UpperCAmelCase		(				nn.Module	):


	'''simple docstring'''





	def __init__(     self      ,      _A      ,      _A					):



		'''simple docstring'''

		super().__init__()
		_SCREAMING_SNAKE_CASE						=nn.Linear(__lowerCAmelCase      ,      out_features * 2      ,      bias=__lowerCAmelCase					)







	def 				UpperCamelCase_  (     self      ,      _A      ,      _A					):



		'''simple docstring'''

		_SCREAMING_SNAKE_CASE						=self.scale_bias(__lowerCAmelCase					)
		_SCREAMING_SNAKE_CASE		,     _SCREAMING_SNAKE_CASE						=torch.chunk(__lowerCAmelCase      ,      2      ,      -1					)
		_SCREAMING_SNAKE_CASE						=x * (1 + scale) + shift
		return x











import unittest

import torch

from diffusers import VQModel
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()

class 							__SCREAMING_SNAKE_CASE       (  _a ,		_a ,		unittest.TestCase							):
							snake_case  :   List[str]      =							VQModel
							snake_case  :   Dict      =							"""sample"""






							@property
							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase=(32, 32)	):
														UpperCamelCase__					  =      4
														UpperCamelCase__					  =      3

														UpperCamelCase__					  =      floats_tensor((batch_size, num_channels) + sizes	).to(__lowerCAmelCase	)

														return {"sample": image}






							@property
							def 					_lowerCamelCase		(    self	):
														return (3, 32, 32)






							@property
							def 					_lowerCamelCase		(    self	):
														return (3, 32, 32)






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      {
														    """block_out_channels""": [32, 64],
														    """in_channels""": 3,
														    """out_channels""": 3,
														    """down_block_types""": ["""DownEncoderBlock2D""", """DownEncoderBlock2D"""],
														    """up_block_types""": ["""UpDecoderBlock2D""", """UpDecoderBlock2D"""],
														    """latent_channels""": 3,
														}
														UpperCamelCase__					  =      self.dummy_input
														return init_dict, inputs_dict






							def 					_lowerCamelCase		(    self	):
														pass






							def 					_lowerCamelCase		(    self	):
														pass






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__  ,  UpperCamelCase__					  =      VQModel.from_pretrained("""fusing/vqgan-dummy"""	,			output_loading_info=__lowerCAmelCase	)
														self.assertIsNotNone(__lowerCAmelCase	)
														self.assertEqual(len(loading_info["""missing_keys"""]	)	,			0	)

														model.to(__lowerCAmelCase	)
														UpperCamelCase__					  =      model(**self.dummy_input	)

														assert image is not None, "Make sure output is not None"







							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      VQModel.from_pretrained("""fusing/vqgan-dummy"""	)
														model.to(__lowerCAmelCase	).eval()

														torch.manual_seed(0	)
														if torch.cuda.is_available():
																					torch.cuda.manual_seed_all(0	)

														UpperCamelCase__					  =      torch.randn(1	,			model.config.in_channels	,			model.config.sample_size	,			model.config.sample_size	)
														UpperCamelCase__					  =      image.to(__lowerCAmelCase	)
														with torch.no_grad():
																					UpperCamelCase__					  =      model(__lowerCAmelCase	).sample

														UpperCamelCase__					  =      output[0, -1, -3:, -3:].flatten().cpu()
														# fmt: off
														UpperCamelCase__					  =      torch.tensor([-0.0153, -0.4044, -0.1880, -0.5161, -0.2418, -0.4072, -0.1612, -0.0633, -0.0143]	)
														# fmt: on
														self.assertTrue(torch.allclose(__lowerCAmelCase	,			__lowerCAmelCase	,			atol=1E-3	)	)




