


'''simple docstring'''



import gc
import math
import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, logging, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


__UpperCAmelCase      :Any			    =						logging.get_logger(__name__)

enable_full_determinism()



class      a    (		_a     ,  _a     ,  unittest.TestCase	):


     """simple docstring"""






     SCREAMING_SNAKE_CASE	:			int							       =       UNetaDModel
     SCREAMING_SNAKE_CASE	:			Optional[int]							       =       "sample"

     @property
     def 			lowerCamelCase__			(    self					:		Tuple      )						->							Optional[Any]:
          __UpperCAmelCase     :			List[Any]  							=      4
          __UpperCAmelCase     :			Optional[Any]  							=      3
          __UpperCAmelCase     :			Optional[int]  							=      (32, 32)

          __UpperCAmelCase     :			int  							=      floats_tensor((batch_size, num_channels) + sizes      ).to(snake_case      )
          __UpperCAmelCase     :			str  							=      torch.tensor([10]      ).to(snake_case      )

          return {"sample": noise, "timestep": time_step}

     @property
     def 			lowerCamelCase__			(    self					:		List[Any]      )						->							Dict:
          return (3, 32, 32)

     @property
     def 			lowerCamelCase__			(    self					:		Optional[int]      )						->							List[str]:
          return (3, 32, 32)






     def 			lowerCamelCase__			(    self					:		Optional[Any]      )						->							List[Any]:
          __UpperCAmelCase     :			Optional[int]  							=      {
              '''block_out_channels''': (32, 64),
              '''down_block_types''': ('''DownBlock2D''', '''AttnDownBlock2D'''),
              '''up_block_types''': ('''AttnUpBlock2D''', '''UpBlock2D'''),
              '''attention_head_dim''': 3,
              '''out_channels''': 3,
              '''in_channels''': 3,
              '''layers_per_block''': 2,
              '''sample_size''': 32,
          }
          __UpperCAmelCase     :			Dict  							=      self.dummy_input
          return init_dict, inputs_dict



class      a    (		_a     ,  _a     ,  unittest.TestCase	):


     """simple docstring"""






     SCREAMING_SNAKE_CASE	:			str							       =       UNetaDModel
     SCREAMING_SNAKE_CASE	:			List[Any]							       =       "sample"

     @property
     def 			lowerCamelCase__			(    self					:		Union[str, Any]      )						->							List[Any]:
          __UpperCAmelCase     :			Tuple  							=      4
          __UpperCAmelCase     :			str  							=      4
          __UpperCAmelCase     :			Union[str, Any]  							=      (32, 32)

          __UpperCAmelCase     :			int  							=      floats_tensor((batch_size, num_channels) + sizes      ).to(snake_case      )
          __UpperCAmelCase     :			Optional[int]  							=      torch.tensor([10]      ).to(snake_case      )

          return {"sample": noise, "timestep": time_step}

     @property
     def 			lowerCamelCase__			(    self					:		Optional[Any]      )						->							str:
          return (4, 32, 32)

     @property
     def 			lowerCamelCase__			(    self					:		List[str]      )						->							Union[str, Any]:
          return (4, 32, 32)

     def 			lowerCamelCase__			(    self					:		str      )						->							int:
          __UpperCAmelCase     :			Optional[Any]  							=      {
              '''sample_size''': 32,
              '''in_channels''': 4,
              '''out_channels''': 4,
              '''layers_per_block''': 2,
              '''block_out_channels''': (32, 64),
              '''attention_head_dim''': 32,
              '''down_block_types''': ('''DownBlock2D''', '''DownBlock2D'''),
              '''up_block_types''': ('''UpBlock2D''', '''UpBlock2D'''),
          }
          __UpperCAmelCase     :			Optional[Any]  							=      self.dummy_input
          return init_dict, inputs_dict

     def 			lowerCamelCase__			(    self					:		Union[str, Any]      )						->							Tuple:
          __UpperCAmelCase     :			Dict  							=      UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''							,			output_loading_info=snake_case      )

          self.assertIsNotNone(snake_case      )
          self.assertEqual(len(loading_info['''missing_keys''']      )							,			0      )

          model.to(snake_case      )
          __UpperCAmelCase     :			Dict  							=      model(**self.dummy_input      ).sample

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

     @unittest.skipIf(torch_device != '''cuda'''							,			'''This test is supposed to run on GPU'''      )
     def 			lowerCamelCase__			(    self					:		Optional[Any]      )						->							Tuple:
          __UpperCAmelCase     :			List[Any]  							=      UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''							,			output_loading_info=snake_case      )
          model.to(snake_case      )
          __UpperCAmelCase     :			Optional[int]  							=      model(**self.dummy_input      ).sample

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

     @unittest.skipIf(torch_device != '''cuda'''							,			'''This test is supposed to run on GPU'''      )
     def 			lowerCamelCase__			(    self					:		List[str]      )						->							Optional[int]:
          # by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
          __UpperCAmelCase     :			Dict  							=      UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''							,			output_loading_info=snake_case      )
          model_accelerate.to(snake_case      )
          model_accelerate.eval()

          __UpperCAmelCase     :			List[Any]  							=      torch.randn(
              1							,			model_accelerate.config.in_channels							,			model_accelerate.config.sample_size							,			model_accelerate.config.sample_size							,			generator=torch.manual_seed(0      )							,			)
          __UpperCAmelCase     :			Any  							=      noise.to(snake_case      )
          __UpperCAmelCase     :			Union[str, Any]  							=      torch.tensor([10] * noise.shape[0]      ).to(snake_case      )

          __UpperCAmelCase     :			Tuple  							=      model_accelerate(snake_case							,			snake_case      )['''sample''']

          # two models don't need to stay in the device at the same time
          del model_accelerate
          torch.cuda.empty_cache()
          gc.collect()

          __UpperCAmelCase     :			List[str]  							=      UNetaDModel.from_pretrained(
              '''fusing/unet-ldm-dummy-update'''							,			output_loading_info=snake_case							,			low_cpu_mem_usage=snake_case      )
          model_normal_load.to(snake_case      )
          model_normal_load.eval()
          __UpperCAmelCase     :			List[Any]  							=      model_normal_load(snake_case							,			snake_case      )['''sample''']

          assert torch_all_close(snake_case							,			snake_case							,			rtol=1E-3      )






     def 			lowerCamelCase__			(    self					:		Any      )						->							int:
          __UpperCAmelCase     :			int  							=      UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''      )
          model.eval()
          model.to(snake_case      )

          __UpperCAmelCase     :			Union[str, Any]  							=      torch.randn(
              1							,			model.config.in_channels							,			model.config.sample_size							,			model.config.sample_size							,			generator=torch.manual_seed(0      )							,			)
          __UpperCAmelCase     :			Optional[int]  							=      noise.to(snake_case      )
          __UpperCAmelCase     :			List[str]  							=      torch.tensor([10] * noise.shape[0]      ).to(snake_case      )

          with torch.no_grad():
               __UpperCAmelCase     :			Optional[Any]  							=      model(snake_case							,			snake_case      ).sample

          __UpperCAmelCase     :			Optional[Any]  							=      output[0, -1, -3:, -3:].flatten().cpu()
          # fmt: off
          __UpperCAmelCase     :			str  							=      torch.tensor([-13.3_258, -20.1_100, -15.9_873, -17.6_617, -23.0_596, -17.9_419, -13.3_675, -16.1_889, -12.3_800]      )
          # fmt: on

          self.assertTrue(torch_all_close(snake_case							,			snake_case							,			rtol=1E-3      )      )



class      a    (		_a     ,  _a     ,  unittest.TestCase	):


     """simple docstring"""






     SCREAMING_SNAKE_CASE	:			Any							       =       UNetaDModel
     SCREAMING_SNAKE_CASE	:			Optional[int]							       =       "sample"

     @property
     def 			lowerCamelCase__			(    self					:		Tuple							,			snake_case					:		Optional[Any]=(32, 32)      )						->							List[Any]:
          __UpperCAmelCase     :			Any  							=      4
          __UpperCAmelCase     :			Optional[int]  							=      3

          __UpperCAmelCase     :			Any  							=      floats_tensor((batch_size, num_channels) + sizes      ).to(snake_case      )
          __UpperCAmelCase     :			Dict  							=      torch.tensor(batch_size * [10]      ).to(dtype=torch.intaa							,			device=snake_case      )

          return {"sample": noise, "timestep": time_step}

     @property
     def 			lowerCamelCase__			(    self					:		List[str]      )						->							int:
          return (3, 32, 32)

     @property
     def 			lowerCamelCase__			(    self					:		int      )						->							Dict:
          return (3, 32, 32)

     def 			lowerCamelCase__			(    self					:		str      )						->							List[str]:
          __UpperCAmelCase     :			Optional[Any]  							=      {
              '''block_out_channels''': [32, 64, 64, 64],
              '''in_channels''': 3,
              '''layers_per_block''': 1,
              '''out_channels''': 3,
              '''time_embedding_type''': '''fourier''',
              '''norm_eps''': 1E-6,
              '''mid_block_scale_factor''': math.sqrt(2.0      ),
              '''norm_num_groups''': None,
              '''down_block_types''': [
                  '''SkipDownBlock2D''',
                  '''AttnSkipDownBlock2D''',
                  '''SkipDownBlock2D''',
                  '''SkipDownBlock2D''',
              ],
              '''up_block_types''': [
                  '''SkipUpBlock2D''',
                  '''SkipUpBlock2D''',
                  '''AttnSkipUpBlock2D''',
                  '''SkipUpBlock2D''',
              ],
          }
          __UpperCAmelCase     :			Any  							=      self.dummy_input
          return init_dict, inputs_dict

     @slow
     def 			lowerCamelCase__			(    self					:		int      )						->							Any:
          __UpperCAmelCase     :			Tuple  							=      UNetaDModel.from_pretrained('''google/ncsnpp-celebahq-256'''							,			output_loading_info=snake_case      )
          self.assertIsNotNone(snake_case      )
          self.assertEqual(len(loading_info['''missing_keys''']      )							,			0      )

          model.to(snake_case      )
          __UpperCAmelCase     :			Dict  							=      self.dummy_input
          __UpperCAmelCase     :			str  							=      floats_tensor((4, 3) + (256, 256)      ).to(snake_case      )
          __UpperCAmelCase     :			Any  							=      noise
          __UpperCAmelCase     :			List[str]  							=      model(**snake_case      )

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

     @slow
     def 			lowerCamelCase__			(    self					:		Optional[int]      )						->							Dict:
          __UpperCAmelCase     :			int  							=      UNetaDModel.from_pretrained('''google/ncsnpp-celebahq-256'''      )
          model.to(snake_case      )

          __UpperCAmelCase     :			Dict  							=      4
          __UpperCAmelCase     :			Optional[int]  							=      3
          __UpperCAmelCase     :			Any  							=      (256, 256)

          __UpperCAmelCase     :			int  							=      torch.ones((batch_size, num_channels) + sizes      ).to(snake_case      )
          __UpperCAmelCase     :			List[str]  							=      torch.tensor(batch_size * [1E-4]      ).to(snake_case      )

          with torch.no_grad():
               __UpperCAmelCase     :			Tuple  							=      model(snake_case							,			snake_case      ).sample

          __UpperCAmelCase     :			Optional[Any]  							=      output[0, -3:, -3:, -1].flatten().cpu()
          # fmt: off
          __UpperCAmelCase     :			Optional[Any]  							=      torch.tensor([-4842.8691, -6499.6631, -3800.1953, -7978.2686, -1_0980.7129, -2_0028.8535, 8148.2822, 2342.2905, 567.7_608]      )
          # fmt: on

          self.assertTrue(torch_all_close(snake_case							,			snake_case							,			rtol=1E-2      )      )

     def 			lowerCamelCase__			(    self					:		str      )						->							Optional[int]:
          __UpperCAmelCase     :			List[Any]  							=      UNetaDModel.from_pretrained('''fusing/ncsnpp-ffhq-ve-dummy-update'''      )
          model.to(snake_case      )

          __UpperCAmelCase     :			Optional[int]  							=      4
          __UpperCAmelCase     :			str  							=      3
          __UpperCAmelCase     :			Optional[Any]  							=      (32, 32)

          __UpperCAmelCase     :			Optional[int]  							=      torch.ones((batch_size, num_channels) + sizes      ).to(snake_case      )
          __UpperCAmelCase     :			Any  							=      torch.tensor(batch_size * [1E-4]      ).to(snake_case      )

          with torch.no_grad():
               __UpperCAmelCase     :			Dict  							=      model(snake_case							,			snake_case      ).sample

          __UpperCAmelCase     :			Optional[Any]  							=      output[0, -3:, -3:, -1].flatten().cpu()
          # fmt: off
          __UpperCAmelCase     :			Union[str, Any]  							=      torch.tensor([-0.0_325, -0.0_900, -0.0_869, -0.0_332, -0.0_725, -0.0_270, -0.0_101, 0.0_227, 0.0_256]      )
          # fmt: on

          self.assertTrue(torch_all_close(snake_case							,			snake_case							,			rtol=1E-2      )      )






     def 			lowerCamelCase__			(    self					:		Dict      )						->							Optional[int]:
          # not required for this model
          pass



'''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      a    (		_a	):


     """simple docstring"""






     SCREAMING_SNAKE_CASE	:			torch.FloatTensor



class      a    (		_a     ,  _a	):


     """simple docstring"""






     @register_to_config
     def __init__(    self					:		str							,			snake_case					:		int = 32							,			snake_case					:		int = 64							,			snake_case					:		int = 20							,			snake_case					:		int = 768							,			snake_case					:		Tuple=77							,			snake_case					:		List[Any]=4							,			snake_case					:		float = 0.0							,			snake_case					:		str = "silu"							,			snake_case					:		Optional[str] = None							,			snake_case					:		Optional[str] = None							,			snake_case					:		Optional[str] = "linear"							,			snake_case					:		Optional[str] = "prd"							,			snake_case					:		Optional[int] = None							,			snake_case					:		Optional[int] = None							,			snake_case					:		Optional[int] = None							,			)						->							List[str]:
          super().__init__()
          __UpperCAmelCase     :			List[Any]  							=      num_attention_heads
          __UpperCAmelCase     :			List[str]  							=      attention_head_dim
          __UpperCAmelCase     :			int  							=      num_attention_heads * attention_head_dim
          __UpperCAmelCase     :			List[Any]  							=      additional_embeddings

          __UpperCAmelCase     :			Any  							=      time_embed_dim or inner_dim
          __UpperCAmelCase     :			Any  							=      embedding_proj_dim or embedding_dim
          __UpperCAmelCase     :			Union[str, Any]  							=      clip_embed_dim or embedding_dim

          __UpperCAmelCase     :			List[Any]  							=      Timesteps(snake_case							,			snake_case							,			0      )
          __UpperCAmelCase     :			Optional[Any]  							=      TimestepEmbedding(snake_case							,			snake_case							,			out_dim=snake_case							,			act_fn=snake_case      )

          __UpperCAmelCase     :			str  							=      nn.Linear(snake_case							,			snake_case      )

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

          __UpperCAmelCase     :			List[Any]  							=      nn.Linear(snake_case							,			snake_case      )

          if encoder_hid_proj_type is None:
               __UpperCAmelCase     :			Union[str, Any]  							=      None
          elif encoder_hid_proj_type == "linear":
               __UpperCAmelCase     :			Any  							=      nn.Linear(snake_case							,			snake_case      )
          else:
               raise ValueError(f'unsupported encoder_hid_proj_type: {encoder_hid_proj_type}'      )

          __UpperCAmelCase     :			Dict  							=      nn.Parameter(torch.zeros(1							,			num_embeddings + additional_embeddings							,			snake_case      )      )

          if added_emb_type == "prd":
               __UpperCAmelCase     :			Any  							=      nn.Parameter(torch.zeros(1							,			1							,			snake_case      )      )
          elif added_emb_type is None:
               __UpperCAmelCase     :			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`.'      )

          __UpperCAmelCase     :			Optional[int]  							=      nn.ModuleList(
              [
                  BasicTransformerBlock(
                      snake_case							,			snake_case							,			snake_case							,			dropout=snake_case							,			activation_fn='''gelu'''							,			attention_bias=snake_case							,			)
                  for d in range(snake_case      )
              ]      )

          if norm_in_type == "layer":
               __UpperCAmelCase     :			Tuple  							=      nn.LayerNorm(snake_case      )
          elif norm_in_type is None:
               __UpperCAmelCase     :			List[Any]  							=      None
          else:
               raise ValueError(f'Unsupported norm_in_type: {norm_in_type}.'      )

          __UpperCAmelCase     :			Dict  							=      nn.LayerNorm(snake_case      )

          __UpperCAmelCase     :			Any  							=      nn.Linear(snake_case							,			snake_case      )

          __UpperCAmelCase     :			Any  							=      torch.full(
              [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]							,			-10_000.0      )
          causal_attention_mask.triu_(1      )
          __UpperCAmelCase     :			str  							=      causal_attention_mask[None, ...]
          self.register_buffer('''causal_attention_mask'''							,			snake_case							,			persistent=snake_case      )

          __UpperCAmelCase     :			Tuple  							=      nn.Parameter(torch.zeros(1							,			snake_case      )      )
          __UpperCAmelCase     :			List[str]  							=      nn.Parameter(torch.zeros(1							,			snake_case      )      )

     @property
     # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
     def 			lowerCamelCase__			(    self					:		int      )						->							Dict[str, AttentionProcessor]:
          __UpperCAmelCase     :			Optional[Any]  							=      {}

          def fn_recursive_add_processors(snake_case					:		str							,			snake_case					:		torch.nn.Module							,			snake_case					:		Dict[str, AttentionProcessor]      ):
               if hasattr(snake_case							,			'''set_processor'''      ):
                    __UpperCAmelCase     :			Union[str, Any]  							=      module.processor

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

               return processors

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

          return processors

     def 			lowerCamelCase__			(    self					:		Optional[Any]							,			snake_case					:		Union[AttentionProcessor, Dict[str, AttentionProcessor]]      )						->							Union[str, Any]:
          __UpperCAmelCase     :			Union[str, Any]  							=      len(self.attn_processors.keys()      )

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

          def fn_recursive_attn_processor(snake_case					:		str							,			snake_case					:		torch.nn.Module							,			snake_case					:		int      ):
               if hasattr(snake_case							,			'''set_processor'''      ):
                    if not isinstance(snake_case							,			snake_case      ):
                         module.set_processor(snake_case      )
                    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}'							,			snake_case							,			snake_case      )

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

     def 			lowerCamelCase__			(    self					:		str      )						->							Tuple:
          self.set_attn_processor(AttnProcessor()      )

     def 			lowerCamelCase__			(    self					:		Optional[Any]							,			snake_case					:		List[Any]							,			snake_case					:		Union[torch.Tensor, float, int]							,			snake_case					:		torch.FloatTensor							,			snake_case					:		Optional[torch.FloatTensor] = None							,			snake_case					:		Optional[torch.BoolTensor] = None							,			snake_case					:		bool = True							,			)						->							List[Any]:
          __UpperCAmelCase     :			Any  							=      hidden_states.shape[0]

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

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

          __UpperCAmelCase     :			Tuple  							=      self.time_proj(snake_case      )

          # 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     :			int  							=      timesteps_projected.to(dtype=self.dtype      )
          __UpperCAmelCase     :			Optional[int]  							=      self.time_embedding(snake_case      )

          if self.embedding_proj_norm is not None:
               __UpperCAmelCase     :			Optional[Any]  							=      self.embedding_proj_norm(snake_case      )

          __UpperCAmelCase     :			str  							=      self.embedding_proj(snake_case      )
          if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
               __UpperCAmelCase     :			Dict  							=      self.encoder_hidden_states_proj(snake_case      )
          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     :			Optional[int]  							=      self.proj_in(snake_case      )

          __UpperCAmelCase     :			Optional[Any]  							=      self.positional_embedding.to(hidden_states.dtype      )

          __UpperCAmelCase     :			Union[str, Any]  							=      []
          __UpperCAmelCase     :			Optional[Any]  							=      0

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

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

          if len(hidden_states.shape      ) == 2:
               __UpperCAmelCase     :			Union[str, Any]  							=      hidden_states[:, None, :]

          __UpperCAmelCase     :			Union[str, Any]  							=      additional_embeds + [
              proj_embeddings,
              time_embeddings[:, None, :],
              hidden_states,
          ]

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

          __UpperCAmelCase     :			Dict  							=      torch.cat(
              snake_case							,			dim=1							,			)

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

          __UpperCAmelCase     :			Optional[int]  							=      hidden_states + positional_embeddings

          if attention_mask is not None:
               __UpperCAmelCase     :			List[Any]  							=      (1 - attention_mask.to(hidden_states.dtype      )) * -10_000.0
               __UpperCAmelCase     :			str  							=      F.pad(snake_case							,			(0, self.additional_embeddings)							,			value=0.0      )
               __UpperCAmelCase     :			Optional[int]  							=      (attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype      )
               __UpperCAmelCase     :			Optional[int]  							=      attention_mask.repeat_interleave(self.config.num_attention_heads							,			dim=0      )

          if self.norm_in is not None:
               __UpperCAmelCase     :			str  							=      self.norm_in(snake_case      )

          for block in self.transformer_blocks:
               __UpperCAmelCase     :			Optional[int]  							=      block(snake_case							,			attention_mask=snake_case      )

          __UpperCAmelCase     :			int  							=      self.norm_out(snake_case      )

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

          __UpperCAmelCase     :			Dict  							=      self.proj_to_clip_embeddings(snake_case      )

          if not return_dict:
               return (predicted_image_embedding,)

          return PriorTransformerOutput(predicted_image_embedding=snake_case      )






     def 			lowerCamelCase__			(    self					:		Optional[int]							,			snake_case					:		List[str]      )						->							str:
          __UpperCAmelCase     :			Dict  							=      (prior_latents * self.clip_std) + self.clip_mean
          return prior_latents