

"""simple docstring"""






import inspect
import unittest

from datasets import load_dataset
from packaging import version

from transformers import BeitConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch
	from torch import nn

	from transformers import (
	    MODEL_MAPPING,
	    BeitForImageClassification,
	    BeitForMaskedImageModeling,
	    BeitForSemanticSegmentation,
	    BeitModel,
	)
	from transformers.models.beit.modeling_beit import BEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	import PIL
	from PIL import Image

	from transformers import BeitImageProcessor







class 	__UpperCamelCase    :

		def __init__(		self    ,_A    ,_A=100    ,_A=13    ,_A=30    ,_A=2    ,_A=3    ,_A=True    ,_A=True    ,_A=32    ,_A=4    ,_A=4    ,_A=37    ,_A="gelu"    ,_A=0.1    ,_A=0.1    ,_A=10    ,_A=0.0_2    ,_A=3    ,_A=None    ,_A=[0, 1, 2, 3]    ,):
						'''simple docstring'''

						_lowerCAmelCase       :   Optional[int]  =    parent
						_lowerCAmelCase       :   str  =    100
						_lowerCAmelCase       :   str  =    batch_size
						_lowerCAmelCase       :   Dict  =    image_size
						_lowerCAmelCase       :   Union[str, Any]  =    patch_size
						_lowerCAmelCase       :   Any  =    num_channels
						_lowerCAmelCase       :   List[str]  =    is_training
						_lowerCAmelCase       :   List[str]  =    use_labels
						_lowerCAmelCase       :   Any  =    hidden_size
						_lowerCAmelCase       :   Tuple  =    num_hidden_layers
						_lowerCAmelCase       :   Dict  =    num_attention_heads
						_lowerCAmelCase       :   str  =    intermediate_size
						_lowerCAmelCase       :   List[Any]  =    hidden_act
						_lowerCAmelCase       :   Union[str, Any]  =    hidden_dropout_prob
						_lowerCAmelCase       :   int  =    attention_probs_dropout_prob
						_lowerCAmelCase       :   Any  =    type_sequence_label_size
						_lowerCAmelCase       :   List[str]  =    initializer_range
						_lowerCAmelCase       :   Optional[int]  =    scope
						_lowerCAmelCase       :   Dict  =    out_indices
						_lowerCAmelCase       :   int  =    num_labels

						# in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
						_lowerCAmelCase       :   List[Any]  =    (image_size // patch_size) ** 2
						_lowerCAmelCase       :   int  =    num_patches + 1

		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						_lowerCAmelCase       :   str  =    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

						_lowerCAmelCase       :   Dict  =    None
						_lowerCAmelCase       :   Union[str, Any]  =    None
						if self.use_labels:
										_lowerCAmelCase       :   str  =    ids_tensor([self.batch_size]    ,self.type_sequence_label_size      )
										_lowerCAmelCase       :   Dict  =    ids_tensor([self.batch_size, self.image_size, self.image_size]    ,self.num_labels      )

						_lowerCAmelCase       :   Tuple  =    self.get_config()

						return config, pixel_values, labels, pixel_labels

		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						return BeitConfig(
						    vocab_size=self.vocab_size    ,image_size=self.image_size    ,patch_size=self.patch_size    ,num_channels=self.num_channels    ,hidden_size=self.hidden_size    ,num_hidden_layers=self.num_hidden_layers    ,num_attention_heads=self.num_attention_heads    ,intermediate_size=self.intermediate_size    ,hidden_act=self.hidden_act    ,hidden_dropout_prob=self.hidden_dropout_prob    ,attention_probs_dropout_prob=self.attention_probs_dropout_prob    ,is_decoder=_lowercase    ,initializer_range=self.initializer_range    ,out_indices=self.out_indices    ,)

		def __lowerCamelCase		(		self    ,_A    ,_A    ,_A    ,_A      ):
						'''simple docstring'''

						_lowerCAmelCase       :   Any  =    BeitModel(config=_lowercase      )
						model.to(_lowercase      )
						model.eval()
						_lowerCAmelCase       :   Optional[int]  =    model(_lowercase      )
						self.parent.assertEqual(result.last_hidden_state.shape    ,(self.batch_size, self.seq_length, self.hidden_size)      )

		def __lowerCamelCase		(		self    ,_A    ,_A    ,_A    ,_A      ):
						'''simple docstring'''

						_lowerCAmelCase       :   Dict  =    BeitForMaskedImageModeling(config=_lowercase      )
						model.to(_lowercase      )
						model.eval()
						_lowerCAmelCase       :   Optional[Any]  =    model(_lowercase      )
						self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.seq_length - 1, self.vocab_size)      )

		def __lowerCamelCase		(		self    ,_A    ,_A    ,_A    ,_A      ):
						'''simple docstring'''

						_lowerCAmelCase       :   Union[str, Any]  =    self.type_sequence_label_size
						_lowerCAmelCase       :   Dict  =    BeitForImageClassification(_lowercase      )
						model.to(_lowercase      )
						model.eval()
						_lowerCAmelCase       :   int  =    model(_lowercase    ,labels=_lowercase      )
						self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.type_sequence_label_size)      )

						# test greyscale images
						_lowerCAmelCase       :   Dict  =    1
						_lowerCAmelCase       :   Tuple  =    BeitForImageClassification(_lowercase      )
						model.to(_lowercase      )
						model.eval()

						_lowerCAmelCase       :   Union[str, Any]  =    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]      )
						_lowerCAmelCase       :   Optional[int]  =    model(_lowercase    ,labels=_lowercase      )
						self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.type_sequence_label_size)      )

		def __lowerCamelCase		(		self    ,_A    ,_A    ,_A    ,_A      ):
						'''simple docstring'''

						_lowerCAmelCase       :   List[Any]  =    self.num_labels
						_lowerCAmelCase       :   Any  =    BeitForSemanticSegmentation(_lowercase      )
						model.to(_lowercase      )
						model.eval()
						_lowerCAmelCase       :   Optional[int]  =    model(_lowercase      )
						self.parent.assertEqual(
						    result.logits.shape    ,(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)      )
						_lowerCAmelCase       :   List[Any]  =    model(_lowercase    ,labels=_lowercase      )
						self.parent.assertEqual(
						    result.logits.shape    ,(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)      )







		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						_lowerCAmelCase       :   Any  =    self.prepare_config_and_inputs()
						_lowerCAmelCase       :   int  =    config_and_inputs
						_lowerCAmelCase       :   int  =    {"""pixel_values""": pixel_values}
						return config, inputs_dict







@require_torch
class 	__UpperCamelCase    (    SCREAMING_SNAKE_CASE__							,							SCREAMING_SNAKE_CASE__							,							unittest.TestCase						):
		_UpperCAmelCase          =		(
		    (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
		    if is_torch_available()
		    else ()
		)
		_UpperCAmelCase          =		(
		    {
		        "feature-extraction": BeitModel,
		        "image-classification": BeitForImageClassification,
		        "image-segmentation": BeitForSemanticSegmentation,
		    }
		    if is_torch_available()
		    else {}
		)

		_UpperCAmelCase          =		False
		_UpperCAmelCase          =		False
		_UpperCAmelCase          =		False

		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						_lowerCAmelCase       :   Optional[int]  =    BeitModelTester(self      )
						_lowerCAmelCase       :   str  =    ConfigTester(self    ,config_class=_lowercase    ,has_text_modality=_lowercase    ,hidden_size=37      )

		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						self.config_tester.run_common_tests()

		@unittest.skip(reason='BEiT does not use inputs_embeds'      )
		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						pass

		@require_torch_multi_gpu
		@unittest.skip(reason='BEiT has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'      )
		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						pass

		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						_lowerCAmelCase       :   Dict  =    self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
										_lowerCAmelCase       :   Tuple  =    model_class(_lowercase      )
										self.assertIsInstance(model.get_input_embeddings()    ,(nn.Module)      )
										_lowerCAmelCase       :   List[str]  =    model.get_output_embeddings()
										self.assertTrue(x is None or isinstance(_lowercase    ,nn.Linear      )      )

		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						_lowerCAmelCase       :   Any  =    self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
										_lowerCAmelCase       :   Any  =    model_class(_lowercase      )
										_lowerCAmelCase       :   Any  =    inspect.signature(model.forward      )
										# signature.parameters is an OrderedDict => so arg_names order is deterministic
										_lowerCAmelCase       :   Optional[int]  =    [*signature.parameters.keys()]

										_lowerCAmelCase       :   Any  =    ["""pixel_values"""]
										self.assertListEqual(arg_names[:1]    ,_lowercase      )

		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						_lowerCAmelCase       :   Any  =    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*_lowercase      )

		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						_lowerCAmelCase       :   Optional[Any]  =    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*_lowercase      )

		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						_lowerCAmelCase       :   List[str]  =    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_image_classification(*_lowercase      )

		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						_lowerCAmelCase       :   Any  =    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_semantic_segmentation(*_lowercase      )

		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						if not self.model_tester.is_training:
										return

						_lowerCAmelCase       :   List[Any]  =    self.model_tester.prepare_config_and_inputs_for_common()
						_lowerCAmelCase       :   Tuple  =    True

						for model_class in self.all_model_classes:
										# we don't test BeitForMaskedImageModeling
										if model_class in [*get_values(_lowercase      ), BeitForMaskedImageModeling]:
														continue

										_lowerCAmelCase       :   Optional[int]  =    model_class(_lowercase      )
										model.to(_lowercase      )
										model.train()
										_lowerCAmelCase       :   Tuple  =    self._prepare_for_class(_lowercase    ,_lowercase    ,return_labels=_lowercase      )
										_lowerCAmelCase       :   str  =    model(**_lowercase      ).loss
										loss.backward()

		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						_lowerCAmelCase       :   int  =    self.model_tester.prepare_config_and_inputs_for_common()
						if not self.model_tester.is_training:
										return

						_lowerCAmelCase       :   Dict  =    False
						_lowerCAmelCase       :   List[str]  =    True

						for model_class in self.all_model_classes:
										# we don't test BeitForMaskedImageModeling
										if (
										    model_class in [*get_values(_lowercase      ), BeitForMaskedImageModeling]
										    or not model_class.supports_gradient_checkpointing
										):
														continue

										_lowerCAmelCase       :   int  =    model_class(_lowercase      )
										model.gradient_checkpointing_enable()
										model.to(_lowercase      )
										model.train()
										_lowerCAmelCase       :   List[Any]  =    self._prepare_for_class(_lowercase    ,_lowercase    ,return_labels=_lowercase      )
										_lowerCAmelCase       :   List[Any]  =    model(**_lowercase      ).loss
										loss.backward()

		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						_lowerCAmelCase       :   List[Any]  =    self.model_tester.prepare_config_and_inputs_for_common()

						_lowerCAmelCase       :   int  =    _config_zero_init(_lowercase      )
						for model_class in self.all_model_classes:
										_lowerCAmelCase       :   Union[str, Any]  =    model_class(config=_lowercase      )
										for name, param in model.named_parameters():
														# we skip lambda parameters as these require special initial values
														# determined by config.layer_scale_init_value
														if "lambda" in name:
																		continue
														if param.requires_grad:
																		self.assertIn(
																		    ((param.data.mean() * 1E9).round() / 1E9).item()    ,[0.0, 1.0]    ,msg=F"""Parameter {name} of model {model_class} seems not properly initialized"""    ,)







		@slow
		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										_lowerCAmelCase       :   Union[str, Any]  =    BeitModel.from_pretrained(_lowercase      )
										self.assertIsNotNone(_lowercase      )





def 		lowerCamelCase__		(					):

				'''simple docstring'''



				_lowerCAmelCase       :   Dict  =    Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'   )
				return image







@require_torch
@require_vision
class 	__UpperCamelCase    (    unittest.TestCase						):

		@cached_property
		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						return BeitImageProcessor.from_pretrained('microsoft/beit-base-patch16-224'      ) if is_vision_available() else None

		@slow
		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						_lowerCAmelCase       :   List[str]  =    BeitForMaskedImageModeling.from_pretrained('microsoft/beit-base-patch16-224-pt22k'      ).to(_lowercase      )

						_lowerCAmelCase       :   str  =    self.default_image_processor
						_lowerCAmelCase       :   int  =    prepare_img()
						_lowerCAmelCase       :   Optional[int]  =    image_processor(images=_lowercase    ,return_tensors='pt'      ).pixel_values.to(_lowercase      )

						# prepare bool_masked_pos
						_lowerCAmelCase       :   Optional[Any]  =    torch.ones((1, 196)    ,dtype=torch.bool      ).to(_lowercase      )

						# forward pass
						with torch.no_grad():
										_lowerCAmelCase       :   Dict  =    model(pixel_values=_lowercase    ,bool_masked_pos=_lowercase      )
						_lowerCAmelCase       :   str  =    outputs.logits

						# verify the logits
						_lowerCAmelCase       :   str  =    torch.Size((1, 196, 8192)      )
						self.assertEqual(logits.shape    ,_lowercase      )

						_lowerCAmelCase       :   int  =    torch.tensor(
						    [[-3.2_4_3_7, 0.5_0_7_2, -1_3.9_1_7_4], [-3.2_4_5_6, 0.4_9_4_8, -1_3.9_4_0_1], [-3.2_0_3_3, 0.5_1_2_1, -1_3.8_5_5_0]]      ).to(_lowercase      )

						self.assertTrue(torch.allclose(logits[bool_masked_pos][:3, :3]    ,_lowercase    ,atol=1E-2      )      )

		@slow
		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						_lowerCAmelCase       :   Tuple  =    BeitForImageClassification.from_pretrained('microsoft/beit-base-patch16-224'      ).to(_lowercase      )

						_lowerCAmelCase       :   Tuple  =    self.default_image_processor
						_lowerCAmelCase       :   Any  =    prepare_img()
						_lowerCAmelCase       :   Dict  =    image_processor(images=_lowercase    ,return_tensors='pt'      ).to(_lowercase      )

						# forward pass
						with torch.no_grad():
										_lowerCAmelCase       :   Union[str, Any]  =    model(**_lowercase      )
						_lowerCAmelCase       :   Dict  =    outputs.logits

						# verify the logits
						_lowerCAmelCase       :   Optional[Any]  =    torch.Size((1, 1000)      )
						self.assertEqual(logits.shape    ,_lowercase      )

						_lowerCAmelCase       :   Any  =    torch.tensor([-1.2_3_8_5, -1.0_9_8_7, -1.0_1_0_8]      ).to(_lowercase      )

						self.assertTrue(torch.allclose(logits[0, :3]    ,_lowercase    ,atol=1E-4      )      )

						_lowerCAmelCase       :   int  =    281
						self.assertEqual(logits.argmax(-1      ).item()    ,_lowercase      )

		@slow
		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						_lowerCAmelCase       :   str  =    BeitForImageClassification.from_pretrained('microsoft/beit-large-patch16-224-pt22k-ft22k'      ).to(
						    _lowercase      )

						_lowerCAmelCase       :   int  =    self.default_image_processor
						_lowerCAmelCase       :   int  =    prepare_img()
						_lowerCAmelCase       :   Union[str, Any]  =    image_processor(images=_lowercase    ,return_tensors='pt'      ).to(_lowercase      )

						# forward pass
						with torch.no_grad():
										_lowerCAmelCase       :   Any  =    model(**_lowercase      )
						_lowerCAmelCase       :   str  =    outputs.logits

						# verify the logits
						_lowerCAmelCase       :   List[str]  =    torch.Size((1, 2_1841)      )
						self.assertEqual(logits.shape    ,_lowercase      )

						_lowerCAmelCase       :   Union[str, Any]  =    torch.tensor([1.6_8_8_1, -0.2_7_8_7, 0.5_9_0_1]      ).to(_lowercase      )

						self.assertTrue(torch.allclose(logits[0, :3]    ,_lowercase    ,atol=1E-4      )      )

						_lowerCAmelCase       :   Dict  =    2396
						self.assertEqual(logits.argmax(-1      ).item()    ,_lowercase      )

		@slow
		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						_lowerCAmelCase       :   str  =    BeitForSemanticSegmentation.from_pretrained('microsoft/beit-base-finetuned-ade-640-640'      )
						_lowerCAmelCase       :   int  =    model.to(_lowercase      )

						_lowerCAmelCase       :   Any  =    BeitImageProcessor(do_resize=_lowercase    ,size=640    ,do_center_crop=_lowercase      )

						_lowerCAmelCase       :   List[str]  =    load_dataset('hf-internal-testing/fixtures_ade20k'    ,split='test'      )
						_lowerCAmelCase       :   Tuple  =    Image.open(ds[0]['file']      )
						_lowerCAmelCase       :   Any  =    image_processor(images=_lowercase    ,return_tensors='pt'      ).to(_lowercase      )

						# forward pass
						with torch.no_grad():
										_lowerCAmelCase       :   List[str]  =    model(**_lowercase      )
						_lowerCAmelCase       :   str  =    outputs.logits

						# verify the logits
						_lowerCAmelCase       :   Optional[Any]  =    torch.Size((1, 150, 160, 160)      )
						self.assertEqual(logits.shape    ,_lowercase      )

						_lowerCAmelCase       :   int  =    version.parse(PIL.__version__      ) < version.parse('9.0.0'      )

						if is_pillow_less_than_a:
										_lowerCAmelCase       :   Dict  =    torch.tensor(
										    [
										        [[-4.9_2_2_5, -2.3_9_5_4, -3.0_5_2_2], [-2.8_8_2_2, -1.0_0_4_6, -1.7_5_6_1], [-2.9_5_4_9, -1.3_2_2_8, -2.1_3_4_7]],
										        [[-5.8_1_6_8, -3.4_1_2_9, -4.0_7_7_8], [-3.8_6_5_1, -2.2_2_1_4, -3.0_2_7_7], [-3.8_3_5_6, -2.4_6_4_3, -3.3_5_3_5]],
										        [[-0.0_0_7_8, 3.9_9_5_2, 4.0_7_5_4], [2.9_8_5_6, 4.6_9_4_4, 5.0_0_3_5], [3.2_4_1_3, 4.7_8_1_3, 4.9_9_6_9]],
										    ]    ,device=_lowercase    ,)
						else:
										_lowerCAmelCase       :   List[str]  =    torch.tensor(
										    [
										        [[-4.8_9_6_0, -2.3_6_8_8, -3.0_3_5_5], [-2.8_4_7_8, -0.9_8_3_6, -1.7_4_1_8], [-2.9_4_4_9, -1.3_3_3_2, -2.1_4_5_6]],
										        [[-5.8_0_8_1, -3.4_1_2_4, -4.1_0_0_6], [-3.8_5_6_1, -2.2_0_8_1, -3.0_3_2_3], [-3.8_3_6_5, -2.4_6_0_1, -3.3_6_6_9]],
										        [[-0.0_3_0_9, 3.9_8_6_8, 4.0_5_4_0], [2.9_6_4_0, 4.6_8_7_7, 4.9_9_7_6], [3.2_0_8_1, 4.7_6_9_0, 4.9_9_4_2]],
										    ]    ,device=_lowercase    ,)

						self.assertTrue(torch.allclose(logits[0, :3, :3, :3]    ,_lowercase    ,atol=1E-4      )      )







		@slow
		def __lowerCamelCase		(		self      ):
						'''simple docstring'''

						_lowerCAmelCase       :   Dict  =    BeitForSemanticSegmentation.from_pretrained('microsoft/beit-base-finetuned-ade-640-640'      )
						_lowerCAmelCase       :   Optional[Any]  =    model.to(_lowercase      )

						_lowerCAmelCase       :   List[Any]  =    BeitImageProcessor(do_resize=_lowercase    ,size=640    ,do_center_crop=_lowercase      )

						_lowerCAmelCase       :   Tuple  =    load_dataset('hf-internal-testing/fixtures_ade20k'    ,split='test'      )
						_lowerCAmelCase       :   Any  =    Image.open(ds[0]['file']      )
						_lowerCAmelCase       :   Optional[int]  =    image_processor(images=_lowercase    ,return_tensors='pt'      ).to(_lowercase      )

						# forward pass
						with torch.no_grad():
										_lowerCAmelCase       :   Optional[Any]  =    model(**_lowercase      )

						_lowerCAmelCase       :   str  =    outputs.logits.detach().cpu()

						_lowerCAmelCase       :   List[str]  =    image_processor.post_process_semantic_segmentation(outputs=_lowercase    ,target_sizes=[(500, 300)]      )
						_lowerCAmelCase       :   int  =    torch.Size((500, 300)      )
						self.assertEqual(segmentation[0].shape    ,_lowercase      )

						_lowerCAmelCase       :   List[str]  =    image_processor.post_process_semantic_segmentation(outputs=_lowercase      )
						_lowerCAmelCase       :   Optional[Any]  =    torch.Size((160, 160)      )
						self.assertEqual(segmentation[0].shape    ,_lowercase      )








"""simple docstring"""






from dataclasses import dataclass
from typing import Dict, Optional, Union

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

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .attention_processor import AttentionProcessor, AttnProcessor
from .embeddings import TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin







@dataclass
class 	__UpperCamelCase    (    a__						):
  _UpperCAmelCase          =		42







class 	__UpperCamelCase    (    a__							,							a__						):

  @register_to_config
  def __init__(		self    ,_A = 32    ,_A = 64    ,_A = 20    ,_A = 768    ,_A=77    ,_A=4    ,_A = 0.0    ,_A = "silu"    ,_A = None    ,_A = None    ,_A = "linear"    ,_A = "prd"    ,_A = None    ,_A = None    ,_A = None    ,):
      '''simple docstring'''

      super().__init__()
      _lowerCAmelCase       :   Any  =    num_attention_heads
      _lowerCAmelCase       :   Optional[int]  =    attention_head_dim
      _lowerCAmelCase       :   Tuple  =    num_attention_heads * attention_head_dim
      _lowerCAmelCase       :   Optional[Any]  =    additional_embeddings

      _lowerCAmelCase       :   Union[str, Any]  =    time_embed_dim or inner_dim
      _lowerCAmelCase       :   Union[str, Any]  =    embedding_proj_dim or embedding_dim
      _lowerCAmelCase       :   Optional[int]  =    clip_embed_dim or embedding_dim

      _lowerCAmelCase       :   int  =    Timesteps(_A    ,_A    ,0      )
      _lowerCAmelCase       :   int  =    TimestepEmbedding(_A    ,_A    ,out_dim=_A    ,act_fn=_A      )

      _lowerCAmelCase       :   List[Any]  =    nn.Linear(_A    ,_A      )

      if embedding_proj_norm_type is None:
          _lowerCAmelCase       :   Optional[Any]  =    None
      elif embedding_proj_norm_type == "layer":
          _lowerCAmelCase       :   List[Any]  =    nn.LayerNorm(_A      )
      else:
          raise ValueError(F"""unsupported embedding_proj_norm_type: {embedding_proj_norm_type}"""      )

      _lowerCAmelCase       :   Tuple  =    nn.Linear(_A    ,_A      )

      if encoder_hid_proj_type is None:
          _lowerCAmelCase       :   int  =    None
      elif encoder_hid_proj_type == "linear":
          _lowerCAmelCase       :   List[Any]  =    nn.Linear(_A    ,_A      )
      else:
          raise ValueError(F"""unsupported encoder_hid_proj_type: {encoder_hid_proj_type}"""      )

      _lowerCAmelCase       :   Dict  =    nn.Parameter(torch.zeros(1    ,num_embeddings + additional_embeddings    ,_A      )      )

      if added_emb_type == "prd":
          _lowerCAmelCase       :   Dict  =    nn.Parameter(torch.zeros(1    ,1    ,_A      )      )
      elif added_emb_type is None:
          _lowerCAmelCase       :   List[Any]  =    None
      else:
          raise ValueError(
              F"""`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `'prd'` or `None`."""      )

      _lowerCAmelCase       :   List[Any]  =    nn.ModuleList(
          [
              BasicTransformerBlock(
                  _A    ,_A    ,_A    ,dropout=_A    ,activation_fn='gelu'    ,attention_bias=_A    ,)
              for d in range(_A      )
          ]      )

      if norm_in_type == "layer":
          _lowerCAmelCase       :   Any  =    nn.LayerNorm(_A      )
      elif norm_in_type is None:
          _lowerCAmelCase       :   Any  =    None
      else:
          raise ValueError(F"""Unsupported norm_in_type: {norm_in_type}."""      )

      _lowerCAmelCase       :   Union[str, Any]  =    nn.LayerNorm(_A      )

      _lowerCAmelCase       :   int  =    nn.Linear(_A    ,_A      )

      _lowerCAmelCase       :   Any  =    torch.full(
          [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]    ,-1_0_0_0_0.0      )
      causal_attention_mask.triu_(1      )
      _lowerCAmelCase       :   Tuple  =    causal_attention_mask[None, ...]
      self.register_buffer('causal_attention_mask'    ,_A    ,persistent=_A      )

      _lowerCAmelCase       :   Tuple  =    nn.Parameter(torch.zeros(1    ,_A      )      )
      _lowerCAmelCase       :   Dict  =    nn.Parameter(torch.zeros(1    ,_A      )      )

  @property
  # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[Any]  =    {}

      def fn_recursive_add_processors(_A    ,_A    ,_A      ):
          if hasattr(_A    ,'set_processor'      ):
              _lowerCAmelCase       :   str  =    module.processor

          for sub_name, child in module.named_children():
              fn_recursive_add_processors(F"""{name}.{sub_name}"""    ,_A    ,_A      )

          return processors

      for name, module in self.named_children():
          fn_recursive_add_processors(_A    ,_A    ,_A      )

      return processors

  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[int]  =    len(self.attn_processors.keys()      )

      if isinstance(_A    ,_A      ) and len(_A      ) != count:
          raise ValueError(
              F"""A dict of processors was passed, but the number of processors {len(_A      )} does not match the"""
              F""" number of attention layers: {count}. Please make sure to pass {count} processor classes."""      )

      def fn_recursive_attn_processor(_A    ,_A    ,_A      ):
          if hasattr(_A    ,'set_processor'      ):
              if not isinstance(_A    ,_A      ):
                  module.set_processor(_A      )
              else:
                  module.set_processor(processor.pop(F"""{name}.processor"""      )      )

          for sub_name, child in module.named_children():
              fn_recursive_attn_processor(F"""{name}.{sub_name}"""    ,_A    ,_A      )

      for name, module in self.named_children():
          fn_recursive_attn_processor(_A    ,_A    ,_A      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      self.set_attn_processor(AttnProcessor()      )

  def __lowerCamelCase		(		self    ,_A    ,_A    ,_A    ,_A = None    ,_A = None    ,_A = True    ,):
      '''simple docstring'''

      _lowerCAmelCase       :   str  =    hidden_states.shape[0]

      _lowerCAmelCase       :   int  =    timestep
      if not torch.is_tensor(_A      ):
          _lowerCAmelCase       :   str  =    torch.tensor([timesteps]    ,dtype=torch.long    ,device=hidden_states.device      )
      elif torch.is_tensor(_A      ) and len(timesteps.shape      ) == 0:
          _lowerCAmelCase       :   Dict  =    timesteps[None].to(hidden_states.device      )

      # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
      _lowerCAmelCase       :   Optional[int]  =    timesteps * torch.ones(_A    ,dtype=timesteps.dtype    ,device=timesteps.device      )

      _lowerCAmelCase       :   Dict  =    self.time_proj(_A      )

      # timesteps does not contain any weights and will always return f32 tensors
      # but time_embedding might be fp16, so we need to cast here.
      _lowerCAmelCase       :   Any  =    timesteps_projected.to(dtype=self.dtype      )
      _lowerCAmelCase       :   Optional[Any]  =    self.time_embedding(_A      )

      if self.embedding_proj_norm is not None:
          _lowerCAmelCase       :   int  =    self.embedding_proj_norm(_A      )

      _lowerCAmelCase       :   str  =    self.embedding_proj(_A      )
      if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
          _lowerCAmelCase       :   str  =    self.encoder_hidden_states_proj(_A      )
      elif self.encoder_hidden_states_proj is not None and encoder_hidden_states is None:
          raise ValueError('`encoder_hidden_states_proj` requires `encoder_hidden_states` to be set'      )

      _lowerCAmelCase       :   Any  =    self.proj_in(_A      )

      _lowerCAmelCase       :   Dict  =    self.positional_embedding.to(hidden_states.dtype      )

      _lowerCAmelCase       :   List[Any]  =    []
      _lowerCAmelCase       :   Optional[Any]  =    0

      if encoder_hidden_states is not None:
          additional_embeds.append(_A      )
          additional_embeddings_len += encoder_hidden_states.shape[1]

      if len(proj_embeddings.shape      ) == 2:
          _lowerCAmelCase       :   int  =    proj_embeddings[:, None, :]

      if len(hidden_states.shape      ) == 2:
          _lowerCAmelCase       :   Any  =    hidden_states[:, None, :]

      _lowerCAmelCase       :   int  =    additional_embeds + [
          proj_embeddings,
          time_embeddings[:, None, :],
          hidden_states,
      ]

      if self.prd_embedding is not None:
          _lowerCAmelCase       :   List[str]  =    self.prd_embedding.to(hidden_states.dtype      ).expand(_A    ,-1    ,-1      )
          additional_embeds.append(_A      )

      _lowerCAmelCase       :   List[str]  =    torch.cat(
          _A    ,dim=1    ,)

      # Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens
      _lowerCAmelCase       :   Tuple  =    additional_embeddings_len + proj_embeddings.shape[1] + 1
      if positional_embeddings.shape[1] < hidden_states.shape[1]:
          _lowerCAmelCase       :   Any  =    F.pad(
              _A    ,(
                  0,
                  0,
                  additional_embeddings_len,
                  self.prd_embedding.shape[1] if self.prd_embedding is not None else 0,
              )    ,value=0.0    ,)

      _lowerCAmelCase       :   int  =    hidden_states + positional_embeddings

      if attention_mask is not None:
          _lowerCAmelCase       :   Optional[Any]  =    (1 - attention_mask.to(hidden_states.dtype      )) * -1_0_0_0_0.0
          _lowerCAmelCase       :   Union[str, Any]  =    F.pad(_A    ,(0, self.additional_embeddings)    ,value=0.0      )
          _lowerCAmelCase       :   Tuple  =    (attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype      )
          _lowerCAmelCase       :   Optional[Any]  =    attention_mask.repeat_interleave(self.config.num_attention_heads    ,dim=0      )

      if self.norm_in is not None:
          _lowerCAmelCase       :   Any  =    self.norm_in(_A      )

      for block in self.transformer_blocks:
          _lowerCAmelCase       :   int  =    block(_A    ,attention_mask=_A      )

      _lowerCAmelCase       :   Union[str, Any]  =    self.norm_out(_A      )

      if self.prd_embedding is not None:
          _lowerCAmelCase       :   Optional[int]  =    hidden_states[:, -1]
      else:
          _lowerCAmelCase       :   Any  =    hidden_states[:, additional_embeddings_len:]

      _lowerCAmelCase       :   Optional[int]  =    self.proj_to_clip_embeddings(_A      )

      if not return_dict:
          return (predicted_image_embedding,)

      return PriorTransformerOutput(predicted_image_embedding=_A      )







  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[Any]  =    (prior_latents * self.clip_std) + self.clip_mean
      return prior_latents





