



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




	'''simple docstring'''





	SCREAMING_SNAKE_CASE        =      42






class      A__  (    __SCREAMING_SNAKE_CASE    ,						__SCREAMING_SNAKE_CASE				):




	'''simple docstring'''





	@register_to_config
	def __init__(	self:						Tuple			,     _SCREAMING_SNAKE_CASE:						int = 32			,     _SCREAMING_SNAKE_CASE:						int = 64			,     _SCREAMING_SNAKE_CASE:						int = 20			,     _SCREAMING_SNAKE_CASE:						int = 768			,     _SCREAMING_SNAKE_CASE:						str=77			,     _SCREAMING_SNAKE_CASE:						int=4			,     _SCREAMING_SNAKE_CASE:						float = 0.0			,     _SCREAMING_SNAKE_CASE:						str = "silu"			,     _SCREAMING_SNAKE_CASE:						Optional[str] = None			,     _SCREAMING_SNAKE_CASE:						Optional[str] = None			,     _SCREAMING_SNAKE_CASE:						Optional[str] = "linear"			,     _SCREAMING_SNAKE_CASE:						Optional[str] = "prd"			,     _SCREAMING_SNAKE_CASE:						Optional[int] = None			,     _SCREAMING_SNAKE_CASE:						Optional[int] = None			,     _SCREAMING_SNAKE_CASE:						Optional[int] = None			,     )							->			Optional[int]:

				"""simple docstring"""

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

				__lowerCAmelCase :    int				=			time_embed_dim or inner_dim
				__lowerCAmelCase :    Optional[int]				=			embedding_proj_dim or embedding_dim
				__lowerCAmelCase :    str				=			clip_embed_dim or embedding_dim

				__lowerCAmelCase :    Dict				=			Timesteps(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE			,     0)
				__lowerCAmelCase :    Optional[int]				=			TimestepEmbedding(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE			,     out_dim=_SCREAMING_SNAKE_CASE			,     act_fn=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    str				=			nn.Linear(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)

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

				__lowerCAmelCase :    Dict				=			nn.Linear(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)

				if encoder_hid_proj_type is None:
							__lowerCAmelCase :    Dict				=			None
				elif encoder_hid_proj_type == "linear":
							__lowerCAmelCase :    List[Any]				=			nn.Linear(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)
				else:
							raise ValueError(F"""unsupported encoder_hid_proj_type: {encoder_hid_proj_type}""")

				__lowerCAmelCase :    Optional[Any]				=			nn.Parameter(torch.zeros(1			,     num_embeddings + additional_embeddings			,     _SCREAMING_SNAKE_CASE))

				if added_emb_type == "prd":
							__lowerCAmelCase :    Union[str, Any]				=			nn.Parameter(torch.zeros(1			,     1			,     _SCREAMING_SNAKE_CASE))
				elif added_emb_type is None:
							__lowerCAmelCase :    List[str]				=			None
				else:
							raise ValueError(
							    F"""`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `'prd'` or `None`.""")

				__lowerCAmelCase :    List[Any]				=			nn.ModuleList(
				    [
				        BasicTransformerBlock(
				            _SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE			,     dropout=_SCREAMING_SNAKE_CASE			,     activation_fn="gelu"			,     attention_bias=_SCREAMING_SNAKE_CASE			,     )
				        for d in range(_SCREAMING_SNAKE_CASE)
				    ])

				if norm_in_type == "layer":
							__lowerCAmelCase :    Tuple				=			nn.LayerNorm(_SCREAMING_SNAKE_CASE)
				elif norm_in_type is None:
							__lowerCAmelCase :    List[Any]				=			None
				else:
							raise ValueError(F"""Unsupported norm_in_type: {norm_in_type}.""")

				__lowerCAmelCase :    str				=			nn.LayerNorm(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    List[Any]				=			nn.Linear(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    List[str]				=			torch.full(
				    [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]			,     -1_0000.0)
				causal_attention_mask.triu_(1)
				__lowerCAmelCase :    Tuple				=			causal_attention_mask[None, ...]
				self.register_buffer("causal_attention_mask"			,     _SCREAMING_SNAKE_CASE			,     persistent=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Optional[int]				=			nn.Parameter(torch.zeros(1			,     _SCREAMING_SNAKE_CASE))
				__lowerCAmelCase :    Dict				=			nn.Parameter(torch.zeros(1			,     _SCREAMING_SNAKE_CASE))

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

				"""simple docstring"""

				__lowerCAmelCase :    str				=			{}

				def fn_recursive_add_processors(_SCREAMING_SNAKE_CASE:						str			,     _SCREAMING_SNAKE_CASE:						torch.nn.Module			,     _SCREAMING_SNAKE_CASE:						Dict[str, AttentionProcessor]):
							if hasattr(_SCREAMING_SNAKE_CASE			,     "set_processor"):
										__lowerCAmelCase :    int				=			module.processor

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

							return processors

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

				return processors

	def 				_SCREAMING_SNAKE_CASE  (	self:						List[Any]			,     _SCREAMING_SNAKE_CASE:						Union[AttentionProcessor, Dict[str, AttentionProcessor]])							->			Union[str, Any]:

				"""simple docstring"""

				__lowerCAmelCase :    Any				=			len(self.attn_processors.keys())

				if isinstance(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE) and len(_SCREAMING_SNAKE_CASE) != count:
							raise ValueError(
							    F"""A dict of processors was passed, but the number of processors {len(_SCREAMING_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(_SCREAMING_SNAKE_CASE:						str			,     _SCREAMING_SNAKE_CASE:						torch.nn.Module			,     _SCREAMING_SNAKE_CASE:						Optional[int]):
							if hasattr(_SCREAMING_SNAKE_CASE			,     "set_processor"):
										if not isinstance(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE):
													module.set_processor(_SCREAMING_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}"""			,     _SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)

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

	def 				_SCREAMING_SNAKE_CASE  (	self:						int)							->			List[str]:

				"""simple docstring"""

				self.set_attn_processor(AttnProcessor())

	def 				_SCREAMING_SNAKE_CASE  (	self:						str			,     _SCREAMING_SNAKE_CASE:						Optional[int]			,     _SCREAMING_SNAKE_CASE:						Union[torch.Tensor, float, int]			,     _SCREAMING_SNAKE_CASE:						torch.FloatTensor			,     _SCREAMING_SNAKE_CASE:						Optional[torch.FloatTensor] = None			,     _SCREAMING_SNAKE_CASE:						Optional[torch.BoolTensor] = None			,     _SCREAMING_SNAKE_CASE:						bool = True			,     )							->			List[Any]:

				"""simple docstring"""

				__lowerCAmelCase :    Dict				=			hidden_states.shape[0]

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

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

				__lowerCAmelCase :    Dict				=			self.time_proj(_SCREAMING_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.
				__lowerCAmelCase :    Dict				=			timesteps_projected.to(dtype=self.dtype)
				__lowerCAmelCase :    Tuple				=			self.time_embedding(_SCREAMING_SNAKE_CASE)

				if self.embedding_proj_norm is not None:
							__lowerCAmelCase :    Any				=			self.embedding_proj_norm(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    str				=			self.embedding_proj(_SCREAMING_SNAKE_CASE)
				if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
							__lowerCAmelCase :    List[Any]				=			self.encoder_hidden_states_proj(_SCREAMING_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")

				__lowerCAmelCase :    Optional[int]				=			self.proj_in(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Tuple				=			self.positional_embedding.to(hidden_states.dtype)

				__lowerCAmelCase :    Union[str, Any]				=			[]
				__lowerCAmelCase :    List[Any]				=			0

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

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

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

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

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

				__lowerCAmelCase :    List[Any]				=			torch.cat(
				    _SCREAMING_SNAKE_CASE			,     dim=1			,     )

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

				__lowerCAmelCase :    List[Any]				=			hidden_states + positional_embeddings

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

				if self.norm_in is not None:
							__lowerCAmelCase :    Optional[int]				=			self.norm_in(_SCREAMING_SNAKE_CASE)

				for block in self.transformer_blocks:
							__lowerCAmelCase :    int				=			block(_SCREAMING_SNAKE_CASE			,     attention_mask=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Tuple				=			self.norm_out(_SCREAMING_SNAKE_CASE)

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

				__lowerCAmelCase :    List[str]				=			self.proj_to_clip_embeddings(_SCREAMING_SNAKE_CASE)

				if not return_dict:
							return (predicted_image_embedding,)

				return PriorTransformerOutput(predicted_image_embedding=_SCREAMING_SNAKE_CASE)







	def 				_SCREAMING_SNAKE_CASE  (	self:						int			,     _SCREAMING_SNAKE_CASE:						str)							->			Optional[Any]:

				"""simple docstring"""

				__lowerCAmelCase :    Optional[Any]				=			(prior_latents * self.clip_std) + self.clip_mean
				return prior_latents




"""simple docstring"""






from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__snake_case				:  Tuple =       {
    'configuration_autoformer': [
        'AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'AutoformerConfig',
    ],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__snake_case				:  Optional[Any] =       [
	    'AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
	    'AutoformerForPrediction',
	    'AutoformerModel',
	    'AutoformerPreTrainedModel',
	]


if TYPE_CHECKING:
	from .configuration_autoformer import (
	    AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
	    AutoformerConfig,
	)

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_autoformer import (
		    AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    AutoformerForPrediction,
		    AutoformerModel,
		    AutoformerPreTrainedModel,
		)

else:
	import sys

	__snake_case				:  Optional[int] =       _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)