




'''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 (				lowerCamelCase__       ):
       lowercase			:					torch.FloatTensor





class 						__UpperCamelCase (				lowerCamelCase__    ,      lowerCamelCase__       ):






       @register_to_config
       def __init__(							self,	lowerCAmelCase = 16,	lowerCAmelCase = 88,	lowerCAmelCase = None,	lowerCAmelCase = None,	lowerCAmelCase = 1,	lowerCAmelCase = 0.0,	lowerCAmelCase = 32,	lowerCAmelCase = None,	lowerCAmelCase = False,	lowerCAmelCase = None,	lowerCAmelCase = None,	lowerCAmelCase = None,	lowerCAmelCase = "geglu",	lowerCAmelCase = None,	lowerCAmelCase = False,	lowerCAmelCase = False,	lowerCAmelCase = False,	lowerCAmelCase = "layer_norm",	lowerCAmelCase = True,	):



              """simple docstring"""







              super().__init__()
              lowerCamelCase_							=use_linear_projection
              lowerCamelCase_							=num_attention_heads
              lowerCamelCase_							=attention_head_dim
              lowerCamelCase_							=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
              lowerCamelCase_							=(in_channels is not None) and (patch_size is None)
              lowerCamelCase_							=num_vector_embeds is not None
              lowerCamelCase_							=in_channels is not None and patch_size is not None

              if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
                     lowerCamelCase_							=(
                         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''',	lowerCAmelCase,	standard_warn=lowerCAmelCase						)
                     lowerCamelCase_							='''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:
                     lowerCamelCase_							=in_channels

                     lowerCamelCase_							=torch.nn.GroupNorm(num_groups=lowerCAmelCase,	num_channels=lowerCAmelCase,	eps=1e-6,	affine=lowerCAmelCase						)
                     if use_linear_projection:
                            lowerCamelCase_							=nn.Linear(lowerCAmelCase,	lowerCAmelCase						)
                     else:
                            lowerCamelCase_							=nn.Convad(lowerCAmelCase,	lowerCAmelCase,	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"

                     lowerCamelCase_							=sample_size
                     lowerCamelCase_							=sample_size
                     lowerCamelCase_							=num_vector_embeds
                     lowerCamelCase_							=self.height * self.width

                     lowerCamelCase_							=ImagePositionalEmbeddings(
                         num_embed=lowerCAmelCase,	embed_dim=lowerCAmelCase,	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"

                     lowerCamelCase_							=sample_size
                     lowerCamelCase_							=sample_size

                     lowerCamelCase_							=patch_size
                     lowerCamelCase_							=PatchEmbed(
                         height=lowerCAmelCase,	width=lowerCAmelCase,	patch_size=lowerCAmelCase,	in_channels=lowerCAmelCase,	embed_dim=lowerCAmelCase,	)

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

              # 4. Define output layers
              lowerCamelCase_							=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:
                            lowerCamelCase_							=nn.Linear(lowerCAmelCase,	lowerCAmelCase						)
                     else:
                            lowerCamelCase_							=nn.Convad(lowerCAmelCase,	lowerCAmelCase,	kernel_size=1,	stride=1,	padding=0						)
              elif self.is_input_vectorized:
                     lowerCamelCase_							=nn.LayerNorm(lowerCAmelCase						)
                     lowerCamelCase_							=nn.Linear(lowerCAmelCase,	self.num_vector_embeds - 1						)
              elif self.is_input_patches:
                     lowerCamelCase_							=nn.LayerNorm(lowerCAmelCase,	elementwise_affine=lowerCAmelCase,	eps=1e-6						)
                     lowerCamelCase_							=nn.Linear(lowerCAmelCase,	2 * inner_dim						)
                     lowerCamelCase_							=nn.Linear(lowerCAmelCase,	patch_size * patch_size * self.out_channels						)






       def   lowercase__  (							self,	lowerCAmelCase,	lowerCAmelCase = None,	lowerCAmelCase = None,	lowerCAmelCase = None,	lowerCAmelCase = None,	lowerCAmelCase = None,	lowerCAmelCase = None,	lowerCAmelCase = True,	):



              """simple docstring"""







              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)
                     lowerCamelCase_							=(1 - attention_mask.to(hidden_states.dtype						)) * -10_000.0
                     lowerCamelCase_							=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:
                     lowerCamelCase_							=(1 - encoder_attention_mask.to(hidden_states.dtype						)) * -10_000.0
                     lowerCamelCase_							=encoder_attention_mask.unsqueeze(1						)

              # 1. Input
              if self.is_input_continuous:
                     lowerCamelCase_,   lowerCamelCase_,   lowerCamelCase_,   lowerCamelCase_							=hidden_states.shape
                     lowerCamelCase_							=hidden_states

                     lowerCamelCase_							=self.norm(lowerCAmelCase						)
                     if not self.use_linear_projection:
                            lowerCamelCase_							=self.proj_in(lowerCAmelCase						)
                            lowerCamelCase_							=hidden_states.shape[1]
                            lowerCamelCase_							=hidden_states.permute(0,	2,	3,	1						).reshape(lowerCAmelCase,	height * width,	lowerCAmelCase						)
                     else:
                            lowerCamelCase_							=hidden_states.shape[1]
                            lowerCamelCase_							=hidden_states.permute(0,	2,	3,	1						).reshape(lowerCAmelCase,	height * width,	lowerCAmelCase						)
                            lowerCamelCase_							=self.proj_in(lowerCAmelCase						)
              elif self.is_input_vectorized:
                     lowerCamelCase_							=self.latent_image_embedding(lowerCAmelCase						)
              elif self.is_input_patches:
                     lowerCamelCase_							=self.pos_embed(lowerCAmelCase						)

              # 2. Blocks
              for block in self.transformer_blocks:
                     lowerCamelCase_							=block(
                         lowerCAmelCase,	attention_mask=lowerCAmelCase,	encoder_hidden_states=lowerCAmelCase,	encoder_attention_mask=lowerCAmelCase,	timestep=lowerCAmelCase,	cross_attention_kwargs=lowerCAmelCase,	class_labels=lowerCAmelCase,	)

              # 3. Output
              if self.is_input_continuous:
                     if not self.use_linear_projection:
                            lowerCamelCase_							=hidden_states.reshape(lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase						).permute(0,	3,	1,	2						).contiguous()
                            lowerCamelCase_							=self.proj_out(lowerCAmelCase						)
                     else:
                            lowerCamelCase_							=self.proj_out(lowerCAmelCase						)
                            lowerCamelCase_							=hidden_states.reshape(lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase						).permute(0,	3,	1,	2						).contiguous()

                     lowerCamelCase_							=hidden_states + residual
              elif self.is_input_vectorized:
                     lowerCamelCase_							=self.norm_out(lowerCAmelCase						)
                     lowerCamelCase_							=self.out(lowerCAmelCase						)
                     # (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
                     lowerCamelCase_							=logits.permute(0,	2,	1						)

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

                     # unpatchify
                     lowerCamelCase_							=lowerCamelCase_							=int(hidden_states.shape[1] ** 0.5						)
                     lowerCamelCase_							=hidden_states.reshape(
                         shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)						)
                     lowerCamelCase_							=torch.einsum('''nhwpqc->nchpwq''',	lowerCAmelCase						)
                     lowerCamelCase_							=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=lowerCAmelCase						)







'''simple docstring'''




from typing import List, Optional, Tuple

from tokenizers import pre_tokenizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_clip import CLIPTokenizer


a_   :	Tuple							   = logging.get_logger(__name__)

a_   :	str							   = {"""vocab_file""": """vocab.json""", """merges_file""": """merges.txt""", """tokenizer_file""": """tokenizer.json"""}

a_   :	Tuple							   = {
    """vocab_file""": {
        """openai/clip-vit-base-patch32""": """https://huggingface.co/openai/clip-vit-base-patch32/resolve/main/vocab.json""",
    },
    """merges_file""": {
        """openai/clip-vit-base-patch32""": """https://huggingface.co/openai/clip-vit-base-patch32/resolve/main/merges.txt""",
    },
    """tokenizer_file""": {
        """openai/clip-vit-base-patch32""": (
            """https://huggingface.co/openai/clip-vit-base-patch32/resolve/main/tokenizer.json"""
        ),
    },
}

a_   :	str							   = {
    """openai/clip-vit-base-patch32""": 77,
}


class 						__UpperCamelCase (				lowerCamelCase__       ):
       lowercase			:					Union[str, Any]   		=VOCAB_FILES_NAMES
       lowercase			:					Optional[int]   		=PRETRAINED_VOCAB_FILES_MAP
       lowercase			:					Any   		=PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       lowercase			:					str   		=['input_ids', 'attention_mask']
       lowercase			:					List[Any]   		=CLIPTokenizer






       def __init__(							self,	lowerCAmelCase=None,	lowerCAmelCase=None,	lowerCAmelCase=None,	lowerCAmelCase="<|endoftext|>",	lowerCAmelCase="<|startoftext|>",	lowerCAmelCase="<|endoftext|>",	lowerCAmelCase="<|endoftext|>",	**lowerCAmelCase,	):



              """simple docstring"""







              super().__init__(
                  lowerCAmelCase,	lowerCAmelCase,	tokenizer_file=lowerCAmelCase,	unk_token=lowerCAmelCase,	bos_token=lowerCAmelCase,	eos_token=lowerCAmelCase,	pad_token=lowerCAmelCase,	**lowerCAmelCase,	)

              if not isinstance(self.backend_tokenizer.pre_tokenizer,	pre_tokenizers.Sequence						):
                     raise ValueError(
                         '''The `backend_tokenizer` provided does not match the expected format. The CLIP tokenizer has been'''
                         ''' heavily modified from transformers version 4.17.0. You need to convert the tokenizer you are using'''
                         ''' to be compatible with this version.The easiest way to do so is'''
                         ''' `CLIPTokenizerFast.from_pretrained("path_to_local_folder_or_hub_repo, from_slow=True)`. If you want'''
                         ''' to use your existing tokenizer, you will have to revert to a version prior to 4.17.0 of'''
                         ''' transformers.'''						)

              self._wrap_decode_method_backend_tokenizer()






       def   lowercase__  (							self						):



              """simple docstring"""







              lowerCamelCase_							=self.backend_tokenizer.decode

              def new_decode_method(*lowerCAmelCase,	**lowerCAmelCase						):
                     lowerCamelCase_							=orig_decode_method(*lowerCAmelCase,	**lowerCAmelCase						)
                     lowerCamelCase_							=text.replace(self.backend_tokenizer.model.end_of_word_suffix,	''' '''						).strip()
                     return text

              lowerCamelCase_							=new_decode_method






       def   lowercase__  (							self,	lowerCAmelCase,	lowerCAmelCase = None						):



              """simple docstring"""







              lowerCamelCase_							=[self.bos_token_id]
              lowerCamelCase_							=[self.eos_token_id]

              if token_ids_a is None:
                     return bos_token + token_ids_a + eos_token
              return bos_token + token_ids_a + eos_token + eos_token + token_ids_a + eos_token






       def   lowercase__  (							self,	lowerCAmelCase,	lowerCAmelCase = None						):



              """simple docstring"""







              lowerCamelCase_							=[self.bos_token_id]
              lowerCamelCase_							=[self.eos_token_id]

              if token_ids_a is None:
                     return len(bos_token + token_ids_a + eos_token						) * [0]
              return len(bos_token + token_ids_a + eos_token + eos_token + token_ids_a + eos_token						) * [0]






       def   lowercase__  (							self,	lowerCAmelCase,	lowerCAmelCase = None						):



              """simple docstring"""







              lowerCamelCase_							=self._tokenizer.model.save(lowerCAmelCase,	name=lowerCAmelCase						)
              return tuple(lowerCAmelCase						)




