





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





   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					,			)				->  Tuple:
      """simple docstring"""





      super().__init__()
      a__	:						Union[str, Any]      		=  only_cross_attention

      a__	:						Optional[int]      		=  (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"
      a__	:						Tuple      		=  (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__	:						Optional[int]      		=  AdaLayerNorm(snake_case					,			snake_case      )
      elif self.use_ada_layer_norm_zero:
         a__	:						Tuple      		=  AdaLayerNormZero(snake_case					,			snake_case      )
      else:
         a__	:						Optional[int]      		=  nn.LayerNorm(snake_case					,			elementwise_affine=snake_case      )
      a__	:						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.
         a__	:						str      		=  (
             AdaLayerNorm(snake_case					,			snake_case      )
             if self.use_ada_layer_norm
             else nn.LayerNorm(snake_case					,			elementwise_affine=snake_case      )
         )
         a__	:						Any      		=  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:
         a__	:						int      		=  None
         a__	:						List[Any]      		=  None

      # 3. Feed-forward
      a__	:						List[Any]      		=  nn.LayerNorm(snake_case					,			elementwise_affine=snake_case      )
      a__	:						Optional[int]      		=  FeedForward(snake_case					,			dropout=snake_case					,			activation_fn=snake_case					,			final_dropout=snake_case      )

      # let chunk size default to None
      a__	:						str      		=  None
      a__	:						Any      		=  0





   def 	_snake_case ( self					,			snake_case					,			snake_case      )				->  int:
      """simple docstring"""





      a__	:						Any      		=  chunk_size
      a__	:						str      		=  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					,			)				->  Optional[int]:
      """simple docstring"""





      if self.use_ada_layer_norm:
         a__	:						List[Any]      		=  self.norma(snake_case					,			snake_case      )
      elif self.use_ada_layer_norm_zero:
         a__		,     a__		,     a__		,     a__		,     a__	:						List[str]      		=  self.norma(
             snake_case					,			snake_case					,			snake_case					,			hidden_dtype=hidden_states.dtype      )
      else:
         a__	:						Optional[int]      		=  self.norma(snake_case      )

      a__	:						List[Any]      		=  cross_attention_kwargs if cross_attention_kwargs is not None else {}

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

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

         a__	:						Dict      		=  self.attna(
             snake_case					,			encoder_hidden_states=snake_case					,			attention_mask=snake_case					,			**snake_case					,			)
         a__	:						List[Any]      		=  attn_output + hidden_states

      # 3. Feed-forward
      a__	:						List[str]      		=  self.norma(snake_case      )

      if self.use_ada_layer_norm_zero:
         a__	:						List[str]      		=  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__	:						List[str]      		=  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:
         a__	:						Optional[int]      		=  self.ff(snake_case      )

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

      a__	:						List[Any]      		=  ff_output + hidden_states

      return hidden_states







class   __lowerCAmelCase			(					nn.Module  ):





   def __init__( self					,			snake_case					,			snake_case = None					,			snake_case = 4					,			snake_case = 0.0					,			snake_case = "geglu"					,			snake_case = False					,			)				->  List[str]:
      """simple docstring"""





      super().__init__()
      a__	:						Union[str, Any]      		=  int(dim * mult      )
      a__	:						Optional[Any]      		=  dim_out if dim_out is not None else dim

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

      a__	:						Union[str, 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:
         a__	:						int      		=  module(snake_case      )
      return hidden_states







class   __lowerCAmelCase			(					nn.Module  ):





   def __init__( self					,			snake_case					,			snake_case					,			snake_case = "none"      )				->  List[str]:
      """simple docstring"""





      super().__init__()
      a__	:						str      		=  nn.Linear(snake_case					,			snake_case      )
      a__	:						List[str]      		=  approximate





   def 	_snake_case ( self					,			snake_case      )				->  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      )				->  str:
      """simple docstring"""





      a__	:						Tuple      		=  self.proj(snake_case      )
      a__	:						List[Any]      		=  self.gelu(snake_case      )
      return hidden_states







class   __lowerCAmelCase			(					nn.Module  ):





   def __init__( self					,			snake_case					,			snake_case      )				->  str:
      """simple docstring"""





      super().__init__()
      a__	:						int      		=  nn.Linear(snake_case					,			dim_out * 2      )





   def 	_snake_case ( self					,			snake_case      )				->  Optional[Any]:
      """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      )				->  Dict:
      """simple docstring"""





      a__		,     a__	:						List[Any]      		=  self.proj(snake_case      ).chunk(2					,			dim=-1      )
      return hidden_states * self.gelu(snake_case      )







class   __lowerCAmelCase			(					nn.Module  ):





   def __init__( self					,			snake_case					,			snake_case      )				->  List[Any]:
      """simple docstring"""





      super().__init__()
      a__	:						List[str]      		=  nn.Linear(snake_case					,			snake_case      )





   def 	_snake_case ( self					,			snake_case      )				->  Optional[int]:
      """simple docstring"""





      a__	:						str      		=  self.proj(snake_case      )
      return x * torch.sigmoid(1.702 * x      )







class   __lowerCAmelCase			(					nn.Module  ):





   def __init__( self					,			snake_case					,			snake_case      )				->  Union[str, Any]:
      """simple docstring"""





      super().__init__()
      a__	:						Any      		=  nn.Embedding(snake_case					,			snake_case      )
      a__	:						Dict      		=  nn.SiLU()
      a__	:						Any      		=  nn.Linear(snake_case					,			embedding_dim * 2      )
      a__	:						Any      		=  nn.LayerNorm(snake_case					,			elementwise_affine=snake_case      )





   def 	_snake_case ( self					,			snake_case					,			snake_case      )				->  List[str]:
      """simple docstring"""





      a__	:						Union[str, Any]      		=  self.linear(self.silu(self.emb(snake_case      )      )      )
      a__		,     a__	:						str      		=  torch.chunk(snake_case					,			2      )
      a__	:						Dict      		=  self.norm(snake_case      ) * (1 + scale) + shift
      return x







class   __lowerCAmelCase			(					nn.Module  ):





   def __init__( self					,			snake_case					,			snake_case      )				->  Optional[int]:
      """simple docstring"""





      super().__init__()

      a__	:						Dict      		=  CombinedTimestepLabelEmbeddings(snake_case					,			snake_case      )

      a__	:						Union[str, Any]      		=  nn.SiLU()
      a__	:						Tuple      		=  nn.Linear(snake_case					,			6 * embedding_dim					,			bias=snake_case      )
      a__	:						List[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      )				->  List[str]:
      """simple docstring"""





      a__	:						Optional[int]      		=  self.linear(self.silu(self.emb(snake_case					,			snake_case					,			hidden_dtype=snake_case      )      )      )
      a__		,     a__		,     a__		,     a__		,     a__		,     a__	:						int      		=  emb.chunk(6					,			dim=1      )
      a__	:						Dict      		=  self.norm(snake_case      ) * (1 + scale_msa[:, None]) + shift_msa[:, None]
      return x, gate_msa, shift_mlp, scale_mlp, gate_mlp







class   __lowerCAmelCase			(					nn.Module  ):





   def __init__( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case = None					,			snake_case = 1E-5      )				->  str:
      """simple docstring"""





      super().__init__()
      a__	:						Tuple      		=  num_groups
      a__	:						Dict      		=  eps

      if act_fn is None:
         a__	:						str      		=  None
      else:
         a__	:						Dict      		=  get_activation(snake_case      )

      a__	:						Dict      		=  nn.Linear(snake_case					,			out_dim * 2      )





   def 	_snake_case ( self					,			snake_case					,			snake_case      )				->  List[Any]:
      """simple docstring"""





      if self.act:
         a__	:						Dict      		=  self.act(snake_case      )
      a__	:						Optional[Any]      		=  self.linear(snake_case      )
      a__	:						int      		=  emb[:, :, None, None]
      a__		,     a__	:						int      		=  emb.chunk(2					,			dim=1      )

      a__	:						Union[str, Any]      		=  F.group_norm(snake_case					,			self.num_groups					,			eps=self.eps      )
      a__	:						List[str]      		=  x * (1 + scale) + shift
      return x








from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


SCREAMING_SNAKE_CASE__       :			Dict			      =						logging.get_logger(__name__)

SCREAMING_SNAKE_CASE__       :			Any			      =						{
    """google/mobilenet_v2_1.4_224""": """https://huggingface.co/google/mobilenet_v2_1.4_224/resolve/main/config.json""",
    """google/mobilenet_v2_1.0_224""": """https://huggingface.co/google/mobilenet_v2_1.0_224/resolve/main/config.json""",
    """google/mobilenet_v2_0.75_160""": """https://huggingface.co/google/mobilenet_v2_0.75_160/resolve/main/config.json""",
    """google/mobilenet_v2_0.35_96""": """https://huggingface.co/google/mobilenet_v2_0.35_96/resolve/main/config.json""",
    # See all MobileNetV2 models at https://huggingface.co/models?filter=mobilenet_v2
}







class   __lowerCAmelCase			(					_UpperCamelCase  ):
   _UpperCamelCase :						Optional[int]     		=				"""mobilenet_v2"""





   def __init__( self					,			snake_case=3					,			snake_case=224					,			snake_case=1.0					,			snake_case=8					,			snake_case=8					,			snake_case=6					,			snake_case=32					,			snake_case=True					,			snake_case=True					,			snake_case="relu6"					,			snake_case=True					,			snake_case=0.8					,			snake_case=0.02					,			snake_case=0.001					,			snake_case=255					,			**snake_case					,			)				->  int:
      """simple docstring"""





      super().__init__(**snake_case      )

      if depth_multiplier <= 0:
         raise ValueError("depth_multiplier must be greater than zero."      )

      a__	:						str      		=  num_channels
      a__	:						Dict      		=  image_size
      a__	:						Any      		=  depth_multiplier
      a__	:						str      		=  depth_divisible_by
      a__	:						Optional[int]      		=  min_depth
      a__	:						Dict      		=  expand_ratio
      a__	:						str      		=  output_stride
      a__	:						Optional[int]      		=  first_layer_is_expansion
      a__	:						Union[str, Any]      		=  finegrained_output
      a__	:						Union[str, Any]      		=  hidden_act
      a__	:						str      		=  tf_padding
      a__	:						List[Any]      		=  classifier_dropout_prob
      a__	:						List[Any]      		=  initializer_range
      a__	:						Optional[Any]      		=  layer_norm_eps
      a__	:						str      		=  semantic_loss_ignore_index







class   __lowerCAmelCase			(					_UpperCamelCase  ):
   _UpperCamelCase :						Any     		=				version.parse("""1.11"""  )





   @property
   def 	_snake_case ( self      )				->  Mapping[str, Mapping[int, str]]:
      """simple docstring"""





      return OrderedDict([("pixel_values", {0: "batch"})]      )





   @property
   def 	_snake_case ( self      )				->  Mapping[str, Mapping[int, str]]:
      """simple docstring"""





      if self.task == "image-classification":
         return OrderedDict([("logits", {0: "batch"})]      )
      else:
         return OrderedDict([("last_hidden_state", {0: "batch"}), ("pooler_output", {0: "batch"})]      )





   @property
   def 	_snake_case ( self      )				->  float:
      """simple docstring"""





      return 1E-4

