import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ConvNextConfig, UperNetConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import 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 transformers import UperNetForSemanticSegmentation
      from transformers.models.upernet.modeling_upernet import UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
      from PIL import Image

      from transformers import AutoImageProcessor



class        _lowerCAmelCase    :


  '''simple docstring'''







  def __init__(      self      :					Optional[int]      ,				UpperCamelCase      :					Dict      ,				UpperCamelCase      :					Any=13      ,				UpperCamelCase      :					List[Any]=32      ,				UpperCamelCase      :					int=3      ,				UpperCamelCase      :					List[str]=4      ,				UpperCamelCase      :					Optional[Any]=[10, 20, 30, 40]      ,				UpperCamelCase      :					Union[str, Any]=[2, 2, 3, 2]      ,				UpperCamelCase      :					Tuple=True      ,				UpperCamelCase      :					int=True      ,				UpperCamelCase      :					Tuple=37      ,				UpperCamelCase      :					Union[str, Any]="gelu"      ,				UpperCamelCase      :					str=10      ,				UpperCamelCase      :					Tuple=0.02      ,				UpperCamelCase      :					Tuple=["stage2", "stage3", "stage4"]      ,				UpperCamelCase      :					Dict=3      ,				UpperCamelCase      :					Optional[Any]=None      ,				):


    '''simple docstring'''





    _snake_case   :				Union[str, Any]														=	parent
    _snake_case   :				Optional[int]														=	batch_size
    _snake_case   :				Optional[int]														=	image_size
    _snake_case   :				Dict														=	num_channels
    _snake_case   :				str														=	num_stages
    _snake_case   :				str														=	hidden_sizes
    _snake_case   :				Optional[Any]														=	depths
    _snake_case   :				List[Any]														=	is_training
    _snake_case   :				Tuple														=	use_labels
    _snake_case   :				Dict														=	intermediate_size
    _snake_case   :				List[Any]														=	hidden_act
    _snake_case   :				List[Any]														=	type_sequence_label_size
    _snake_case   :				Any														=	initializer_range
    _snake_case   :				List[Any]														=	out_features
    _snake_case   :				Optional[Any]														=	num_labels
    _snake_case   :				int														=	scope
    _snake_case   :				Union[str, Any]														=	num_stages







  def UpperCamelCase_ (      self      :					List[str]							):


    '''simple docstring'''





    _snake_case   :				Optional[Any]														=	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

    _snake_case   :				str														=	None
    if self.use_labels:
      _snake_case   :				Dict														=	ids_tensor([self.batch_size]      ,				self.type_sequence_label_size							)

    _snake_case   :				Union[str, Any]														=	self.get_config()

    return config, pixel_values, labels







  def UpperCamelCase_ (      self      :					Optional[int]							):


    '''simple docstring'''





    return ConvNextConfig(
        num_channels=self.num_channels      ,				num_stages=self.num_stages      ,				hidden_sizes=self.hidden_sizes      ,				depths=self.depths      ,				is_training=self.is_training      ,				intermediate_size=self.intermediate_size      ,				hidden_act=self.hidden_act      ,				out_features=self.out_features      ,				)







  def UpperCamelCase_ (      self      :					Dict							):


    '''simple docstring'''





    return UperNetConfig(
        backbone_config=self.get_backbone_config()      ,				hidden_size=5_12      ,				pool_scales=[1, 2, 3, 6]      ,				use_auxiliary_head=UpperCamelCase      ,				auxiliary_loss_weight=0.4      ,				auxiliary_in_channels=40      ,				auxiliary_channels=2_56      ,				auxiliary_num_convs=1      ,				auxiliary_concat_input=UpperCamelCase      ,				loss_ignore_index=2_55      ,				num_labels=self.num_labels      ,				)







  def UpperCamelCase_ (      self      :					Union[str, Any]      ,				UpperCamelCase      :					Optional[Any]      ,				UpperCamelCase      :					List[str]      ,				UpperCamelCase      :					Union[str, Any]							):


    '''simple docstring'''





    _snake_case   :				Optional[Any]														=	UperNetForSemanticSegmentation(config=UpperCamelCase							)
    model.to(UpperCamelCase							)
    model.eval()
    _snake_case   :				Any														=	model(UpperCamelCase							)
    self.parent.assertEqual(
        result.logits.shape      ,				(self.batch_size, self.num_labels, self.image_size, self.image_size)							)







  def UpperCamelCase_ (      self      :					int							):


    '''simple docstring'''





    _snake_case   :				Tuple														=	self.prepare_config_and_inputs()
    (
        _snake_case					
    )   :				Tuple														=	config_and_inputs
    _snake_case   :				List[str]														=	{'pixel_values': pixel_values}
    return config, inputs_dict




@require_torch
class        _lowerCAmelCase    (			UpperCAmelCase_							,    UpperCAmelCase_							,    unittest.TestCase							):


  '''simple docstring'''
  a_							:       Union[str, Any]      					=(UperNetForSemanticSegmentation,) if is_torch_available() else ()
  a_							:       Union[str, Any]      					={"""image-segmentation""": UperNetForSemanticSegmentation} if is_torch_available() else {}
  a_							:       int      					=False
  a_							:       List[str]      					=False
  a_							:       List[Any]      					=False
  a_							:       Any      					=False
  a_							:       Optional[int]      					=False
  a_							:       Dict      					=False







  def UpperCamelCase_ (      self      :					Any							):


    '''simple docstring'''





    _snake_case   :				Union[str, Any]														=	UperNetModelTester(self							)
    _snake_case   :				Optional[Any]														=	ConfigTester(self      ,				config_class=UpperCamelCase      ,				has_text_modality=UpperCamelCase      ,				hidden_size=37							)







  def UpperCamelCase_ (      self      :					str							):


    '''simple docstring'''





    self.create_and_test_config_common_properties()
    self.config_tester.create_and_test_config_to_json_string()
    self.config_tester.create_and_test_config_to_json_file()
    self.config_tester.create_and_test_config_from_and_save_pretrained()
    self.config_tester.create_and_test_config_with_num_labels()
    self.config_tester.check_config_can_be_init_without_params()
    self.config_tester.check_config_arguments_init()







  def UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





    return







  def UpperCamelCase_ (      self      :					Tuple							):


    '''simple docstring'''





    _snake_case   :				Any														=	self.model_tester.prepare_config_and_inputs_for_common()

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

      _snake_case   :				str														=	['pixel_values']
      self.assertListEqual(arg_names[:1]      ,				UpperCamelCase							)







  def UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





    _snake_case   :				Dict														=	self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_semantic_segmentation(*UpperCamelCase							)







  @unittest.skip(reason='UperNet does not use inputs_embeds'							)
  def UpperCamelCase_ (      self      :					Tuple							):


    '''simple docstring'''





    pass







  @unittest.skip(reason='UperNet does not support input and output embeddings'							)
  def UpperCamelCase_ (      self      :					int							):


    '''simple docstring'''





    pass







  @unittest.skip(reason='UperNet does not have a base model'							)
  def UpperCamelCase_ (      self      :					int							):


    '''simple docstring'''





    pass







  @unittest.skip(reason='UperNet does not have a base model'							)
  def UpperCamelCase_ (      self      :					Tuple							):


    '''simple docstring'''





    pass







  @require_torch_multi_gpu
  @unittest.skip(reason='UperNet has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'							)
  def UpperCamelCase_ (      self      :					List[str]							):


    '''simple docstring'''





    pass







  @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'							)
  def UpperCamelCase_ (      self      :					Union[str, Any]							):


    '''simple docstring'''





    pass







  def UpperCamelCase_ (      self      :					Union[str, Any]							):


    '''simple docstring'''





    def check_hidden_states_output(UpperCamelCase      :					Union[str, Any]      ,				UpperCamelCase      :					List[str]      ,				UpperCamelCase      :					int							):
      _snake_case   :				List[Any]														=	model_class(UpperCamelCase							)
      model.to(UpperCamelCase							)
      model.eval()

      with torch.no_grad():
        _snake_case   :				List[Any]														=	model(**self._prepare_for_class(UpperCamelCase      ,				UpperCamelCase							)							)

      _snake_case   :				Any														=	outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

      _snake_case   :				Any														=	self.model_tester.num_stages
      self.assertEqual(len(UpperCamelCase							)      ,				expected_num_stages + 1							)

      # ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
      self.assertListEqual(
          list(hidden_states[0].shape[-2:]							)      ,				[self.model_tester.image_size // 4, self.model_tester.image_size // 4]      ,				)

    _snake_case   :				Tuple														=	self.model_tester.prepare_config_and_inputs_for_common()

    for model_class in self.all_model_classes:
      _snake_case   :				List[Any]														=	True
      check_hidden_states_output(UpperCamelCase      ,				UpperCamelCase      ,				UpperCamelCase							)

      # check that output_hidden_states also work using config
      del inputs_dict["output_hidden_states"]
      _snake_case   :				Optional[Any]														=	True

      check_hidden_states_output(UpperCamelCase      ,				UpperCamelCase      ,				UpperCamelCase							)







  def UpperCamelCase_ (      self      :					Dict							):


    '''simple docstring'''





    _snake_case   :				Optional[Any]														=	self.model_tester.prepare_config_and_inputs_for_common()

    _snake_case   :				List[Any]														=	_config_zero_init(UpperCamelCase							)
    _snake_case   :				int														=	_config_zero_init(configs_no_init.backbone_config							)
    for model_class in self.all_model_classes:
      _snake_case   :				List[str]														=	model_class(config=UpperCamelCase							)
      for name, param in model.named_parameters():
        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"""      ,				)







  @unittest.skip(reason='UperNet does not have tied weights'							)
  def UpperCamelCase_ (      self      :					Union[str, Any]							):


    '''simple docstring'''





    pass







  @slow
  def UpperCamelCase_ (      self      :					int							):


    '''simple docstring'''





    for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      _snake_case   :				int														=	UperNetForSemanticSegmentation.from_pretrained(UpperCamelCase							)
      self.assertIsNotNone(UpperCamelCase							)




def     lowerCamelCase_	(      )->		Optional[Any]:
  _snake_case   :				List[str]														=	hf_hub_download(
      repo_id='hf-internal-testing/fixtures_ade20k'    ,							repo_type='dataset'    ,							filename='ADE_val_00000001.jpg'				)
  _snake_case   :				int														=	Image.open(lowerCAmelCase				).convert('RGB'				)
  return image







@require_torch
@require_vision
@slow
class        _lowerCAmelCase    (			unittest.TestCase							):


  '''simple docstring'''







  def UpperCamelCase_ (      self      :					Any							):


    '''simple docstring'''





    _snake_case   :				Tuple														=	AutoImageProcessor.from_pretrained('openmmlab/upernet-swin-tiny'							)
    _snake_case   :				Tuple														=	UperNetForSemanticSegmentation.from_pretrained('openmmlab/upernet-swin-tiny'							).to(UpperCamelCase							)

    _snake_case   :				int														=	prepare_img()
    _snake_case   :				List[Any]														=	processor(images=UpperCamelCase      ,				return_tensors='pt'							).to(UpperCamelCase							)

    with torch.no_grad():
      _snake_case   :				Dict														=	model(**UpperCamelCase							)

    _snake_case   :				Union[str, Any]														=	torch.Size((1, model.config.num_labels, 5_12, 5_12)							)
    self.assertEqual(outputs.logits.shape      ,				UpperCamelCase							)

    _snake_case   :				Optional[Any]														=	torch.tensor(
        [[-7.59_58, -7.59_58, -7.43_02], [-7.59_58, -7.59_58, -7.43_02], [-7.47_97, -7.47_97, -7.30_68]]							).to(UpperCamelCase							)
    self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]      ,				UpperCamelCase      ,				atol=1e-4							)							)







  def UpperCamelCase_ (      self      :					Any							):


    '''simple docstring'''





    _snake_case   :				List[str]														=	AutoImageProcessor.from_pretrained('openmmlab/upernet-convnext-tiny'							)
    _snake_case   :				Union[str, Any]														=	UperNetForSemanticSegmentation.from_pretrained('openmmlab/upernet-convnext-tiny'							).to(UpperCamelCase							)

    _snake_case   :				Optional[Any]														=	prepare_img()
    _snake_case   :				Tuple														=	processor(images=UpperCamelCase      ,				return_tensors='pt'							).to(UpperCamelCase							)

    with torch.no_grad():
      _snake_case   :				Dict														=	model(**UpperCamelCase							)

    _snake_case   :				int														=	torch.Size((1, model.config.num_labels, 5_12, 5_12)							)
    self.assertEqual(outputs.logits.shape      ,				UpperCamelCase							)

    _snake_case   :				List[str]														=	torch.tensor(
        [[-8.81_10, -8.81_10, -8.65_21], [-8.81_10, -8.81_10, -8.65_21], [-8.77_46, -8.77_46, -8.61_30]]							).to(UpperCamelCase							)
    self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]      ,				UpperCamelCase      ,				atol=1e-4							)							)











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    (			UpperCAmelCase_							):


  '''simple docstring'''
  a_							:       torch.FloatTensor




class        _lowerCAmelCase    (			UpperCAmelCase_							,    UpperCAmelCase_							):


  '''simple docstring'''







  @register_to_config
  def __init__(      self      :					str      ,				UpperCamelCase      :					int = 32      ,				UpperCamelCase      :					int = 64      ,				UpperCamelCase      :					int = 20      ,				UpperCamelCase      :					int = 7_68      ,				UpperCamelCase      :					Optional[int]=77      ,				UpperCamelCase      :					int=4      ,				UpperCamelCase      :					float = 0.0      ,				UpperCamelCase      :					str = "silu"      ,				UpperCamelCase      :					Optional[str] = None      ,				UpperCamelCase      :					Optional[str] = None      ,				UpperCamelCase      :					Optional[str] = "linear"      ,				UpperCamelCase      :					Optional[str] = "prd"      ,				UpperCamelCase      :					Optional[int] = None      ,				UpperCamelCase      :					Optional[int] = None      ,				UpperCamelCase      :					Optional[int] = None      ,				):


    '''simple docstring'''





    super().__init__()
    _snake_case   :				str														=	num_attention_heads
    _snake_case   :				Optional[int]														=	attention_head_dim
    _snake_case   :				Any														=	num_attention_heads * attention_head_dim
    _snake_case   :				List[Any]														=	additional_embeddings

    _snake_case   :				List[str]														=	time_embed_dim or inner_dim
    _snake_case   :				int														=	embedding_proj_dim or embedding_dim
    _snake_case   :				List[Any]														=	clip_embed_dim or embedding_dim

    _snake_case   :				Optional[Any]														=	Timesteps(UpperCamelCase      ,				UpperCamelCase      ,				0							)
    _snake_case   :				List[Any]														=	TimestepEmbedding(UpperCamelCase      ,				UpperCamelCase      ,				out_dim=UpperCamelCase      ,				act_fn=UpperCamelCase							)

    _snake_case   :				Optional[int]														=	nn.Linear(UpperCamelCase      ,				UpperCamelCase							)

    if embedding_proj_norm_type is None:
      _snake_case   :				str														=	None
    elif embedding_proj_norm_type == "layer":
      _snake_case   :				List[Any]														=	nn.LayerNorm(UpperCamelCase							)
    else:
      raise ValueError(f"""unsupported embedding_proj_norm_type: {embedding_proj_norm_type}"""							)

    _snake_case   :				str														=	nn.Linear(UpperCamelCase      ,				UpperCamelCase							)

    if encoder_hid_proj_type is None:
      _snake_case   :				Any														=	None
    elif encoder_hid_proj_type == "linear":
      _snake_case   :				Optional[int]														=	nn.Linear(UpperCamelCase      ,				UpperCamelCase							)
    else:
      raise ValueError(f"""unsupported encoder_hid_proj_type: {encoder_hid_proj_type}"""							)

    _snake_case   :				List[str]														=	nn.Parameter(torch.zeros(1      ,				num_embeddings + additional_embeddings      ,				UpperCamelCase							)							)

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

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

    if norm_in_type == "layer":
      _snake_case   :				Optional[int]														=	nn.LayerNorm(UpperCamelCase							)
    elif norm_in_type is None:
      _snake_case   :				Optional[Any]														=	None
    else:
      raise ValueError(f"""Unsupported norm_in_type: {norm_in_type}."""							)

    _snake_case   :				Optional[Any]														=	nn.LayerNorm(UpperCamelCase							)

    _snake_case   :				Union[str, Any]														=	nn.Linear(UpperCamelCase      ,				UpperCamelCase							)

    _snake_case   :				List[Any]														=	torch.full(
        [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]      ,				-1_00_00.0							)
    causal_attention_mask.triu_(1							)
    _snake_case   :				Optional[Any]														=	causal_attention_mask[None, ...]
    self.register_buffer('causal_attention_mask'      ,				UpperCamelCase      ,				persistent=UpperCamelCase							)

    _snake_case   :				str														=	nn.Parameter(torch.zeros(1      ,				UpperCamelCase							)							)
    _snake_case   :				List[str]														=	nn.Parameter(torch.zeros(1      ,				UpperCamelCase							)							)







  @property
  # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
  def UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





    _snake_case   :				Optional[Any]														=	{}

    def fn_recursive_add_processors(UpperCamelCase      :					str      ,				UpperCamelCase      :					torch.nn.Module      ,				UpperCamelCase      :					Dict[str, AttentionProcessor]							):
      if hasattr(UpperCamelCase      ,				'set_processor'							):
        _snake_case   :				Tuple														=	module.processor

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

      return processors

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

    return processors







  def UpperCamelCase_ (      self      :					List[Any]      ,				UpperCamelCase      :					Union[AttentionProcessor, Dict[str, AttentionProcessor]]							):


    '''simple docstring'''





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

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

    def fn_recursive_attn_processor(UpperCamelCase      :					str      ,				UpperCamelCase      :					torch.nn.Module      ,				UpperCamelCase      :					Union[str, Any]							):
      if hasattr(UpperCamelCase      ,				'set_processor'							):
        if not isinstance(UpperCamelCase      ,				UpperCamelCase							):
          module.set_processor(UpperCamelCase							)
        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}"""      ,				UpperCamelCase      ,				UpperCamelCase							)

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







  def UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





    self.set_attn_processor(AttnProcessor()							)







  def UpperCamelCase_ (      self      :					Optional[Any]      ,				UpperCamelCase      :					Any      ,				UpperCamelCase      :					Union[torch.Tensor, float, int]      ,				UpperCamelCase      :					torch.FloatTensor      ,				UpperCamelCase      :					Optional[torch.FloatTensor] = None      ,				UpperCamelCase      :					Optional[torch.BoolTensor] = None      ,				UpperCamelCase      :					bool = True      ,				):


    '''simple docstring'''





    _snake_case   :				Dict														=	hidden_states.shape[0]

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

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

    _snake_case   :				Union[str, Any]														=	self.time_proj(UpperCamelCase							)

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

    if self.embedding_proj_norm is not None:
      _snake_case   :				Optional[Any]														=	self.embedding_proj_norm(UpperCamelCase							)

    _snake_case   :				Union[str, Any]														=	self.embedding_proj(UpperCamelCase							)
    if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
      _snake_case   :				Dict														=	self.encoder_hidden_states_proj(UpperCamelCase							)
    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'							)

    _snake_case   :				str														=	self.proj_in(UpperCamelCase							)

    _snake_case   :				int														=	self.positional_embedding.to(hidden_states.dtype							)

    _snake_case   :				Optional[int]														=	[]
    _snake_case   :				List[Any]														=	0

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

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

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

    _snake_case   :				str														=	additional_embeds + [
        proj_embeddings,
        time_embeddings[:, None, :],
        hidden_states,
    ]

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

    _snake_case   :				Optional[int]														=	torch.cat(
        UpperCamelCase      ,				dim=1      ,				)

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

    _snake_case   :				Optional[Any]														=	hidden_states + positional_embeddings

    if attention_mask is not None:
      _snake_case   :				Any														=	(1 - attention_mask.to(hidden_states.dtype							)) * -1_00_00.0
      _snake_case   :				Tuple														=	F.pad(UpperCamelCase      ,				(0, self.additional_embeddings)      ,				value=0.0							)
      _snake_case   :				int														=	(attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype							)
      _snake_case   :				str														=	attention_mask.repeat_interleave(self.config.num_attention_heads      ,				dim=0							)

    if self.norm_in is not None:
      _snake_case   :				Tuple														=	self.norm_in(UpperCamelCase							)

    for block in self.transformer_blocks:
      _snake_case   :				Any														=	block(UpperCamelCase      ,				attention_mask=UpperCamelCase							)

    _snake_case   :				Dict														=	self.norm_out(UpperCamelCase							)

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

    _snake_case   :				List[Any]														=	self.proj_to_clip_embeddings(UpperCamelCase							)

    if not return_dict:
      return (predicted_image_embedding,)

    return PriorTransformerOutput(predicted_image_embedding=UpperCamelCase							)







  def UpperCamelCase_ (      self      :					Tuple      ,				UpperCamelCase      :					Union[str, Any]							):


    '''simple docstring'''





    _snake_case   :				List[Any]														=	(prior_latents * self.clip_std) + self.clip_mean
    return prior_latents




