






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





	def __init__(   self      :				Any			,			lowerCAmelCase_      :				int			,			lowerCAmelCase_      :				int			,			lowerCAmelCase_      :				int			,			lowerCAmelCase_      :				List[Any]=0.0			,			lowerCAmelCase_      :				Optional[int] = None			,			lowerCAmelCase_      :				str = "geglu"			,			lowerCAmelCase_      :				Optional[int] = None			,			lowerCAmelCase_      :				bool = False			,			lowerCAmelCase_      :				bool = False			,			lowerCAmelCase_      :				bool = False			,			lowerCAmelCase_      :				bool = False			,			lowerCAmelCase_      :				bool = True			,			lowerCAmelCase_      :				str = "layer_norm"			,			lowerCAmelCase_      :				bool = False			,			)		->    Dict:
				super().__init__()
				UpperCAmelCase_  :				Dict      =      only_cross_attention

				UpperCAmelCase_  :				Union[str, Any]      =      (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"
				UpperCAmelCase_  :				int      =      (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:
							UpperCAmelCase_  :				Union[str, Any]      =      AdaLayerNorm(lowerCAmelCase_			,			lowerCAmelCase_		)
				elif self.use_ada_layer_norm_zero:
							UpperCAmelCase_  :				Optional[Any]      =      AdaLayerNormZero(lowerCAmelCase_			,			lowerCAmelCase_		)
				else:
							UpperCAmelCase_  :				Optional[Any]      =      nn.LayerNorm(lowerCAmelCase_			,			elementwise_affine=lowerCAmelCase_		)
				UpperCAmelCase_  :				Any      =      Attention(
				    query_dim=lowerCAmelCase_			,			heads=lowerCAmelCase_			,			dim_head=lowerCAmelCase_			,			dropout=lowerCAmelCase_			,			bias=lowerCAmelCase_			,			cross_attention_dim=cross_attention_dim if only_cross_attention else None			,			upcast_attention=lowerCAmelCase_			,			)

				# 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.
							UpperCAmelCase_  :				List[Any]      =      (
							    AdaLayerNorm(lowerCAmelCase_			,			lowerCAmelCase_		)
							    if self.use_ada_layer_norm
							    else nn.LayerNorm(lowerCAmelCase_			,			elementwise_affine=lowerCAmelCase_		)
							)
							UpperCAmelCase_  :				List[str]      =      Attention(
							    query_dim=lowerCAmelCase_			,			cross_attention_dim=cross_attention_dim if not double_self_attention else None			,			heads=lowerCAmelCase_			,			dim_head=lowerCAmelCase_			,			dropout=lowerCAmelCase_			,			bias=lowerCAmelCase_			,			upcast_attention=lowerCAmelCase_			,			)  # is self-attn if encoder_hidden_states is none
				else:
							UpperCAmelCase_  :				Union[str, Any]      =      None
							UpperCAmelCase_  :				Tuple      =      None

				# 3. Feed-forward
				UpperCAmelCase_  :				List[Any]      =      nn.LayerNorm(lowerCAmelCase_			,			elementwise_affine=lowerCAmelCase_		)
				UpperCAmelCase_  :				List[str]      =      FeedForward(lowerCAmelCase_			,			dropout=lowerCAmelCase_			,			activation_fn=lowerCAmelCase_			,			final_dropout=lowerCAmelCase_		)

				# let chunk size default to None
				UpperCAmelCase_  :				Optional[int]      =      None
				UpperCAmelCase_  :				List[str]      =      0





	def   _SCREAMING_SNAKE_CASE   (   self      :				int			,			lowerCAmelCase_      :				Optional[int]			,			lowerCAmelCase_      :				int		)		->    List[str]:
				# Sets chunk feed-forward
				UpperCAmelCase_  :				Dict      =      chunk_size
				UpperCAmelCase_  :				Union[str, Any]      =      dim





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

				UpperCAmelCase_  :				Tuple      =      cross_attention_kwargs if cross_attention_kwargs is not None else {}

				UpperCAmelCase_  :				Optional[Any]      =      self.attna(
				    lowerCAmelCase_			,			encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None			,			attention_mask=lowerCAmelCase_			,			**lowerCAmelCase_			,			)
				if self.use_ada_layer_norm_zero:
							UpperCAmelCase_  :				Union[str, Any]      =      gate_msa.unsqueeze(1		) * attn_output
				UpperCAmelCase_  :				Tuple      =      attn_output + hidden_states

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

							UpperCAmelCase_  :				List[str]      =      self.attna(
							    lowerCAmelCase_			,			encoder_hidden_states=lowerCAmelCase_			,			attention_mask=lowerCAmelCase_			,			**lowerCAmelCase_			,			)
							UpperCAmelCase_  :				Dict      =      attn_output + hidden_states

				# 3. Feed-forward
				UpperCAmelCase_  :				Optional[int]      =      self.norma(lowerCAmelCase_		)

				if self.use_ada_layer_norm_zero:
							UpperCAmelCase_  :				Optional[int]      =      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`."""		)

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

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

				UpperCAmelCase_  :				Optional[Any]      =      ff_output + hidden_states

				return hidden_states







class     UpperCamelCase_ (nn.Module  ):





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

				if activation_fn == "gelu":
							UpperCAmelCase_  :				Dict      =      GELU(lowerCAmelCase_			,			lowerCAmelCase_		)
				if activation_fn == "gelu-approximate":
							UpperCAmelCase_  :				Tuple      =      GELU(lowerCAmelCase_			,			lowerCAmelCase_			,			approximate="tanh"		)
				elif activation_fn == "geglu":
							UpperCAmelCase_  :				Tuple      =      GEGLU(lowerCAmelCase_			,			lowerCAmelCase_		)
				elif activation_fn == "geglu-approximate":
							UpperCAmelCase_  :				Optional[int]      =      ApproximateGELU(lowerCAmelCase_			,			lowerCAmelCase_		)

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





	def   _SCREAMING_SNAKE_CASE   (   self      :				Union[str, Any]			,			lowerCAmelCase_      :				Optional[Any]		)		->    Optional[Any]:
				for module in self.net:
							UpperCAmelCase_  :				Optional[Any]      =      module(lowerCAmelCase_		)
				return hidden_states







class     UpperCamelCase_ (nn.Module  ):





	def __init__(   self      :				Optional[int]			,			lowerCAmelCase_      :				int			,			lowerCAmelCase_      :				int			,			lowerCAmelCase_      :				str = "none"		)		->    Optional[Any]:
				super().__init__()
				UpperCAmelCase_  :				Union[str, Any]      =      nn.Linear(lowerCAmelCase_			,			lowerCAmelCase_		)
				UpperCAmelCase_  :				int      =      approximate





	def   _SCREAMING_SNAKE_CASE   (   self      :				List[Any]			,			lowerCAmelCase_      :				Optional[int]		)		->    List[str]:
				if gate.device.type != "mps":
							return F.gelu(lowerCAmelCase_			,			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   _SCREAMING_SNAKE_CASE   (   self      :				List[Any]			,			lowerCAmelCase_      :				Tuple		)		->    List[str]:
				UpperCAmelCase_  :				List[Any]      =      self.proj(lowerCAmelCase_		)
				UpperCAmelCase_  :				Union[str, Any]      =      self.gelu(lowerCAmelCase_		)
				return hidden_states







class     UpperCamelCase_ (nn.Module  ):





	def __init__(   self      :				List[Any]			,			lowerCAmelCase_      :				int			,			lowerCAmelCase_      :				int		)		->    Dict:
				super().__init__()
				UpperCAmelCase_  :				str      =      nn.Linear(lowerCAmelCase_			,			dim_out * 2		)





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





	def   _SCREAMING_SNAKE_CASE   (   self      :				Union[str, Any]			,			lowerCAmelCase_      :				List[Any]		)		->    Tuple:
				UpperCAmelCase_  :				Tuple      =      self.proj(lowerCAmelCase_		).chunk(2			,			dim=-1		)
				return hidden_states * self.gelu(lowerCAmelCase_		)







class     UpperCamelCase_ (nn.Module  ):





	def __init__(   self      :				str			,			lowerCAmelCase_      :				int			,			lowerCAmelCase_      :				int		)		->    Dict:
				super().__init__()
				UpperCAmelCase_  :				List[str]      =      nn.Linear(lowerCAmelCase_			,			lowerCAmelCase_		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				int			,			lowerCAmelCase_      :				Dict		)		->    Union[str, Any]:
				UpperCAmelCase_  :				str      =      self.proj(lowerCAmelCase_		)
				return x * torch.sigmoid(1.7_0_2 * x		)







class     UpperCamelCase_ (nn.Module  ):





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





	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[Any]			,			lowerCAmelCase_      :				Union[str, Any]			,			lowerCAmelCase_      :				Tuple		)		->    Optional[int]:
				UpperCAmelCase_  :				Optional[Any]      =      self.linear(self.silu(self.emb(lowerCAmelCase_		)		)		)
				UpperCAmelCase_  :				Optional[int]      =      torch.chunk(lowerCAmelCase_			,			2		)
				UpperCAmelCase_  :				Optional[Any]      =      self.norm(lowerCAmelCase_		) * (1 + scale) + shift
				return x







class     UpperCamelCase_ (nn.Module  ):





	def __init__(   self      :				List[str]			,			lowerCAmelCase_      :				Union[str, Any]			,			lowerCAmelCase_      :				Tuple		)		->    List[Any]:
				super().__init__()

				UpperCAmelCase_  :				Any      =      CombinedTimestepLabelEmbeddings(lowerCAmelCase_			,			lowerCAmelCase_		)

				UpperCAmelCase_  :				Dict      =      nn.SiLU()
				UpperCAmelCase_  :				Any      =      nn.Linear(lowerCAmelCase_			,			6 * embedding_dim			,			bias=lowerCAmelCase_		)
				UpperCAmelCase_  :				Optional[Any]      =      nn.LayerNorm(lowerCAmelCase_			,			elementwise_affine=lowerCAmelCase_			,			eps=1e-6		)





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







class     UpperCamelCase_ (nn.Module  ):





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

				if act_fn is None:
							UpperCAmelCase_  :				List[Any]      =      None
				else:
							UpperCAmelCase_  :				Dict      =      get_activation(lowerCAmelCase_		)

				UpperCAmelCase_  :				Dict      =      nn.Linear(lowerCAmelCase_			,			out_dim * 2		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				List[str]			,			lowerCAmelCase_      :				Dict			,			lowerCAmelCase_      :				List[str]		)		->    List[str]:
				if self.act:
							UpperCAmelCase_  :				List[str]      =      self.act(lowerCAmelCase_		)
				UpperCAmelCase_  :				Any      =      self.linear(lowerCAmelCase_		)
				UpperCAmelCase_  :				List[str]      =      emb[:, :, None, None]
				UpperCAmelCase_  :				Union[str, Any]      =      emb.chunk(2			,			dim=1		)

				UpperCAmelCase_  :				Optional[Any]      =      F.group_norm(lowerCAmelCase_			,			self.num_groups			,			eps=self.eps		)
				UpperCAmelCase_  :				List[Any]      =      x * (1 + scale) + shift
				return x











"""simple docstring"""
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
    is_vision_available,
)


lowerCamelCase_       			=   {
    '''configuration_clip''': [
        '''CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''CLIPConfig''',
        '''CLIPOnnxConfig''',
        '''CLIPTextConfig''',
        '''CLIPVisionConfig''',
    ],
    '''processing_clip''': ['''CLIPProcessor'''],
    '''tokenization_clip''': ['''CLIPTokenizer'''],
}

try:
			if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCamelCase_       			=   ['''CLIPTokenizerFast''']

try:
			if not is_vision_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCamelCase_       			=   ['''CLIPFeatureExtractor''']
			lowerCamelCase_       			=   ['''CLIPImageProcessor''']

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCamelCase_       			=   [
			    '''CLIP_PRETRAINED_MODEL_ARCHIVE_LIST''',
			    '''CLIPModel''',
			    '''CLIPPreTrainedModel''',
			    '''CLIPTextModel''',
			    '''CLIPTextModelWithProjection''',
			    '''CLIPVisionModel''',
			    '''CLIPVisionModelWithProjection''',
			]

try:
			if not is_tf_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCamelCase_       			=   [
			    '''TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST''',
			    '''TFCLIPModel''',
			    '''TFCLIPPreTrainedModel''',
			    '''TFCLIPTextModel''',
			    '''TFCLIPVisionModel''',
			]

try:
			if not is_flax_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCamelCase_       			=   [
			    '''FlaxCLIPModel''',
			    '''FlaxCLIPPreTrainedModel''',
			    '''FlaxCLIPTextModel''',
			    '''FlaxCLIPTextPreTrainedModel''',
			    '''FlaxCLIPVisionModel''',
			    '''FlaxCLIPVisionPreTrainedModel''',
			]


if TYPE_CHECKING:
			from .configuration_clip import (
			    CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP,
			    CLIPConfig,
			    CLIPOnnxConfig,
			    CLIPTextConfig,
			    CLIPVisionConfig,
			)
			from .processing_clip import CLIPProcessor
			from .tokenization_clip import CLIPTokenizer

			try:
						if not is_tokenizers_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .tokenization_clip_fast import CLIPTokenizerFast

			try:
						if not is_vision_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .feature_extraction_clip import CLIPFeatureExtractor
						from .image_processing_clip import CLIPImageProcessor

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_clip import (
						    CLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
						    CLIPModel,
						    CLIPPreTrainedModel,
						    CLIPTextModel,
						    CLIPTextModelWithProjection,
						    CLIPVisionModel,
						    CLIPVisionModelWithProjection,
						)

			try:
						if not is_tf_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_tf_clip import (
						    TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
						    TFCLIPModel,
						    TFCLIPPreTrainedModel,
						    TFCLIPTextModel,
						    TFCLIPVisionModel,
						)

			try:
						if not is_flax_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_flax_clip import (
						    FlaxCLIPModel,
						    FlaxCLIPPreTrainedModel,
						    FlaxCLIPTextModel,
						    FlaxCLIPTextPreTrainedModel,
						    FlaxCLIPVisionModel,
						    FlaxCLIPVisionPreTrainedModel,
						)


else:
			import sys

			lowerCamelCase_       			=   _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



