






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      UpperCamelCase (				nn.Module ):







		"""simple docstring"""



		def __init__(  self		:       Union[str, Any]						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       Any=0.0						,_SCREAMING_SNAKE_CASE		:       Optional[int] = None						,_SCREAMING_SNAKE_CASE		:       str = "geglu"						,_SCREAMING_SNAKE_CASE		:       Optional[int] = None						,_SCREAMING_SNAKE_CASE		:       bool = False						,_SCREAMING_SNAKE_CASE		:       bool = False						,_SCREAMING_SNAKE_CASE		:       bool = False						,_SCREAMING_SNAKE_CASE		:       bool = False						,_SCREAMING_SNAKE_CASE		:       bool = True						,_SCREAMING_SNAKE_CASE		:       str = "layer_norm"						,_SCREAMING_SNAKE_CASE		:       bool = False						,)    ->   Optional[int]:


						'''simple docstring'''
						super().__init__()
						A	 =      only_cross_attention

						A	 =      (num_embeds_ada_norm is not None) and norm_type == 'ada_norm_zero'
						A	 =      (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:
										A	 =      AdaLayerNorm(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		)
						elif self.use_ada_layer_norm_zero:
										A	 =      AdaLayerNormZero(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		)
						else:
										A	 =      nn.LayerNorm(_SCREAMING_SNAKE_CASE						,elementwise_affine=_SCREAMING_SNAKE_CASE		)
						A	 =      Attention(
						    query_dim=_SCREAMING_SNAKE_CASE						,heads=_SCREAMING_SNAKE_CASE						,dim_head=_SCREAMING_SNAKE_CASE						,dropout=_SCREAMING_SNAKE_CASE						,bias=_SCREAMING_SNAKE_CASE						,cross_attention_dim=cross_attention_dim if only_cross_attention else None						,upcast_attention=_SCREAMING_SNAKE_CASE						,)

						# 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.
										A	 =      (
										    AdaLayerNorm(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		)
										    if self.use_ada_layer_norm
										    else nn.LayerNorm(_SCREAMING_SNAKE_CASE						,elementwise_affine=_SCREAMING_SNAKE_CASE		)
										)
										A	 =      Attention(
										    query_dim=_SCREAMING_SNAKE_CASE						,cross_attention_dim=cross_attention_dim if not double_self_attention else None						,heads=_SCREAMING_SNAKE_CASE						,dim_head=_SCREAMING_SNAKE_CASE						,dropout=_SCREAMING_SNAKE_CASE						,bias=_SCREAMING_SNAKE_CASE						,upcast_attention=_SCREAMING_SNAKE_CASE						,)  # is self-attn if encoder_hidden_states is none
						else:
										A	 =      None
										A	 =      None

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

						# let chunk size default to None
						A	 =      None
						A	 =      0
		def       A(  self		:       Tuple						,_SCREAMING_SNAKE_CASE		:       Optional[int]						,_SCREAMING_SNAKE_CASE		:       int		)    ->   Optional[Any]:


						'''simple docstring'''
						# Sets chunk feed-forward
						A	 =      chunk_size
						A	 =      dim







		def       A(  self		:       Optional[Any]						,_SCREAMING_SNAKE_CASE		:       torch.FloatTensor						,_SCREAMING_SNAKE_CASE		:       Optional[torch.FloatTensor] = None						,_SCREAMING_SNAKE_CASE		:       Optional[torch.FloatTensor] = None						,_SCREAMING_SNAKE_CASE		:       Optional[torch.FloatTensor] = None						,_SCREAMING_SNAKE_CASE		:       Optional[torch.LongTensor] = None						,_SCREAMING_SNAKE_CASE		:       Dict[str, Any] = None						,_SCREAMING_SNAKE_CASE		:       Optional[torch.LongTensor] = None						,)    ->   str:


						'''simple docstring'''
						# Notice that normalization is always applied before the real computation in the following blocks.
						# 1. Self-Attention
						if self.use_ada_layer_norm:
										A	 =      self.norma(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		)
						elif self.use_ada_layer_norm_zero:
										A ,		A ,		A ,		A ,		A	 =      self.norma(
										    _SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE						,hidden_dtype=hidden_states.dtype		)
						else:
										A	 =      self.norma(_SCREAMING_SNAKE_CASE		)

						A	 =      cross_attention_kwargs if cross_attention_kwargs is not None else {}

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

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

										A	 =      self.attna(
										    _SCREAMING_SNAKE_CASE						,encoder_hidden_states=_SCREAMING_SNAKE_CASE						,attention_mask=_SCREAMING_SNAKE_CASE						,**_SCREAMING_SNAKE_CASE						,)
										A	 =      attn_output + hidden_states

						# 3. Feed-forward
						A	 =      self.norma(_SCREAMING_SNAKE_CASE		)

						if self.use_ada_layer_norm_zero:
										A	 =      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`.'		)

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

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

						A	 =      ff_output + hidden_states

						return hidden_states


class      UpperCamelCase (				nn.Module ):







		"""simple docstring"""



		def __init__(  self		:       int						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       Optional[int] = None						,_SCREAMING_SNAKE_CASE		:       int = 4						,_SCREAMING_SNAKE_CASE		:       float = 0.0						,_SCREAMING_SNAKE_CASE		:       str = "geglu"						,_SCREAMING_SNAKE_CASE		:       bool = False						,)    ->   str:


						'''simple docstring'''
						super().__init__()
						A	 =      int(dim * mult		)
						A	 =      dim_out if dim_out is not None else dim

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

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







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


						'''simple docstring'''
						for module in self.net:
										A	 =      module(_SCREAMING_SNAKE_CASE		)
						return hidden_states


class      UpperCamelCase (				nn.Module ):







		"""simple docstring"""



		def __init__(  self		:       Tuple						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       str = "none"		)    ->   str:


						'''simple docstring'''
						super().__init__()
						A	 =      nn.Linear(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		)
						A	 =      approximate
		def       A(  self		:       str						,_SCREAMING_SNAKE_CASE		:       Union[str, Any]		)    ->   Optional[int]:


						'''simple docstring'''
						if gate.device.type != "mps":
										return F.gelu(_SCREAMING_SNAKE_CASE						,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       A(  self		:       str						,_SCREAMING_SNAKE_CASE		:       int		)    ->   List[Any]:


						'''simple docstring'''
						A	 =      self.proj(_SCREAMING_SNAKE_CASE		)
						A	 =      self.gelu(_SCREAMING_SNAKE_CASE		)
						return hidden_states


class      UpperCamelCase (				nn.Module ):







		"""simple docstring"""



		def __init__(  self		:       Dict						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       int		)    ->   Tuple:


						'''simple docstring'''
						super().__init__()
						A	 =      nn.Linear(_SCREAMING_SNAKE_CASE						,dim_out * 2		)
		def       A(  self		:       Any						,_SCREAMING_SNAKE_CASE		:       Any		)    ->   Union[str, Any]:


						'''simple docstring'''
						if gate.device.type != "mps":
										return F.gelu(_SCREAMING_SNAKE_CASE		)
						# mps: gelu is not implemented for float16
						return F.gelu(gate.to(dtype=torch.floataa		)		).to(dtype=gate.dtype		)







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


						'''simple docstring'''
						A ,		A	 =      self.proj(_SCREAMING_SNAKE_CASE		).chunk(2						,dim=-1		)
						return hidden_states * self.gelu(_SCREAMING_SNAKE_CASE		)


class      UpperCamelCase (				nn.Module ):







		"""simple docstring"""



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


						'''simple docstring'''
						super().__init__()
						A	 =      nn.Linear(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		)







		def       A(  self		:       str						,_SCREAMING_SNAKE_CASE		:       str		)    ->   List[Any]:


						'''simple docstring'''
						A	 =      self.proj(_SCREAMING_SNAKE_CASE		)
						return x * torch.sigmoid(1.702 * x		)


class      UpperCamelCase (				nn.Module ):







		"""simple docstring"""



		def __init__(  self		:       Optional[int]						,_SCREAMING_SNAKE_CASE		:       List[str]						,_SCREAMING_SNAKE_CASE		:       Tuple		)    ->   Dict:


						'''simple docstring'''
						super().__init__()
						A	 =      nn.Embedding(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		)
						A	 =      nn.SiLU()
						A	 =      nn.Linear(_SCREAMING_SNAKE_CASE						,embedding_dim * 2		)
						A	 =      nn.LayerNorm(_SCREAMING_SNAKE_CASE						,elementwise_affine=_SCREAMING_SNAKE_CASE		)







		def       A(  self		:       Dict						,_SCREAMING_SNAKE_CASE		:       str						,_SCREAMING_SNAKE_CASE		:       int		)    ->   Any:


						'''simple docstring'''
						A	 =      self.linear(self.silu(self.emb(_SCREAMING_SNAKE_CASE		)		)		)
						A ,		A	 =      torch.chunk(_SCREAMING_SNAKE_CASE						,2		)
						A	 =      self.norm(_SCREAMING_SNAKE_CASE		) * (1 + scale) + shift
						return x


class      UpperCamelCase (				nn.Module ):







		"""simple docstring"""



		def __init__(  self		:       List[str]						,_SCREAMING_SNAKE_CASE		:       Any						,_SCREAMING_SNAKE_CASE		:       Tuple		)    ->   Any:


						'''simple docstring'''
						super().__init__()

						A	 =      CombinedTimestepLabelEmbeddings(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		)

						A	 =      nn.SiLU()
						A	 =      nn.Linear(_SCREAMING_SNAKE_CASE						,6 * embedding_dim						,bias=_SCREAMING_SNAKE_CASE		)
						A	 =      nn.LayerNorm(_SCREAMING_SNAKE_CASE						,elementwise_affine=_SCREAMING_SNAKE_CASE						,eps=1E-6		)







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


						'''simple docstring'''
						A	 =      self.linear(self.silu(self.emb(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE						,hidden_dtype=_SCREAMING_SNAKE_CASE		)		)		)
						A ,		A ,		A ,		A ,		A ,		A	 =      emb.chunk(6						,dim=1		)
						A	 =      self.norm(_SCREAMING_SNAKE_CASE		) * (1 + scale_msa[:, None]) + shift_msa[:, None]
						return x, gate_msa, shift_mlp, scale_mlp, gate_mlp


class      UpperCamelCase (				nn.Module ):







		"""simple docstring"""



		def __init__(  self		:       List[str]						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       Optional[str] = None						,_SCREAMING_SNAKE_CASE		:       float = 1E-5		)    ->   Dict:


						'''simple docstring'''
						super().__init__()
						A	 =      num_groups
						A	 =      eps

						if act_fn is None:
										A	 =      None
						else:
										A	 =      get_activation(_SCREAMING_SNAKE_CASE		)

						A	 =      nn.Linear(_SCREAMING_SNAKE_CASE						,out_dim * 2		)







		def       A(  self		:       int						,_SCREAMING_SNAKE_CASE		:       str						,_SCREAMING_SNAKE_CASE		:       Dict		)    ->   str:


						'''simple docstring'''
						if self.act:
										A	 =      self.act(_SCREAMING_SNAKE_CASE		)
						A	 =      self.linear(_SCREAMING_SNAKE_CASE		)
						A	 =      emb[:, :, None, None]
						A ,		A	 =      emb.chunk(2						,dim=1		)

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













from collections.abc import Callable
class      UpperCamelCase :







		"""simple docstring"""



		def __init__(  self		:       Tuple						,_SCREAMING_SNAKE_CASE		:       Callable | None = None		)    ->   None:


						'''simple docstring'''
						# Stores actual heap items.
						A	 =      []
						# Stores indexes of each item for supporting updates and deletion.
						A	 =      {}
						# Stores current size of heap.
						A	 =      0
						# Stores function used to evaluate the score of an item on which basis ordering
						# will be done.
						A	 =      key or (lambda _SCREAMING_SNAKE_CASE		: x)
		def       A(  self		:       Tuple						,_SCREAMING_SNAKE_CASE		:       int		)    ->   int | None:


						'''simple docstring'''
						return int((i - 1) / 2		) if i > 0 else None
		def       A(  self		:       Union[str, Any]						,_SCREAMING_SNAKE_CASE		:       int		)    ->   int | None:


						'''simple docstring'''
						A	 =      int(2 * i + 1		)
						return left if 0 < left < self.size else None
		def       A(  self		:       Union[str, Any]						,_SCREAMING_SNAKE_CASE		:       int		)    ->   int | None:


						'''simple docstring'''
						A	 =      int(2 * i + 2		)
						return right if 0 < right < self.size else None
		def       A(  self		:       List[str]						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       int		)    ->   None:


						'''simple docstring'''
						A ,		A	 =      (
						    self.pos_map[self.arr[j][0]],
						    self.pos_map[self.arr[i][0]],
						)
						# Then swap the items in the list.
						A ,		A	 =      self.arr[j], self.arr[i]
		def       A(  self		:       List[str]						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       int		)    ->   bool:


						'''simple docstring'''
						return self.arr[i][1] < self.arr[j][1]
		def       A(  self		:       int						,_SCREAMING_SNAKE_CASE		:       int		)    ->   int:


						'''simple docstring'''
						A	 =      self._left(_SCREAMING_SNAKE_CASE		)
						A	 =      self._right(_SCREAMING_SNAKE_CASE		)
						A	 =      i

						if left is not None and not self._cmp(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		):
										A	 =      left
						if right is not None and not self._cmp(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		):
										A	 =      right

						return valid_parent
		def       A(  self		:       Any						,_SCREAMING_SNAKE_CASE		:       int		)    ->   None:


						'''simple docstring'''
						A	 =      self._parent(_SCREAMING_SNAKE_CASE		)
						while parent is not None and not self._cmp(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		):
										self._swap(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		)
										A ,		A	 =      parent, self._parent(_SCREAMING_SNAKE_CASE		)
		def       A(  self		:       List[Any]						,_SCREAMING_SNAKE_CASE		:       int		)    ->   None:


						'''simple docstring'''
						A	 =      self._get_valid_parent(_SCREAMING_SNAKE_CASE		)
						while valid_parent != index:
										self._swap(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		)
										A ,		A	 =      valid_parent, self._get_valid_parent(_SCREAMING_SNAKE_CASE		)
		def       A(  self		:       Optional[Any]						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       int		)    ->   None:


						'''simple docstring'''
						if item not in self.pos_map:
										return
						A	 =      self.pos_map[item]
						A	 =      [item, self.key(_SCREAMING_SNAKE_CASE		)]
						# Make sure heap is right in both up and down direction.
						# Ideally only one of them will make any change.
						self._heapify_up(_SCREAMING_SNAKE_CASE		)
						self._heapify_down(_SCREAMING_SNAKE_CASE		)
		def       A(  self		:       int						,_SCREAMING_SNAKE_CASE		:       int		)    ->   None:


						'''simple docstring'''
						if item not in self.pos_map:
										return
						A	 =      self.pos_map[item]
						del self.pos_map[item]
						A	 =      self.arr[self.size - 1]
						A	 =      index
						self.size -= 1
						# Make sure heap is right in both up and down direction. Ideally only one
						# of them will make any change- so no performance loss in calling both.
						if self.size > index:
										self._heapify_up(_SCREAMING_SNAKE_CASE		)
										self._heapify_down(_SCREAMING_SNAKE_CASE		)
		def       A(  self		:       Optional[Any]						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       int		)    ->   None:


						'''simple docstring'''
						A	 =      len(self.arr		)
						if arr_len == self.size:
										self.arr.append([item, self.key(_SCREAMING_SNAKE_CASE		)]		)
						else:
										A	 =      [item, self.key(_SCREAMING_SNAKE_CASE		)]
						A	 =      self.size
						self.size += 1
						self._heapify_up(self.size - 1		)
		def       A(  self		:       str		)    ->   tuple | None:


						'''simple docstring'''
						return self.arr[0] if self.size else None







		def       A(  self		:       Any		)    ->   tuple | None:


						'''simple docstring'''
						A	 =      self.get_top()
						if top_item_tuple:
										self.delete_item(top_item_tuple[0]		)
						return top_item_tuple







def      snake_case		(							):
				pass


if __name__ == "__main__":
					import doctest

					doctest.testmod()





