





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  , __lowerCAmelCase  , __lowerCAmelCase  , __lowerCAmelCase  , __lowerCAmelCase=0.0  , __lowerCAmelCase = None  , __lowerCAmelCase = "geglu"  , __lowerCAmelCase = None  , __lowerCAmelCase = False  , __lowerCAmelCase = False  , __lowerCAmelCase = False  , __lowerCAmelCase = False  , __lowerCAmelCase = True  , __lowerCAmelCase = "layer_norm"  , __lowerCAmelCase = False  , ):
														super().__init__()
														UpperCamelCase__				     =			only_cross_attention

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

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

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






							def 				_lowerCamelCase		(						self  , __lowerCAmelCase  , __lowerCAmelCase  ):
														# Sets chunk feed-forward
														UpperCamelCase__				     =			chunk_size
														UpperCamelCase__				     =			dim







							def 				_lowerCamelCase		(						self  , __lowerCAmelCase  , __lowerCAmelCase = None  , __lowerCAmelCase = None  , __lowerCAmelCase = None  , __lowerCAmelCase = None  , __lowerCAmelCase = None  , __lowerCAmelCase = None  , ):
														# Notice that normalization is always applied before the real computation in the following blocks.
														# 1. Self-Attention
														if self.use_ada_layer_norm:
																					UpperCamelCase__				     =			self.norma(__lowerCAmelCase  , __lowerCAmelCase  )
														elif self.use_ada_layer_norm_zero:
																					UpperCamelCase__     ,     UpperCamelCase__     ,     UpperCamelCase__     ,     UpperCamelCase__     ,     UpperCamelCase__				     =			self.norma(
																					    __lowerCAmelCase  , __lowerCAmelCase  , __lowerCAmelCase  , hidden_dtype=hidden_states.dtype  )
														else:
																					UpperCamelCase__				     =			self.norma(__lowerCAmelCase  )

														UpperCamelCase__				     =			cross_attention_kwargs if cross_attention_kwargs is not None else {}

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

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

																					UpperCamelCase__				     =			self.attna(
																					    __lowerCAmelCase  , encoder_hidden_states=__lowerCAmelCase  , attention_mask=__lowerCAmelCase  , **__lowerCAmelCase  , )
																					UpperCamelCase__				     =			attn_output + hidden_states

														# 3. Feed-forward
														UpperCamelCase__				     =			self.norma(__lowerCAmelCase  )

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

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

														UpperCamelCase__				     =			ff_output + hidden_states

														return hidden_states






class        __SCREAMING_SNAKE_CASE					(     nn.Module       ):






							def __init__(						self  , __lowerCAmelCase  , __lowerCAmelCase = None  , __lowerCAmelCase = 4  , __lowerCAmelCase = 0.0  , __lowerCAmelCase = "geglu"  , __lowerCAmelCase = False  , ):
														super().__init__()
														UpperCamelCase__				     =			int(dim * mult  )
														UpperCamelCase__				     =			dim_out if dim_out is not None else dim

														if activation_fn == "gelu":
																					UpperCamelCase__				     =			GELU(__lowerCAmelCase  , __lowerCAmelCase  )
														if activation_fn == "gelu-approximate":
																					UpperCamelCase__				     =			GELU(__lowerCAmelCase  , __lowerCAmelCase  , approximate="""tanh"""  )
														elif activation_fn == "geglu":
																					UpperCamelCase__				     =			GEGLU(__lowerCAmelCase  , __lowerCAmelCase  )
														elif activation_fn == "geglu-approximate":
																					UpperCamelCase__				     =			ApproximateGELU(__lowerCAmelCase  , __lowerCAmelCase  )

														UpperCamelCase__				     =			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 				_lowerCamelCase		(						self  , __lowerCAmelCase  ):
														for module in self.net:
																					UpperCamelCase__				     =			module(__lowerCAmelCase  )
														return hidden_states






class        __SCREAMING_SNAKE_CASE					(     nn.Module       ):






							def __init__(						self  , __lowerCAmelCase  , __lowerCAmelCase  , __lowerCAmelCase = "none"  ):
														super().__init__()
														UpperCamelCase__				     =			nn.Linear(__lowerCAmelCase  , __lowerCAmelCase  )
														UpperCamelCase__				     =			approximate






							def 				_lowerCamelCase		(						self  , __lowerCAmelCase  ):
														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 				_lowerCamelCase		(						self  , __lowerCAmelCase  ):
														UpperCamelCase__				     =			self.proj(__lowerCAmelCase  )
														UpperCamelCase__				     =			self.gelu(__lowerCAmelCase  )
														return hidden_states






class        __SCREAMING_SNAKE_CASE					(     nn.Module       ):






							def __init__(						self  , __lowerCAmelCase  , __lowerCAmelCase  ):
														super().__init__()
														UpperCamelCase__				     =			nn.Linear(__lowerCAmelCase  , dim_out * 2  )






							def 				_lowerCamelCase		(						self  , __lowerCAmelCase  ):
														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 				_lowerCamelCase		(						self  , __lowerCAmelCase  ):
														UpperCamelCase__     ,     UpperCamelCase__				     =			self.proj(__lowerCAmelCase  ).chunk(2  , dim=-1  )
														return hidden_states * self.gelu(__lowerCAmelCase  )






class        __SCREAMING_SNAKE_CASE					(     nn.Module       ):






							def __init__(						self  , __lowerCAmelCase  , __lowerCAmelCase  ):
														super().__init__()
														UpperCamelCase__				     =			nn.Linear(__lowerCAmelCase  , __lowerCAmelCase  )







							def 				_lowerCamelCase		(						self  , __lowerCAmelCase  ):
														UpperCamelCase__				     =			self.proj(__lowerCAmelCase  )
														return x * torch.sigmoid(1.702 * x  )






class        __SCREAMING_SNAKE_CASE					(     nn.Module       ):






							def __init__(						self  , __lowerCAmelCase  , __lowerCAmelCase  ):
														super().__init__()
														UpperCamelCase__				     =			nn.Embedding(__lowerCAmelCase  , __lowerCAmelCase  )
														UpperCamelCase__				     =			nn.SiLU()
														UpperCamelCase__				     =			nn.Linear(__lowerCAmelCase  , embedding_dim * 2  )
														UpperCamelCase__				     =			nn.LayerNorm(__lowerCAmelCase  , elementwise_affine=__lowerCAmelCase  )







							def 				_lowerCamelCase		(						self  , __lowerCAmelCase  , __lowerCAmelCase  ):
														UpperCamelCase__				     =			self.linear(self.silu(self.emb(__lowerCAmelCase  )  )  )
														UpperCamelCase__     ,     UpperCamelCase__				     =			torch.chunk(__lowerCAmelCase  , 2  )
														UpperCamelCase__				     =			self.norm(__lowerCAmelCase  ) * (1 + scale) + shift
														return x






class        __SCREAMING_SNAKE_CASE					(     nn.Module       ):






							def __init__(						self  , __lowerCAmelCase  , __lowerCAmelCase  ):
														super().__init__()

														UpperCamelCase__				     =			CombinedTimestepLabelEmbeddings(__lowerCAmelCase  , __lowerCAmelCase  )

														UpperCamelCase__				     =			nn.SiLU()
														UpperCamelCase__				     =			nn.Linear(__lowerCAmelCase  , 6 * embedding_dim  , bias=__lowerCAmelCase  )
														UpperCamelCase__				     =			nn.LayerNorm(__lowerCAmelCase  , elementwise_affine=__lowerCAmelCase  , eps=1E-6  )







							def 				_lowerCamelCase		(						self  , __lowerCAmelCase  , __lowerCAmelCase  , __lowerCAmelCase  , __lowerCAmelCase=None  ):
														UpperCamelCase__				     =			self.linear(self.silu(self.emb(__lowerCAmelCase  , __lowerCAmelCase  , hidden_dtype=__lowerCAmelCase  )  )  )
														UpperCamelCase__     ,     UpperCamelCase__     ,     UpperCamelCase__     ,     UpperCamelCase__     ,     UpperCamelCase__     ,     UpperCamelCase__				     =			emb.chunk(6  , dim=1  )
														UpperCamelCase__				     =			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  , __lowerCAmelCase  , __lowerCAmelCase  , __lowerCAmelCase  , __lowerCAmelCase = None  , __lowerCAmelCase = 1E-5  ):
														super().__init__()
														UpperCamelCase__				     =			num_groups
														UpperCamelCase__				     =			eps

														if act_fn is None:
																					UpperCamelCase__				     =			None
														else:
																					UpperCamelCase__				     =			get_activation(__lowerCAmelCase  )

														UpperCamelCase__				     =			nn.Linear(__lowerCAmelCase  , out_dim * 2  )







							def 				_lowerCamelCase		(						self  , __lowerCAmelCase  , __lowerCAmelCase  ):
														if self.act:
																					UpperCamelCase__				     =			self.act(__lowerCAmelCase  )
														UpperCamelCase__				     =			self.linear(__lowerCAmelCase  )
														UpperCamelCase__				     =			emb[:, :, None, None]
														UpperCamelCase__     ,     UpperCamelCase__				     =			emb.chunk(2  , dim=1  )

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











import argparse
import datetime





def     _UpperCamelCase     (a__   :str       ):






							"""simple docstring"""





							UpperCamelCase__				     =			{
							    """0""": """Sunday""",
							    """1""": """Monday""",
							    """2""": """Tuesday""",
							    """3""": """Wednesday""",
							    """4""": """Thursday""",
							    """5""": """Friday""",
							    """6""": """Saturday""",
							}

							UpperCamelCase__				     =			{0: 1, 1: 2, 2: 3, 3: 4, 4: 5, 5: 6, 6: 0}

							# Validate
							if not 0 < len(a__       ) < 11:
														raise ValueError("""Must be 10 characters long"""       )

							# Get month
							UpperCamelCase__				     =			int(date_input[0] + date_input[1]       )
							# Validate
							if not 0 < m < 13:
														raise ValueError("""Month must be between 1 - 12"""       )

							UpperCamelCase__				     =			date_input[2]
							# Validate
							if sep_a not in ["-", "/"]:
														raise ValueError("""Date separator must be '-' or '/'"""       )

							# Get day
							UpperCamelCase__				     =			int(date_input[3] + date_input[4]       )
							# Validate
							if not 0 < d < 32:
														raise ValueError("""Date must be between 1 - 31"""       )

							# Get second separator
							UpperCamelCase__				     =			date_input[5]
							# Validate
							if sep_a not in ["-", "/"]:
														raise ValueError("""Date separator must be '-' or '/'"""       )

							# Get year
							UpperCamelCase__				     =			int(date_input[6] + date_input[7] + date_input[8] + date_input[9]       )
							# Arbitrary year range
							if not 45 < y < 8500:
														raise ValueError(
														    """Year out of range. There has to be some sort of limit...right?"""       )

							# Get datetime obj for validation
							UpperCamelCase__				     =			datetime.date(int(a__       )   ,			int(a__       )   ,			int(a__       )       )

							# Start math
							if m <= 2:
														UpperCamelCase__				     =			y - 1
														UpperCamelCase__				     =			m + 12
							# maths var
							UpperCamelCase__				     =			int(str(a__       )[:2]       )
							UpperCamelCase__				     =			int(str(a__       )[2:]       )
							UpperCamelCase__				     =			int(2.6 * m - 5.39       )
							UpperCamelCase__				     =			int(c / 4       )
							UpperCamelCase__				     =			int(k / 4       )
							UpperCamelCase__				     =			int(d + k       )
							UpperCamelCase__				     =			int(t + u + v + x       )
							UpperCamelCase__				     =			int(z - (2 * c)       )
							UpperCamelCase__				     =			round(w % 7       )
							# End math

							# Validate math
							if f != convert_datetime_days[dt_ck.weekday()]:
														raise AssertionError("""The date was evaluated incorrectly. Contact developer."""       )

							# Response
							UpperCamelCase__				     =			f"""Your date {date_input}, is a {days[str(a__       )]}!"""
							return response


if __name__ == "__main__":
					import doctest

					doctest.testmod()
					UpperCamelCase__											= argparse.ArgumentParser(
					    description=(
					        "Find out what day of the week nearly any date is or was. Enter "
					        "date as a string in the mm-dd-yyyy or mm/dd/yyyy format"
					    )
					)
					parser.add_argument(
					    "date_input", type=str, help="Date as a string (mm-dd-yyyy or mm/dd/yyyy)"
					)
					UpperCamelCase__											= parser.parse_args()
					zeller(args.date_input)




