


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

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase=0.0    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = "geglu"    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = False    ,				__UpperCAmelCase = False    ,				__UpperCAmelCase = False    ,				__UpperCAmelCase = False    ,				__UpperCAmelCase = True    ,				__UpperCAmelCase = "layer_norm"    ,				__UpperCAmelCase = False    ,				)       ->							Optional[Any]:
          super().__init__()
          _a    	=				only_cross_attention

          _a    	=				(num_embeds_ada_norm is not None) and norm_type == '''ada_norm_zero'''
          _a    	=				(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    	=				AdaLayerNorm(__UpperCAmelCase    ,				__UpperCAmelCase							)
          elif self.use_ada_layer_norm_zero:
               _a    	=				AdaLayerNormZero(__UpperCAmelCase    ,				__UpperCAmelCase							)
          else:
               _a    	=				nn.LayerNorm(__UpperCAmelCase    ,				elementwise_affine=__UpperCAmelCase							)
          _a    	=				Attention(
              query_dim=__UpperCAmelCase    ,				heads=__UpperCAmelCase    ,				dim_head=__UpperCAmelCase    ,				dropout=__UpperCAmelCase    ,				bias=__UpperCAmelCase    ,				cross_attention_dim=cross_attention_dim if only_cross_attention else None    ,				upcast_attention=__UpperCAmelCase    ,				)

          # 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    	=				(
                   AdaLayerNorm(__UpperCAmelCase    ,				__UpperCAmelCase							)
                   if self.use_ada_layer_norm
                   else nn.LayerNorm(__UpperCAmelCase    ,				elementwise_affine=__UpperCAmelCase							)
               )
               _a    	=				Attention(
                   query_dim=__UpperCAmelCase    ,				cross_attention_dim=cross_attention_dim if not double_self_attention else None    ,				heads=__UpperCAmelCase    ,				dim_head=__UpperCAmelCase    ,				dropout=__UpperCAmelCase    ,				bias=__UpperCAmelCase    ,				upcast_attention=__UpperCAmelCase    ,				)  # is self-attn if encoder_hidden_states is none
          else:
               _a    	=				None
               _a    	=				None

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

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


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)       ->							Dict:
          # Sets chunk feed-forward
          _a    	=				chunk_size
          _a    	=				dim





     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = 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    	=				self.norma(__UpperCAmelCase    ,				__UpperCAmelCase							)
          elif self.use_ada_layer_norm_zero:
               _a	,      _a	,      _a	,      _a	,      _a    	=				self.norma(
                   __UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				hidden_dtype=hidden_states.dtype							)
          else:
               _a    	=				self.norma(__UpperCAmelCase							)

          _a    	=				cross_attention_kwargs if cross_attention_kwargs is not None else {}

          _a    	=				self.attna(
              __UpperCAmelCase    ,				encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None    ,				attention_mask=__UpperCAmelCase    ,				**__UpperCAmelCase    ,				)
          if self.use_ada_layer_norm_zero:
               _a    	=				gate_msa.unsqueeze(1							) * attn_output
          _a    	=				attn_output + hidden_states

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

               _a    	=				self.attna(
                   __UpperCAmelCase    ,				encoder_hidden_states=__UpperCAmelCase    ,				attention_mask=__UpperCAmelCase    ,				**__UpperCAmelCase    ,				)
               _a    	=				attn_output + hidden_states

          # 3. Feed-forward
          _a    	=				self.norma(__UpperCAmelCase							)

          if self.use_ada_layer_norm_zero:
               _a    	=				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    	=				norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
               _a    	=				torch.cat(
                   [self.ff(__UpperCAmelCase							) for hid_slice in norm_hidden_states.chunk(__UpperCAmelCase    ,				dim=self._chunk_dim							)]    ,				dim=self._chunk_dim    ,				)
          else:
               _a    	=				self.ff(__UpperCAmelCase							)

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

          _a    	=				ff_output + hidden_states

          return hidden_states







class     __lowerCamelCase  (					nn.Module				):

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = 4    ,				__UpperCAmelCase = 0.0    ,				__UpperCAmelCase = "geglu"    ,				__UpperCAmelCase = False    ,				)       ->							List[Any]:
          super().__init__()
          _a    	=				int(dim * mult							)
          _a    	=				dim_out if dim_out is not None else dim

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

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





     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							Optional[int]:
          for module in self.net:
               _a    	=				module(__UpperCAmelCase							)
          return hidden_states







class     __lowerCamelCase  (					nn.Module				):

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase = "none"							)       ->							Union[str, Any]:
          super().__init__()
          _a    	=				nn.Linear(__UpperCAmelCase    ,				__UpperCAmelCase							)
          _a    	=				approximate


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							Optional[Any]:
          if gate.device.type != "mps":
               return F.gelu(__UpperCAmelCase    ,				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 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							Tuple:
          _a    	=				self.proj(__UpperCAmelCase							)
          _a    	=				self.gelu(__UpperCAmelCase							)
          return hidden_states







class     __lowerCamelCase  (					nn.Module				):

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)       ->							Optional[Any]:
          super().__init__()
          _a    	=				nn.Linear(__UpperCAmelCase    ,				dim_out * 2							)


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





     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							Optional[Any]:
          _a	,      _a    	=				self.proj(__UpperCAmelCase							).chunk(2    ,				dim=-1							)
          return hidden_states * self.gelu(__UpperCAmelCase							)







class     __lowerCamelCase  (					nn.Module				):

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)       ->							List[Any]:
          super().__init__()
          _a    	=				nn.Linear(__UpperCAmelCase    ,				__UpperCAmelCase							)





     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							str:
          _a    	=				self.proj(__UpperCAmelCase							)
          return x * torch.sigmoid(1.702 * x							)







class     __lowerCamelCase  (					nn.Module				):

     '''simple docstring'''


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





     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)       ->							int:
          _a    	=				self.linear(self.silu(self.emb(__UpperCAmelCase							)							)							)
          _a	,      _a    	=				torch.chunk(__UpperCAmelCase    ,				2							)
          _a    	=				self.norm(__UpperCAmelCase							) * (1 + scale) + shift
          return x







class     __lowerCamelCase  (					nn.Module				):

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)       ->							List[str]:
          super().__init__()

          _a    	=				CombinedTimestepLabelEmbeddings(__UpperCAmelCase    ,				__UpperCAmelCase							)

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





     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase=None							)       ->							Optional[Any]:
          _a    	=				self.linear(self.silu(self.emb(__UpperCAmelCase    ,				__UpperCAmelCase    ,				hidden_dtype=__UpperCAmelCase							)							)							)
          _a	,      _a	,      _a	,      _a	,      _a	,      _a    	=				emb.chunk(6    ,				dim=1							)
          _a    	=				self.norm(__UpperCAmelCase							) * (1 + scale_msa[:, None]) + shift_msa[:, None]
          return x, gate_msa, shift_mlp, scale_mlp, gate_mlp







class     __lowerCamelCase  (					nn.Module				):

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = 1e-5							)       ->							Any:
          super().__init__()
          _a    	=				num_groups
          _a    	=				eps

          if act_fn is None:
               _a    	=				None
          else:
               _a    	=				get_activation(__UpperCAmelCase							)

          _a    	=				nn.Linear(__UpperCAmelCase    ,				out_dim * 2							)





     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)       ->							int:
          if self.act:
               _a    	=				self.act(__UpperCAmelCase							)
          _a    	=				self.linear(__UpperCAmelCase							)
          _a    	=				emb[:, :, None, None]
          _a	,      _a    	=				emb.chunk(2    ,				dim=1							)

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



"""simple docstring"""






def       A_   (   _lowerCAmelCase     :					int = 10**12       ):






     """simple docstring"""

     _a    	=				1
     _a    	=				0

     _a    	=				1
     _a    	=				1

     while numerator <= 2 * min_total - 1:
          prev_numerator += 2 * numerator
          numerator += 2 * prev_numerator

          prev_denominator += 2 * denominator
          denominator += 2 * prev_denominator

     return (denominator + 1) // 2


if __name__ == "__main__":
   print(f'{solution() = }')