


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 ,						__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 ,						)			->				Optional[int]:
            super().__init__()
            _A      :       Optional[Any]				    = only_cross_attention

            _A      :       int				    = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"
            _A      :       List[Any]				    = (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      :       Any				    = AdaLayerNorm(__lowerCamelCase ,						__lowerCamelCase)
            elif self.use_ada_layer_norm_zero:
                   _A      :       Optional[int]				    = AdaLayerNormZero(__lowerCamelCase ,						__lowerCamelCase)
            else:
                   _A      :       str				    = nn.LayerNorm(__lowerCamelCase ,						elementwise_affine=__lowerCamelCase)
            _A      :       int				    = 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      :       int				    = (
                       AdaLayerNorm(__lowerCamelCase ,						__lowerCamelCase)
                       if self.use_ada_layer_norm
                       else nn.LayerNorm(__lowerCamelCase ,						elementwise_affine=__lowerCamelCase)
                   )
                   _A      :       Optional[int]				    = 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      :       Tuple				    = None
                   _A      :       List[Any]				    = None

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

            # let chunk size default to None
            _A      :       Optional[int]				    = None
            _A      :       List[Any]				    = 0


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


     def  _lowerCamelCase (						self ,						__lowerCamelCase ,						__lowerCamelCase = None ,						__lowerCamelCase = None ,						__lowerCamelCase = None ,						__lowerCamelCase = None ,						__lowerCamelCase = None ,						__lowerCamelCase = None ,						)			->				int:
            # Notice that normalization is always applied before the real computation in the following blocks.
            # 1. Self-Attention
            if self.use_ada_layer_norm:
                   _A      :       Union[str, Any]				    = self.norma(__lowerCamelCase ,						__lowerCamelCase)
            elif self.use_ada_layer_norm_zero:
                   _A    ,			_A    ,			_A    ,			_A    ,			_A      :       Tuple				    = self.norma(
                       __lowerCamelCase ,						__lowerCamelCase ,						__lowerCamelCase ,						hidden_dtype=hidden_states.dtype)
            else:
                   _A      :       Tuple				    = self.norma(__lowerCamelCase)

            _A      :       Dict				    = cross_attention_kwargs if cross_attention_kwargs is not None else {}

            _A      :       Optional[Any]				    = 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      :       int				    = gate_msa.unsqueeze(1) * attn_output
            _A      :       str				    = attn_output + hidden_states

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

                   _A      :       Union[str, Any]				    = self.attna(
                       __lowerCamelCase ,						encoder_hidden_states=__lowerCamelCase ,						attention_mask=__lowerCamelCase ,						**__lowerCamelCase ,						)
                   _A      :       Any				    = attn_output + hidden_states

            # 3. Feed-forward
            _A      :       str				    = self.norma(__lowerCamelCase)

            if self.use_ada_layer_norm_zero:
                   _A      :       Optional[Any]				    = 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      :       Optional[Any]				    = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
                   _A      :       int				    = torch.cat(
                       [self.ff(__lowerCamelCase) for hid_slice in norm_hidden_states.chunk(__lowerCamelCase ,						dim=self._chunk_dim)] ,						dim=self._chunk_dim ,						)
            else:
                   _A      :       List[str]				    = self.ff(__lowerCamelCase)

            if self.use_ada_layer_norm_zero:
                   _A      :       Optional[int]				    = gate_mlp.unsqueeze(1) * ff_output

            _A      :       Optional[int]				    = ff_output + hidden_states

            return hidden_states







class      lowerCAmelCase__       (				nn.Module):




     '''simple docstring'''


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

            if activation_fn == "gelu":
                   _A      :       Tuple				    = GELU(__lowerCamelCase ,						__lowerCamelCase)
            if activation_fn == "gelu-approximate":
                   _A      :       List[Any]				    = GELU(__lowerCamelCase ,						__lowerCamelCase ,						approximate="tanh")
            elif activation_fn == "geglu":
                   _A      :       Optional[Any]				    = GEGLU(__lowerCamelCase ,						__lowerCamelCase)
            elif activation_fn == "geglu-approximate":
                   _A      :       Any				    = ApproximateGELU(__lowerCamelCase ,						__lowerCamelCase)

            _A      :       List[str]				    = 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)			->				List[str]:
            for module in self.net:
                   _A      :       List[Any]				    = module(__lowerCamelCase)
            return hidden_states







class      lowerCAmelCase__       (				nn.Module):




     '''simple docstring'''


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


     def  _lowerCamelCase (						self ,						__lowerCamelCase)			->				Optional[int]:
            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)			->				Union[str, Any]:
            _A      :       Dict				    = self.proj(__lowerCamelCase)
            _A      :       int				    = self.gelu(__lowerCamelCase)
            return hidden_states







class      lowerCAmelCase__       (				nn.Module):




     '''simple docstring'''


     def __init__(						self ,						__lowerCamelCase ,						__lowerCamelCase)			->				Union[str, Any]:
            super().__init__()
            _A      :       int				    = nn.Linear(__lowerCamelCase ,						dim_out * 2)


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







class      lowerCAmelCase__       (				nn.Module):




     '''simple docstring'''


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


     def  _lowerCamelCase (						self ,						__lowerCamelCase)			->				Any:
            _A      :       Any				    = self.proj(__lowerCamelCase)
            return x * torch.sigmoid(1.7_0_2 * x)







class      lowerCAmelCase__       (				nn.Module):




     '''simple docstring'''


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


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







class      lowerCAmelCase__       (				nn.Module):




     '''simple docstring'''


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

            _A      :       str				    = CombinedTimestepLabelEmbeddings(__lowerCamelCase ,						__lowerCamelCase)

            _A      :       Optional[int]				    = nn.SiLU()
            _A      :       Dict				    = nn.Linear(__lowerCamelCase ,						6 * embedding_dim ,						bias=__lowerCamelCase)
            _A      :       int				    = nn.LayerNorm(__lowerCamelCase ,						elementwise_affine=__lowerCamelCase ,						eps=1e-6)


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

            if act_fn is None:
                   _A      :       Tuple				    = None
            else:
                   _A      :       str				    = get_activation(__lowerCamelCase)

            _A      :       List[Any]				    = nn.Linear(__lowerCamelCase ,						out_dim * 2)


     def  _lowerCamelCase (						self ,						__lowerCamelCase ,						__lowerCamelCase)			->				Optional[Any]:
            if self.act:
                   _A      :       Tuple				    = self.act(__lowerCamelCase)
            _A      :       Optional[Any]				    = self.linear(__lowerCamelCase)
            _A      :       Tuple				    = emb[:, :, None, None]
            _A    ,			_A      :       Union[str, Any]				    = emb.chunk(2 ,						dim=1)

            _A      :       List[str]				    = F.group_norm(__lowerCamelCase ,						self.num_groups ,						eps=self.eps)
            _A      :       int				    = x * (1 + scale) + shift
            return x




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      lowerCAmelCase__       (				a     ,    a     ,    unittest.TestCase):




     '''simple docstring'''
     __SCREAMING_SNAKE_CASE  			= VQModel
     __SCREAMING_SNAKE_CASE  			= "sample"


     @property
     def  _lowerCamelCase (						self ,						__lowerCamelCase=(3_2, 3_2))			->				Optional[Any]:
            _A      :       Optional[int]				    = 4
            _A      :       Tuple				    = 3

            _A      :       List[Any]				    = floats_tensor((batch_size, num_channels) + sizes).to(__lowerCamelCase)

            return {"sample": image}


     @property
     def  _lowerCamelCase (						self)			->				int:
            return (3, 3_2, 3_2)


     @property
     def  _lowerCamelCase (						self)			->				List[Any]:
            return (3, 3_2, 3_2)


     def  _lowerCamelCase (						self)			->				Union[str, Any]:
            _A      :       List[Any]				    = {
                "block_out_channels": [3_2, 6_4],
                "in_channels": 3,
                "out_channels": 3,
                "down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D"],
                "up_block_types": ["UpDecoderBlock2D", "UpDecoderBlock2D"],
                "latent_channels": 3,
            }
            _A      :       int				    = self.dummy_input
            return init_dict, inputs_dict


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


     def  _lowerCamelCase (						self)			->				Any:
            pass


     def  _lowerCamelCase (						self)			->				Any:
            _A    ,			_A      :       List[Any]				    = VQModel.from_pretrained("fusing/vqgan-dummy" ,						output_loading_info=__lowerCamelCase)
            self.assertIsNotNone(__lowerCamelCase)
            self.assertEqual(len(loading_info["missing_keys"]) ,						0)

            model.to(__lowerCamelCase)
            _A      :       str				    = model(**self.dummy_input)

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


     def  _lowerCamelCase (						self)			->				Union[str, Any]:
            _A      :       Optional[Any]				    = VQModel.from_pretrained("fusing/vqgan-dummy")
            model.to(__lowerCamelCase).eval()

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

            _A      :       Tuple				    = torch.randn(1 ,						model.config.in_channels ,						model.config.sample_size ,						model.config.sample_size)
            _A      :       Optional[int]				    = image.to(__lowerCamelCase)
            with torch.no_grad():
                   _A      :       List[str]				    = model(__lowerCamelCase).sample

            _A      :       int				    = output[0, -1, -3:, -3:].flatten().cpu()
            # fmt: off
            _A      :       Optional[Any]				    = torch.tensor([-0.0_1_5_3, -0.4_0_4_4, -0.1_8_8_0, -0.5_1_6_1, -0.2_4_1_8, -0.4_0_7_2, -0.1_6_1_2, -0.0_6_3_3, -0.0_1_4_3])
            # fmt: on
            self.assertTrue(torch.allclose(__lowerCamelCase ,						__lowerCamelCase ,						atol=1e-3))
