






"""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       snake_case_       (    SCREAMING_SNAKE_CASE__       ):



	"""simple docstring"""

	SCREAMING_SNAKE_CASE_:							Optional[int]   			=					42





class       snake_case_       (    SCREAMING_SNAKE_CASE__    ,							SCREAMING_SNAKE_CASE__       ):



	"""simple docstring"""

	@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__()
					A__							 =       num_attention_heads
					A__							 =       attention_head_dim
					A__							 =       num_attention_heads * attention_head_dim
					A__							 =       additional_embeddings

					A__							 =       time_embed_dim or inner_dim
					A__							 =       embedding_proj_dim or embedding_dim
					A__							 =       clip_embed_dim or embedding_dim

					A__							 =       Timesteps(_lowercase  ,					_lowercase  ,					0						)
					A__							 =       TimestepEmbedding(_lowercase  ,					_lowercase  ,					out_dim=_lowercase  ,					act_fn=_lowercase						)

					A__							 =       nn.Linear(_lowercase  ,					_lowercase						)

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

					A__							 =       nn.Linear(_lowercase  ,					_lowercase						)

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

					A__							 =       nn.Parameter(torch.zeros(1  ,					num_embeddings + additional_embeddings  ,					_lowercase						)						)

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

					A__							 =       nn.ModuleList(
					    [
					        BasicTransformerBlock(
					            _lowercase  ,					_lowercase  ,					_lowercase  ,					dropout=_lowercase  ,					activation_fn='gelu'  ,					attention_bias=_lowercase  ,					)
					        for d in range(_lowercase						)
					    ]						)

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

					A__							 =       nn.LayerNorm(_lowercase						)

					A__							 =       nn.Linear(_lowercase  ,					_lowercase						)

					A__							 =       torch.full(
					    [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]  ,					-1_0000.0						)
					causal_attention_mask.triu_(1						)
					A__							 =       causal_attention_mask[None, ...]
					self.register_buffer('causal_attention_mask'  ,					_lowercase  ,					persistent=_lowercase						)

					A__							 =       nn.Parameter(torch.zeros(1  ,					_lowercase						)						)
					A__							 =       nn.Parameter(torch.zeros(1  ,					_lowercase						)						)


	@property
	# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       {}

					def fn_recursive_add_processors(__a  ,					__a  ,					__a						):
									if hasattr(_lowercase  ,					'set_processor'						):
													A__							 =       module.processor

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

									return processors

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

					return processors


	def       _UpperCAmelCase						(      self  ,					__a						):


					"""simple docstring"""

					A__							 =       len(self.attn_processors.keys()						)

					if isinstance(_lowercase  ,					_lowercase						) and len(_lowercase						) != count:
									raise ValueError(
									    f'''A dict of processors was passed, but the number of processors {len(_lowercase						)} 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(_lowercase  ,					'set_processor'						):
													if not isinstance(_lowercase  ,					_lowercase						):
																	module.set_processor(_lowercase						)
													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}'''  ,					_lowercase  ,					_lowercase						)

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


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					self.set_attn_processor(AttnProcessor()						)


	def       _UpperCAmelCase						(      self  ,					__a  ,					__a  ,					__a  ,					__a = None  ,					__a = None  ,					__a = True  ,					):


					"""simple docstring"""

					A__							 =       hidden_states.shape[0]

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

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

					A__							 =       self.time_proj(_lowercase						)

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

					if self.embedding_proj_norm is not None:
									A__							 =       self.embedding_proj_norm(_lowercase						)

					A__							 =       self.embedding_proj(_lowercase						)
					if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
									A__							 =       self.encoder_hidden_states_proj(_lowercase						)
					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'						)

					A__							 =       self.proj_in(_lowercase						)

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

					A__							 =       []
					A__							 =       0

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

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

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

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

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

					A__							 =       torch.cat(
					    _lowercase  ,					dim=1  ,					)

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

					A__							 =       hidden_states + positional_embeddings

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

					if self.norm_in is not None:
									A__							 =       self.norm_in(_lowercase						)

					for block in self.transformer_blocks:
									A__							 =       block(_lowercase  ,					attention_mask=_lowercase						)

					A__							 =       self.norm_out(_lowercase						)

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

					A__							 =       self.proj_to_clip_embeddings(_lowercase						)

					if not return_dict:
									return (predicted_image_embedding,)

					return PriorTransformerOutput(predicted_image_embedding=_lowercase						)






	def       _UpperCAmelCase						(      self  ,					__a						):


					"""simple docstring"""

					A__							 =       (prior_latents * self.clip_std) + self.clip_mean
					return prior_latents







"""simple docstring"""






import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
					import torch

if is_vision_available():
					from PIL import Image

					from transformers import PoolFormerImageProcessor





class       snake_case_       (    unittest.TestCase       ):



	"""simple docstring"""

	def __init__(      self  ,					__a  ,					__a=7  ,					__a=3  ,					__a=30  ,					__a=400  ,					__a=True  ,					__a=None  ,					__a=0.9  ,					__a=None  ,					__a=True  ,					__a=[0.5, 0.5, 0.5]  ,					__a=[0.5, 0.5, 0.5]  ,					):


					"""simple docstring"""

					A__							 =       size if size is not None else {'shortest_edge': 30}
					A__							 =       crop_size if crop_size is not None else {'height': 30, 'width': 30}
					A__							 =       parent
					A__							 =       batch_size
					A__							 =       num_channels
					A__							 =       min_resolution
					A__							 =       max_resolution
					A__							 =       do_resize_and_center_crop
					A__							 =       size
					A__							 =       crop_pct
					A__							 =       crop_size
					A__							 =       do_normalize
					A__							 =       image_mean
					A__							 =       image_std






	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					return {
					    "size": self.size,
					    "do_resize_and_center_crop": self.do_resize_and_center_crop,
					    "crop_pct": self.crop_pct,
					    "crop_size": self.crop_size,
					    "do_normalize": self.do_normalize,
					    "image_mean": self.image_mean,
					    "image_std": self.image_std,
					}





@require_torch
@require_vision
class       snake_case_       (    _lowerCamelCase    ,							unittest.TestCase       ):



	"""simple docstring"""

	SCREAMING_SNAKE_CASE_:							List[Any]   			=					PoolFormerImageProcessor if is_vision_available() else None


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       PoolFormerImageProcessingTester(self						)


	@property
	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					return self.image_processor_tester.prepare_image_processor_dict()


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       self.image_processing_class(**self.image_processor_dict						)
					self.assertTrue(hasattr(__a  ,					'do_resize_and_center_crop'						)						)
					self.assertTrue(hasattr(__a  ,					'size'						)						)
					self.assertTrue(hasattr(__a  ,					'crop_pct'						)						)
					self.assertTrue(hasattr(__a  ,					'do_normalize'						)						)
					self.assertTrue(hasattr(__a  ,					'image_mean'						)						)
					self.assertTrue(hasattr(__a  ,					'image_std'						)						)


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       self.image_processing_class.from_dict(self.image_processor_dict						)
					self.assertEqual(image_processor.size  ,					{'shortest_edge': 30}						)
					self.assertEqual(image_processor.crop_size  ,					{'height': 30, 'width': 30}						)

					A__							 =       self.image_processing_class.from_dict(self.image_processor_dict  ,					size=42  ,					crop_size=84						)
					self.assertEqual(image_processor.size  ,					{'shortest_edge': 42}						)
					self.assertEqual(image_processor.crop_size  ,					{'height': 84, 'width': 84}						)


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					pass


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       self.image_processing_class(**self.image_processor_dict						)
					# create random PIL images
					A__							 =       prepare_image_inputs(self.image_processor_tester  ,					equal_resolution=__a						)
					for image in image_inputs:
									self.assertIsInstance(__a  ,					Image.Image						)

					# Test not batched input
					A__							 =       image_processing(image_inputs[0]  ,					return_tensors='pt'						).pixel_values

					self.assertEqual(
					    encoded_images.shape  ,					(
					        1,
					        self.image_processor_tester.num_channels,
					        self.image_processor_tester.crop_size['height'],
					        self.image_processor_tester.crop_size['width'],
					    )  ,					)

					# Test batched
					A__							 =       image_processing(__a  ,					return_tensors='pt'						).pixel_values
					self.assertEqual(
					    encoded_images.shape  ,					(
					        self.image_processor_tester.batch_size,
					        self.image_processor_tester.num_channels,
					        self.image_processor_tester.crop_size['height'],
					        self.image_processor_tester.crop_size['width'],
					    )  ,					)


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       self.image_processing_class(**self.image_processor_dict						)
					# create random numpy tensors
					A__							 =       prepare_image_inputs(self.image_processor_tester  ,					equal_resolution=__a  ,					numpify=__a						)
					for image in image_inputs:
									self.assertIsInstance(__a  ,					np.ndarray						)

					# Test not batched input
					A__							 =       image_processing(image_inputs[0]  ,					return_tensors='pt'						).pixel_values
					self.assertEqual(
					    encoded_images.shape  ,					(
					        1,
					        self.image_processor_tester.num_channels,
					        self.image_processor_tester.crop_size['height'],
					        self.image_processor_tester.crop_size['width'],
					    )  ,					)

					# Test batched
					A__							 =       image_processing(__a  ,					return_tensors='pt'						).pixel_values
					self.assertEqual(
					    encoded_images.shape  ,					(
					        self.image_processor_tester.batch_size,
					        self.image_processor_tester.num_channels,
					        self.image_processor_tester.crop_size['height'],
					        self.image_processor_tester.crop_size['width'],
					    )  ,					)






	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       self.image_processing_class(**self.image_processor_dict						)
					# create random PyTorch tensors
					A__							 =       prepare_image_inputs(self.image_processor_tester  ,					equal_resolution=__a  ,					torchify=__a						)
					for image in image_inputs:
									self.assertIsInstance(__a  ,					torch.Tensor						)

					# Test not batched input
					A__							 =       image_processing(image_inputs[0]  ,					return_tensors='pt'						).pixel_values
					self.assertEqual(
					    encoded_images.shape  ,					(
					        1,
					        self.image_processor_tester.num_channels,
					        self.image_processor_tester.crop_size['height'],
					        self.image_processor_tester.crop_size['width'],
					    )  ,					)

					# Test batched
					A__							 =       image_processing(__a  ,					return_tensors='pt'						).pixel_values
					self.assertEqual(
					    encoded_images.shape  ,					(
					        self.image_processor_tester.batch_size,
					        self.image_processor_tester.num_channels,
					        self.image_processor_tester.crop_size['height'],
					        self.image_processor_tester.crop_size['width'],
					    )  ,					)


