






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






			def __init__( self      :			Tuple  ,						__SCREAMING_SNAKE_CASE      :			int  ,						__SCREAMING_SNAKE_CASE      :			int  ,						__SCREAMING_SNAKE_CASE      :			int  ,						__SCREAMING_SNAKE_CASE      :			Tuple=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  ,						)			->			List[Any]:
							super().__init__()
							a_		:  List[Any]											=      only_cross_attention

							a_		:  Optional[Any]											=      (num_embeds_ada_norm is not None) and norm_type == '''ada_norm_zero'''
							a_		:  Dict											=      (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_		:  List[Any]											=      AdaLayerNorm(_lowerCamelCase  ,						_lowerCamelCase					)
							elif self.use_ada_layer_norm_zero:
											a_		:  List[str]											=      AdaLayerNormZero(_lowerCamelCase  ,						_lowerCamelCase					)
							else:
											a_		:  Optional[int]											=      nn.LayerNorm(_lowerCamelCase  ,						elementwise_affine=_lowerCamelCase					)
							a_		:  List[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.
											a_		:  int											=      (
											    AdaLayerNorm(_lowerCamelCase  ,						_lowerCamelCase					)
											    if self.use_ada_layer_norm
											    else nn.LayerNorm(_lowerCamelCase  ,						elementwise_affine=_lowerCamelCase					)
											)
											a_		:  Optional[int]											=      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:
											a_		:  str											=      None
											a_		:  List[Any]											=      None

							# 3. Feed-forward
							a_		:  Optional[Any]											=      nn.LayerNorm(_lowerCamelCase  ,						elementwise_affine=_lowerCamelCase					)
							a_		:  Dict											=      FeedForward(_lowerCamelCase  ,						dropout=_lowerCamelCase  ,						activation_fn=_lowerCamelCase  ,						final_dropout=_lowerCamelCase					)

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






			def 			SCREAMING_SNAKE_CASE      ( self      :			List[str]  ,						__SCREAMING_SNAKE_CASE      :			Optional[int]  ,						__SCREAMING_SNAKE_CASE      :			int					)			->			Optional[Any]:
							# Sets chunk feed-forward
							a_		:  List[Any]											=      chunk_size
							a_		:  Dict											=      dim







			def 			SCREAMING_SNAKE_CASE      ( self      :			Dict  ,						__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  ,						)			->			Tuple:
							# Notice that normalization is always applied before the real computation in the following blocks.
							# 1. Self-Attention
							if self.use_ada_layer_norm:
											a_		:  List[str]											=      self.norma(_lowerCamelCase  ,						_lowerCamelCase					)
							elif self.use_ada_layer_norm_zero:
											a_				,			a_				,			a_				,			a_				,			a_		:  Optional[int]											=      self.norma(
											    _lowerCamelCase  ,						_lowerCamelCase  ,						_lowerCamelCase  ,						hidden_dtype=hidden_states.dtype					)
							else:
											a_		:  Optional[Any]											=      self.norma(_lowerCamelCase					)

							a_		:  List[str]											=      cross_attention_kwargs if cross_attention_kwargs is not None else {}

							a_		:  Tuple											=      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:
											a_		:  Optional[int]											=      gate_msa.unsqueeze(1					) * attn_output
							a_		:  str											=      attn_output + hidden_states

							# 2. Cross-Attention
							if self.attna is not None:
											a_		:  List[str]											=      (
											    self.norma(_lowerCamelCase  ,						_lowerCamelCase					) if self.use_ada_layer_norm else self.norma(_lowerCamelCase					)
											)

											a_		:  Optional[int]											=      self.attna(
											    _lowerCamelCase  ,						encoder_hidden_states=_lowerCamelCase  ,						attention_mask=_lowerCamelCase  ,						**_lowerCamelCase  ,						)
											a_		:  List[str]											=      attn_output + hidden_states

							# 3. Feed-forward
							a_		:  str											=      self.norma(_lowerCamelCase					)

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

											a_		:  Dict											=      norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
											a_		:  Optional[int]											=      torch.cat(
											    [self.ff(_lowerCamelCase					) for hid_slice in norm_hidden_states.chunk(_lowerCamelCase  ,						dim=self._chunk_dim					)]  ,						dim=self._chunk_dim  ,						)
							else:
											a_		:  Optional[int]											=      self.ff(_lowerCamelCase					)

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

							a_		:  Optional[int]											=      ff_output + hidden_states

							return hidden_states





class 	SCREAMING_SNAKE_CASE					(					nn.Module       ):






			def __init__( self      :			Any  ,						__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  ,						)			->			Optional[int]:
							super().__init__()
							a_		:  Any											=      int(dim * mult					)
							a_		:  int											=      dim_out if dim_out is not None else dim

							if activation_fn == "gelu":
											a_		:  Optional[int]											=      GELU(_lowerCamelCase  ,						_lowerCamelCase					)
							if activation_fn == "gelu-approximate":
											a_		:  Tuple											=      GELU(_lowerCamelCase  ,						_lowerCamelCase  ,						approximate='''tanh'''					)
							elif activation_fn == "geglu":
											a_		:  Union[str, Any]											=      GEGLU(_lowerCamelCase  ,						_lowerCamelCase					)
							elif activation_fn == "geglu-approximate":
											a_		:  Any											=      ApproximateGELU(_lowerCamelCase  ,						_lowerCamelCase					)

							a_		:  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      :			int  ,						__SCREAMING_SNAKE_CASE      :			Any					)			->			Any:
							for module in self.net:
											a_		:  Union[str, Any]											=      module(_lowerCamelCase					)
							return hidden_states





class 	SCREAMING_SNAKE_CASE					(					nn.Module       ):






			def __init__( self      :			List[Any]  ,						__SCREAMING_SNAKE_CASE      :			int  ,						__SCREAMING_SNAKE_CASE      :			int  ,						__SCREAMING_SNAKE_CASE      :			str = "none"					)			->			int:
							super().__init__()
							a_		:  List[Any]											=      nn.Linear(_lowerCamelCase  ,						_lowerCamelCase					)
							a_		:  Optional[int]											=      approximate






			def 			SCREAMING_SNAKE_CASE      ( self      :			List[Any]  ,						__SCREAMING_SNAKE_CASE      :			List[str]					)			->			List[Any]:
							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      :			Tuple  ,						__SCREAMING_SNAKE_CASE      :			Union[str, Any]					)			->			Union[str, Any]:
							a_		:  str											=      self.proj(_lowerCamelCase					)
							a_		:  str											=      self.gelu(_lowerCamelCase					)
							return hidden_states





class 	SCREAMING_SNAKE_CASE					(					nn.Module       ):






			def __init__( self      :			Any  ,						__SCREAMING_SNAKE_CASE      :			int  ,						__SCREAMING_SNAKE_CASE      :			int					)			->			str:
							super().__init__()
							a_		:  List[Any]											=      nn.Linear(_lowerCamelCase  ,						dim_out * 2					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Any  ,						__SCREAMING_SNAKE_CASE      :			Optional[Any]					)			->			Dict:
							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      :			int  ,						__SCREAMING_SNAKE_CASE      :			Union[str, Any]					)			->			Dict:
							a_				,			a_		:  List[str]											=      self.proj(_lowerCamelCase					).chunk(2  ,						dim=-1					)
							return hidden_states * self.gelu(_lowerCamelCase					)





class 	SCREAMING_SNAKE_CASE					(					nn.Module       ):






			def __init__( self      :			Optional[int]  ,						__SCREAMING_SNAKE_CASE      :			int  ,						__SCREAMING_SNAKE_CASE      :			int					)			->			Union[str, Any]:
							super().__init__()
							a_		:  Any											=      nn.Linear(_lowerCamelCase  ,						_lowerCamelCase					)







			def 			SCREAMING_SNAKE_CASE      ( self      :			Tuple  ,						__SCREAMING_SNAKE_CASE      :			List[Any]					)			->			List[Any]:
							a_		:  Dict											=      self.proj(_lowerCamelCase					)
							return x * torch.sigmoid(1.702 * x					)





class 	SCREAMING_SNAKE_CASE					(					nn.Module       ):






			def __init__( self      :			Dict  ,						__SCREAMING_SNAKE_CASE      :			Dict  ,						__SCREAMING_SNAKE_CASE      :			Tuple					)			->			Optional[int]:
							super().__init__()
							a_		:  int											=      nn.Embedding(_lowerCamelCase  ,						_lowerCamelCase					)
							a_		:  int											=      nn.SiLU()
							a_		:  int											=      nn.Linear(_lowerCamelCase  ,						embedding_dim * 2					)
							a_		:  List[Any]											=      nn.LayerNorm(_lowerCamelCase  ,						elementwise_affine=_lowerCamelCase					)







			def 			SCREAMING_SNAKE_CASE      ( self      :			Tuple  ,						__SCREAMING_SNAKE_CASE      :			int  ,						__SCREAMING_SNAKE_CASE      :			Optional[Any]					)			->			Optional[Any]:
							a_		:  List[Any]											=      self.linear(self.silu(self.emb(_lowerCamelCase					)					)					)
							a_				,			a_		:  Tuple											=      torch.chunk(_lowerCamelCase  ,						2					)
							a_		:  List[Any]											=      self.norm(_lowerCamelCase					) * (1 + scale) + shift
							return x





class 	SCREAMING_SNAKE_CASE					(					nn.Module       ):






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

							a_		:  int											=      CombinedTimestepLabelEmbeddings(_lowerCamelCase  ,						_lowerCamelCase					)

							a_		:  str											=      nn.SiLU()
							a_		:  Any											=      nn.Linear(_lowerCamelCase  ,						6 * embedding_dim  ,						bias=_lowerCamelCase					)
							a_		:  Optional[int]											=      nn.LayerNorm(_lowerCamelCase  ,						elementwise_affine=_lowerCamelCase  ,						eps=1e-6					)







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





class 	SCREAMING_SNAKE_CASE					(					nn.Module       ):






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

							if act_fn is None:
											a_		:  Tuple											=      None
							else:
											a_		:  List[str]											=      get_activation(_lowerCamelCase					)

							a_		:  Dict											=      nn.Linear(_lowerCamelCase  ,						out_dim * 2					)







			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[Any]  ,						__SCREAMING_SNAKE_CASE      :			Union[str, Any]  ,						__SCREAMING_SNAKE_CASE      :			Union[str, Any]					)			->			Tuple:
							if self.act:
											a_		:  Union[str, Any]											=      self.act(_lowerCamelCase					)
							a_		:  List[Any]											=      self.linear(_lowerCamelCase					)
							a_		:  str											=      emb[:, :, None, None]
							a_				,			a_		:  int											=      emb.chunk(2  ,						dim=1					)

							a_		:  Any											=      F.group_norm(_lowerCamelCase  ,						self.num_groups  ,						eps=self.eps					)
							a_		:  int											=      x * (1 + scale) + shift
							return x












'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__lowerCAmelCase       		=      {
    'configuration_git': ['GIT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'GitConfig', 'GitVisionConfig'],
    'processing_git': ['GitProcessor'],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__lowerCAmelCase       		=      [
		    'GIT_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'GitForCausalLM',
		    'GitModel',
		    'GitPreTrainedModel',
		    'GitVisionModel',
		]

if TYPE_CHECKING:
		from .configuration_git import GIT_PRETRAINED_CONFIG_ARCHIVE_MAP, GitConfig, GitVisionConfig
		from .processing_git import GitProcessor

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_git import (
				    GIT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    GitForCausalLM,
				    GitModel,
				    GitPreTrainedModel,
				    GitVisionModel,
				)

else:
		import sys

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




