





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		:	int			,     __lowerCamelCase		:	int			,     __lowerCamelCase		:	int			,     __lowerCamelCase		:	int			,     __lowerCamelCase		:	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			,     )	->							str:
       super().__init__()
       SCREAMING_SNAKE_CASE__					=					only_cross_attention

       SCREAMING_SNAKE_CASE__					=					(num_embeds_ada_norm is not None) and norm_type == '''ada_norm_zero'''
       SCREAMING_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:
            SCREAMING_SNAKE_CASE__					=					AdaLayerNorm(__lowerCamelCase			,     __lowerCamelCase      )
       elif self.use_ada_layer_norm_zero:
            SCREAMING_SNAKE_CASE__					=					AdaLayerNormZero(__lowerCamelCase			,     __lowerCamelCase      )
       else:
            SCREAMING_SNAKE_CASE__					=					nn.LayerNorm(__lowerCamelCase			,     elementwise_affine=__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					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.
            SCREAMING_SNAKE_CASE__					=					(
                AdaLayerNorm(__lowerCamelCase			,     __lowerCamelCase      )
                if self.use_ada_layer_norm
                else nn.LayerNorm(__lowerCamelCase			,     elementwise_affine=__lowerCamelCase      )
            )
            SCREAMING_SNAKE_CASE__					=					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:
            SCREAMING_SNAKE_CASE__					=					None
            SCREAMING_SNAKE_CASE__					=					None

       # 3. Feed-forward
       SCREAMING_SNAKE_CASE__					=					nn.LayerNorm(__lowerCamelCase			,     elementwise_affine=__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					FeedForward(__lowerCamelCase			,     dropout=__lowerCamelCase			,     activation_fn=__lowerCamelCase			,     final_dropout=__lowerCamelCase      )

       # let chunk size default to None
       SCREAMING_SNAKE_CASE__					=					None
       SCREAMING_SNAKE_CASE__					=					0






  def    lowercase_	(							self		:	int			,     __lowerCamelCase		:	Optional[int]			,     __lowerCamelCase		:	int      )	->							Union[str, Any]:
       # Sets chunk feed-forward
       SCREAMING_SNAKE_CASE__					=					chunk_size
       SCREAMING_SNAKE_CASE__					=					dim







  def    lowercase_	(							self		:	Union[str, Any]			,     __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			,     )	->							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:
            SCREAMING_SNAKE_CASE__					=					self.norma(__lowerCamelCase			,     __lowerCamelCase      )
       elif self.use_ada_layer_norm_zero:
            SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__					=					self.norma(
                __lowerCamelCase			,     __lowerCamelCase			,     __lowerCamelCase			,     hidden_dtype=hidden_states.dtype      )
       else:
            SCREAMING_SNAKE_CASE__					=					self.norma(__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					cross_attention_kwargs if cross_attention_kwargs is not None else {}

       SCREAMING_SNAKE_CASE__					=					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:
            SCREAMING_SNAKE_CASE__					=					gate_msa.unsqueeze(1      ) * attn_output
       SCREAMING_SNAKE_CASE__					=					attn_output + hidden_states

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

            SCREAMING_SNAKE_CASE__					=					self.attna(
                __lowerCamelCase			,     encoder_hidden_states=__lowerCamelCase			,     attention_mask=__lowerCamelCase			,     **__lowerCamelCase			,     )
            SCREAMING_SNAKE_CASE__					=					attn_output + hidden_states

       # 3. Feed-forward
       SCREAMING_SNAKE_CASE__					=					self.norma(__lowerCamelCase      )

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

            SCREAMING_SNAKE_CASE__					=					norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
            SCREAMING_SNAKE_CASE__					=					torch.cat(
                [self.ff(__lowerCamelCase      ) for hid_slice in norm_hidden_states.chunk(__lowerCamelCase			,     dim=self._chunk_dim      )]			,     dim=self._chunk_dim			,     )
       else:
            SCREAMING_SNAKE_CASE__					=					self.ff(__lowerCamelCase      )

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

       SCREAMING_SNAKE_CASE__					=					ff_output + hidden_states

       return hidden_states






class  UpperCAmelCase__	( nn.Module ):






  """simple docstring"""


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

       if activation_fn == "gelu":
            SCREAMING_SNAKE_CASE__					=					GELU(__lowerCamelCase			,     __lowerCamelCase      )
       if activation_fn == "gelu-approximate":
            SCREAMING_SNAKE_CASE__					=					GELU(__lowerCamelCase			,     __lowerCamelCase			,     approximate='''tanh'''      )
       elif activation_fn == "geglu":
            SCREAMING_SNAKE_CASE__					=					GEGLU(__lowerCamelCase			,     __lowerCamelCase      )
       elif activation_fn == "geglu-approximate":
            SCREAMING_SNAKE_CASE__					=					ApproximateGELU(__lowerCamelCase			,     __lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					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    lowercase_	(							self		:	Optional[Any]			,     __lowerCamelCase		:	int      )	->							Optional[Any]:
       for module in self.net:
            SCREAMING_SNAKE_CASE__					=					module(__lowerCamelCase      )
       return hidden_states






class  UpperCAmelCase__	( nn.Module ):






  """simple docstring"""


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






  def    lowercase_	(							self		:	Union[str, Any]			,     __lowerCamelCase		:	Any      )	->							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    lowercase_	(							self		:	Union[str, Any]			,     __lowerCamelCase		:	List[str]      )	->							int:
       SCREAMING_SNAKE_CASE__					=					self.proj(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					self.gelu(__lowerCamelCase      )
       return hidden_states






class  UpperCAmelCase__	( nn.Module ):






  """simple docstring"""


  def __init__(							self		:	Optional[int]			,     __lowerCamelCase		:	int			,     __lowerCamelCase		:	int      )	->							Tuple:
       super().__init__()
       SCREAMING_SNAKE_CASE__					=					nn.Linear(__lowerCamelCase			,     dim_out * 2      )






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






class  UpperCAmelCase__	( nn.Module ):






  """simple docstring"""


  def __init__(							self		:	List[str]			,     __lowerCamelCase		:	int			,     __lowerCamelCase		:	int      )	->							Union[str, Any]:
       super().__init__()
       SCREAMING_SNAKE_CASE__					=					nn.Linear(__lowerCamelCase			,     __lowerCamelCase      )







  def    lowercase_	(							self		:	str			,     __lowerCamelCase		:	List[Any]      )	->							Dict:
       SCREAMING_SNAKE_CASE__					=					self.proj(__lowerCamelCase      )
       return x * torch.sigmoid(1.702 * x      )






class  UpperCAmelCase__	( nn.Module ):






  """simple docstring"""


  def __init__(							self		:	str			,     __lowerCamelCase		:	Any			,     __lowerCamelCase		:	Optional[Any]      )	->							Any:
       super().__init__()
       SCREAMING_SNAKE_CASE__					=					nn.Embedding(__lowerCamelCase			,     __lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					nn.SiLU()
       SCREAMING_SNAKE_CASE__					=					nn.Linear(__lowerCamelCase			,     embedding_dim * 2      )
       SCREAMING_SNAKE_CASE__					=					nn.LayerNorm(__lowerCamelCase			,     elementwise_affine=__lowerCamelCase      )







  def    lowercase_	(							self		:	Optional[int]			,     __lowerCamelCase		:	Optional[int]			,     __lowerCamelCase		:	int      )	->							Any:
       SCREAMING_SNAKE_CASE__					=					self.linear(self.silu(self.emb(__lowerCamelCase      )      )      )
       SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__					=					torch.chunk(__lowerCamelCase			,     2      )
       SCREAMING_SNAKE_CASE__					=					self.norm(__lowerCamelCase      ) * (1 + scale) + shift
       return x






class  UpperCAmelCase__	( nn.Module ):






  """simple docstring"""


  def __init__(							self		:	Optional[Any]			,     __lowerCamelCase		:	Optional[Any]			,     __lowerCamelCase		:	int      )	->							List[Any]:
       super().__init__()

       SCREAMING_SNAKE_CASE__					=					CombinedTimestepLabelEmbeddings(__lowerCamelCase			,     __lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					nn.SiLU()
       SCREAMING_SNAKE_CASE__					=					nn.Linear(__lowerCamelCase			,     6 * embedding_dim			,     bias=__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					nn.LayerNorm(__lowerCamelCase			,     elementwise_affine=__lowerCamelCase			,     eps=1e-6      )







  def    lowercase_	(							self		:	Optional[int]			,     __lowerCamelCase		:	Union[str, Any]			,     __lowerCamelCase		:	Optional[int]			,     __lowerCamelCase		:	List[Any]			,     __lowerCamelCase		:	Tuple=None      )	->							List[Any]:
       SCREAMING_SNAKE_CASE__					=					self.linear(self.silu(self.emb(__lowerCamelCase			,     __lowerCamelCase			,     hidden_dtype=__lowerCamelCase      )      )      )
       SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__					=					emb.chunk(6			,     dim=1      )
       SCREAMING_SNAKE_CASE__					=					self.norm(__lowerCamelCase      ) * (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		:	str			,     __lowerCamelCase		:	int			,     __lowerCamelCase		:	int			,     __lowerCamelCase		:	int			,     __lowerCamelCase		:	Optional[str] = None			,     __lowerCamelCase		:	float = 1e-5      )	->							int:
       super().__init__()
       SCREAMING_SNAKE_CASE__					=					num_groups
       SCREAMING_SNAKE_CASE__					=					eps

       if act_fn is None:
            SCREAMING_SNAKE_CASE__					=					None
       else:
            SCREAMING_SNAKE_CASE__					=					get_activation(__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					nn.Linear(__lowerCamelCase			,     out_dim * 2      )







  def    lowercase_	(							self		:	Union[str, Any]			,     __lowerCamelCase		:	Optional[Any]			,     __lowerCamelCase		:	List[str]      )	->							Optional[int]:
       if self.act:
            SCREAMING_SNAKE_CASE__					=					self.act(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					self.linear(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					emb[:, :, None, None]
       SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__					=					emb.chunk(2			,     dim=1      )

       SCREAMING_SNAKE_CASE__					=					F.group_norm(__lowerCamelCase			,     self.num_groups			,     eps=self.eps      )
       SCREAMING_SNAKE_CASE__					=					x * (1 + scale) + shift
       return x





from collections.abc import Callable

import numpy as np
def        UpperCAmelCase_     (		_A   ,	_A   ,	_A   ,	_A   ,	_A       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					int(np.ceil((x_end - xa) / step_size       )       )
     SCREAMING_SNAKE_CASE__					=					np.zeros((n + 1,)       )
     SCREAMING_SNAKE_CASE__					=					ya
     SCREAMING_SNAKE_CASE__					=					xa

     for k in range(_A       ):
          SCREAMING_SNAKE_CASE__					=					y[k] + step_size * ode_func(_A   ,	y[k]       )
          SCREAMING_SNAKE_CASE__					=					y[k] + (
              (step_size / 2) * (ode_func(_A   ,	y[k]       ) + ode_func(x + step_size   ,	_A       ))
          )
          x += step_size

     return y


if __name__ == "__main__":
       import doctest

       doctest.testmod()


