


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			:  Dict     =logging.get_logger(__name__)

_UpperCAmelCase			:  Optional[int]     ={
    """microsoft/layoutlmv3-base""": """https://huggingface.co/microsoft/layoutlmv3-base/resolve/main/config.json""",
}







class 				snake_case__(							UpperCAmelCase__							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__      :  Any					=							"""layoutlmv3"""




						def __init__(				self   ,						__lowercase=5_0_2_6_5   ,						__lowercase=7_6_8   ,						__lowercase=1_2   ,						__lowercase=1_2   ,						__lowercase=3_0_7_2   ,						__lowercase="gelu"   ,						__lowercase=0.1   ,						__lowercase=0.1   ,						__lowercase=5_1_2   ,						__lowercase=2   ,						__lowercase=0.02   ,						__lowercase=1e-5   ,						__lowercase=1   ,						__lowercase=0   ,						__lowercase=2   ,						__lowercase=1_0_2_4   ,						__lowercase=1_2_8   ,						__lowercase=1_2_8   ,						__lowercase=True   ,						__lowercase=3_2   ,						__lowercase=1_2_8   ,						__lowercase=6_4   ,						__lowercase=2_5_6   ,						__lowercase=True   ,						__lowercase=True   ,						__lowercase=True   ,						__lowercase=2_2_4   ,						__lowercase=3   ,						__lowercase=1_6   ,						__lowercase=None   ,						**__lowercase   ,						)					->    Dict:
										super().__init__(
										    vocab_size=__lowercase   ,						hidden_size=__lowercase   ,						num_hidden_layers=__lowercase   ,						num_attention_heads=__lowercase   ,						intermediate_size=__lowercase   ,						hidden_act=__lowercase   ,						hidden_dropout_prob=__lowercase   ,						attention_probs_dropout_prob=__lowercase   ,						max_position_embeddings=__lowercase   ,						type_vocab_size=__lowercase   ,						initializer_range=__lowercase   ,						layer_norm_eps=__lowercase   ,						pad_token_id=__lowercase   ,						bos_token_id=__lowercase   ,						eos_token_id=__lowercase   ,						**__lowercase   ,						)
										lowerCAmelCase_ :      Tuple    				=				max_ad_position_embeddings
										lowerCAmelCase_ :      List[str]    				=				coordinate_size
										lowerCAmelCase_ :      Optional[Any]    				=				shape_size
										lowerCAmelCase_ :      Optional[Any]    				=				has_relative_attention_bias
										lowerCAmelCase_ :      List[str]    				=				rel_pos_bins
										lowerCAmelCase_ :      str    				=				max_rel_pos
										lowerCAmelCase_ :      Optional[Any]    				=				has_spatial_attention_bias
										lowerCAmelCase_ :      Optional[int]    				=				rel_ad_pos_bins
										lowerCAmelCase_ :      List[str]    				=				max_rel_ad_pos
										lowerCAmelCase_ :      List[Any]    				=				text_embed
										lowerCAmelCase_ :      Tuple    				=				visual_embed
										lowerCAmelCase_ :      List[str]    				=				input_size
										lowerCAmelCase_ :      int    				=				num_channels
										lowerCAmelCase_ :      Optional[int]    				=				patch_size
										lowerCAmelCase_ :      Any    				=				classifier_dropout







class 				snake_case__(							UpperCAmelCase__							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__      :  Optional[Any]					=							version.parse("""1.12"""							)




						@property
						def 						lowercase_ (				self							)					->    Mapping[str, Mapping[int, str]]:
										# The order of inputs is different for question answering and sequence classification
										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 						lowercase_ (				self							)					->    float:
										return 1e-5




						@property
						def 						lowercase_ (				self							)					->    int:
										return 1_2




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase = -1   ,						__lowercase = -1   ,						__lowercase = False   ,						__lowercase = None   ,						__lowercase = 3   ,						__lowercase = 4_0   ,						__lowercase = 4_0   ,						)					->    Mapping[str, Any]:
										setattr(processor.image_processor   ,						'''apply_ocr'''   ,						__lowercase							)

										# If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
										lowerCAmelCase_ :      Any    				=				compute_effective_axis_dimension(
										    __lowercase   ,						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_ :      Optional[Any]    				=				processor.tokenizer.num_special_tokens_to_add(__lowercase							)
										lowerCAmelCase_ :      Union[str, Any]    				=				compute_effective_axis_dimension(
										    __lowercase   ,						fixed_dimension=OnnxConfig.default_fixed_sequence   ,						num_token_to_add=__lowercase							)
										# Generate dummy inputs according to compute batch and sequence
										lowerCAmelCase_ :      Any    				=				[[''' '''.join([processor.tokenizer.unk_token]							) * seq_length]] * batch_size

										# Generate dummy bounding boxes
										lowerCAmelCase_ :      List[Any]    				=				[[[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_ :      Optional[Any]    				=				self._generate_dummy_images(__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase							)

										lowerCAmelCase_ :      str    				=				dict(
										    processor(
										        __lowercase   ,						text=__lowercase   ,						boxes=__lowercase   ,						return_tensors=__lowercase   ,						)							)

										return inputs



import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


_UpperCAmelCase			:  Any     =False







class 				snake_case__(							unittest.TestCase							):







						'''simple docstring'''




						def 						lowercase_ (				self							)					->    Optional[int]:
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()




						@property
						def 						lowercase_ (				self							)					->    Union[str, Any]:
										return 1_2




						@property
						def 						lowercase_ (				self							)					->    Any:
										return 1_2




						@property
						def 						lowercase_ (				self							)					->    Optional[Any]:
										return 3_2




						@property
						def 						lowercase_ (				self							)					->    int:
										torch.manual_seed(0							)
										lowerCAmelCase_ :      Any    				=				VQModel(
										    block_out_channels=[3_2, 6_4]   ,						in_channels=3   ,						out_channels=3   ,						down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']   ,						up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']   ,						latent_channels=3   ,						num_vq_embeddings=self.num_embed   ,						vq_embed_dim=3   ,						)
										return model




						@property
						def 						lowercase_ (				self							)					->    Dict:
										lowerCAmelCase_ :      Optional[Any]    				=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)
										return tokenizer




						@property
						def 						lowercase_ (				self							)					->    int:
										torch.manual_seed(0							)
										lowerCAmelCase_ :      Union[str, Any]    				=				CLIPTextConfig(
										    bos_token_id=0   ,						eos_token_id=2   ,						hidden_size=self.text_embedder_hidden_size   ,						intermediate_size=3_7   ,						layer_norm_eps=1e-05   ,						num_attention_heads=4   ,						num_hidden_layers=5   ,						pad_token_id=1   ,						vocab_size=1_0_0_0   ,						)
										return CLIPTextModel(__lowercase							)




						@property
						def 						lowercase_ (				self							)					->    List[str]:
										torch.manual_seed(0							)

										lowerCAmelCase_ :      Union[str, Any]    				=				1_2
										lowerCAmelCase_ :      int    				=				1_2

										lowerCAmelCase_ :      Union[str, Any]    				=				{
										    '''attention_bias''': True,
										    '''cross_attention_dim''': 3_2,
										    '''attention_head_dim''': height * width,
										    '''num_attention_heads''': 1,
										    '''num_vector_embeds''': self.num_embed,
										    '''num_embeds_ada_norm''': self.num_embeds_ada_norm,
										    '''norm_num_groups''': 3_2,
										    '''sample_size''': width,
										    '''activation_fn''': '''geglu-approximate''',
										}

										lowerCAmelCase_ :      List[str]    				=				TransformeraDModel(**__lowercase							)
										return model




						def 						lowercase_ (				self							)					->    str:
										lowerCAmelCase_ :      List[Any]    				=				'''cpu'''

										lowerCAmelCase_ :      Any    				=				self.dummy_vqvae
										lowerCAmelCase_ :      str    				=				self.dummy_text_encoder
										lowerCAmelCase_ :      Union[str, Any]    				=				self.dummy_tokenizer
										lowerCAmelCase_ :      int    				=				self.dummy_transformer
										lowerCAmelCase_ :      List[str]    				=				VQDiffusionScheduler(self.num_embed							)
										lowerCAmelCase_ :      Union[str, Any]    				=				LearnedClassifierFreeSamplingEmbeddings(learnable=__lowercase							)

										lowerCAmelCase_ :      Dict    				=				VQDiffusionPipeline(
										    vqvae=__lowercase   ,						text_encoder=__lowercase   ,						tokenizer=__lowercase   ,						transformer=__lowercase   ,						scheduler=__lowercase   ,						learned_classifier_free_sampling_embeddings=__lowercase   ,						)
										lowerCAmelCase_ :      int    				=				pipe.to(__lowercase							)
										pipe.set_progress_bar_config(disable=__lowercase							)

										lowerCAmelCase_ :      Any    				=				'''teddy bear playing in the pool'''

										lowerCAmelCase_ :      int    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      Tuple    				=				pipe([prompt]   ,						generator=__lowercase   ,						num_inference_steps=2   ,						output_type='''np'''							)
										lowerCAmelCase_ :      Union[str, Any]    				=				output.images

										lowerCAmelCase_ :      List[Any]    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      List[Any]    				=				pipe(
										    [prompt]   ,						generator=__lowercase   ,						output_type='''np'''   ,						return_dict=__lowercase   ,						num_inference_steps=2							)[0]

										lowerCAmelCase_ :      List[str]    				=				image[0, -3:, -3:, -1]
										lowerCAmelCase_ :      str    				=				image_from_tuple[0, -3:, -3:, -1]

										assert image.shape == (1, 2_4, 2_4, 3)

										lowerCAmelCase_ :      Optional[int]    				=				np.array([0.65_51, 0.61_68, 0.50_08, 0.56_76, 0.56_59, 0.42_95, 0.60_73, 0.55_99, 0.49_92]							)

										assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-2
										assert np.abs(image_from_tuple_slice.flatten() - expected_slice							).max() < 1e-2




						def 						lowercase_ (				self							)					->    List[str]:
										lowerCAmelCase_ :      Optional[Any]    				=				'''cpu'''

										lowerCAmelCase_ :      str    				=				self.dummy_vqvae
										lowerCAmelCase_ :      Dict    				=				self.dummy_text_encoder
										lowerCAmelCase_ :      List[Any]    				=				self.dummy_tokenizer
										lowerCAmelCase_ :      Union[str, Any]    				=				self.dummy_transformer
										lowerCAmelCase_ :      Tuple    				=				VQDiffusionScheduler(self.num_embed							)
										lowerCAmelCase_ :      str    				=				LearnedClassifierFreeSamplingEmbeddings(
										    learnable=__lowercase   ,						hidden_size=self.text_embedder_hidden_size   ,						length=tokenizer.model_max_length							)

										lowerCAmelCase_ :      List[str]    				=				VQDiffusionPipeline(
										    vqvae=__lowercase   ,						text_encoder=__lowercase   ,						tokenizer=__lowercase   ,						transformer=__lowercase   ,						scheduler=__lowercase   ,						learned_classifier_free_sampling_embeddings=__lowercase   ,						)
										lowerCAmelCase_ :      Union[str, Any]    				=				pipe.to(__lowercase							)
										pipe.set_progress_bar_config(disable=__lowercase							)

										lowerCAmelCase_ :      Any    				=				'''teddy bear playing in the pool'''

										lowerCAmelCase_ :      List[str]    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      Dict    				=				pipe([prompt]   ,						generator=__lowercase   ,						num_inference_steps=2   ,						output_type='''np'''							)
										lowerCAmelCase_ :      str    				=				output.images

										lowerCAmelCase_ :      List[Any]    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      Union[str, Any]    				=				pipe(
										    [prompt]   ,						generator=__lowercase   ,						output_type='''np'''   ,						return_dict=__lowercase   ,						num_inference_steps=2							)[0]

										lowerCAmelCase_ :      List[str]    				=				image[0, -3:, -3:, -1]
										lowerCAmelCase_ :      str    				=				image_from_tuple[0, -3:, -3:, -1]

										assert image.shape == (1, 2_4, 2_4, 3)

										lowerCAmelCase_ :      Union[str, Any]    				=				np.array([0.66_93, 0.60_75, 0.49_59, 0.57_01, 0.55_83, 0.43_33, 0.61_71, 0.56_84, 0.49_88]							)

										assert np.abs(image_slice.flatten() - expected_slice							).max() < 2.0
										assert np.abs(image_from_tuple_slice.flatten() - expected_slice							).max() < 1e-2







@slow
@require_torch_gpu
class 				snake_case__(							unittest.TestCase							):







						'''simple docstring'''




						def 						lowercase_ (				self							)					->    Optional[int]:
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()




						def 						lowercase_ (				self							)					->    int:
										lowerCAmelCase_ :      Tuple    				=				load_numpy(
										    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
										    '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy'''							)

										lowerCAmelCase_ :      str    				=				VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq'''							)
										lowerCAmelCase_ :      List[Any]    				=				pipeline.to(__lowercase							)
										pipeline.set_progress_bar_config(disable=__lowercase							)

										# requires GPU generator for gumbel softmax
										# don't use GPU generator in tests though
										lowerCAmelCase_ :      List[Any]    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      Optional[int]    				=				pipeline(
										    '''teddy bear playing in the pool'''   ,						num_images_per_prompt=1   ,						generator=__lowercase   ,						output_type='''np'''   ,						)

										lowerCAmelCase_ :      Union[str, Any]    				=				output.images[0]

										assert image.shape == (2_5_6, 2_5_6, 3)
										assert np.abs(expected_image - image							).max() < 2.0