






"""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  lowerCAmelCase__ (						__magic_name__      ):
   SCREAMING_SNAKE_CASE_         =42




class  lowerCAmelCase__ (						__magic_name__       ,	__magic_name__      ):
   @register_to_config
   def __init__(		self			:  Dict       , snake_case__			:  int = 1_6       , snake_case__			:  int = 8_8       , snake_case__			:  Optional[int] = None       , snake_case__			:  Optional[int] = None       , snake_case__			:  int = 1       , snake_case__			:  float = 0.0       , snake_case__			:  int = 3_2       , snake_case__			:  Optional[int] = None       , snake_case__			:  bool = False       , snake_case__			:  Optional[int] = None       , snake_case__			:  Optional[int] = None       , snake_case__			:  Optional[int] = None       , snake_case__			:  str = "geglu"       , snake_case__			:  Optional[int] = None       , snake_case__			:  bool = False       , snake_case__			:  bool = False       , snake_case__			:  bool = False       , snake_case__			:  str = "layer_norm"       , snake_case__			:  bool = True       , ):


        '''simple docstring'''
        super().__init__()
        UpperCAmelCase__	:						Optional[int]       =       use_linear_projection
        UpperCAmelCase__	:						Optional[int]       =       num_attention_heads
        UpperCAmelCase__	:						Optional[Any]       =       attention_head_dim
        UpperCAmelCase__	:						str       =       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__	:						Optional[Any]       =       (in_channels is not None) and (patch_size is None)
        UpperCAmelCase__	:						Union[str, Any]       =       num_vector_embeds is not None
        UpperCAmelCase__	:						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__	:						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"       , snake_case__       , standard_warn=snake_case__       )
             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__	:						str       =       in_channels

             UpperCAmelCase__	:						Dict       =       torch.nn.GroupNorm(num_groups=snake_case__       , num_channels=snake_case__       , eps=1e-6       , affine=snake_case__       )
             if use_linear_projection:
                  UpperCAmelCase__	:						Any       =       nn.Linear(snake_case__       , snake_case__       )
             else:
                  UpperCAmelCase__	:						Any       =       nn.Convad(snake_case__       , snake_case__       , 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__	:						str       =       sample_size
             UpperCAmelCase__	:						Optional[int]       =       sample_size
             UpperCAmelCase__	:						List[Any]       =       num_vector_embeds
             UpperCAmelCase__	:						List[str]       =       self.height * self.width

             UpperCAmelCase__	:						str       =       ImagePositionalEmbeddings(
                 num_embed=snake_case__       , embed_dim=snake_case__       , 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__	:						str       =       sample_size
             UpperCAmelCase__	:						Any       =       sample_size

             UpperCAmelCase__	:						Union[str, Any]       =       patch_size
             UpperCAmelCase__	:						List[str]       =       PatchEmbed(
                 height=snake_case__       , width=snake_case__       , patch_size=snake_case__       , in_channels=snake_case__       , embed_dim=snake_case__       , )

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

        # 4. Define output layers
        UpperCAmelCase__	:						int       =       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__	:						List[Any]       =       nn.Linear(snake_case__       , snake_case__       )
             else:
                  UpperCAmelCase__	:						List[Any]       =       nn.Convad(snake_case__       , snake_case__       , kernel_size=1       , stride=1       , padding=0       )
        elif self.is_input_vectorized:
             UpperCAmelCase__	:						Optional[Any]       =       nn.LayerNorm(snake_case__       )
             UpperCAmelCase__	:						Tuple       =       nn.Linear(snake_case__       , self.num_vector_embeds - 1       )
        elif self.is_input_patches:
             UpperCAmelCase__	:						Optional[Any]       =       nn.LayerNorm(snake_case__       , elementwise_affine=snake_case__       , eps=1e-6       )
             UpperCAmelCase__	:						int       =       nn.Linear(snake_case__       , 2 * inner_dim       )
             UpperCAmelCase__	:						int       =       nn.Linear(snake_case__       , patch_size * patch_size * self.out_channels       )




   def 			__a       (		self			:  Any       , snake_case__			:  torch.Tensor       , snake_case__			:  Optional[torch.Tensor] = None       , snake_case__			:  Optional[torch.LongTensor] = None       , snake_case__			:  Optional[torch.LongTensor] = None       , snake_case__			:  Dict[str, Any] = None       , snake_case__			:  Optional[torch.Tensor] = None       , snake_case__			:  Optional[torch.Tensor] = None       , snake_case__			:  bool = True       , ):


        '''simple docstring'''
        # 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__	:						List[Any]       =       (1 - attention_mask.to(hidden_states.dtype       )) * -1_0000.0
             UpperCAmelCase__	:						Tuple       =       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_0000.0
             UpperCAmelCase__	:						Any       =       encoder_attention_mask.unsqueeze(1       )

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

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

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

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

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

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

             # unpatchify
             UpperCAmelCase__	:						int       =       int(hidden_states.shape[1] ** 0.5       )
             UpperCAmelCase__	:						Any       =       hidden_states.reshape(
                 shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)       )
             UpperCAmelCase__	:						Dict       =       torch.einsum("nhwpqc->nchpwq"       , snake_case__       )
             UpperCAmelCase__	:						Any       =       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=snake_case__       )










"""simple docstring"""



from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


_lowerCAmelCase  :  List[Any]   =			{
    """configuration_markuplm""": ["""MARKUPLM_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MarkupLMConfig"""],
    """feature_extraction_markuplm""": ["""MarkupLMFeatureExtractor"""],
    """processing_markuplm""": ["""MarkupLMProcessor"""],
    """tokenization_markuplm""": ["""MarkupLMTokenizer"""],
}

try:
 if not is_tokenizers_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _lowerCAmelCase  :  str   =			["""MarkupLMTokenizerFast"""]

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _lowerCAmelCase  :  List[str]   =			[
     """MARKUPLM_PRETRAINED_MODEL_ARCHIVE_LIST""",
     """MarkupLMForQuestionAnswering""",
     """MarkupLMForSequenceClassification""",
     """MarkupLMForTokenClassification""",
     """MarkupLMModel""",
     """MarkupLMPreTrainedModel""",
 ]


if TYPE_CHECKING:
 from .configuration_markuplm import MARKUPLM_PRETRAINED_CONFIG_ARCHIVE_MAP, MarkupLMConfig
 from .feature_extraction_markuplm import MarkupLMFeatureExtractor
 from .processing_markuplm import MarkupLMProcessor
 from .tokenization_markuplm import MarkupLMTokenizer

 try:
  if not is_tokenizers_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_markuplm_fast import MarkupLMTokenizerFast

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_markuplm import (
      MARKUPLM_PRETRAINED_MODEL_ARCHIVE_LIST,
      MarkupLMForQuestionAnswering,
      MarkupLMForSequenceClassification,
      MarkupLMForTokenClassification,
      MarkupLMModel,
      MarkupLMPreTrainedModel,
  )


else:
 import sys

 _lowerCAmelCase  :  Dict   =			_LazyModule(__name__, globals()["""__file__"""], _import_structure)


