




from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Mapping, Optional

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import logging


if TYPE_CHECKING:
					from ...processing_utils import ProcessorMixin
					from ...utils import TensorType


UpperCamelCase__         =      logging.get_logger(__name__)

UpperCamelCase__         =      {
    """microsoft/layoutlmv3-base""": """https://huggingface.co/microsoft/layoutlmv3-base/resolve/main/config.json""",
}
class       a__   ( snake_case__				):
						_a							:   str =   """layoutlmv3"""






						def __init__(       self		,      _A=5_0_2_6_5		,      _A=7_6_8		,      _A=1_2		,      _A=1_2		,      _A=3_0_7_2		,      _A="gelu"		,      _A=0.1		,      _A=0.1		,      _A=5_1_2		,      _A=2		,      _A=0.02		,      _A=1E-5		,      _A=1		,      _A=0		,      _A=2		,      _A=1_0_2_4		,      _A=1_2_8		,      _A=1_2_8		,      _A=True		,      _A=3_2		,      _A=1_2_8		,      _A=6_4		,      _A=2_5_6		,      _A=True		,      _A=True		,      _A=True		,      _A=2_2_4		,      _A=3		,      _A=1_6		,      _A=None		,      **_A		,      ):




											"""simple docstring"""



											super().__init__(
											    vocab_size=_A		,      hidden_size=_A		,      num_hidden_layers=_A		,      num_attention_heads=_A		,      intermediate_size=_A		,      hidden_act=_A		,      hidden_dropout_prob=_A		,      attention_probs_dropout_prob=_A		,      max_position_embeddings=_A		,      type_vocab_size=_A		,      initializer_range=_A		,      layer_norm_eps=_A		,      pad_token_id=_A		,      bos_token_id=_A		,      eos_token_id=_A		,      **_A		,      )
											__lowerCAmelCase           =	max_ad_position_embeddings
											__lowerCAmelCase           =	coordinate_size
											__lowerCAmelCase           =	shape_size
											__lowerCAmelCase           =	has_relative_attention_bias
											__lowerCAmelCase           =	rel_pos_bins
											__lowerCAmelCase           =	max_rel_pos
											__lowerCAmelCase           =	has_spatial_attention_bias
											__lowerCAmelCase           =	rel_ad_pos_bins
											__lowerCAmelCase           =	max_rel_ad_pos
											__lowerCAmelCase           =	text_embed
											__lowerCAmelCase           =	visual_embed
											__lowerCAmelCase           =	input_size
											__lowerCAmelCase           =	num_channels
											__lowerCAmelCase           =	patch_size
											__lowerCAmelCase           =	classifier_dropout






class       a__   ( snake_case__				):
						_a							:   Tuple =   version.parse("""1.12"""				)






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											if self.task in ["question-answering", "sequence-classification"]:
																return OrderedDict(
																    [
																        ("input_ids", {0: "batch", 1: "sequence"}),
																        ("attention_mask", {0: "batch", 1: "sequence"}),
																        ("bbox", {0: "batch", 1: "sequence"}),
																        ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
																    ]      )
											else:
																return OrderedDict(
																    [
																        ("input_ids", {0: "batch", 1: "sequence"}),
																        ("bbox", {0: "batch", 1: "sequence"}),
																        ("attention_mask", {0: "batch", 1: "sequence"}),
																        ("pixel_values", {0: "batch", 1: "num_channels"}),
																    ]      )






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											return 1E-5






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											return 1_2






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A = -1		,      _A = -1		,      _A = False		,      _A = None		,      _A = 3		,      _A = 4_0		,      _A = 4_0		,      ):




											"""simple docstring"""



											setattr(processor.image_processor		,      "apply_ocr"		,      _A      )

											# If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
											__lowerCAmelCase           =	compute_effective_axis_dimension(
											    _A		,      fixed_dimension=OnnxConfig.default_fixed_batch		,      num_token_to_add=0      )
											# If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
											__lowerCAmelCase           =	processor.tokenizer.num_special_tokens_to_add(_A      )
											__lowerCAmelCase           =	compute_effective_axis_dimension(
											    _A		,      fixed_dimension=OnnxConfig.default_fixed_sequence		,      num_token_to_add=_A      )
											# Generate dummy inputs according to compute batch and sequence
											__lowerCAmelCase           =	[[" ".join([processor.tokenizer.unk_token]      ) * seq_length]] * batch_size

											# Generate dummy bounding boxes
											__lowerCAmelCase           =	[[[4_8, 8_4, 7_3, 1_2_8]]] * batch_size

											# If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
											# batch_size = compute_effective_axis_dimension(batch_size, fixed_dimension=OnnxConfig.default_fixed_batch)
											__lowerCAmelCase           =	self._generate_dummy_images(_A		,      _A		,      _A		,      _A      )

											__lowerCAmelCase           =	dict(
											    processor(
											        _A		,      text=_A		,      boxes=_A		,      return_tensors=_A		,      )      )

											return inputs










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				):






						def __init__(       self		,      _A		,      _A		,      _A		,      _A=0.0		,      _A = None		,      _A = "geglu"		,      _A = None		,      _A = False		,      _A = False		,      _A = False		,      _A = False		,      _A = True		,      _A = "layer_norm"		,      _A = False		,      ):




											"""simple docstring"""



											super().__init__()
											__lowerCAmelCase           =	only_cross_attention

											__lowerCAmelCase           =	(num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"
											__lowerCAmelCase           =	(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:
																__lowerCAmelCase           =	AdaLayerNorm(_A		,      _A      )
											elif self.use_ada_layer_norm_zero:
																__lowerCAmelCase           =	AdaLayerNormZero(_A		,      _A      )
											else:
																__lowerCAmelCase           =	nn.LayerNorm(_A		,      elementwise_affine=_A      )
											__lowerCAmelCase           =	Attention(
											    query_dim=_A		,      heads=_A		,      dim_head=_A		,      dropout=_A		,      bias=_A		,      cross_attention_dim=cross_attention_dim if only_cross_attention else None		,      upcast_attention=_A		,      )

											# 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.
																__lowerCAmelCase           =	(
																    AdaLayerNorm(_A		,      _A      )
																    if self.use_ada_layer_norm
																    else nn.LayerNorm(_A		,      elementwise_affine=_A      )
																)
																__lowerCAmelCase           =	Attention(
																    query_dim=_A		,      cross_attention_dim=cross_attention_dim if not double_self_attention else None		,      heads=_A		,      dim_head=_A		,      dropout=_A		,      bias=_A		,      upcast_attention=_A		,      )  # is self-attn if encoder_hidden_states is none
											else:
																__lowerCAmelCase           =	None
																__lowerCAmelCase           =	None

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

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






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	chunk_size
											__lowerCAmelCase           =	dim






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A = None		,      _A = None		,      _A = None		,      _A = None		,      _A = None		,      _A = None		,      ):




											"""simple docstring"""



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

											__lowerCAmelCase           =	cross_attention_kwargs if cross_attention_kwargs is not None else {}

											__lowerCAmelCase           =	self.attna(
											    _A		,      encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None		,      attention_mask=_A		,      **_A		,      )
											if self.use_ada_layer_norm_zero:
																__lowerCAmelCase           =	gate_msa.unsqueeze(1      ) * attn_output
											__lowerCAmelCase           =	attn_output + hidden_states

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

																__lowerCAmelCase           =	self.attna(
																    _A		,      encoder_hidden_states=_A		,      attention_mask=_A		,      **_A		,      )
																__lowerCAmelCase           =	attn_output + hidden_states

											# 3. Feed-forward
											__lowerCAmelCase           =	self.norma(_A      )

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

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

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

											__lowerCAmelCase           =	ff_output + hidden_states

											return hidden_states






class       a__   ( nn.Module				):






						def __init__(       self		,      _A		,      _A = None		,      _A = 4		,      _A = 0.0		,      _A = "geglu"		,      _A = False		,      ):




											"""simple docstring"""



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

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

											__lowerCAmelCase           =	nn.ModuleList([]      )
											# project in
											self.net.append(_A      )
											# project dropout
											self.net.append(nn.Dropout(_A      )      )
											# project out
											self.net.append(nn.Linear(_A		,      _A      )      )
											# FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
											if final_dropout:
																self.net.append(nn.Dropout(_A      )      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											for module in self.net:
																__lowerCAmelCase           =	module(_A      )
											return hidden_states






class       a__   ( nn.Module				):






						def __init__(       self		,      _A		,      _A		,      _A = "none"      ):




											"""simple docstring"""



											super().__init__()
											__lowerCAmelCase           =	nn.Linear(_A		,      _A      )
											__lowerCAmelCase           =	approximate






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											if gate.device.type != "mps":
																return F.gelu(_A		,      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     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.proj(_A      )
											__lowerCAmelCase           =	self.gelu(_A      )
											return hidden_states






class       a__   ( nn.Module				):






						def __init__(       self		,      _A		,      _A      ):




											"""simple docstring"""



											super().__init__()
											__lowerCAmelCase           =	nn.Linear(_A		,      dim_out * 2      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											if gate.device.type != "mps":
																return F.gelu(_A      )
											# mps: gelu is not implemented for float16
											return F.gelu(gate.to(dtype=torch.floataa      )      ).to(dtype=gate.dtype      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



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






class       a__   ( nn.Module				):






						def __init__(       self		,      _A		,      _A      ):




											"""simple docstring"""



											super().__init__()
											__lowerCAmelCase           =	nn.Linear(_A		,      _A      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.proj(_A      )
											return x * torch.sigmoid(1.7_02 * x      )






class       a__   ( nn.Module				):






						def __init__(       self		,      _A		,      _A      ):




											"""simple docstring"""



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






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A      ):




											"""simple docstring"""



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






class       a__   ( nn.Module				):






						def __init__(       self		,      _A		,      _A      ):




											"""simple docstring"""



											super().__init__()

											__lowerCAmelCase           =	CombinedTimestepLabelEmbeddings(_A		,      _A      )

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






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A		,      _A		,      _A=None      ):




											"""simple docstring"""



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






class       a__   ( nn.Module				):






						def __init__(       self		,      _A		,      _A		,      _A		,      _A = None		,      _A = 1E-5      ):




											"""simple docstring"""



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

											if act_fn is None:
																__lowerCAmelCase           =	None
											else:
																__lowerCAmelCase           =	get_activation(_A      )

											__lowerCAmelCase           =	nn.Linear(_A		,      out_dim * 2      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A      ):




											"""simple docstring"""



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

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




