


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





							'''simple docstring'''

							def __init__(				self							:      List[Any]		,       UpperCAmelCase_							:      int		,       UpperCAmelCase_							:      int		,       UpperCAmelCase_							:      int		,       UpperCAmelCase_							:      Optional[int]=0.0		,       UpperCAmelCase_							:      Optional[int] = None		,       UpperCAmelCase_							:      str = "geglu"		,       UpperCAmelCase_							:      Optional[int] = None		,       UpperCAmelCase_							:      bool = False		,       UpperCAmelCase_							:      bool = False		,       UpperCAmelCase_							:      bool = False		,       UpperCAmelCase_							:      bool = False		,       UpperCAmelCase_							:      bool = True		,       UpperCAmelCase_							:      str = "layer_norm"		,       UpperCAmelCase_							:      bool = False		,       )					->Tuple:



														"""simple docstring"""



														super().__init__()
														snake_case_           =		only_cross_attention

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

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

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



							def 		lowerCAmelCase       (				self							:      Any		,       UpperCAmelCase_							:      Optional[int]		,       UpperCAmelCase_							:      int       )					->str:



														"""simple docstring"""



														snake_case_           =		chunk_size
														snake_case_           =		dim



							def 		lowerCAmelCase       (				self							:      int		,       UpperCAmelCase_							:      torch.FloatTensor		,       UpperCAmelCase_							:      Optional[torch.FloatTensor] = None		,       UpperCAmelCase_							:      Optional[torch.FloatTensor] = None		,       UpperCAmelCase_							:      Optional[torch.FloatTensor] = None		,       UpperCAmelCase_							:      Optional[torch.LongTensor] = None		,       UpperCAmelCase_							:      Dict[str, Any] = None		,       UpperCAmelCase_							:      Optional[torch.LongTensor] = None		,       )					->Union[str, Any]:



														"""simple docstring"""



														if self.use_ada_layer_norm:
																					snake_case_           =		self.norma(UpperCAmelCase_		,       UpperCAmelCase_       )
														elif self.use_ada_layer_norm_zero:
																					snake_case_       ,						snake_case_       ,						snake_case_       ,						snake_case_       ,						snake_case_           =		self.norma(
																					    UpperCAmelCase_		,       UpperCAmelCase_		,       UpperCAmelCase_		,       hidden_dtype=hidden_states.dtype       )
														else:
																					snake_case_           =		self.norma(UpperCAmelCase_       )

														snake_case_           =		cross_attention_kwargs if cross_attention_kwargs is not None else {}

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

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

																					snake_case_           =		self.attna(
																					    UpperCAmelCase_		,       encoder_hidden_states=UpperCAmelCase_		,       attention_mask=UpperCAmelCase_		,       **UpperCAmelCase_		,       )
																					snake_case_           =		attn_output + hidden_states

														# 3. Feed-forward
														snake_case_           =		self.norma(UpperCAmelCase_       )

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

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

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

														snake_case_           =		ff_output + hidden_states

														return hidden_states





class      __A     (nn.Module):





							'''simple docstring'''

							def __init__(				self							:      Tuple		,       UpperCAmelCase_							:      int		,       UpperCAmelCase_							:      Optional[int] = None		,       UpperCAmelCase_							:      int = 4		,       UpperCAmelCase_							:      float = 0.0		,       UpperCAmelCase_							:      str = "geglu"		,       UpperCAmelCase_							:      bool = False		,       )					->List[str]:



														"""simple docstring"""



														super().__init__()
														snake_case_           =		int(dim * mult       )
														snake_case_           =		dim_out if dim_out is not None else dim

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

														snake_case_           =		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 		lowerCAmelCase       (				self							:      Optional[int]		,       UpperCAmelCase_							:      int       )					->Any:



														"""simple docstring"""



														for module in self.net:
																					snake_case_           =		module(UpperCAmelCase_       )
														return hidden_states





class      __A     (nn.Module):





							'''simple docstring'''

							def __init__(				self							:      Union[str, Any]		,       UpperCAmelCase_							:      int		,       UpperCAmelCase_							:      int		,       UpperCAmelCase_							:      str = "none"       )					->Dict:



														"""simple docstring"""



														super().__init__()
														snake_case_           =		nn.Linear(UpperCAmelCase_		,       UpperCAmelCase_       )
														snake_case_           =		approximate



							def 		lowerCAmelCase       (				self							:      Optional[int]		,       UpperCAmelCase_							:      Any       )					->int:



														"""simple docstring"""



														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 		lowerCAmelCase       (				self							:      Dict		,       UpperCAmelCase_							:      str       )					->int:



														"""simple docstring"""



														snake_case_           =		self.proj(UpperCAmelCase_       )
														snake_case_           =		self.gelu(UpperCAmelCase_       )
														return hidden_states





class      __A     (nn.Module):





							'''simple docstring'''

							def __init__(				self							:      List[Any]		,       UpperCAmelCase_							:      int		,       UpperCAmelCase_							:      int       )					->Union[str, Any]:



														"""simple docstring"""



														super().__init__()
														snake_case_           =		nn.Linear(UpperCAmelCase_		,       dim_out * 2       )



							def 		lowerCAmelCase       (				self							:      Union[str, Any]		,       UpperCAmelCase_							:      Tuple       )					->Optional[Any]:



														"""simple docstring"""



														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 		lowerCAmelCase       (				self							:      List[Any]		,       UpperCAmelCase_							:      int       )					->str:



														"""simple docstring"""



														snake_case_       ,						snake_case_           =		self.proj(UpperCAmelCase_       ).chunk(2		,       dim=-1       )
														return hidden_states * self.gelu(UpperCAmelCase_       )





class      __A     (nn.Module):





							'''simple docstring'''

							def __init__(				self							:      Optional[int]		,       UpperCAmelCase_							:      int		,       UpperCAmelCase_							:      int       )					->Union[str, Any]:



														"""simple docstring"""



														super().__init__()
														snake_case_           =		nn.Linear(UpperCAmelCase_		,       UpperCAmelCase_       )



							def 		lowerCAmelCase       (				self							:      Any		,       UpperCAmelCase_							:      Union[str, Any]       )					->Union[str, Any]:



														"""simple docstring"""



														snake_case_           =		self.proj(UpperCAmelCase_       )
														return x * torch.sigmoid(1.702 * x       )





class      __A     (nn.Module):





							'''simple docstring'''

							def __init__(				self							:      Optional[int]		,       UpperCAmelCase_							:      List[Any]		,       UpperCAmelCase_							:      int       )					->Union[str, Any]:



														"""simple docstring"""



														super().__init__()
														snake_case_           =		nn.Embedding(UpperCAmelCase_		,       UpperCAmelCase_       )
														snake_case_           =		nn.SiLU()
														snake_case_           =		nn.Linear(UpperCAmelCase_		,       embedding_dim * 2       )
														snake_case_           =		nn.LayerNorm(UpperCAmelCase_		,       elementwise_affine=UpperCAmelCase_       )



							def 		lowerCAmelCase       (				self							:      Tuple		,       UpperCAmelCase_							:      Optional[int]		,       UpperCAmelCase_							:      Tuple       )					->int:



														"""simple docstring"""



														snake_case_           =		self.linear(self.silu(self.emb(UpperCAmelCase_       )       )       )
														snake_case_       ,						snake_case_           =		torch.chunk(UpperCAmelCase_		,       2       )
														snake_case_           =		self.norm(UpperCAmelCase_       ) * (1 + scale) + shift
														return x





class      __A     (nn.Module):





							'''simple docstring'''

							def __init__(				self							:      List[str]		,       UpperCAmelCase_							:      int		,       UpperCAmelCase_							:      List[str]       )					->Tuple:



														"""simple docstring"""



														super().__init__()

														snake_case_           =		CombinedTimestepLabelEmbeddings(UpperCAmelCase_		,       UpperCAmelCase_       )

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



							def 		lowerCAmelCase       (				self							:      List[str]		,       UpperCAmelCase_							:      int		,       UpperCAmelCase_							:      Optional[Any]		,       UpperCAmelCase_							:      Tuple		,       UpperCAmelCase_							:      Any=None       )					->Any:



														"""simple docstring"""



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





class      __A     (nn.Module):





							'''simple docstring'''

							def __init__(				self							:      int		,       UpperCAmelCase_							:      int		,       UpperCAmelCase_							:      int		,       UpperCAmelCase_							:      int		,       UpperCAmelCase_							:      Optional[str] = None		,       UpperCAmelCase_							:      float = 1E-5       )					->str:



														"""simple docstring"""



														super().__init__()
														snake_case_           =		num_groups
														snake_case_           =		eps

														if act_fn is None:
																					snake_case_           =		None
														else:
																					snake_case_           =		get_activation(UpperCAmelCase_       )

														snake_case_           =		nn.Linear(UpperCAmelCase_		,       out_dim * 2       )



							def 		lowerCAmelCase       (				self							:      List[str]		,       UpperCAmelCase_							:      Any		,       UpperCAmelCase_							:      List[str]       )					->Dict:



														"""simple docstring"""



														if self.act:
																					snake_case_           =		self.act(UpperCAmelCase_       )
														snake_case_           =		self.linear(UpperCAmelCase_       )
														snake_case_           =		emb[:, :, None, None]
														snake_case_       ,						snake_case_           =		emb.chunk(2		,       dim=1       )

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







"""simple docstring"""
import unittest

import torch

from diffusers import VQModel
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()





class      __A     (snake_case__						,		snake_case__						,		unittest.TestCase):





							'''simple docstring'''

							__lowercase:		str							      =    VQModel
							__lowercase:		Union[str, Any]							      =    """sample"""



							@property
							def 		lowerCAmelCase       (				self							:      List[str]		,       UpperCAmelCase_							:      List[str]=(32, 32)       )					->Tuple:



														"""simple docstring"""



														snake_case_           =		4
														snake_case_           =		3

														snake_case_           =		floats_tensor((batch_size, num_channels) + sizes       ).to(UpperCAmelCase_       )

														return {"sample": image}



							@property
							def 		lowerCAmelCase       (				self							:      Tuple       )					->str:



														"""simple docstring"""



														return (3, 32, 32)



							@property
							def 		lowerCAmelCase       (				self							:      List[Any]       )					->Any:



														"""simple docstring"""



														return (3, 32, 32)



							def 		lowerCAmelCase       (				self							:      Optional[int]       )					->Dict:



														"""simple docstring"""



														snake_case_           =		{
														    """block_out_channels""": [32, 64],
														    """in_channels""": 3,
														    """out_channels""": 3,
														    """down_block_types""": ["""DownEncoderBlock2D""", """DownEncoderBlock2D"""],
														    """up_block_types""": ["""UpDecoderBlock2D""", """UpDecoderBlock2D"""],
														    """latent_channels""": 3,
														}
														snake_case_           =		self.dummy_input
														return init_dict, inputs_dict



							def 		lowerCAmelCase       (				self							:      List[str]       )					->Dict:



														"""simple docstring"""



														pass



							def 		lowerCAmelCase       (				self							:      Optional[Any]       )					->List[str]:



														"""simple docstring"""



														pass



							def 		lowerCAmelCase       (				self							:      Any       )					->Union[str, Any]:



														"""simple docstring"""



														snake_case_       ,						snake_case_           =		VQModel.from_pretrained("""fusing/vqgan-dummy"""		,       output_loading_info=UpperCAmelCase_       )
														self.assertIsNotNone(UpperCAmelCase_       )
														self.assertEqual(len(loading_info["""missing_keys"""]       )		,       0       )

														model.to(UpperCAmelCase_       )
														snake_case_           =		model(**self.dummy_input       )

														assert image is not None, "Make sure output is not None"



							def 		lowerCAmelCase       (				self							:      Tuple       )					->Optional[Any]:



														"""simple docstring"""



														snake_case_           =		VQModel.from_pretrained("""fusing/vqgan-dummy"""       )
														model.to(UpperCAmelCase_       ).eval()

														torch.manual_seed(0       )
														if torch.cuda.is_available():
																					torch.cuda.manual_seed_all(0       )

														snake_case_           =		torch.randn(1		,       model.config.in_channels		,       model.config.sample_size		,       model.config.sample_size       )
														snake_case_           =		image.to(UpperCAmelCase_       )
														with torch.no_grad():
																					snake_case_           =		model(UpperCAmelCase_       ).sample

														snake_case_           =		output[0, -1, -3:, -3:].flatten().cpu()
														# fmt: off
														snake_case_           =		torch.tensor([-0.0_153, -0.4_044, -0.1_880, -0.5_161, -0.2_418, -0.4_072, -0.1_612, -0.0_633, -0.0_143]       )
														# fmt: on
														self.assertTrue(torch.allclose(UpperCAmelCase_		,       UpperCAmelCase_		,       atol=1E-3       )       )



