






'''simple docstring'''







from dataclasses import dataclass
from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..models.embeddings import ImagePositionalEmbeddings
from ..utils import BaseOutput, deprecate
from .attention import BasicTransformerBlock
from .embeddings import PatchEmbed
from .modeling_utils import ModelMixin




@dataclass
class 					__UpperCAmelCase (		UpperCamelCase__	):

      '''simple docstring'''





      SCREAMING_SNAKE_CASE				:    torch.FloatTensor




class 					__UpperCAmelCase (		UpperCamelCase__      ,				UpperCamelCase__	):

      '''simple docstring'''





      @register_to_config
      def __init__(					self  ,		_UpperCAmelCase = 16  ,		_UpperCAmelCase = 88  ,		_UpperCAmelCase = None  ,		_UpperCAmelCase = None  ,		_UpperCAmelCase = 1  ,		_UpperCAmelCase = 0.0  ,		_UpperCAmelCase = 32  ,		_UpperCAmelCase = None  ,		_UpperCAmelCase = False  ,		_UpperCAmelCase = None  ,		_UpperCAmelCase = None  ,		_UpperCAmelCase = None  ,		_UpperCAmelCase = "geglu"  ,		_UpperCAmelCase = None  ,		_UpperCAmelCase = False  ,		_UpperCAmelCase = False  ,		_UpperCAmelCase = False  ,		_UpperCAmelCase = "layer_norm"  ,		_UpperCAmelCase = True  ,		):
             super().__init__()
             UpperCAmelCase__	:   List[Any]												=	use_linear_projection
             UpperCAmelCase__	:   str												=	num_attention_heads
             UpperCAmelCase__	:   int												=	attention_head_dim
             UpperCAmelCase__	:   Union[str, Any]												=	num_attention_heads * attention_head_dim

             # 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
             # Define whether input is continuous or discrete depending on configuration
             UpperCAmelCase__	:   Dict												=	(in_channels is not None) and (patch_size is None)
             UpperCAmelCase__	:   Optional[Any]												=	num_vector_embeds is not None
             UpperCAmelCase__	:   Optional[int]												=	in_channels is not None and patch_size is not None

             if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
                    UpperCAmelCase__	:   Union[str, Any]												=	(
                        F"""The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"""
                        ''' incorrectly set to `\'layer_norm\'`.Make sure to set `norm_type` to `\'ada_norm\'` in the config.'''
                        ''' Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect'''
                        ''' results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it'''
                        ''' would be very nice if you could open a Pull request for the `transformer/config.json` file'''
                    )
                    deprecate('''norm_type!=num_embeds_ada_norm'''  ,		'''1.0.0'''  ,		_UpperCAmelCase  ,		standard_warn=_UpperCAmelCase						)
                    UpperCAmelCase__	:   Union[str, Any]												=	'''ada_norm'''

             if self.is_input_continuous and self.is_input_vectorized:
                    raise ValueError(
                        F"""Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make"""
                        ''' sure that either `in_channels` or `num_vector_embeds` is None.'''						)
             elif self.is_input_vectorized and self.is_input_patches:
                    raise ValueError(
                        F"""Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make"""
                        ''' sure that either `num_vector_embeds` or `num_patches` is None.'''						)
             elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:
                    raise ValueError(
                        F"""Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:"""
                        F""" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None."""						)

             # 2. Define input layers
             if self.is_input_continuous:
                    UpperCAmelCase__	:   Tuple												=	in_channels

                    UpperCAmelCase__	:   Optional[int]												=	torch.nn.GroupNorm(num_groups=_UpperCAmelCase  ,		num_channels=_UpperCAmelCase  ,		eps=1E-6  ,		affine=_UpperCAmelCase						)
                    if use_linear_projection:
                           UpperCAmelCase__	:   List[str]												=	nn.Linear(_UpperCAmelCase  ,		_UpperCAmelCase						)
                    else:
                           UpperCAmelCase__	:   Union[str, Any]												=	nn.Convad(_UpperCAmelCase  ,		_UpperCAmelCase  ,		kernel_size=1  ,		stride=1  ,		padding=0						)
             elif self.is_input_vectorized:
                    assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
                    assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"

                    UpperCAmelCase__	:   List[str]												=	sample_size
                    UpperCAmelCase__	:   int												=	sample_size
                    UpperCAmelCase__	:   Dict												=	num_vector_embeds
                    UpperCAmelCase__	:   str												=	self.height * self.width

                    UpperCAmelCase__	:   Union[str, Any]												=	ImagePositionalEmbeddings(
                        num_embed=_UpperCAmelCase  ,		embed_dim=_UpperCAmelCase  ,		height=self.height  ,		width=self.width						)
             elif self.is_input_patches:
                    assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"

                    UpperCAmelCase__	:   Union[str, Any]												=	sample_size
                    UpperCAmelCase__	:   List[Any]												=	sample_size

                    UpperCAmelCase__	:   Optional[Any]												=	patch_size
                    UpperCAmelCase__	:   Optional[Any]												=	PatchEmbed(
                        height=_UpperCAmelCase  ,		width=_UpperCAmelCase  ,		patch_size=_UpperCAmelCase  ,		in_channels=_UpperCAmelCase  ,		embed_dim=_UpperCAmelCase  ,		)

             # 3. Define transformers blocks
             UpperCAmelCase__	:   List[str]												=	nn.ModuleList(
                 [
                     BasicTransformerBlock(
                         _UpperCAmelCase  ,		_UpperCAmelCase  ,		_UpperCAmelCase  ,		dropout=_UpperCAmelCase  ,		cross_attention_dim=_UpperCAmelCase  ,		activation_fn=_UpperCAmelCase  ,		num_embeds_ada_norm=_UpperCAmelCase  ,		attention_bias=_UpperCAmelCase  ,		only_cross_attention=_UpperCAmelCase  ,		upcast_attention=_UpperCAmelCase  ,		norm_type=_UpperCAmelCase  ,		norm_elementwise_affine=_UpperCAmelCase  ,		)
                     for d in range(_UpperCAmelCase						)
                 ]						)

             # 4. Define output layers
             UpperCAmelCase__	:   Dict												=	in_channels if out_channels is None else out_channels
             if self.is_input_continuous:
                    # TODO: should use out_channels for continuous projections
                    if use_linear_projection:
                           UpperCAmelCase__	:   int												=	nn.Linear(_UpperCAmelCase  ,		_UpperCAmelCase						)
                    else:
                           UpperCAmelCase__	:   Any												=	nn.Convad(_UpperCAmelCase  ,		_UpperCAmelCase  ,		kernel_size=1  ,		stride=1  ,		padding=0						)
             elif self.is_input_vectorized:
                    UpperCAmelCase__	:   Optional[int]												=	nn.LayerNorm(_UpperCAmelCase						)
                    UpperCAmelCase__	:   str												=	nn.Linear(_UpperCAmelCase  ,		self.num_vector_embeds - 1						)
             elif self.is_input_patches:
                    UpperCAmelCase__	:   Dict												=	nn.LayerNorm(_UpperCAmelCase  ,		elementwise_affine=_UpperCAmelCase  ,		eps=1E-6						)
                    UpperCAmelCase__	:   Dict												=	nn.Linear(_UpperCAmelCase  ,		2 * inner_dim						)
                    UpperCAmelCase__	:   str												=	nn.Linear(_UpperCAmelCase  ,		patch_size * patch_size * self.out_channels						)




      def 							lowerCamelCase     (					self  ,		_UpperCAmelCase  ,		_UpperCAmelCase = None  ,		_UpperCAmelCase = None  ,		_UpperCAmelCase = None  ,		_UpperCAmelCase = None  ,		_UpperCAmelCase = None  ,		_UpperCAmelCase = None  ,		_UpperCAmelCase = True  ,		):
             # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
             #   we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
             #   we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
             # expects mask of shape:
             #   [batch, key_tokens]
             # adds singleton query_tokens dimension:
             #   [batch,                    1, key_tokens]
             # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
             #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
             #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
             if attention_mask is not None and attention_mask.ndim == 2:
                    # assume that mask is expressed as:
                    #   (1 = keep,      0 = discard)
                    # convert mask into a bias that can be added to attention scores:
                    #       (keep = +0,     discard = -10000.0)
                    UpperCAmelCase__	:   str												=	(1 - attention_mask.to(hidden_states.dtype						)) * -1_00_00.0
                    UpperCAmelCase__	:   Dict												=	attention_mask.unsqueeze(1						)

             # convert encoder_attention_mask to a bias the same way we do for attention_mask
             if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
                    UpperCAmelCase__	:   Any												=	(1 - encoder_attention_mask.to(hidden_states.dtype						)) * -1_00_00.0
                    UpperCAmelCase__	:   List[Any]												=	encoder_attention_mask.unsqueeze(1						)

             # 1. Input
             if self.is_input_continuous:
                    UpperCAmelCase__							,							UpperCAmelCase__							,							UpperCAmelCase__							,							UpperCAmelCase__	:   Optional[int]												=	hidden_states.shape
                    UpperCAmelCase__	:   int												=	hidden_states

                    UpperCAmelCase__	:   Dict												=	self.norm(_UpperCAmelCase						)
                    if not self.use_linear_projection:
                           UpperCAmelCase__	:   List[Any]												=	self.proj_in(_UpperCAmelCase						)
                           UpperCAmelCase__	:   Any												=	hidden_states.shape[1]
                           UpperCAmelCase__	:   Tuple												=	hidden_states.permute(0  ,		2  ,		3  ,		1						).reshape(_UpperCAmelCase  ,		height * width  ,		_UpperCAmelCase						)
                    else:
                           UpperCAmelCase__	:   Union[str, Any]												=	hidden_states.shape[1]
                           UpperCAmelCase__	:   Dict												=	hidden_states.permute(0  ,		2  ,		3  ,		1						).reshape(_UpperCAmelCase  ,		height * width  ,		_UpperCAmelCase						)
                           UpperCAmelCase__	:   List[str]												=	self.proj_in(_UpperCAmelCase						)
             elif self.is_input_vectorized:
                    UpperCAmelCase__	:   int												=	self.latent_image_embedding(_UpperCAmelCase						)
             elif self.is_input_patches:
                    UpperCAmelCase__	:   Optional[int]												=	self.pos_embed(_UpperCAmelCase						)

             # 2. Blocks
             for block in self.transformer_blocks:
                    UpperCAmelCase__	:   Optional[int]												=	block(
                        _UpperCAmelCase  ,		attention_mask=_UpperCAmelCase  ,		encoder_hidden_states=_UpperCAmelCase  ,		encoder_attention_mask=_UpperCAmelCase  ,		timestep=_UpperCAmelCase  ,		cross_attention_kwargs=_UpperCAmelCase  ,		class_labels=_UpperCAmelCase  ,		)

             # 3. Output
             if self.is_input_continuous:
                    if not self.use_linear_projection:
                           UpperCAmelCase__	:   Any												=	hidden_states.reshape(_UpperCAmelCase  ,		_UpperCAmelCase  ,		_UpperCAmelCase  ,		_UpperCAmelCase						).permute(0  ,		3  ,		1  ,		2						).contiguous()
                           UpperCAmelCase__	:   int												=	self.proj_out(_UpperCAmelCase						)
                    else:
                           UpperCAmelCase__	:   Tuple												=	self.proj_out(_UpperCAmelCase						)
                           UpperCAmelCase__	:   Tuple												=	hidden_states.reshape(_UpperCAmelCase  ,		_UpperCAmelCase  ,		_UpperCAmelCase  ,		_UpperCAmelCase						).permute(0  ,		3  ,		1  ,		2						).contiguous()

                    UpperCAmelCase__	:   Tuple												=	hidden_states + residual
             elif self.is_input_vectorized:
                    UpperCAmelCase__	:   str												=	self.norm_out(_UpperCAmelCase						)
                    UpperCAmelCase__	:   List[Any]												=	self.out(_UpperCAmelCase						)
                    # (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
                    UpperCAmelCase__	:   List[Any]												=	logits.permute(0  ,		2  ,		1						)

                    # log(p(x_0))
                    UpperCAmelCase__	:   List[str]												=	F.log_softmax(logits.double()  ,		dim=1						).float()
             elif self.is_input_patches:
                    # TODO: cleanup!
                    UpperCAmelCase__	:   str												=	self.transformer_blocks[0].norma.emb(
                        _UpperCAmelCase  ,		_UpperCAmelCase  ,		hidden_dtype=hidden_states.dtype						)
                    UpperCAmelCase__							,							UpperCAmelCase__	:   int												=	self.proj_out_a(F.silu(_UpperCAmelCase						)						).chunk(2  ,		dim=1						)
                    UpperCAmelCase__	:   Optional[Any]												=	self.norm_out(_UpperCAmelCase						) * (1 + scale[:, None]) + shift[:, None]
                    UpperCAmelCase__	:   Tuple												=	self.proj_out_a(_UpperCAmelCase						)

                    # unpatchify
                    UpperCAmelCase__	:   Optional[int]												=	int(hidden_states.shape[1] ** 0.5						)
                    UpperCAmelCase__	:   Dict												=	hidden_states.reshape(
                        shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)						)
                    UpperCAmelCase__	:   Optional[int]												=	torch.einsum('''nhwpqc->nchpwq'''  ,		_UpperCAmelCase						)
                    UpperCAmelCase__	:   int												=	hidden_states.reshape(
                        shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)						)

             if not return_dict:
                    return (output,)

             return TransformeraDModelOutput(sample=_UpperCAmelCase						)







'''simple docstring'''







import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
  import torch

  from transformers import CamembertModel




@require_torch
@require_sentencepiece
@require_tokenizers
class 					__UpperCAmelCase (		unittest.TestCase	):

      '''simple docstring'''





      @slow
      def 							lowerCamelCase     (					self						):
             UpperCAmelCase__	:   Optional[int]												=	CamembertModel.from_pretrained('''camembert-base'''						)
             model.to(_UpperCAmelCase						)

             UpperCAmelCase__	:   Optional[int]												=	torch.tensor(
                 [[5, 121, 11, 660, 16, 730, 2_5543, 110, 83, 6]]  ,		device=_UpperCAmelCase  ,		dtype=torch.long  ,		)  # J'aime le camembert !
             with torch.no_grad():
                    UpperCAmelCase__	:   Tuple												=	model(_UpperCAmelCase						)['''last_hidden_state''']
             UpperCAmelCase__	:   List[str]												=	torch.Size((1, 10, 768)						)
             self.assertEqual(output.shape  ,		_UpperCAmelCase						)
             # compare the actual values for a slice.
             UpperCAmelCase__	:   List[Any]												=	torch.tensor(
                 [[[-0.02_54, 0.02_35, 0.10_27], [0.06_06, -0.18_11, -0.04_18], [-0.15_61, -0.11_27, 0.26_87]]]  ,		device=_UpperCAmelCase  ,		dtype=torch.float  ,		)
             # camembert = torch.hub.load('pytorch/fairseq', 'camembert.v0')
             # camembert.eval()
             # expected_slice = roberta.model.forward(input_ids)[0][:, :3, :3].detach()

             self.assertTrue(torch.allclose(output[:, :3, :3]  ,		_UpperCAmelCase  ,		atol=1E-4						)						)