





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,   snake_case__,   snake_case__,   snake_case__,   snake_case__=0.0,   snake_case__ = None,   snake_case__ = "geglu",   snake_case__ = None,   snake_case__ = False,   snake_case__ = False,   snake_case__ = False,   snake_case__ = False,   snake_case__ = True,   snake_case__ = "layer_norm",   snake_case__ = False,   )       -> Union[str, Any]:






									"""simple docstring"""
									super().__init__()
									lowercase_	:   List[Any]  				=     only_cross_attention

									lowercase_	:   Optional[Any]  				=     (num_embeds_ada_norm is not None) and norm_type == """ada_norm_zero"""
									lowercase_	:   List[str]  				=     (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:
																lowercase_	:   List[str]  				=     AdaLayerNorm(snake_case__,   snake_case__  )
									elif self.use_ada_layer_norm_zero:
																lowercase_	:   Dict  				=     AdaLayerNormZero(snake_case__,   snake_case__  )
									else:
																lowercase_	:   Any  				=     nn.LayerNorm(snake_case__,   elementwise_affine=snake_case__  )
									lowercase_	:   str  				=     Attention(
									    query_dim=snake_case__,   heads=snake_case__,   dim_head=snake_case__,   dropout=snake_case__,   bias=snake_case__,   cross_attention_dim=cross_attention_dim if only_cross_attention else None,   upcast_attention=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.
																lowercase_	:   List[str]  				=     (
																    AdaLayerNorm(snake_case__,   snake_case__  )
																    if self.use_ada_layer_norm
																    else nn.LayerNorm(snake_case__,   elementwise_affine=snake_case__  )
																)
																lowercase_	:   Tuple  				=     Attention(
																    query_dim=snake_case__,   cross_attention_dim=cross_attention_dim if not double_self_attention else None,   heads=snake_case__,   dim_head=snake_case__,   dropout=snake_case__,   bias=snake_case__,   upcast_attention=snake_case__,   )  # is self-attn if encoder_hidden_states is none
									else:
																lowercase_	:   Optional[Any]  				=     None
																lowercase_	:   str  				=     None

									# 3. Feed-forward
									lowercase_	:   Optional[int]  				=     nn.LayerNorm(snake_case__,   elementwise_affine=snake_case__  )
									lowercase_	:   Union[str, Any]  				=     FeedForward(snake_case__,   dropout=snake_case__,   activation_fn=snake_case__,   final_dropout=snake_case__  )

									# let chunk size default to None
									lowercase_	:   Any  				=     None
									lowercase_	:   Optional[int]  				=     0






		def 						snake_case__				(							self,   snake_case__,   snake_case__  )       -> List[str]:






									"""simple docstring"""
									# Sets chunk feed-forward
									lowercase_	:   int  				=     chunk_size
									lowercase_	:   Tuple  				=     dim






		def 						snake_case__				(							self,   snake_case__,   snake_case__ = None,   snake_case__ = None,   snake_case__ = None,   snake_case__ = None,   snake_case__ = None,   snake_case__ = None,   )       -> Union[str, Any]:






									"""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:
																lowercase_	:   Dict  				=     self.norma(snake_case__,   snake_case__  )
									elif self.use_ada_layer_norm_zero:
																lowercase_				,		lowercase_				,		lowercase_				,		lowercase_				,		lowercase_	:   Any  				=     self.norma(
																    snake_case__,   snake_case__,   snake_case__,   hidden_dtype=hidden_states.dtype  )
									else:
																lowercase_	:   Optional[Any]  				=     self.norma(snake_case__  )

									lowercase_	:   Optional[int]  				=     cross_attention_kwargs if cross_attention_kwargs is not None else {}

									lowercase_	:   int  				=     self.attna(
									    snake_case__,   encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,   attention_mask=snake_case__,   **snake_case__,   )
									if self.use_ada_layer_norm_zero:
																lowercase_	:   Optional[int]  				=     gate_msa.unsqueeze(1  ) * attn_output
									lowercase_	:   List[Any]  				=     attn_output + hidden_states

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

																lowercase_	:   Union[str, Any]  				=     self.attna(
																    snake_case__,   encoder_hidden_states=snake_case__,   attention_mask=snake_case__,   **snake_case__,   )
																lowercase_	:   List[str]  				=     attn_output + hidden_states

									# 3. Feed-forward
									lowercase_	:   Optional[Any]  				=     self.norma(snake_case__  )

									if self.use_ada_layer_norm_zero:
																lowercase_	:   Tuple  				=     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`."""  )

																lowercase_	:   List[Any]  				=     norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
																lowercase_	:   Union[str, Any]  				=     torch.cat(
																    [self.ff(snake_case__  ) for hid_slice in norm_hidden_states.chunk(snake_case__,   dim=self._chunk_dim  )],   dim=self._chunk_dim,   )
									else:
																lowercase_	:   Tuple  				=     self.ff(snake_case__  )

									if self.use_ada_layer_norm_zero:
																lowercase_	:   List[Any]  				=     gate_mlp.unsqueeze(1  ) * ff_output

									lowercase_	:   Dict  				=     ff_output + hidden_states

									return hidden_states




class     UpperCamelCase__					(     nn.Module				):






		'''simple docstring'''






		def __init__(							self,   snake_case__,   snake_case__ = None,   snake_case__ = 4,   snake_case__ = 0.0,   snake_case__ = "geglu",   snake_case__ = False,   )       -> List[Any]:






									"""simple docstring"""
									super().__init__()
									lowercase_	:   Tuple  				=     int(dim * mult  )
									lowercase_	:   List[str]  				=     dim_out if dim_out is not None else dim

									if activation_fn == "gelu":
																lowercase_	:   Optional[Any]  				=     GELU(snake_case__,   snake_case__  )
									if activation_fn == "gelu-approximate":
																lowercase_	:   Tuple  				=     GELU(snake_case__,   snake_case__,   approximate="""tanh"""  )
									elif activation_fn == "geglu":
																lowercase_	:   Dict  				=     GEGLU(snake_case__,   snake_case__  )
									elif activation_fn == "geglu-approximate":
																lowercase_	:   Union[str, Any]  				=     ApproximateGELU(snake_case__,   snake_case__  )

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






		def 						snake_case__				(							self,   snake_case__  )       -> List[str]:






									"""simple docstring"""
									for module in self.net:
																lowercase_	:   Dict  				=     module(snake_case__  )
									return hidden_states




class     UpperCamelCase__					(     nn.Module				):






		'''simple docstring'''






		def __init__(							self,   snake_case__,   snake_case__,   snake_case__ = "none"  )       -> Tuple:






									"""simple docstring"""
									super().__init__()
									lowercase_	:   Tuple  				=     nn.Linear(snake_case__,   snake_case__  )
									lowercase_	:   Optional[Any]  				=     approximate






		def 						snake_case__				(							self,   snake_case__  )       -> List[Any]:






									"""simple docstring"""
									if gate.device.type != "mps":
																return F.gelu(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 						snake_case__				(							self,   snake_case__  )       -> Optional[Any]:






									"""simple docstring"""
									lowercase_	:   str  				=     self.proj(snake_case__  )
									lowercase_	:   int  				=     self.gelu(snake_case__  )
									return hidden_states




class     UpperCamelCase__					(     nn.Module				):






		'''simple docstring'''






		def __init__(							self,   snake_case__,   snake_case__  )       -> Optional[Any]:






									"""simple docstring"""
									super().__init__()
									lowercase_	:   Union[str, Any]  				=     nn.Linear(snake_case__,   dim_out * 2  )






		def 						snake_case__				(							self,   snake_case__  )       -> str:






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






		def 						snake_case__				(							self,   snake_case__  )       -> Optional[Any]:






									"""simple docstring"""
									lowercase_				,		lowercase_	:   Any  				=     self.proj(snake_case__  ).chunk(2,   dim=-1  )
									return hidden_states * self.gelu(snake_case__  )




class     UpperCamelCase__					(     nn.Module				):






		'''simple docstring'''






		def __init__(							self,   snake_case__,   snake_case__  )       -> Union[str, Any]:






									"""simple docstring"""
									super().__init__()
									lowercase_	:   int  				=     nn.Linear(snake_case__,   snake_case__  )






		def 						snake_case__				(							self,   snake_case__  )       -> Optional[Any]:






									"""simple docstring"""
									lowercase_	:   Optional[Any]  				=     self.proj(snake_case__  )
									return x * torch.sigmoid(1.702 * x  )




class     UpperCamelCase__					(     nn.Module				):






		'''simple docstring'''






		def __init__(							self,   snake_case__,   snake_case__  )       -> List[Any]:






									"""simple docstring"""
									super().__init__()
									lowercase_	:   int  				=     nn.Embedding(snake_case__,   snake_case__  )
									lowercase_	:   Tuple  				=     nn.SiLU()
									lowercase_	:   Tuple  				=     nn.Linear(snake_case__,   embedding_dim * 2  )
									lowercase_	:   Tuple  				=     nn.LayerNorm(snake_case__,   elementwise_affine=snake_case__  )






		def 						snake_case__				(							self,   snake_case__,   snake_case__  )       -> List[Any]:






									"""simple docstring"""
									lowercase_	:   Optional[Any]  				=     self.linear(self.silu(self.emb(snake_case__  )  )  )
									lowercase_				,		lowercase_	:   Dict  				=     torch.chunk(snake_case__,   2  )
									lowercase_	:   Tuple  				=     self.norm(snake_case__  ) * (1 + scale) + shift
									return x




class     UpperCamelCase__					(     nn.Module				):






		'''simple docstring'''






		def __init__(							self,   snake_case__,   snake_case__  )       -> Tuple:






									"""simple docstring"""
									super().__init__()

									lowercase_	:   List[Any]  				=     CombinedTimestepLabelEmbeddings(snake_case__,   snake_case__  )

									lowercase_	:   int  				=     nn.SiLU()
									lowercase_	:   Dict  				=     nn.Linear(snake_case__,   6 * embedding_dim,   bias=snake_case__  )
									lowercase_	:   Optional[Any]  				=     nn.LayerNorm(snake_case__,   elementwise_affine=snake_case__,   eps=1E-6  )






		def 						snake_case__				(							self,   snake_case__,   snake_case__,   snake_case__,   snake_case__=None  )       -> Union[str, Any]:






									"""simple docstring"""
									lowercase_	:   Optional[Any]  				=     self.linear(self.silu(self.emb(snake_case__,   snake_case__,   hidden_dtype=snake_case__  )  )  )
									lowercase_				,		lowercase_				,		lowercase_				,		lowercase_				,		lowercase_				,		lowercase_	:   List[str]  				=     emb.chunk(6,   dim=1  )
									lowercase_	:   int  				=     self.norm(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,   snake_case__,   snake_case__,   snake_case__,   snake_case__ = None,   snake_case__ = 1E-5  )       -> Dict:






									"""simple docstring"""
									super().__init__()
									lowercase_	:   str  				=     num_groups
									lowercase_	:   Any  				=     eps

									if act_fn is None:
																lowercase_	:   Union[str, Any]  				=     None
									else:
																lowercase_	:   Dict  				=     get_activation(snake_case__  )

									lowercase_	:   Union[str, Any]  				=     nn.Linear(snake_case__,   out_dim * 2  )






		def 						snake_case__				(							self,   snake_case__,   snake_case__  )       -> List[Any]:






									"""simple docstring"""
									if self.act:
																lowercase_	:   Optional[int]  				=     self.act(snake_case__  )
									lowercase_	:   int  				=     self.linear(snake_case__  )
									lowercase_	:   int  				=     emb[:, :, None, None]
									lowercase_				,		lowercase_	:   Optional[int]  				=     emb.chunk(2,   dim=1  )

									lowercase_	:   Dict  				=     F.group_norm(snake_case__,   self.num_groups,   eps=self.eps  )
									lowercase_	:   Tuple  				=     x * (1 + scale) + shift
									return x






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


UpperCAmelCase_							=				{
    """configuration_maskformer""": ["""MASKFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MaskFormerConfig"""],
    """configuration_maskformer_swin""": ["""MaskFormerSwinConfig"""],
}

try:
							if not is_vision_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCAmelCase_							=				["""MaskFormerFeatureExtractor"""]
							UpperCAmelCase_							=				["""MaskFormerImageProcessor"""]


try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCAmelCase_							=				[
							    """MASKFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
							    """MaskFormerForInstanceSegmentation""",
							    """MaskFormerModel""",
							    """MaskFormerPreTrainedModel""",
							]
							UpperCAmelCase_							=				[
							    """MaskFormerSwinBackbone""",
							    """MaskFormerSwinModel""",
							    """MaskFormerSwinPreTrainedModel""",
							]

if TYPE_CHECKING:
							from .configuration_maskformer import MASKFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, MaskFormerConfig
							from .configuration_maskformer_swin import MaskFormerSwinConfig

							try:
														if not is_vision_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .feature_extraction_maskformer import MaskFormerFeatureExtractor
														from .image_processing_maskformer import MaskFormerImageProcessor
							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_maskformer import (
														    MASKFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
														    MaskFormerForInstanceSegmentation,
														    MaskFormerModel,
														    MaskFormerPreTrainedModel,
														)
														from .modeling_maskformer_swin import (
														    MaskFormerSwinBackbone,
														    MaskFormerSwinModel,
														    MaskFormerSwinPreTrainedModel,
														)


else:
							import sys

							UpperCAmelCase_							=				_LazyModule(__name__, globals()["""__file__"""], _import_structure)