
"""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     (							snake_case__   ):
							UpperCAmelCase__						     =      42







class  snake_case     (							snake_case__       ,			snake_case__   ):

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

									"""simple docstring"""







									super().__init__()
									SCREAMING_SNAKE_CASE_						 = num_attention_heads
									SCREAMING_SNAKE_CASE_						 = attention_head_dim
									SCREAMING_SNAKE_CASE_						 = num_attention_heads * attention_head_dim
									SCREAMING_SNAKE_CASE_						 = additional_embeddings

									SCREAMING_SNAKE_CASE_						 = time_embed_dim or inner_dim
									SCREAMING_SNAKE_CASE_						 = embedding_proj_dim or embedding_dim
									SCREAMING_SNAKE_CASE_						 = clip_embed_dim or embedding_dim

									SCREAMING_SNAKE_CASE_						 = Timesteps(_A       ,       _A       ,       0    )
									SCREAMING_SNAKE_CASE_						 = TimestepEmbedding(_A       ,       _A       ,       out_dim=_A       ,       act_fn=_A    )

									SCREAMING_SNAKE_CASE_						 = nn.Linear(_A       ,       _A    )

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

									SCREAMING_SNAKE_CASE_						 = nn.Linear(_A       ,       _A    )

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

									SCREAMING_SNAKE_CASE_						 = nn.Parameter(torch.zeros(1       ,       num_embeddings + additional_embeddings       ,       _A    )    )

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

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

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

									SCREAMING_SNAKE_CASE_						 = nn.LayerNorm(_A    )

									SCREAMING_SNAKE_CASE_						 = nn.Linear(_A       ,       _A    )

									SCREAMING_SNAKE_CASE_						 = torch.full(
									    [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]       ,       -1_00_00.0    )
									causal_attention_mask.triu_(1    )
									SCREAMING_SNAKE_CASE_						 = causal_attention_mask[None, ...]
									self.register_buffer('''causal_attention_mask'''       ,       _A       ,       persistent=_A    )

									SCREAMING_SNAKE_CASE_						 = nn.Parameter(torch.zeros(1       ,       _A    )    )
									SCREAMING_SNAKE_CASE_						 = nn.Parameter(torch.zeros(1       ,       _A    )    )

							@property
							# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = {}

									def fn_recursive_add_processors(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    ):
											if hasattr(_A       ,       '''set_processor'''    ):
													SCREAMING_SNAKE_CASE_						 = 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 _lowercase					(self       ,       SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = 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(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    ):
											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 _lowercase					(self    ):

									"""simple docstring"""







									self.set_attn_processor(AttnProcessor()    )

							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_ = None       ,       SCREAMING_SNAKE_CASE_ = None       ,       SCREAMING_SNAKE_CASE_ = True       ,       ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = hidden_states.shape[0]

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

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

									SCREAMING_SNAKE_CASE_						 = 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.
									SCREAMING_SNAKE_CASE_						 = timesteps_projected.to(dtype=self.dtype    )
									SCREAMING_SNAKE_CASE_						 = self.time_embedding(_A    )

									if self.embedding_proj_norm is not None:
											SCREAMING_SNAKE_CASE_						 = self.embedding_proj_norm(_A    )

									SCREAMING_SNAKE_CASE_						 = self.embedding_proj(_A    )
									if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
											SCREAMING_SNAKE_CASE_						 = 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'''    )

									SCREAMING_SNAKE_CASE_						 = self.proj_in(_A    )

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

									SCREAMING_SNAKE_CASE_						 = []
									SCREAMING_SNAKE_CASE_						 = 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:
											SCREAMING_SNAKE_CASE_						 = proj_embeddings[:, None, :]

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

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

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

									SCREAMING_SNAKE_CASE_						 = torch.cat(
									    _A       ,       dim=1       ,       )

									# Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens
									SCREAMING_SNAKE_CASE_						 = additional_embeddings_len + proj_embeddings.shape[1] + 1
									if positional_embeddings.shape[1] < hidden_states.shape[1]:
											SCREAMING_SNAKE_CASE_						 = 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       ,       )

									SCREAMING_SNAKE_CASE_						 = hidden_states + positional_embeddings

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

									if self.norm_in is not None:
											SCREAMING_SNAKE_CASE_						 = self.norm_in(_A    )

									for block in self.transformer_blocks:
											SCREAMING_SNAKE_CASE_						 = block(_A       ,       attention_mask=_A    )

									SCREAMING_SNAKE_CASE_						 = self.norm_out(_A    )

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

									SCREAMING_SNAKE_CASE_						 = self.proj_to_clip_embeddings(_A    )

									if not return_dict:
											return (predicted_image_embedding,)

									return PriorTransformerOutput(predicted_image_embedding=_A    )







							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = (prior_latents * self.clip_std) + self.clip_mean
									return prior_latents

"""simple docstring"""




def        _lowerCamelCase ( __a    ):

		if not isinstance(__a,							__a    ):
				SCREAMING_SNAKE_CASE_						 = F'Input value of [number={number}] must be an integer'
				raise TypeError(__a    )

		if number < 1:
				SCREAMING_SNAKE_CASE_						 = F'Input value of [number={number}] must be > 0'
				raise ValueError(__a    )

		SCREAMING_SNAKE_CASE_						 = 1

		for i in range(1,							__a    ):
				current_number *= 4 * i - 2
				current_number //= i + 1

		return current_number


if __name__ == "__main__":
						import doctest

						doctest.testmod()