





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







	def __init__(				self			:   Any , lowerCamelCase			:   int , lowerCamelCase			:   int , lowerCamelCase			:   int , lowerCamelCase			:   Optional[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 , ):




				'''simple docstring'''
				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(lowerCamelCase , lowerCamelCase		)
				elif self.use_ada_layer_norm_zero:
							a__     = AdaLayerNormZero(lowerCamelCase , lowerCamelCase		)
				else:
							a__     = nn.LayerNorm(lowerCamelCase , elementwise_affine=lowerCamelCase		)
				a__     = 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__     = (
							    AdaLayerNorm(lowerCamelCase , lowerCamelCase		)
							    if self.use_ada_layer_norm
							    else nn.LayerNorm(lowerCamelCase , elementwise_affine=lowerCamelCase		)
							)
							a__     = 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__     = None
							a__     = None

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

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







	def 						__a   (				self			:   Tuple , lowerCamelCase			:   Optional[int] , lowerCamelCase			:   int		):




				'''simple docstring'''
				# Sets chunk feed-forward
				a__     = chunk_size
				a__     = dim







	def 						__a   (				self			:   List[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 , ):




				'''simple docstring'''
				# 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(lowerCamelCase , lowerCamelCase		)
				elif self.use_ada_layer_norm_zero:
							a__   ,				a__   ,				a__   ,				a__   ,				a__     = self.norma(
							    lowerCamelCase , lowerCamelCase , lowerCamelCase , hidden_dtype=hidden_states.dtype		)
				else:
							a__     = self.norma(lowerCamelCase		)

				a__     = cross_attention_kwargs if cross_attention_kwargs is not None else {}

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

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

							a__     = self.attna(
							    lowerCamelCase , encoder_hidden_states=lowerCamelCase , attention_mask=lowerCamelCase , **lowerCamelCase , )
							a__     = attn_output + hidden_states

				# 3. Feed-forward
				a__     = self.norma(lowerCamelCase		)

				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(lowerCamelCase		) for hid_slice in norm_hidden_states.chunk(lowerCamelCase , dim=self._chunk_dim		)] , dim=self._chunk_dim , )
				else:
							a__     = self.ff(lowerCamelCase		)

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

				a__     = ff_output + hidden_states

				return hidden_states






class       UpperCamelCase__  (		nn.Module     ):







	def __init__(				self			:   str , lowerCamelCase			:   int , lowerCamelCase			:   Optional[int] = None , lowerCamelCase			:   int = 4 , lowerCamelCase			:   float = 0.0 , lowerCamelCase			:   str = "geglu" , lowerCamelCase			:   bool = False , ):




				'''simple docstring'''
				super().__init__()
				a__     = int(dim * mult		)
				a__     = dim_out if dim_out is not None else dim

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

				a__     = 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 						__a   (				self			:   List[Any] , lowerCamelCase			:   str		):




				'''simple docstring'''
				for module in self.net:
							a__     = module(lowerCamelCase		)
				return hidden_states






class       UpperCamelCase__  (		nn.Module     ):







	def __init__(				self			:   Tuple , lowerCamelCase			:   int , lowerCamelCase			:   int , lowerCamelCase			:   str = "none"		):




				'''simple docstring'''
				super().__init__()
				a__     = nn.Linear(lowerCamelCase , lowerCamelCase		)
				a__     = approximate







	def 						__a   (				self			:   Any , lowerCamelCase			:   Optional[int]		):




				'''simple docstring'''
				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 						__a   (				self			:   Tuple , lowerCamelCase			:   List[str]		):




				'''simple docstring'''
				a__     = self.proj(lowerCamelCase		)
				a__     = self.gelu(lowerCamelCase		)
				return hidden_states






class       UpperCamelCase__  (		nn.Module     ):







	def __init__(				self			:   Union[str, Any] , lowerCamelCase			:   int , lowerCamelCase			:   int		):




				'''simple docstring'''
				super().__init__()
				a__     = nn.Linear(lowerCamelCase , dim_out * 2		)







	def 						__a   (				self			:   Optional[int] , lowerCamelCase			:   Union[str, Any]		):




				'''simple docstring'''
				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 						__a   (				self			:   Union[str, Any] , lowerCamelCase			:   Optional[Any]		):




				'''simple docstring'''
				a__   ,				a__     = self.proj(lowerCamelCase		).chunk(2 , dim=-1		)
				return hidden_states * self.gelu(lowerCamelCase		)






class       UpperCamelCase__  (		nn.Module     ):







	def __init__(				self			:   Any , lowerCamelCase			:   int , lowerCamelCase			:   int		):




				'''simple docstring'''
				super().__init__()
				a__     = nn.Linear(lowerCamelCase , lowerCamelCase		)







	def 						__a   (				self			:   Optional[Any] , lowerCamelCase			:   List[Any]		):




				'''simple docstring'''
				a__     = self.proj(lowerCamelCase		)
				return x * torch.sigmoid(1.702 * x		)






class       UpperCamelCase__  (		nn.Module     ):







	def __init__(				self			:   Optional[Any] , lowerCamelCase			:   Optional[Any] , lowerCamelCase			:   List[str]		):




				'''simple docstring'''
				super().__init__()
				a__     = nn.Embedding(lowerCamelCase , lowerCamelCase		)
				a__     = nn.SiLU()
				a__     = nn.Linear(lowerCamelCase , embedding_dim * 2		)
				a__     = nn.LayerNorm(lowerCamelCase , elementwise_affine=lowerCamelCase		)







	def 						__a   (				self			:   Optional[Any] , lowerCamelCase			:   int , lowerCamelCase			:   List[str]		):




				'''simple docstring'''
				a__     = self.linear(self.silu(self.emb(lowerCamelCase		)		)		)
				a__   ,				a__     = torch.chunk(lowerCamelCase , 2		)
				a__     = self.norm(lowerCamelCase		) * (1 + scale) + shift
				return x






class       UpperCamelCase__  (		nn.Module     ):







	def __init__(				self			:   Dict , lowerCamelCase			:   Dict , lowerCamelCase			:   List[str]		):




				'''simple docstring'''
				super().__init__()

				a__     = CombinedTimestepLabelEmbeddings(lowerCamelCase , lowerCamelCase		)

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







	def 						__a   (				self			:   str , lowerCamelCase			:   List[Any] , lowerCamelCase			:   Dict , lowerCamelCase			:   int , lowerCamelCase			:   str=None		):




				'''simple docstring'''
				a__     = self.linear(self.silu(self.emb(lowerCamelCase , lowerCamelCase , hidden_dtype=lowerCamelCase		)		)		)
				a__   ,				a__   ,				a__   ,				a__   ,				a__   ,				a__     = emb.chunk(6 , dim=1		)
				a__     = self.norm(lowerCamelCase		) * (1 + scale_msa[:, None]) + shift_msa[:, None]
				return x, gate_msa, shift_mlp, scale_mlp, gate_mlp






class       UpperCamelCase__  (		nn.Module     ):







	def __init__(				self			:   List[str] , lowerCamelCase			:   int , lowerCamelCase			:   int , lowerCamelCase			:   int , lowerCamelCase			:   Optional[str] = None , lowerCamelCase			:   float = 1e-5		):




				'''simple docstring'''
				super().__init__()
				a__     = num_groups
				a__     = eps

				if act_fn is None:
							a__     = None
				else:
							a__     = get_activation(lowerCamelCase		)

				a__     = nn.Linear(lowerCamelCase , out_dim * 2		)







	def 						__a   (				self			:   Optional[Any] , lowerCamelCase			:   Any , lowerCamelCase			:   Optional[Any]		):




				'''simple docstring'''
				if self.act:
							a__     = self.act(lowerCamelCase		)
				a__     = self.linear(lowerCamelCase		)
				a__     = emb[:, :, None, None]
				a__   ,				a__     = emb.chunk(2 , dim=1		)

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










'''simple docstring'''


import importlib
import os
import sys


# This is required to make the module import works (when the python process is running from the root of the repo)
sys.path.append(".")







def    _lowerCamelCase  (__lowerCamelCase  :     List[str]     )   ->    List[str]:
			a__     = test_file.split(os.path.sep     )
			if components[0:2] != ["tests", "models"]:
						raise ValueError(
						    "`test_file` should start with `tests/models/` (with `/` being the OS specific path separator). Got "
						    f'''{test_file} instead.'''     )
			a__     = components[-1]
			if not test_fn.endswith("py"     ):
						raise ValueError(f'''`test_file` should be a python file. Got {test_fn} instead.'''     )
			if not test_fn.startswith("test_modeling_"     ):
						raise ValueError(
						    f'''`test_file` should point to a file name of the form `test_modeling_*.py`. Got {test_fn} instead.'''     )

			a__     = components[:-1] + [test_fn.replace(".py"					,  ""     )]
			a__     = ".".join(__lowerCamelCase     )

			return test_module_path







def    _lowerCamelCase  (__lowerCamelCase  :     Union[str, Any]     )   ->    Optional[Any]:
			a__     = get_module_path(__lowerCamelCase     )
			a__     = importlib.import_module(__lowerCamelCase     )

			return test_module







def    _lowerCamelCase  (__lowerCamelCase  :     Optional[int]     )   ->    Union[str, Any]:
			a__     = []
			a__     = get_test_module(__lowerCamelCase     )
			for attr in dir(__lowerCamelCase     ):
						if attr.endswith("ModelTester"     ):
									tester_classes.append(getattr(__lowerCamelCase					,  __lowerCamelCase     )     )

    # sort with class names
			return sorted(__lowerCamelCase					,  key=lambda __lowerCamelCase     : x.__name__     )







def    _lowerCamelCase  (__lowerCamelCase  :     Tuple     )   ->    Tuple:
			a__     = []
			a__     = get_test_module(__lowerCamelCase     )
			for attr in dir(__lowerCamelCase     ):
						a__     = getattr(__lowerCamelCase					,  __lowerCamelCase     )
						# (TF/Flax)ModelTesterMixin is also an attribute in specific model test module. Let's exclude them by checking
						# `all_model_classes` is not empty (which also excludes other special classes).
						a__     = getattr(__lowerCamelCase					,  "all_model_classes"					,  []     )
						if len(__lowerCamelCase     ) > 0:
									test_classes.append(__lowerCamelCase     )

    # sort with class names
			return sorted(__lowerCamelCase					,  key=lambda __lowerCamelCase     : x.__name__     )







def    _lowerCamelCase  (__lowerCamelCase  :     List[Any]     )   ->    List[Any]:
			a__     = get_test_classes(__lowerCamelCase     )
			a__     = set()
			for test_class in test_classes:
						model_classes.update(test_class.all_model_classes     )

			# sort with class names
			return sorted(__lowerCamelCase					,  key=lambda __lowerCamelCase     : x.__name__     )







def    _lowerCamelCase  (__lowerCamelCase  :     str     )   ->    Dict:
			a__     = test_class()
			if hasattr(__lowerCamelCase					,  "setUp"     ):
						test.setUp()

			a__     = None
			if hasattr(__lowerCamelCase					,  "model_tester"     ):
						# `(TF/Flax)ModelTesterMixin` has this attribute default to `None`. Let's skip this case.
						if test.model_tester is not None:
									a__     = test.model_tester.__class__

			return model_tester







def    _lowerCamelCase  (__lowerCamelCase  :     Dict					,  __lowerCamelCase  :     Optional[int]     )   ->    List[Any]:
			a__     = get_test_classes(__lowerCamelCase     )

			a__     = []
			for test_class in test_classes:
						if model_class in test_class.all_model_classes:
									target_test_classes.append(__lowerCamelCase     )

    # sort with class names
			return sorted(__lowerCamelCase					,  key=lambda __lowerCamelCase     : x.__name__     )







def    _lowerCamelCase  (__lowerCamelCase  :     Any					,  __lowerCamelCase  :     Tuple     )   ->    Dict:
			a__     = get_test_classes_for_model(__lowerCamelCase					,  __lowerCamelCase     )

			a__     = []
			for test_class in test_classes:
						a__     = get_model_tester_from_test_class(__lowerCamelCase     )
						if tester_class is not None:
									tester_classes.append(__lowerCamelCase     )

    # sort with class names
			return sorted(__lowerCamelCase					,  key=lambda __lowerCamelCase     : x.__name__     )







def    _lowerCamelCase  (__lowerCamelCase  :     List[str]     )   ->    int:
			a__     = get_test_classes(__lowerCamelCase     )
			a__     = {test_class: get_model_tester_from_test_class(__lowerCamelCase     ) for test_class in test_classes}
			return test_tester_mapping







def    _lowerCamelCase  (__lowerCamelCase  :     Any     )   ->    List[Any]:
			a__     = get_model_classes(__lowerCamelCase     )
			a__     = {
			    model_class: get_test_classes_for_model(__lowerCamelCase					,  __lowerCamelCase     ) for model_class in model_classes
			}
			return model_test_mapping







def    _lowerCamelCase  (__lowerCamelCase  :     Tuple     )   ->    Tuple:
			a__     = get_model_classes(__lowerCamelCase     )
			a__     = {
			    model_class: get_tester_classes_for_model(__lowerCamelCase					,  __lowerCamelCase     ) for model_class in model_classes
			}
			return model_to_tester_mapping







def    _lowerCamelCase  (__lowerCamelCase  :     str     )   ->    Dict:
			if isinstance(__lowerCamelCase					,  __lowerCamelCase     ):
						return o
			elif isinstance(__lowerCamelCase					,  __lowerCamelCase     ):
						return o.__name__
			elif isinstance(__lowerCamelCase					,  (list, tuple)     ):
						return [to_json(__lowerCamelCase     ) for x in o]
			elif isinstance(__lowerCamelCase					,  __lowerCamelCase     ):
						return {to_json(__lowerCamelCase     ): to_json(__lowerCamelCase     ) for k, v in o.items()}
			else:
						return o



