



import inspect
import unittest
import warnings

from transformers import DeiTConfig
from transformers.models.auto import get_values
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_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, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
			import torch
			from torch import nn

			from transformers import (
			    MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
			    MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
			    MODEL_MAPPING,
			    DeiTForImageClassification,
			    DeiTForImageClassificationWithTeacher,
			    DeiTForMaskedImageModeling,
			    DeiTModel,
			)
			from transformers.models.deit.modeling_deit import DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
			from PIL import Image

			from transformers import DeiTImageProcessor



class 	lowerCamelCase      :







					'''simple docstring'''







					def __init__(			self    ,	lowerCAmelCase    ,	lowerCAmelCase=13    ,	lowerCAmelCase=30    ,	lowerCAmelCase=2    ,	lowerCAmelCase=3    ,	lowerCAmelCase=True    ,	lowerCAmelCase=True    ,	lowerCAmelCase=32    ,	lowerCAmelCase=5    ,	lowerCAmelCase=4    ,	lowerCAmelCase=37    ,	lowerCAmelCase="gelu"    ,	lowerCAmelCase=0.1    ,	lowerCAmelCase=0.1    ,	lowerCAmelCase=10    ,	lowerCAmelCase=0.02    ,	lowerCAmelCase=3    ,	lowerCAmelCase=None    ,	lowerCAmelCase=2    ,	):
							UpperCAmelCase_						=				parent
							UpperCAmelCase_						=				batch_size
							UpperCAmelCase_						=				image_size
							UpperCAmelCase_						=				patch_size
							UpperCAmelCase_						=				num_channels
							UpperCAmelCase_						=				is_training
							UpperCAmelCase_						=				use_labels
							UpperCAmelCase_						=				hidden_size
							UpperCAmelCase_						=				num_hidden_layers
							UpperCAmelCase_						=				num_attention_heads
							UpperCAmelCase_						=				intermediate_size
							UpperCAmelCase_						=				hidden_act
							UpperCAmelCase_						=				hidden_dropout_prob
							UpperCAmelCase_						=				attention_probs_dropout_prob
							UpperCAmelCase_						=				type_sequence_label_size
							UpperCAmelCase_						=				initializer_range
							UpperCAmelCase_						=				scope
							UpperCAmelCase_						=				encoder_stride

							# in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
							UpperCAmelCase_						=				(image_size // patch_size) ** 2
							UpperCAmelCase_						=				num_patches + 2



					def    A__		(			self     ):
							UpperCAmelCase_						=				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

							UpperCAmelCase_						=				None
							if self.use_labels:
									UpperCAmelCase_						=				ids_tensor([self.batch_size]    ,	self.type_sequence_label_size     )

							UpperCAmelCase_						=				self.get_config()

							return config, pixel_values, labels



					def    A__		(			self     ):
							return DeiTConfig(
							    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=lowerCAmelCase    ,	initializer_range=self.initializer_range    ,	encoder_stride=self.encoder_stride    ,	)



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase     ):
							UpperCAmelCase_						=				DeiTModel(config=lowerCAmelCase     )
							model.to(lowerCAmelCase     )
							model.eval()
							UpperCAmelCase_						=				model(lowerCAmelCase     )
							self.parent.assertEqual(result.last_hidden_state.shape    ,	(self.batch_size, self.seq_length, self.hidden_size)     )



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase     ):
							UpperCAmelCase_						=				DeiTForMaskedImageModeling(config=lowerCAmelCase     )
							model.to(lowerCAmelCase     )
							model.eval()
							UpperCAmelCase_						=				model(lowerCAmelCase     )
							self.parent.assertEqual(
							    result.reconstruction.shape    ,	(self.batch_size, self.num_channels, self.image_size, self.image_size)     )

							# test greyscale images
							UpperCAmelCase_						=				1
							UpperCAmelCase_						=				DeiTForMaskedImageModeling(lowerCAmelCase     )
							model.to(lowerCAmelCase     )
							model.eval()

							UpperCAmelCase_						=				floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
							UpperCAmelCase_						=				model(lowerCAmelCase     )
							self.parent.assertEqual(result.reconstruction.shape    ,	(self.batch_size, 1, self.image_size, self.image_size)     )



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase     ):
							UpperCAmelCase_						=				self.type_sequence_label_size
							UpperCAmelCase_						=				DeiTForImageClassification(lowerCAmelCase     )
							model.to(lowerCAmelCase     )
							model.eval()
							UpperCAmelCase_						=				model(lowerCAmelCase    ,	labels=lowerCAmelCase     )
							self.parent.assertEqual(result.logits.shape    ,	(self.batch_size, self.type_sequence_label_size)     )

							# test greyscale images
							UpperCAmelCase_						=				1
							UpperCAmelCase_						=				DeiTForImageClassification(lowerCAmelCase     )
							model.to(lowerCAmelCase     )
							model.eval()

							UpperCAmelCase_						=				floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
							UpperCAmelCase_						=				model(lowerCAmelCase    ,	labels=lowerCAmelCase     )
							self.parent.assertEqual(result.logits.shape    ,	(self.batch_size, self.type_sequence_label_size)     )



					def    A__		(			self     ):
							UpperCAmelCase_						=				self.prepare_config_and_inputs()
							(
							    (
							    UpperCAmelCase_
							)					,						(
							    UpperCAmelCase_
							)					,						(
							    UpperCAmelCase_
							)					,						
							)						=				config_and_inputs
							UpperCAmelCase_						=				{"pixel_values": pixel_values}
							return config, inputs_dict





@require_torch
class 	lowerCamelCase      (						lowercase__,     lowercase__,     unittest.TestCase  ):







					'''simple docstring'''







					lowerCAmelCase_				:    Optional[int]									=  (
					    (
					        DeiTModel,
					        DeiTForImageClassification,
					        DeiTForImageClassificationWithTeacher,
					        DeiTForMaskedImageModeling,
					    )
					    if is_torch_available()
					    else ()
					)
					lowerCAmelCase_				:    Optional[Any]									=  (
					    {
					        'feature-extraction': DeiTModel,
					        'image-classification': (DeiTForImageClassification, DeiTForImageClassificationWithTeacher),
					    }
					    if is_torch_available()
					    else {}
					)

					lowerCAmelCase_				:    Union[str, Any]									=  False
					lowerCAmelCase_				:    int									=  False
					lowerCAmelCase_				:    Dict									=  False



					def    A__		(			self     ):
							UpperCAmelCase_						=				DeiTModelTester(self     )
							UpperCAmelCase_						=				ConfigTester(self    ,	config_class=lowerCAmelCase    ,	has_text_modality=lowerCAmelCase    ,	hidden_size=37     )



					def    A__		(			self     ):
							self.config_tester.run_common_tests()



					@unittest.skip(reason="DeiT does not use inputs_embeds"     )
					def    A__		(			self     ):
							pass



					def    A__		(			self     ):
							UpperCAmelCase_					,						UpperCAmelCase_						=				self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
									UpperCAmelCase_						=				model_class(lowerCAmelCase     )
									self.assertIsInstance(model.get_input_embeddings()    ,	(nn.Module)     )
									UpperCAmelCase_						=				model.get_output_embeddings()
									self.assertTrue(x is None or isinstance(lowerCAmelCase    ,	nn.Linear     )     )



					def    A__		(			self     ):
							UpperCAmelCase_					,						UpperCAmelCase_						=				self.model_tester.prepare_config_and_inputs_for_common()

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

									UpperCAmelCase_						=				["pixel_values"]
									self.assertListEqual(arg_names[:1]    ,	lowerCAmelCase     )



					def    A__		(			self     ):
							UpperCAmelCase_						=				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*lowerCAmelCase     )



					def    A__		(			self     ):
							UpperCAmelCase_						=				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_image_modeling(*lowerCAmelCase     )



					def    A__		(			self     ):
							UpperCAmelCase_						=				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*lowerCAmelCase     )



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase=False     ):
							UpperCAmelCase_						=				super()._prepare_for_class(lowerCAmelCase    ,	lowerCAmelCase    ,	return_labels=lowerCAmelCase     )

							if return_labels:
									if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
											del inputs_dict["labels"]

							return inputs_dict



					def    A__		(			self     ):
							if not self.model_tester.is_training:
									return

							UpperCAmelCase_					,						UpperCAmelCase_						=				self.model_tester.prepare_config_and_inputs_for_common()
							UpperCAmelCase_						=				True

							for model_class in self.all_model_classes:
									# DeiTForImageClassificationWithTeacher supports inference-only
									if (
									    model_class in get_values(lowerCAmelCase     )
									    or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
									):
											continue
									UpperCAmelCase_						=				model_class(lowerCAmelCase     )
									model.to(lowerCAmelCase     )
									model.train()
									UpperCAmelCase_						=				self._prepare_for_class(lowerCAmelCase    ,	lowerCAmelCase    ,	return_labels=lowerCAmelCase     )
									UpperCAmelCase_						=				model(**lowerCAmelCase     ).loss
									loss.backward()



					def    A__		(			self     ):
							UpperCAmelCase_					,						UpperCAmelCase_						=				self.model_tester.prepare_config_and_inputs_for_common()
							if not self.model_tester.is_training:
									return

							UpperCAmelCase_						=				False
							UpperCAmelCase_						=				True

							for model_class in self.all_model_classes:
									if model_class in get_values(lowerCAmelCase     ) or not model_class.supports_gradient_checkpointing:
											continue
									# DeiTForImageClassificationWithTeacher supports inference-only
									if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
											continue
									UpperCAmelCase_						=				model_class(lowerCAmelCase     )
									model.gradient_checkpointing_enable()
									model.to(lowerCAmelCase     )
									model.train()
									UpperCAmelCase_						=				self._prepare_for_class(lowerCAmelCase    ,	lowerCAmelCase    ,	return_labels=lowerCAmelCase     )
									UpperCAmelCase_						=				model(**lowerCAmelCase     ).loss
									loss.backward()



					def    A__		(			self     ):
							UpperCAmelCase_					,						UpperCAmelCase_						=				self.model_tester.prepare_config_and_inputs_for_common()

							UpperCAmelCase_						=				[
							    {"title": "multi_label_classification", "num_labels": 2, "dtype": torch.float},
							    {"title": "single_label_classification", "num_labels": 1, "dtype": torch.long},
							    {"title": "regression", "num_labels": 1, "dtype": torch.float},
							]

							for model_class in self.all_model_classes:
									if (
									    model_class
									    not in [
									        *get_values(lowerCAmelCase     ),
									        *get_values(lowerCAmelCase     ),
									    ]
									    or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
									):
											continue

									for problem_type in problem_types:
											with self.subTest(msg=f'''Testing {model_class} with {problem_type['title']}'''     ):
													UpperCAmelCase_						=				problem_type["title"]
													UpperCAmelCase_						=				problem_type["num_labels"]

													UpperCAmelCase_						=				model_class(lowerCAmelCase     )
													model.to(lowerCAmelCase     )
													model.train()

													UpperCAmelCase_						=				self._prepare_for_class(lowerCAmelCase    ,	lowerCAmelCase    ,	return_labels=lowerCAmelCase     )

													if problem_type["num_labels"] > 1:
															UpperCAmelCase_						=				inputs["labels"].unsqueeze(1     ).repeat(1    ,	problem_type["num_labels"]     )

													UpperCAmelCase_						=				inputs["labels"].to(problem_type["dtype"]     )

													# This tests that we do not trigger the warning form PyTorch "Using a target size that is different
													# to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
													# they have the same size." which is a symptom something in wrong for the regression problem.
													# See https://github.com/huggingface/transformers/issues/11780
													with warnings.catch_warnings(record=lowerCAmelCase     ) as warning_list:
															UpperCAmelCase_						=				model(**lowerCAmelCase     ).loss
													for w in warning_list:
															if "Using a target size that is different to the input size" in str(w.message     ):
																	raise ValueError(
																	    f'''Something is going wrong in the regression problem: intercepted {w.message}'''     )

													loss.backward()



					@slow
					def    A__		(			self     ):
							for model_name in DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									UpperCAmelCase_						=				DeiTModel.from_pretrained(lowerCAmelCase     )
									self.assertIsNotNone(lowerCAmelCase     )







def snake_case__  (			)		->							Any:
		UpperCAmelCase_						=				Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png" )
		return image



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







					'''simple docstring'''







					@cached_property
					def    A__		(			self     ):
							return (
							    DeiTImageProcessor.from_pretrained("facebook/deit-base-distilled-patch16-224"     )
							    if is_vision_available()
							    else None
							)



					@slow
					def    A__		(			self     ):
							UpperCAmelCase_						=				DeiTForImageClassificationWithTeacher.from_pretrained("facebook/deit-base-distilled-patch16-224"     ).to(
							    lowerCAmelCase     )

							UpperCAmelCase_						=				self.default_image_processor
							UpperCAmelCase_						=				prepare_img()
							UpperCAmelCase_						=				image_processor(images=lowerCAmelCase    ,	return_tensors="pt"     ).to(lowerCAmelCase     )

							# forward pass
							with torch.no_grad():
									UpperCAmelCase_						=				model(**lowerCAmelCase     )

							# verify the logits
							UpperCAmelCase_						=				torch.Size((1, 1000)     )
							self.assertEqual(outputs.logits.shape    ,	lowerCAmelCase     )

							UpperCAmelCase_						=				torch.tensor([-1.0266, 0.1912, -1.2861]     ).to(lowerCAmelCase     )

							self.assertTrue(torch.allclose(outputs.logits[0, :3]    ,	lowerCAmelCase    ,	atol=1e-4     )     )



					@slow
					@require_accelerate
					@require_torch_gpu
					def    A__		(			self     ):
							UpperCAmelCase_						=				DeiTModel.from_pretrained(
							    "facebook/deit-base-distilled-patch16-224"    ,	torch_dtype=torch.floataa    ,	device_map="auto"     )
							UpperCAmelCase_						=				self.default_image_processor

							UpperCAmelCase_						=				prepare_img()
							UpperCAmelCase_						=				image_processor(images=lowerCAmelCase    ,	return_tensors="pt"     )
							UpperCAmelCase_						=				inputs.pixel_values.to(lowerCAmelCase     )

							# forward pass to make sure inference works in fp16
							with torch.no_grad():
									UpperCAmelCase_						=				model(lowerCAmelCase     )






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 	lowerCamelCase      (						lowercase__  ):







					'''simple docstring'''







					lowerCAmelCase_				:    torch.FloatTensor





class 	lowerCamelCase      (						lowercase__,     lowercase__  ):







					'''simple docstring'''







					@register_to_config
					def __init__(			self    ,	lowerCAmelCase = 32    ,	lowerCAmelCase = 64    ,	lowerCAmelCase = 20    ,	lowerCAmelCase = 768    ,	lowerCAmelCase=77    ,	lowerCAmelCase=4    ,	lowerCAmelCase = 0.0    ,	lowerCAmelCase = "silu"    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = "linear"    ,	lowerCAmelCase = "prd"    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	):
							super().__init__()
							UpperCAmelCase_						=				num_attention_heads
							UpperCAmelCase_						=				attention_head_dim
							UpperCAmelCase_						=				num_attention_heads * attention_head_dim
							UpperCAmelCase_						=				additional_embeddings

							UpperCAmelCase_						=				time_embed_dim or inner_dim
							UpperCAmelCase_						=				embedding_proj_dim or embedding_dim
							UpperCAmelCase_						=				clip_embed_dim or embedding_dim

							UpperCAmelCase_						=				Timesteps(lowerCAmelCase    ,	lowerCAmelCase    ,	0     )
							UpperCAmelCase_						=				TimestepEmbedding(lowerCAmelCase    ,	lowerCAmelCase    ,	out_dim=lowerCAmelCase    ,	act_fn=lowerCAmelCase     )

							UpperCAmelCase_						=				nn.Linear(lowerCAmelCase    ,	lowerCAmelCase     )

							if embedding_proj_norm_type is None:
									UpperCAmelCase_						=				None
							elif embedding_proj_norm_type == "layer":
									UpperCAmelCase_						=				nn.LayerNorm(lowerCAmelCase     )
							else:
									raise ValueError(f'''unsupported embedding_proj_norm_type: {embedding_proj_norm_type}'''     )

							UpperCAmelCase_						=				nn.Linear(lowerCAmelCase    ,	lowerCAmelCase     )

							if encoder_hid_proj_type is None:
									UpperCAmelCase_						=				None
							elif encoder_hid_proj_type == "linear":
									UpperCAmelCase_						=				nn.Linear(lowerCAmelCase    ,	lowerCAmelCase     )
							else:
									raise ValueError(f'''unsupported encoder_hid_proj_type: {encoder_hid_proj_type}'''     )

							UpperCAmelCase_						=				nn.Parameter(torch.zeros(1    ,	num_embeddings + additional_embeddings    ,	lowerCAmelCase     )     )

							if added_emb_type == "prd":
									UpperCAmelCase_						=				nn.Parameter(torch.zeros(1    ,	1    ,	lowerCAmelCase     )     )
							elif added_emb_type is None:
									UpperCAmelCase_						=				None
							else:
									raise ValueError(
									    f'''`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `\'prd\'` or `None`.'''     )

							UpperCAmelCase_						=				nn.ModuleList(
							    [
							        BasicTransformerBlock(
							            lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	dropout=lowerCAmelCase    ,	activation_fn="gelu"    ,	attention_bias=lowerCAmelCase    ,	)
							        for d in range(lowerCAmelCase     )
							    ]     )

							if norm_in_type == "layer":
									UpperCAmelCase_						=				nn.LayerNorm(lowerCAmelCase     )
							elif norm_in_type is None:
									UpperCAmelCase_						=				None
							else:
									raise ValueError(f'''Unsupported norm_in_type: {norm_in_type}.'''     )

							UpperCAmelCase_						=				nn.LayerNorm(lowerCAmelCase     )

							UpperCAmelCase_						=				nn.Linear(lowerCAmelCase    ,	lowerCAmelCase     )

							UpperCAmelCase_						=				torch.full(
							    [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]    ,	-10000.0     )
							causal_attention_mask.triu_(1     )
							UpperCAmelCase_						=				causal_attention_mask[None, ...]
							self.register_buffer("causal_attention_mask"    ,	lowerCAmelCase    ,	persistent=lowerCAmelCase     )

							UpperCAmelCase_						=				nn.Parameter(torch.zeros(1    ,	lowerCAmelCase     )     )
							UpperCAmelCase_						=				nn.Parameter(torch.zeros(1    ,	lowerCAmelCase     )     )



					@property
					# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
					def    A__		(			self     ):
							UpperCAmelCase_						=				{}

							def fn_recursive_add_processors(lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase     ):
									if hasattr(lowerCAmelCase    ,	"set_processor"     ):
											UpperCAmelCase_						=				module.processor

									for sub_name, child in module.named_children():
											fn_recursive_add_processors(f'''{name}.{sub_name}'''    ,	lowerCAmelCase    ,	lowerCAmelCase     )

									return processors

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

							return processors



					def    A__		(			self    ,	lowerCAmelCase     ):
							UpperCAmelCase_						=				len(self.attn_processors.keys()     )

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

							def fn_recursive_attn_processor(lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase     ):
									if hasattr(lowerCAmelCase    ,	"set_processor"     ):
											if not isinstance(lowerCAmelCase    ,	lowerCAmelCase     ):
													module.set_processor(lowerCAmelCase     )
											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}'''    ,	lowerCAmelCase    ,	lowerCAmelCase     )

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



					def    A__		(			self     ):
							self.set_attn_processor(AttnProcessor()     )



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = True    ,	):
							UpperCAmelCase_						=				hidden_states.shape[0]

							UpperCAmelCase_						=				timestep
							if not torch.is_tensor(lowerCAmelCase     ):
									UpperCAmelCase_						=				torch.tensor([timesteps]    ,	dtype=torch.long    ,	device=hidden_states.device     )
							elif torch.is_tensor(lowerCAmelCase     ) and len(timesteps.shape     ) == 0:
									UpperCAmelCase_						=				timesteps[None].to(hidden_states.device     )

							# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
							UpperCAmelCase_						=				timesteps * torch.ones(lowerCAmelCase    ,	dtype=timesteps.dtype    ,	device=timesteps.device     )

							UpperCAmelCase_						=				self.time_proj(lowerCAmelCase     )

							# timesteps does not contain any weights and will always return f32 tensors
							# but time_embedding might be fp16, so we need to cast here.
							UpperCAmelCase_						=				timesteps_projected.to(dtype=self.dtype     )
							UpperCAmelCase_						=				self.time_embedding(lowerCAmelCase     )

							if self.embedding_proj_norm is not None:
									UpperCAmelCase_						=				self.embedding_proj_norm(lowerCAmelCase     )

							UpperCAmelCase_						=				self.embedding_proj(lowerCAmelCase     )
							if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
									UpperCAmelCase_						=				self.encoder_hidden_states_proj(lowerCAmelCase     )
							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"     )

							UpperCAmelCase_						=				self.proj_in(lowerCAmelCase     )

							UpperCAmelCase_						=				self.positional_embedding.to(hidden_states.dtype     )

							UpperCAmelCase_						=				[]
							UpperCAmelCase_						=				0

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

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

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

							UpperCAmelCase_						=				additional_embeds + [
							    proj_embeddings,
							    time_embeddings[:, None, :],
							    hidden_states,
							]

							if self.prd_embedding is not None:
									UpperCAmelCase_						=				self.prd_embedding.to(hidden_states.dtype     ).expand(lowerCAmelCase    ,	-1    ,	-1     )
									additional_embeds.append(lowerCAmelCase     )

							UpperCAmelCase_						=				torch.cat(
							    lowerCAmelCase    ,	dim=1    ,	)

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

							UpperCAmelCase_						=				hidden_states + positional_embeddings

							if attention_mask is not None:
									UpperCAmelCase_						=				(1 - attention_mask.to(hidden_states.dtype     )) * -10000.0
									UpperCAmelCase_						=				F.pad(lowerCAmelCase    ,	(0, self.additional_embeddings)    ,	value=0.0     )
									UpperCAmelCase_						=				(attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype     )
									UpperCAmelCase_						=				attention_mask.repeat_interleave(self.config.num_attention_heads    ,	dim=0     )

							if self.norm_in is not None:
									UpperCAmelCase_						=				self.norm_in(lowerCAmelCase     )

							for block in self.transformer_blocks:
									UpperCAmelCase_						=				block(lowerCAmelCase    ,	attention_mask=lowerCAmelCase     )

							UpperCAmelCase_						=				self.norm_out(lowerCAmelCase     )

							if self.prd_embedding is not None:
									UpperCAmelCase_						=				hidden_states[:, -1]
							else:
									UpperCAmelCase_						=				hidden_states[:, additional_embeddings_len:]

							UpperCAmelCase_						=				self.proj_to_clip_embeddings(lowerCAmelCase     )

							if not return_dict:
									return (predicted_image_embedding,)

							return PriorTransformerOutput(predicted_image_embedding=lowerCAmelCase     )



					def    A__		(			self    ,	lowerCAmelCase     ):
							UpperCAmelCase_						=				(prior_latents * self.clip_std) + self.clip_mean
							return prior_latents

