
"""simple docstring"""





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

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

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






@dataclass
class    _UpperCAmelCase(   lowerCamelCase     ):
		lowercase__      =				42






class    _UpperCAmelCase(   lowerCamelCase				,					lowerCamelCase     ):



		@register_to_config
		def __init__( self    ,		__a = 32    ,		__a = 64    ,		__a = 20    ,		__a = 7_68    ,		__a=77    ,		__a=4    ,		__a = 0.0    ,		__a = "silu"    ,		__a = None    ,		__a = None    ,		__a = "linear"    ,		__a = "prd"    ,		__a = None    ,		__a = None    ,		__a = None    ,		)	->	str:






							'''simple docstring'''
							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(__a    ,		__a    ,		0)
							_UpperCamelCase										=  TimestepEmbedding(__a    ,		__a    ,		out_dim=__a    ,		act_fn=__a)

							_UpperCamelCase										=  nn.Linear(__a    ,		__a)

							if embedding_proj_norm_type is None:
												_UpperCamelCase										=  None
							elif embedding_proj_norm_type == "layer":
												_UpperCamelCase										=  nn.LayerNorm(__a)
							else:
												raise ValueError(F'''unsupported embedding_proj_norm_type: {embedding_proj_norm_type}''')

							_UpperCamelCase										=  nn.Linear(__a    ,		__a)

							if encoder_hid_proj_type is None:
												_UpperCamelCase										=  None
							elif encoder_hid_proj_type == "linear":
												_UpperCamelCase										=  nn.Linear(__a    ,		__a)
							else:
												raise ValueError(F'''unsupported encoder_hid_proj_type: {encoder_hid_proj_type}''')

							_UpperCamelCase										=  nn.Parameter(torch.zeros(1    ,		num_embeddings + additional_embeddings    ,		__a))

							if added_emb_type == "prd":
												_UpperCamelCase										=  nn.Parameter(torch.zeros(1    ,		1    ,		__a))
							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(
							            __a    ,		__a    ,		__a    ,		dropout=__a    ,		activation_fn='''gelu'''    ,		attention_bias=__a    ,		)
							        for d in range(__a)
							    ])

							if norm_in_type == "layer":
												_UpperCamelCase										=  nn.LayerNorm(__a)
							elif norm_in_type is None:
												_UpperCamelCase										=  None
							else:
												raise ValueError(F'''Unsupported norm_in_type: {norm_in_type}.''')

							_UpperCamelCase										=  nn.LayerNorm(__a)

							_UpperCamelCase										=  nn.Linear(__a    ,		__a)

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

							_UpperCamelCase										=  nn.Parameter(torch.zeros(1    ,		__a))
							_UpperCamelCase										=  nn.Parameter(torch.zeros(1    ,		__a))



		@property
		# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
		def 						UpperCAmelCase							( self)	->	Dict[str, AttentionProcessor]:






							'''simple docstring'''
							_UpperCamelCase										=  {}

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

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

												return processors

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

							return processors



		def 						UpperCAmelCase							( self    ,		__a)	->	Tuple:






							'''simple docstring'''
							_UpperCamelCase										=  len(self.attn_processors.keys())

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

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

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

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



		def 						UpperCAmelCase							( self)	->	Optional[int]:






							'''simple docstring'''
							self.set_attn_processor(AttnProcessor())



		def 						UpperCAmelCase							( self    ,		__a    ,		__a    ,		__a    ,		__a = None    ,		__a = None    ,		__a = True    ,		)	->	Union[str, Any]:






							'''simple docstring'''
							_UpperCamelCase										=  hidden_states.shape[0]

							_UpperCamelCase										=  timestep
							if not torch.is_tensor(__a):
												_UpperCamelCase										=  torch.tensor([timesteps]    ,		dtype=torch.long    ,		device=hidden_states.device)
							elif torch.is_tensor(__a) 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(__a    ,		dtype=timesteps.dtype    ,		device=timesteps.device)

							_UpperCamelCase										=  self.time_proj(__a)

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

							if self.embedding_proj_norm is not None:
												_UpperCamelCase										=  self.embedding_proj_norm(__a)

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

							_UpperCamelCase										=  self.proj_in(__a)

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

							_UpperCamelCase										=  []
							_UpperCamelCase										=  0

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

							if len(proj_embeddings.shape) == 2:
												_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(__a    ,		-1    ,		-1)
												additional_embeds.append(__a)

							_UpperCamelCase										=  torch.cat(
							    __a    ,		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(
												    __a    ,		(
												        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)) * -1_0000.0
												_UpperCamelCase										=  F.pad(__a    ,		(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(__a)

							for block in self.transformer_blocks:
												_UpperCamelCase										=  block(__a    ,		attention_mask=__a)

							_UpperCamelCase										=  self.norm_out(__a)

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

							_UpperCamelCase										=  self.proj_to_clip_embeddings(__a)

							if not return_dict:
												return (predicted_image_embedding,)

							return PriorTransformerOutput(predicted_image_embedding=__a)




		def 						UpperCAmelCase							( self    ,		__a)	->	Union[str, Any]:






							'''simple docstring'''
							_UpperCamelCase										=  (prior_latents * self.clip_std) + self.clip_mean
							return prior_latents








"""simple docstring"""





from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


_a														=		logging.get_logger(__name__)

_a														=		{
    """bert-base-uncased""": """https://huggingface.co/bert-base-uncased/resolve/main/config.json""",
    """bert-large-uncased""": """https://huggingface.co/bert-large-uncased/resolve/main/config.json""",
    """bert-base-cased""": """https://huggingface.co/bert-base-cased/resolve/main/config.json""",
    """bert-large-cased""": """https://huggingface.co/bert-large-cased/resolve/main/config.json""",
    """bert-base-multilingual-uncased""": """https://huggingface.co/bert-base-multilingual-uncased/resolve/main/config.json""",
    """bert-base-multilingual-cased""": """https://huggingface.co/bert-base-multilingual-cased/resolve/main/config.json""",
    """bert-base-chinese""": """https://huggingface.co/bert-base-chinese/resolve/main/config.json""",
    """bert-base-german-cased""": """https://huggingface.co/bert-base-german-cased/resolve/main/config.json""",
    """bert-large-uncased-whole-word-masking""": (
        """https://huggingface.co/bert-large-uncased-whole-word-masking/resolve/main/config.json"""
    ),
    """bert-large-cased-whole-word-masking""": (
        """https://huggingface.co/bert-large-cased-whole-word-masking/resolve/main/config.json"""
    ),
    """bert-large-uncased-whole-word-masking-finetuned-squad""": (
        """https://huggingface.co/bert-large-uncased-whole-word-masking-finetuned-squad/resolve/main/config.json"""
    ),
    """bert-large-cased-whole-word-masking-finetuned-squad""": (
        """https://huggingface.co/bert-large-cased-whole-word-masking-finetuned-squad/resolve/main/config.json"""
    ),
    """bert-base-cased-finetuned-mrpc""": """https://huggingface.co/bert-base-cased-finetuned-mrpc/resolve/main/config.json""",
    """bert-base-german-dbmdz-cased""": """https://huggingface.co/bert-base-german-dbmdz-cased/resolve/main/config.json""",
    """bert-base-german-dbmdz-uncased""": """https://huggingface.co/bert-base-german-dbmdz-uncased/resolve/main/config.json""",
    """cl-tohoku/bert-base-japanese""": """https://huggingface.co/cl-tohoku/bert-base-japanese/resolve/main/config.json""",
    """cl-tohoku/bert-base-japanese-whole-word-masking""": (
        """https://huggingface.co/cl-tohoku/bert-base-japanese-whole-word-masking/resolve/main/config.json"""
    ),
    """cl-tohoku/bert-base-japanese-char""": (
        """https://huggingface.co/cl-tohoku/bert-base-japanese-char/resolve/main/config.json"""
    ),
    """cl-tohoku/bert-base-japanese-char-whole-word-masking""": (
        """https://huggingface.co/cl-tohoku/bert-base-japanese-char-whole-word-masking/resolve/main/config.json"""
    ),
    """TurkuNLP/bert-base-finnish-cased-v1""": (
        """https://huggingface.co/TurkuNLP/bert-base-finnish-cased-v1/resolve/main/config.json"""
    ),
    """TurkuNLP/bert-base-finnish-uncased-v1""": (
        """https://huggingface.co/TurkuNLP/bert-base-finnish-uncased-v1/resolve/main/config.json"""
    ),
    """wietsedv/bert-base-dutch-cased""": """https://huggingface.co/wietsedv/bert-base-dutch-cased/resolve/main/config.json""",
    # See all BERT models at https://huggingface.co/models?filter=bert
}






class    _UpperCAmelCase(   lowerCamelCase     ):
		lowercase__      =				'bert'



		def __init__( self    ,		__a=3_05_22    ,		__a=7_68    ,		__a=12    ,		__a=12    ,		__a=30_72    ,		__a="gelu"    ,		__a=0.1    ,		__a=0.1    ,		__a=5_12    ,		__a=2    ,		__a=0.02    ,		__a=1e-12    ,		__a=0    ,		__a="absolute"    ,		__a=True    ,		__a=None    ,		**__a    ,		)	->	Optional[Any]:






							'''simple docstring'''
							super().__init__(pad_token_id=__a    ,		**__a)

							_UpperCamelCase										=  vocab_size
							_UpperCamelCase										=  hidden_size
							_UpperCamelCase										=  num_hidden_layers
							_UpperCamelCase										=  num_attention_heads
							_UpperCamelCase										=  hidden_act
							_UpperCamelCase										=  intermediate_size
							_UpperCamelCase										=  hidden_dropout_prob
							_UpperCamelCase										=  attention_probs_dropout_prob
							_UpperCamelCase										=  max_position_embeddings
							_UpperCamelCase										=  type_vocab_size
							_UpperCamelCase										=  initializer_range
							_UpperCamelCase										=  layer_norm_eps
							_UpperCamelCase										=  position_embedding_type
							_UpperCamelCase										=  use_cache
							_UpperCamelCase										=  classifier_dropout






class    _UpperCAmelCase(   lowerCamelCase     ):



		@property
		def 						UpperCAmelCase							( self)	->	Mapping[str, Mapping[int, str]]:






							'''simple docstring'''
							if self.task == "multiple-choice":
												_UpperCamelCase										=  {0: '''batch''', 1: '''choice''', 2: '''sequence'''}
							else:
												_UpperCamelCase										=  {0: '''batch''', 1: '''sequence'''}
							return OrderedDict(
							    [
							        ('''input_ids''', dynamic_axis),
							        ('''attention_mask''', dynamic_axis),
							        ('''token_type_ids''', dynamic_axis),
							    ])






