



'''simple docstring'''

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






			def __init__(     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase=0.0					,   _UpperCAmelCase = None					,   _UpperCAmelCase = "geglu"					,   _UpperCAmelCase = None					,   _UpperCAmelCase = False					,   _UpperCAmelCase = False					,   _UpperCAmelCase = False					,   _UpperCAmelCase = False					,   _UpperCAmelCase = True					,   _UpperCAmelCase = "layer_norm"					,   _UpperCAmelCase = 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 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							List[Any]:
					# Sets chunk feed-forward
					snake_case__	     =chunk_size
					snake_case__	     =dim






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase = None					,   _UpperCAmelCase = None					,   _UpperCAmelCase = None					,   _UpperCAmelCase = None					,   _UpperCAmelCase = None					,   _UpperCAmelCase = 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 							a__( nn.Module     ):






			def __init__(     self					,   _UpperCAmelCase					,   _UpperCAmelCase = None					,   _UpperCAmelCase = 4					,   _UpperCAmelCase = 0.0					,   _UpperCAmelCase = "geglu"					,   _UpperCAmelCase = False					,   )							->							Optional[Any]:
					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 						_lowercase  (     self					,   _UpperCAmelCase						)							->							List[Any]:
					for module in self.net:
							snake_case__	     =module(_UpperCAmelCase						)
					return hidden_states





class 							a__( nn.Module     ):






			def __init__(     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase = "none"						)							->							Any:
					super().__init__()
					snake_case__	     =nn.Linear(_UpperCAmelCase					,   _UpperCAmelCase						)
					snake_case__	     =approximate






			def 						_lowercase  (     self					,   _UpperCAmelCase						)							->							Any:
					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 						_lowercase  (     self					,   _UpperCAmelCase						)							->							Union[str, Any]:
					snake_case__	     =self.proj(_UpperCAmelCase						)
					snake_case__	     =self.gelu(_UpperCAmelCase						)
					return hidden_states





class 							a__( nn.Module     ):






			def __init__(     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							Any:
					super().__init__()
					snake_case__	     =nn.Linear(_UpperCAmelCase					,   dim_out * 2						)






			def 						_lowercase  (     self					,   _UpperCAmelCase						)							->							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 						_lowercase  (     self					,   _UpperCAmelCase						)							->							Any:
					snake_case__     ,	snake_case__	     =self.proj(_UpperCAmelCase						).chunk(2					,   dim=-1						)
					return hidden_states * self.gelu(_UpperCAmelCase						)





class 							a__( nn.Module     ):






			def __init__(     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							str:
					super().__init__()
					snake_case__	     =nn.Linear(_UpperCAmelCase					,   _UpperCAmelCase						)






			def 						_lowercase  (     self					,   _UpperCAmelCase						)							->							Dict:
					snake_case__	     =self.proj(_UpperCAmelCase						)
					return x * torch.sigmoid(1.702 * x						)





class 							a__( nn.Module     ):






			def __init__(     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							Optional[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 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							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 							a__( nn.Module     ):






			def __init__(     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							Any:
					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 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase=None						)							->							int:
					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 							a__( nn.Module     ):






			def __init__(     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase = None					,   _UpperCAmelCase = 1E-5						)							->							Tuple:
					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 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							Optional[int]:
					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









'''simple docstring'''

import warnings

from ...utils import logging
from .image_processing_glpn import GLPNImageProcessor


SCREAMING_SNAKE_CASE__     :						Optional[Any]	=  logging.get_logger(__name__)

class 							a__( snake_case__     ):






			def __init__(     self					,   *_UpperCAmelCase					,   **_UpperCAmelCase						)							->							None:
					warnings.warn(
					    'The class GLPNFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please'
					    ' use GLPNImageProcessor instead.'					,   _UpperCAmelCase					,   )
					super().__init__(*_UpperCAmelCase					,   **_UpperCAmelCase						)




