

from typing import Any, Dict, Optional

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

from ..utils import maybe_allow_in_graph
from .activations import get_activation
from .attention_processor import Attention
from .embeddings import CombinedTimestepLabelEmbeddings





@maybe_allow_in_graph
class 	snake_case_   (		nn.Module					):







			'''simple docstring'''

			def __init__(     self :				str							,   _UpperCamelCase :				int							,   _UpperCamelCase :				int							,   _UpperCamelCase :				int							,   _UpperCamelCase :				str=0.0							,   _UpperCamelCase :				Optional[int] = None							,   _UpperCamelCase :				str = "geglu"							,   _UpperCamelCase :				Optional[int] = None							,   _UpperCamelCase :				bool = False							,   _UpperCamelCase :				bool = False							,   _UpperCamelCase :				bool = False							,   _UpperCamelCase :				bool = False							,   _UpperCamelCase :				bool = True							,   _UpperCamelCase :				str = "layer_norm"							,   _UpperCamelCase :				bool = False							,   )      ->str:
									super().__init__()
									snake_case_   =	only_cross_attention

									snake_case_   =	(num_embeds_ada_norm is not None) and norm_type == '''ada_norm_zero'''
									snake_case_   =	(num_embeds_ada_norm is not None) and norm_type == '''ada_norm'''

									if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
															raise ValueError(
															    f'''`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to'''
															    f''' define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}.'''							)

									# Define 3 blocks. Each block has its own normalization layer.
									# 1. Self-Attn
									if self.use_ada_layer_norm:
															snake_case_   =	AdaLayerNorm(_UpperCamelCase							,   _UpperCamelCase							)
									elif self.use_ada_layer_norm_zero:
															snake_case_   =	AdaLayerNormZero(_UpperCamelCase							,   _UpperCamelCase							)
									else:
															snake_case_   =	nn.LayerNorm(_UpperCamelCase							,   elementwise_affine=_UpperCamelCase							)
									snake_case_   =	Attention(
									    query_dim=_UpperCamelCase							,   heads=_UpperCamelCase							,   dim_head=_UpperCamelCase							,   dropout=_UpperCamelCase							,   bias=_UpperCamelCase							,   cross_attention_dim=cross_attention_dim if only_cross_attention else None							,   upcast_attention=_UpperCamelCase							,   )

									# 2. Cross-Attn
									if cross_attention_dim is not None or double_self_attention:
															# We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
															# I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
															# the second cross attention block.
															snake_case_   =	(
															    AdaLayerNorm(_UpperCamelCase							,   _UpperCamelCase							)
															    if self.use_ada_layer_norm
															    else nn.LayerNorm(_UpperCamelCase							,   elementwise_affine=_UpperCamelCase							)
															)
															snake_case_   =	Attention(
															    query_dim=_UpperCamelCase							,   cross_attention_dim=cross_attention_dim if not double_self_attention else None							,   heads=_UpperCamelCase							,   dim_head=_UpperCamelCase							,   dropout=_UpperCamelCase							,   bias=_UpperCamelCase							,   upcast_attention=_UpperCamelCase							,   )  # is self-attn if encoder_hidden_states is none
									else:
															snake_case_   =	None
															snake_case_   =	None

									# 3. Feed-forward
									snake_case_   =	nn.LayerNorm(_UpperCamelCase							,   elementwise_affine=_UpperCamelCase							)
									snake_case_   =	FeedForward(_UpperCamelCase							,   dropout=_UpperCamelCase							,   activation_fn=_UpperCamelCase							,   final_dropout=_UpperCamelCase							)

									# let chunk size default to None
									snake_case_   =	None
									snake_case_   =	0







			def 				snake_case__(     self :				str							,   _UpperCamelCase :				Optional[int]							,   _UpperCamelCase :				int							)      ->Any:
									# Sets chunk feed-forward
									snake_case_   =	chunk_size
									snake_case_   =	dim







			def 				snake_case__(     self :				Optional[Any]							,   _UpperCamelCase :				torch.FloatTensor							,   _UpperCamelCase :				Optional[torch.FloatTensor] = None							,   _UpperCamelCase :				Optional[torch.FloatTensor] = None							,   _UpperCamelCase :				Optional[torch.FloatTensor] = None							,   _UpperCamelCase :				Optional[torch.LongTensor] = None							,   _UpperCamelCase :				Dict[str, Any] = None							,   _UpperCamelCase :				Optional[torch.LongTensor] = None							,   )      ->Optional[Any]:
									# Notice that normalization is always applied before the real computation in the following blocks.
									# 1. Self-Attention
									if self.use_ada_layer_norm:
															snake_case_   =	self.norma(_UpperCamelCase							,   _UpperCamelCase							)
									elif self.use_ada_layer_norm_zero:
															snake_case_,    snake_case_,    snake_case_,    snake_case_,    snake_case_   =	self.norma(
															    _UpperCamelCase							,   _UpperCamelCase							,   _UpperCamelCase							,   hidden_dtype=hidden_states.dtype							)
									else:
															snake_case_   =	self.norma(_UpperCamelCase							)

									snake_case_   =	cross_attention_kwargs if cross_attention_kwargs is not None else {}

									snake_case_   =	self.attna(
									    _UpperCamelCase							,   encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None							,   attention_mask=_UpperCamelCase							,   **_UpperCamelCase							,   )
									if self.use_ada_layer_norm_zero:
															snake_case_   =	gate_msa.unsqueeze(1							) * attn_output
									snake_case_   =	attn_output + hidden_states

									# 2. Cross-Attention
									if self.attna is not None:
															snake_case_   =	(
															    self.norma(_UpperCamelCase							,   _UpperCamelCase							) if self.use_ada_layer_norm else self.norma(_UpperCamelCase							)
															)

															snake_case_   =	self.attna(
															    _UpperCamelCase							,   encoder_hidden_states=_UpperCamelCase							,   attention_mask=_UpperCamelCase							,   **_UpperCamelCase							,   )
															snake_case_   =	attn_output + hidden_states

									# 3. Feed-forward
									snake_case_   =	self.norma(_UpperCamelCase							)

									if self.use_ada_layer_norm_zero:
															snake_case_   =	norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]

									if self._chunk_size is not None:
															# "feed_forward_chunk_size" can be used to save memory
															if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
																					raise ValueError(
																					    f'''`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`.'''							)

															snake_case_   =	norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
															snake_case_   =	torch.cat(
															    [self.ff(_UpperCamelCase							) for hid_slice in norm_hidden_states.chunk(_UpperCamelCase							,   dim=self._chunk_dim							)]							,   dim=self._chunk_dim							,   )
									else:
															snake_case_   =	self.ff(_UpperCamelCase							)

									if self.use_ada_layer_norm_zero:
															snake_case_   =	gate_mlp.unsqueeze(1							) * ff_output

									snake_case_   =	ff_output + hidden_states

									return hidden_states






class 	snake_case_   (		nn.Module					):







			'''simple docstring'''

			def __init__(     self :				Optional[Any]							,   _UpperCamelCase :				int							,   _UpperCamelCase :				Optional[int] = None							,   _UpperCamelCase :				int = 4							,   _UpperCamelCase :				float = 0.0							,   _UpperCamelCase :				str = "geglu"							,   _UpperCamelCase :				bool = False							,   )      ->List[str]:
									super().__init__()
									snake_case_   =	int(dim * mult							)
									snake_case_   =	dim_out if dim_out is not None else dim

									if activation_fn == "gelu":
															snake_case_   =	GELU(_UpperCamelCase							,   _UpperCamelCase							)
									if activation_fn == "gelu-approximate":
															snake_case_   =	GELU(_UpperCamelCase							,   _UpperCamelCase							,   approximate='''tanh'''							)
									elif activation_fn == "geglu":
															snake_case_   =	GEGLU(_UpperCamelCase							,   _UpperCamelCase							)
									elif activation_fn == "geglu-approximate":
															snake_case_   =	ApproximateGELU(_UpperCamelCase							,   _UpperCamelCase							)

									snake_case_   =	nn.ModuleList([]							)
									# project in
									self.net.append(_UpperCamelCase							)
									# project dropout
									self.net.append(nn.Dropout(_UpperCamelCase							)							)
									# project out
									self.net.append(nn.Linear(_UpperCamelCase							,   _UpperCamelCase							)							)
									# FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
									if final_dropout:
															self.net.append(nn.Dropout(_UpperCamelCase							)							)







			def 				snake_case__(     self :				Optional[Any]							,   _UpperCamelCase :				Union[str, Any]							)      ->Tuple:
									for module in self.net:
															snake_case_   =	module(_UpperCamelCase							)
									return hidden_states






class 	snake_case_   (		nn.Module					):







			'''simple docstring'''

			def __init__(     self :				List[str]							,   _UpperCamelCase :				int							,   _UpperCamelCase :				int							,   _UpperCamelCase :				str = "none"							)      ->int:
									super().__init__()
									snake_case_   =	nn.Linear(_UpperCamelCase							,   _UpperCamelCase							)
									snake_case_   =	approximate







			def 				snake_case__(     self :				Tuple							,   _UpperCamelCase :				int							)      ->Dict:
									if gate.device.type != "mps":
															return F.gelu(_UpperCamelCase							,   approximate=self.approximate							)
									# mps: gelu is not implemented for float16
									return F.gelu(gate.to(dtype=torch.floataa							)							,   approximate=self.approximate							).to(dtype=gate.dtype							)







			def 				snake_case__(     self :				Any							,   _UpperCamelCase :				List[str]							)      ->List[Any]:
									snake_case_   =	self.proj(_UpperCamelCase							)
									snake_case_   =	self.gelu(_UpperCamelCase							)
									return hidden_states






class 	snake_case_   (		nn.Module					):







			'''simple docstring'''

			def __init__(     self :				Tuple							,   _UpperCamelCase :				int							,   _UpperCamelCase :				int							)      ->Dict:
									super().__init__()
									snake_case_   =	nn.Linear(_UpperCamelCase							,   dim_out * 2							)







			def 				snake_case__(     self :				Union[str, Any]							,   _UpperCamelCase :				Dict							)      ->Optional[int]:
									if gate.device.type != "mps":
															return F.gelu(_UpperCamelCase							)
									# mps: gelu is not implemented for float16
									return F.gelu(gate.to(dtype=torch.floataa							)							).to(dtype=gate.dtype							)







			def 				snake_case__(     self :				Optional[Any]							,   _UpperCamelCase :				Dict							)      ->List[str]:
									snake_case_,    snake_case_   =	self.proj(_UpperCamelCase							).chunk(2							,   dim=-1							)
									return hidden_states * self.gelu(_UpperCamelCase							)






class 	snake_case_   (		nn.Module					):







			'''simple docstring'''

			def __init__(     self :				int							,   _UpperCamelCase :				int							,   _UpperCamelCase :				int							)      ->Union[str, Any]:
									super().__init__()
									snake_case_   =	nn.Linear(_UpperCamelCase							,   _UpperCamelCase							)







			def 				snake_case__(     self :				str							,   _UpperCamelCase :				Optional[int]							)      ->int:
									snake_case_   =	self.proj(_UpperCamelCase							)
									return x * torch.sigmoid(1.702 * x							)






class 	snake_case_   (		nn.Module					):







			'''simple docstring'''

			def __init__(     self :				str							,   _UpperCamelCase :				Optional[int]							,   _UpperCamelCase :				Tuple							)      ->Union[str, Any]:
									super().__init__()
									snake_case_   =	nn.Embedding(_UpperCamelCase							,   _UpperCamelCase							)
									snake_case_   =	nn.SiLU()
									snake_case_   =	nn.Linear(_UpperCamelCase							,   embedding_dim * 2							)
									snake_case_   =	nn.LayerNorm(_UpperCamelCase							,   elementwise_affine=_UpperCamelCase							)







			def 				snake_case__(     self :				Optional[Any]							,   _UpperCamelCase :				Any							,   _UpperCamelCase :				List[str]							)      ->Union[str, Any]:
									snake_case_   =	self.linear(self.silu(self.emb(_UpperCamelCase							)							)							)
									snake_case_,    snake_case_   =	torch.chunk(_UpperCamelCase							,   2							)
									snake_case_   =	self.norm(_UpperCamelCase							) * (1 + scale) + shift
									return x






class 	snake_case_   (		nn.Module					):







			'''simple docstring'''

			def __init__(     self :				int							,   _UpperCamelCase :				int							,   _UpperCamelCase :				Any							)      ->str:
									super().__init__()

									snake_case_   =	CombinedTimestepLabelEmbeddings(_UpperCamelCase							,   _UpperCamelCase							)

									snake_case_   =	nn.SiLU()
									snake_case_   =	nn.Linear(_UpperCamelCase							,   6 * embedding_dim							,   bias=_UpperCamelCase							)
									snake_case_   =	nn.LayerNorm(_UpperCamelCase							,   elementwise_affine=_UpperCamelCase							,   eps=1e-6							)







			def 				snake_case__(     self :				Tuple							,   _UpperCamelCase :				Optional[Any]							,   _UpperCamelCase :				Union[str, Any]							,   _UpperCamelCase :				Optional[int]							,   _UpperCamelCase :				Optional[Any]=None							)      ->Optional[Any]:
									snake_case_   =	self.linear(self.silu(self.emb(_UpperCamelCase							,   _UpperCamelCase							,   hidden_dtype=_UpperCamelCase							)							)							)
									snake_case_,    snake_case_,    snake_case_,    snake_case_,    snake_case_,    snake_case_   =	emb.chunk(6							,   dim=1							)
									snake_case_   =	self.norm(_UpperCamelCase							) * (1 + scale_msa[:, None]) + shift_msa[:, None]
									return x, gate_msa, shift_mlp, scale_mlp, gate_mlp






class 	snake_case_   (		nn.Module					):







			'''simple docstring'''

			def __init__(     self :				List[str]							,   _UpperCamelCase :				int							,   _UpperCamelCase :				int							,   _UpperCamelCase :				int							,   _UpperCamelCase :				Optional[str] = None							,   _UpperCamelCase :				float = 1e-5							)      ->List[str]:
									super().__init__()
									snake_case_   =	num_groups
									snake_case_   =	eps

									if act_fn is None:
															snake_case_   =	None
									else:
															snake_case_   =	get_activation(_UpperCamelCase							)

									snake_case_   =	nn.Linear(_UpperCamelCase							,   out_dim * 2							)







			def 				snake_case__(     self :				List[Any]							,   _UpperCamelCase :				Dict							,   _UpperCamelCase :				Any							)      ->Any:
									if self.act:
															snake_case_   =	self.act(_UpperCamelCase							)
									snake_case_   =	self.linear(_UpperCamelCase							)
									snake_case_   =	emb[:, :, None, None]
									snake_case_,    snake_case_   =	emb.chunk(2							,   dim=1							)

									snake_case_   =	F.group_norm(_UpperCamelCase							,   self.num_groups							,   eps=self.eps							)
									snake_case_   =	x * (1 + scale) + shift
									return x


def   __SCREAMING_SNAKE_CASE					(SCREAMING_SNAKE_CASE__       ):
						snake_case_   =	[]
						if len(SCREAMING_SNAKE_CASE__       ) == 1:
												return [nums.copy()]
						for _ in range(len(SCREAMING_SNAKE_CASE__       )       ):
												snake_case_   =	nums.pop(0       )
												snake_case_   =	permute(SCREAMING_SNAKE_CASE__       )
												for perm in permutations:
																		perm.append(SCREAMING_SNAKE_CASE__       )
												result.extend(SCREAMING_SNAKE_CASE__       )
												nums.append(SCREAMING_SNAKE_CASE__       )
						return result


def   __SCREAMING_SNAKE_CASE					(SCREAMING_SNAKE_CASE__       ):

						def backtrack(SCREAMING_SNAKE_CASE__       ):
												if start == len(SCREAMING_SNAKE_CASE__       ) - 1:
																		output.append(nums[:]       )
												else:
																		for i in range(SCREAMING_SNAKE_CASE__  ,		len(SCREAMING_SNAKE_CASE__       )       ):
																								snake_case_,    snake_case_   =	nums[i], nums[start]
																								backtrack(start + 1       )
																								snake_case_,    snake_case_   =	nums[i], nums[start]  # backtrack

						snake_case_   =	[]
						backtrack(0       )
						return output


if __name__ == "__main__":
					import doctest

					# use res to print the data in permute2 function
					lowerCAmelCase_       	=	permutea([1, 2, 3])
					print(res)
					doctest.testmod()