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




					"""simple docstring"""

					snake_case__       					=				42






class        snake_case (				__UpperCAmelCase   ,					__UpperCAmelCase       ):




					"""simple docstring"""

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

									UpperCAmelCase__    =							time_embed_dim or inner_dim
									UpperCAmelCase__    =							embedding_proj_dim or embedding_dim
									UpperCAmelCase__    =							clip_embed_dim or embedding_dim

									UpperCAmelCase__    =							Timesteps(lowerCamelCase__					,lowerCamelCase__					,0    )
									UpperCAmelCase__    =							TimestepEmbedding(lowerCamelCase__					,lowerCamelCase__					,out_dim=lowerCamelCase__					,act_fn=lowerCamelCase__    )

									UpperCAmelCase__    =							nn.Linear(lowerCamelCase__					,lowerCamelCase__    )

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

									UpperCAmelCase__    =							nn.Linear(lowerCamelCase__					,lowerCamelCase__    )

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

									UpperCAmelCase__    =							nn.Parameter(torch.zeros(1					,num_embeddings + additional_embeddings					,lowerCamelCase__    )    )

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

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

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

									UpperCAmelCase__    =							nn.LayerNorm(lowerCamelCase__    )

									UpperCAmelCase__    =							nn.Linear(lowerCamelCase__					,lowerCamelCase__    )

									UpperCAmelCase__    =							torch.full(
									    [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]					,-10_000.0    )
									causal_attention_mask.triu_(1    )
									UpperCAmelCase__    =							causal_attention_mask[None, ...]
									self.register_buffer('causal_attention_mask'					,lowerCamelCase__					,persistent=lowerCamelCase__    )

									UpperCAmelCase__    =							nn.Parameter(torch.zeros(1					,lowerCamelCase__    )    )
									UpperCAmelCase__    =							nn.Parameter(torch.zeros(1					,lowerCamelCase__    )    )







					@property
					# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
					def 	__lowerCAmelCase		(       self     :			Optional[Any]    ):
									UpperCAmelCase__    =							{}

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

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

													return processors

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

									return processors







					def 	__lowerCAmelCase		(       self     :			List[Any]					,lowerCamelCase__     :			Union[AttentionProcessor, Dict[str, AttentionProcessor]]    ):
									UpperCAmelCase__    =							len(self.attn_processors.keys()    )

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

									def fn_recursive_attn_processor(lowerCamelCase__     :			str					,lowerCamelCase__     :			torch.nn.Module					,lowerCamelCase__     :			List[str]    ):
													if hasattr(lowerCamelCase__					,'set_processor'    ):
																	if not isinstance(lowerCamelCase__					,lowerCamelCase__    ):
																					module.set_processor(lowerCamelCase__    )
																	else:
																					module.set_processor(processor.pop(f'''{name}.processor'''    )    )

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

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







					def 	__lowerCAmelCase		(       self     :			int    ):
									self.set_attn_processor(AttnProcessor()    )







					def 	__lowerCAmelCase		(       self     :			str					,lowerCamelCase__     :			List[Any]					,lowerCamelCase__     :			Union[torch.Tensor, float, int]					,lowerCamelCase__     :			torch.FloatTensor					,lowerCamelCase__     :			Optional[torch.FloatTensor] = None					,lowerCamelCase__     :			Optional[torch.BoolTensor] = None					,lowerCamelCase__     :			bool = True					,):
									UpperCAmelCase__    =							hidden_states.shape[0]

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

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

									UpperCAmelCase__    =							self.time_proj(lowerCamelCase__    )

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

									if self.embedding_proj_norm is not None:
													UpperCAmelCase__    =							self.embedding_proj_norm(lowerCamelCase__    )

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

									UpperCAmelCase__    =							self.proj_in(lowerCamelCase__    )

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

									UpperCAmelCase__    =							[]
									UpperCAmelCase__    =							0

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

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

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

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

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

									UpperCAmelCase__    =							torch.cat(
									    lowerCamelCase__					,dim=1					,)

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

									UpperCAmelCase__    =							hidden_states + positional_embeddings

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

									if self.norm_in is not None:
													UpperCAmelCase__    =							self.norm_in(lowerCamelCase__    )

									for block in self.transformer_blocks:
													UpperCAmelCase__    =							block(lowerCamelCase__					,attention_mask=lowerCamelCase__    )

									UpperCAmelCase__    =							self.norm_out(lowerCamelCase__    )

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

									UpperCAmelCase__    =							self.proj_to_clip_embeddings(lowerCamelCase__    )

									if not return_dict:
													return (predicted_image_embedding,)

									return PriorTransformerOutput(predicted_image_embedding=lowerCamelCase__    )







					def 	__lowerCAmelCase		(       self     :			List[Any]					,lowerCamelCase__     :			Optional[Any]    ):
									UpperCAmelCase__    =							(prior_latents * self.clip_std) + self.clip_mean
									return prior_latents




"""simple docstring"""






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


lowerCAmelCase__						:						int		  =			{
    'configuration_tapas': ['TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP', 'TapasConfig'],
    'tokenization_tapas': ['TapasTokenizer'],
}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						lowerCAmelCase__						:						Union[str, Any]		  =			[
						    'TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST',
						    'TapasForMaskedLM',
						    'TapasForQuestionAnswering',
						    'TapasForSequenceClassification',
						    'TapasModel',
						    'TapasPreTrainedModel',
						    'load_tf_weights_in_tapas',
						]
try:
						if not is_tf_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						lowerCAmelCase__						:						Optional[Any]		  =			[
						    'TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST',
						    'TFTapasForMaskedLM',
						    'TFTapasForQuestionAnswering',
						    'TFTapasForSequenceClassification',
						    'TFTapasModel',
						    'TFTapasPreTrainedModel',
						]


if TYPE_CHECKING:
						from .configuration_tapas import TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP, TapasConfig
						from .tokenization_tapas import TapasTokenizer

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_tapas import (
												    TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST,
												    TapasForMaskedLM,
												    TapasForQuestionAnswering,
												    TapasForSequenceClassification,
												    TapasModel,
												    TapasPreTrainedModel,
												    load_tf_weights_in_tapas,
												)

						try:
												if not is_tf_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_tf_tapas import (
												    TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST,
												    TFTapasForMaskedLM,
												    TFTapasForQuestionAnswering,
												    TFTapasForSequenceClassification,
												    TFTapasModel,
												    TFTapasPreTrainedModel,
												)


else:
						import sys

						lowerCAmelCase__						:						Dict		  =			_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



