



from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block




@dataclass
class   UpperCamelCase_       (     UpperCamelCase__    ):
							lowerCamelCase_           =					42




class   UpperCamelCase_       (     nn.Module    ):


							def __init__(							self    :int							,				__A    :str=3							,				__A    :Optional[Any]=3							,				__A    :str=("DownEncoderBlock2D",)							,				__A    :Tuple=(64,)							,				__A    :List[Any]=2							,				__A    :List[str]=32							,				__A    :str="silu"							,				__A    :Optional[int]=True							,				)   ->   List[str]:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__			=      layers_per_block

									SCREAMING_SNAKE_CASE__			=      torch.nn.Convad(
									    __A							,				block_out_channels[0]							,				kernel_size=3							,				stride=1							,				padding=1							,				)

									SCREAMING_SNAKE_CASE__			=      None
									SCREAMING_SNAKE_CASE__			=      nn.ModuleList([]		)

									# down
									SCREAMING_SNAKE_CASE__			=      block_out_channels[0]
									for i, down_block_type in enumerate(__A		):
											SCREAMING_SNAKE_CASE__			=      output_channel
											SCREAMING_SNAKE_CASE__			=      block_out_channels[i]
											SCREAMING_SNAKE_CASE__			=      i == len(__A		) - 1

											SCREAMING_SNAKE_CASE__			=      get_down_block(
											    __A							,				num_layers=self.layers_per_block							,				in_channels=__A							,				out_channels=__A							,				add_downsample=not is_final_block							,				resnet_eps=1E-6							,				downsample_padding=0							,				resnet_act_fn=__A							,				resnet_groups=__A							,				attention_head_dim=__A							,				temb_channels=__A							,				)
											self.down_blocks.append(__A		)

									# mid
									SCREAMING_SNAKE_CASE__			=      UNetMidBlockaD(
									    in_channels=block_out_channels[-1]							,				resnet_eps=1E-6							,				resnet_act_fn=__A							,				output_scale_factor=1							,				resnet_time_scale_shift="""default"""							,				attention_head_dim=block_out_channels[-1]							,				resnet_groups=__A							,				temb_channels=__A							,				)

									# out
									SCREAMING_SNAKE_CASE__			=      nn.GroupNorm(num_channels=block_out_channels[-1]							,				num_groups=__A							,				eps=1E-6		)
									SCREAMING_SNAKE_CASE__			=      nn.SiLU()

									SCREAMING_SNAKE_CASE__			=      2 * out_channels if double_z else out_channels
									SCREAMING_SNAKE_CASE__			=      nn.Convad(block_out_channels[-1]							,				__A							,				3							,				padding=1		)

									SCREAMING_SNAKE_CASE__			=      False


							def      _snake_case  (							self    :str							,				__A    :int		)   ->   Optional[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      x
									SCREAMING_SNAKE_CASE__			=      self.conv_in(__A		)

									if self.training and self.gradient_checkpointing:

											def create_custom_forward(__A    :str		):
													def custom_forward(*__A    :str		):
															return module(*__A		)

													return custom_forward

											# down
											if is_torch_version(""">="""							,				"""1.11.0"""		):
													for down_block in self.down_blocks:
															SCREAMING_SNAKE_CASE__			=      torch.utils.checkpoint.checkpoint(
															    create_custom_forward(__A		)							,				__A							,				use_reentrant=__A		)
													# middle
													SCREAMING_SNAKE_CASE__			=      torch.utils.checkpoint.checkpoint(
													    create_custom_forward(self.mid_block		)							,				__A							,				use_reentrant=__A		)
											else:
													for down_block in self.down_blocks:
															SCREAMING_SNAKE_CASE__			=      torch.utils.checkpoint.checkpoint(create_custom_forward(__A		)							,				__A		)
													# middle
													SCREAMING_SNAKE_CASE__			=      torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block		)							,				__A		)

									else:
											# down
											for down_block in self.down_blocks:
													SCREAMING_SNAKE_CASE__			=      down_block(__A		)

											# middle
											SCREAMING_SNAKE_CASE__			=      self.mid_block(__A		)

									# post-process
									SCREAMING_SNAKE_CASE__			=      self.conv_norm_out(__A		)
									SCREAMING_SNAKE_CASE__			=      self.conv_act(__A		)
									SCREAMING_SNAKE_CASE__			=      self.conv_out(__A		)

									return sample




class   UpperCamelCase_       (     nn.Module    ):


							def __init__(							self    :Dict							,				__A    :Optional[Any]=3							,				__A    :Tuple=3							,				__A    :Optional[int]=("UpDecoderBlock2D",)							,				__A    :int=(64,)							,				__A    :int=2							,				__A    :Dict=32							,				__A    :Dict="silu"							,				__A    :str="group"							,				)   ->   Optional[Any]:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__			=      layers_per_block

									SCREAMING_SNAKE_CASE__			=      nn.Convad(
									    __A							,				block_out_channels[-1]							,				kernel_size=3							,				stride=1							,				padding=1							,				)

									SCREAMING_SNAKE_CASE__			=      None
									SCREAMING_SNAKE_CASE__			=      nn.ModuleList([]		)

									SCREAMING_SNAKE_CASE__			=      in_channels if norm_type == """spatial""" else None

									# mid
									SCREAMING_SNAKE_CASE__			=      UNetMidBlockaD(
									    in_channels=block_out_channels[-1]							,				resnet_eps=1E-6							,				resnet_act_fn=__A							,				output_scale_factor=1							,				resnet_time_scale_shift="""default""" if norm_type == """group""" else norm_type							,				attention_head_dim=block_out_channels[-1]							,				resnet_groups=__A							,				temb_channels=__A							,				)

									# up
									SCREAMING_SNAKE_CASE__			=      list(reversed(__A		)		)
									SCREAMING_SNAKE_CASE__			=      reversed_block_out_channels[0]
									for i, up_block_type in enumerate(__A		):
											SCREAMING_SNAKE_CASE__			=      output_channel
											SCREAMING_SNAKE_CASE__			=      reversed_block_out_channels[i]

											SCREAMING_SNAKE_CASE__			=      i == len(__A		) - 1

											SCREAMING_SNAKE_CASE__			=      get_up_block(
											    __A							,				num_layers=self.layers_per_block + 1							,				in_channels=__A							,				out_channels=__A							,				prev_output_channel=__A							,				add_upsample=not is_final_block							,				resnet_eps=1E-6							,				resnet_act_fn=__A							,				resnet_groups=__A							,				attention_head_dim=__A							,				temb_channels=__A							,				resnet_time_scale_shift=__A							,				)
											self.up_blocks.append(__A		)
											SCREAMING_SNAKE_CASE__			=      output_channel

									# out
									if norm_type == "spatial":
											SCREAMING_SNAKE_CASE__			=      SpatialNorm(block_out_channels[0]							,				__A		)
									else:
											SCREAMING_SNAKE_CASE__			=      nn.GroupNorm(num_channels=block_out_channels[0]							,				num_groups=__A							,				eps=1E-6		)
									SCREAMING_SNAKE_CASE__			=      nn.SiLU()
									SCREAMING_SNAKE_CASE__			=      nn.Convad(block_out_channels[0]							,				__A							,				3							,				padding=1		)

									SCREAMING_SNAKE_CASE__			=      False


							def      _snake_case  (							self    :List[str]							,				__A    :Any							,				__A    :Dict=None		)   ->   Optional[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      z
									SCREAMING_SNAKE_CASE__			=      self.conv_in(__A		)

									SCREAMING_SNAKE_CASE__			=      next(iter(self.up_blocks.parameters()		)		).dtype
									if self.training and self.gradient_checkpointing:

											def create_custom_forward(__A    :str		):
													def custom_forward(*__A    :Dict		):
															return module(*__A		)

													return custom_forward

											if is_torch_version(""">="""							,				"""1.11.0"""		):
													# middle
													SCREAMING_SNAKE_CASE__			=      torch.utils.checkpoint.checkpoint(
													    create_custom_forward(self.mid_block		)							,				__A							,				__A							,				use_reentrant=__A		)
													SCREAMING_SNAKE_CASE__			=      sample.to(__A		)

													# up
													for up_block in self.up_blocks:
															SCREAMING_SNAKE_CASE__			=      torch.utils.checkpoint.checkpoint(
															    create_custom_forward(__A		)							,				__A							,				__A							,				use_reentrant=__A		)
											else:
													# middle
													SCREAMING_SNAKE_CASE__			=      torch.utils.checkpoint.checkpoint(
													    create_custom_forward(self.mid_block		)							,				__A							,				__A		)
													SCREAMING_SNAKE_CASE__			=      sample.to(__A		)

													# up
													for up_block in self.up_blocks:
															SCREAMING_SNAKE_CASE__			=      torch.utils.checkpoint.checkpoint(create_custom_forward(__A		)							,				__A							,				__A		)
									else:
											# middle
											SCREAMING_SNAKE_CASE__			=      self.mid_block(__A							,				__A		)
											SCREAMING_SNAKE_CASE__			=      sample.to(__A		)

											# up
											for up_block in self.up_blocks:
													SCREAMING_SNAKE_CASE__			=      up_block(__A							,				__A		)

        # post-process
									if latent_embeds is None:
											SCREAMING_SNAKE_CASE__			=      self.conv_norm_out(__A		)
									else:
											SCREAMING_SNAKE_CASE__			=      self.conv_norm_out(__A							,				__A		)
									SCREAMING_SNAKE_CASE__			=      self.conv_act(__A		)
									SCREAMING_SNAKE_CASE__			=      self.conv_out(__A		)

									return sample




class   UpperCamelCase_       (     nn.Module    ):


							def __init__(							self    :List[str]							,				__A    :Optional[Any]							,				__A    :Tuple							,				__A    :Optional[int]							,				__A    :int=None							,				__A    :Optional[Any]="random"							,				__A    :int=False							,				__A    :Any=True		)   ->   Any:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__			=      n_e
									SCREAMING_SNAKE_CASE__			=      vq_embed_dim
									SCREAMING_SNAKE_CASE__			=      beta
									SCREAMING_SNAKE_CASE__			=      legacy

									SCREAMING_SNAKE_CASE__			=      nn.Embedding(self.n_e							,				self.vq_embed_dim		)
									self.embedding.weight.data.uniform_(-1.0 / self.n_e							,				1.0 / self.n_e		)

									SCREAMING_SNAKE_CASE__			=      remap
									if self.remap is not None:
											self.register_buffer("""used"""							,				torch.tensor(np.load(self.remap		)		)		)
											SCREAMING_SNAKE_CASE__			=      self.used.shape[0]
											SCREAMING_SNAKE_CASE__			=      unknown_index  # "random" or "extra" or integer
											if self.unknown_index == "extra":
													SCREAMING_SNAKE_CASE__			=      self.re_embed
													SCREAMING_SNAKE_CASE__			=      self.re_embed + 1
											print(
											    f'''Remapping {self.n_e} indices to {self.re_embed} indices. '''
											    f'''Using {self.unknown_index} for unknown indices.'''		)
									else:
											SCREAMING_SNAKE_CASE__			=      n_e

									SCREAMING_SNAKE_CASE__			=      sane_index_shape


							def      _snake_case  (							self    :str							,				__A    :List[str]		)   ->   Union[str, Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      inds.shape
									assert len(__A		) > 1
									SCREAMING_SNAKE_CASE__			=      inds.reshape(ishape[0]							,				-1		)
									SCREAMING_SNAKE_CASE__			=      self.used.to(__A		)
									SCREAMING_SNAKE_CASE__			=      (inds[:, :, None] == used[None, None, ...]).long()
									SCREAMING_SNAKE_CASE__			=      match.argmax(-1		)
									SCREAMING_SNAKE_CASE__			=      match.sum(2		) < 1
									if self.unknown_index == "random":
											SCREAMING_SNAKE_CASE__			=      torch.randint(0							,				self.re_embed							,				size=new[unknown].shape		).to(device=new.device		)
									else:
											SCREAMING_SNAKE_CASE__			=      self.unknown_index
									return new.reshape(__A		)


							def      _snake_case  (							self    :Dict							,				__A    :str		)   ->   List[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      inds.shape
									assert len(__A		) > 1
									SCREAMING_SNAKE_CASE__			=      inds.reshape(ishape[0]							,				-1		)
									SCREAMING_SNAKE_CASE__			=      self.used.to(__A		)
									if self.re_embed > self.used.shape[0]:  # extra token
											SCREAMING_SNAKE_CASE__			=      0  # simply set to zero
									SCREAMING_SNAKE_CASE__			=      torch.gather(used[None, :][inds.shape[0] * [0], :]							,				1							,				__A		)
									return back.reshape(__A		)


							def      _snake_case  (							self    :Tuple							,				__A    :Optional[int]		)   ->   Dict:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      z.permute(0							,				2							,				3							,				1		).contiguous()
									SCREAMING_SNAKE_CASE__			=      z.view(-1							,				self.vq_embed_dim		)

									# distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
									SCREAMING_SNAKE_CASE__			=      torch.argmin(torch.cdist(__A							,				self.embedding.weight		)							,				dim=1		)

									SCREAMING_SNAKE_CASE__			=      self.embedding(__A		).view(z.shape		)
									SCREAMING_SNAKE_CASE__			=      None
									SCREAMING_SNAKE_CASE__			=      None

									# compute loss for embedding
									if not self.legacy:
											SCREAMING_SNAKE_CASE__			=      self.beta * torch.mean((z_q.detach() - z) ** 2		) + torch.mean((z_q - z.detach()) ** 2		)
									else:
											SCREAMING_SNAKE_CASE__			=      torch.mean((z_q.detach() - z) ** 2		) + self.beta * torch.mean((z_q - z.detach()) ** 2		)

									# preserve gradients
									SCREAMING_SNAKE_CASE__			=      z + (z_q - z).detach()

									# reshape back to match original input shape
									SCREAMING_SNAKE_CASE__			=      z_q.permute(0							,				3							,				1							,				2		).contiguous()

									if self.remap is not None:
											SCREAMING_SNAKE_CASE__			=      min_encoding_indices.reshape(z.shape[0]							,				-1		)  # add batch axis
											SCREAMING_SNAKE_CASE__			=      self.remap_to_used(__A		)
											SCREAMING_SNAKE_CASE__			=      min_encoding_indices.reshape(-1							,				1		)  # flatten

									if self.sane_index_shape:
											SCREAMING_SNAKE_CASE__			=      min_encoding_indices.reshape(z_q.shape[0]							,				z_q.shape[2]							,				z_q.shape[3]		)

									return z_q, loss, (perplexity, min_encodings, min_encoding_indices)


							def      _snake_case  (							self    :List[Any]							,				__A    :Tuple							,				__A    :List[str]		)   ->   List[Any]:



									"""simple docstring"""



									if self.remap is not None:
											SCREAMING_SNAKE_CASE__			=      indices.reshape(shape[0]							,				-1		)  # add batch axis
											SCREAMING_SNAKE_CASE__			=      self.unmap_to_all(__A		)
											SCREAMING_SNAKE_CASE__			=      indices.reshape(-1		)  # flatten again

									# get quantized latent vectors
									SCREAMING_SNAKE_CASE__			=      self.embedding(__A		)

									if shape is not None:
											SCREAMING_SNAKE_CASE__			=      z_q.view(__A		)
											# reshape back to match original input shape
											SCREAMING_SNAKE_CASE__			=      z_q.permute(0							,				3							,				1							,				2		).contiguous()

									return z_q




class   UpperCamelCase_       (     UpperCamelCase__    ):


							def __init__(							self    :List[Any]							,				__A    :List[Any]							,				__A    :int=False		)   ->   Any:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      parameters
									SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__			=      torch.chunk(__A							,				2							,				dim=1		)
									SCREAMING_SNAKE_CASE__			=      torch.clamp(self.logvar							,				-30.0							,				20.0		)
									SCREAMING_SNAKE_CASE__			=      deterministic
									SCREAMING_SNAKE_CASE__			=      torch.exp(0.5 * self.logvar		)
									SCREAMING_SNAKE_CASE__			=      torch.exp(self.logvar		)
									if self.deterministic:
											SCREAMING_SNAKE_CASE__			=      SCREAMING_SNAKE_CASE__			=      torch.zeros_like(
											    self.mean							,				device=self.parameters.device							,				dtype=self.parameters.dtype		)


							def      _snake_case  (							self    :Any							,				__A    :Optional[torch.Generator] = None		)   ->   torch.FloatTensor:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      randn_tensor(
									    self.mean.shape							,				generator=__A							,				device=self.parameters.device							,				dtype=self.parameters.dtype		)
									SCREAMING_SNAKE_CASE__			=      self.mean + self.std * sample
									return x


							def      _snake_case  (							self    :Union[str, Any]							,				__A    :int=None		)   ->   int:



									"""simple docstring"""



									if self.deterministic:
											return torch.Tensor([0.0]		)
									else:
											if other is None:
													return 0.5 * torch.sum(torch.pow(self.mean							,				2		) + self.var - 1.0 - self.logvar							,				dim=[1, 2, 3]		)
											else:
													return 0.5 * torch.sum(
													    torch.pow(self.mean - other.mean							,				2		) / other.var
													    + self.var / other.var
													    - 1.0
													    - self.logvar
													    + other.logvar							,				dim=[1, 2, 3]							,				)


							def      _snake_case  (							self    :str							,				__A    :Tuple							,				__A    :int=[1, 2, 3]		)   ->   int:



									"""simple docstring"""



									if self.deterministic:
											return torch.Tensor([0.0]		)
									SCREAMING_SNAKE_CASE__			=      np.log(2.0 * np.pi		)
									return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean							,				2		) / self.var							,				dim=__A		)


							def      _snake_case  (							self    :Tuple		)   ->   Optional[int]:



									"""simple docstring"""



									return self.mean






import warnings

from ...utils import logging
from .image_processing_layoutlmva import LayoutLMvaImageProcessor


_lowerCamelCase   				=			logging.get_logger(__name__)




class   UpperCamelCase_       (     UpperCamelCase__    ):


							def __init__(							self    :List[Any]							,				*__A    :Tuple							,				**__A    :Dict		)   ->   None:



									"""simple docstring"""



									warnings.warn(
									    """The class LayoutLMv2FeatureExtractor is deprecated and will be removed in version 5 of Transformers."""
									    """ Please use LayoutLMv2ImageProcessor instead."""							,				__A							,				)
									super().__init__(*__A							,				**__A		)