






import gc
import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, KandinskyVaaPipeline, KandinskyVaaPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()
class     lowerCamelCase_					(   lowerCamelCase      ,					unittest.TestCase			):
				a__       							=	KandinskyVaaPipeline
				a__       							=	[
				    '''image_embeds''',
				    '''negative_image_embeds''',
				]
				a__       							=	['''image_embeds''', '''negative_image_embeds''']
				a__       							=	[
				    '''generator''',
				    '''height''',
				    '''width''',
				    '''latents''',
				    '''guidance_scale''',
				    '''num_inference_steps''',
				    '''return_dict''',
				    '''guidance_scale''',
				    '''num_images_per_prompt''',
				    '''output_type''',
				    '''return_dict''',
				]
				a__       							=	False







				@property
				def       A						(    self  ):
								"""simple docstring"""
								return 3_2







				@property
				def       A						(    self  ):
								"""simple docstring"""
								return 3_2







				@property
				def       A						(    self  ):
								"""simple docstring"""
								return self.time_input_dim







				@property
				def       A						(    self  ):
								"""simple docstring"""
								return self.time_input_dim * 4







				@property
				def       A						(    self  ):
								"""simple docstring"""
								return 1_0_0







				@property
				def       A						(    self  ):
								"""simple docstring"""
								torch.manual_seed(0  )

								__magic_name__		:str   =   {
								    '''in_channels''': 4,
								    # Out channels is double in channels because predicts mean and variance
								    '''out_channels''': 8,
								    '''addition_embed_type''': '''image''',
								    '''down_block_types''': ('''ResnetDownsampleBlock2D''', '''SimpleCrossAttnDownBlock2D'''),
								    '''up_block_types''': ('''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D'''),
								    '''mid_block_type''': '''UNetMidBlock2DSimpleCrossAttn''',
								    '''block_out_channels''': (self.block_out_channels_a, self.block_out_channels_a * 2),
								    '''layers_per_block''': 1,
								    '''encoder_hid_dim''': self.text_embedder_hidden_size,
								    '''encoder_hid_dim_type''': '''image_proj''',
								    '''cross_attention_dim''': self.cross_attention_dim,
								    '''attention_head_dim''': 4,
								    '''resnet_time_scale_shift''': '''scale_shift''',
								    '''class_embed_type''': None,
								}

								__magic_name__		:Optional[int]   =   UNetaDConditionModel(**__lowerCAmelCase  )
								return model







				@property
				def       A						(    self  ):
								"""simple docstring"""
								return {
								    "block_out_channels": [3_2, 6_4],
								    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
								    "in_channels": 3,
								    "latent_channels": 4,
								    "layers_per_block": 1,
								    "norm_num_groups": 8,
								    "norm_type": "spatial",
								    "num_vq_embeddings": 1_2,
								    "out_channels": 3,
								    "up_block_types": [
								        "AttnUpDecoderBlock2D",
								        "UpDecoderBlock2D",
								    ],
								    "vq_embed_dim": 4,
								}







				@property
				def       A						(    self  ):
								"""simple docstring"""
								torch.manual_seed(0  )
								__magic_name__		:int   =   VQModel(**self.dummy_movq_kwargs  )
								return model







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Dict   =   self.dummy_unet
								__magic_name__		:str   =   self.dummy_movq

								__magic_name__		:Tuple   =   DDIMScheduler(
								    num_train_timesteps=1_0_0_0							,       beta_schedule='''linear'''							,       beta_start=0.00085							,       beta_end=0.012							,       clip_sample=__lowerCAmelCase							,       set_alpha_to_one=__lowerCAmelCase							,       steps_offset=1							,       prediction_type='''epsilon'''							,       thresholding=__lowerCAmelCase							,       )

								__magic_name__		:Dict   =   {
								    '''unet''': unet,
								    '''scheduler''': scheduler,
								    '''movq''': movq,
								}
								return components







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase=0  ):
								"""simple docstring"""
								__magic_name__		:int   =   floats_tensor((1, self.text_embedder_hidden_size)							,       rng=random.Random(__lowerCAmelCase  )  ).to(__lowerCAmelCase  )
								__magic_name__		:Optional[int]   =   floats_tensor((1, self.text_embedder_hidden_size)							,       rng=random.Random(seed + 1  )  ).to(
								    __lowerCAmelCase  )
								if str(__lowerCAmelCase  ).startswith('''mps'''  ):
												__magic_name__		:Dict   =   torch.manual_seed(__lowerCAmelCase  )
								else:
												__magic_name__		:Any   =   torch.Generator(device=__lowerCAmelCase  ).manual_seed(__lowerCAmelCase  )
								__magic_name__		:Union[str, Any]   =   {
								    '''image_embeds''': image_embeds,
								    '''negative_image_embeds''': negative_image_embeds,
								    '''generator''': generator,
								    '''height''': 6_4,
								    '''width''': 6_4,
								    '''guidance_scale''': 4.0,
								    '''num_inference_steps''': 2,
								    '''output_type''': '''np''',
								}
								return inputs







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Dict   =   '''cpu'''

								__magic_name__		:int   =   self.get_dummy_components()

								__magic_name__		:List[Any]   =   self.pipeline_class(**__lowerCAmelCase  )
								__magic_name__		:Optional[int]   =   pipe.to(__lowerCAmelCase  )

								pipe.set_progress_bar_config(disable=__lowerCAmelCase  )

								__magic_name__		:Optional[int]   =   pipe(**self.get_dummy_inputs(__lowerCAmelCase  )  )
								__magic_name__		:Optional[int]   =   output.images

								__magic_name__		:Optional[int]   =   pipe(
								    **self.get_dummy_inputs(__lowerCAmelCase  )							,       return_dict=__lowerCAmelCase							,       )[0]

								__magic_name__		:List[Any]   =   image[0, -3:, -3:, -1]
								__magic_name__		:Dict   =   image_from_tuple[0, -3:, -3:, -1]

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

								__magic_name__		:str   =   np.array(
								    [0.6237976, 1.0, 0.36441332, 1.0, 0.70639634, 0.29877186, 0.85652125, 0.5216843, 0.54454046]  )

								assert (
								    np.abs(image_slice.flatten() - expected_slice  ).max() < 1E-2
								), F''' expected_slice {expected_slice}, but got {image_slice.flatten()}'''

								assert (
								    np.abs(image_from_tuple_slice.flatten() - expected_slice  ).max() < 1E-2
								), F''' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'''







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







				def       A						(    self  ):
								"""simple docstring"""
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Tuple   =   load_numpy(
								    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
								    '''/kandinskyv22/kandinskyv22_text2img_cat_fp16.npy'''  )

								__magic_name__		:Dict   =   KandinskyVaaPriorPipeline.from_pretrained(
								    '''kandinsky-community/kandinsky-2-2-prior'''							,       torch_dtype=torch.floataa  )
								pipe_prior.to(__lowerCAmelCase  )

								__magic_name__		:Any   =   KandinskyVaaPipeline.from_pretrained(
								    '''kandinsky-community/kandinsky-2-2-decoder'''							,       torch_dtype=torch.floataa  )
								__magic_name__		:List[str]   =   pipeline.to(__lowerCAmelCase  )
								pipeline.set_progress_bar_config(disable=__lowerCAmelCase  )

								__magic_name__		:int   =   '''red cat, 4k photo'''

								__magic_name__		:int   =   torch.Generator(device='''cuda'''  ).manual_seed(0  )
								__magic_name__  ,					__magic_name__		:int   =   pipe_prior(
								    __lowerCAmelCase							,       generator=__lowerCAmelCase							,       num_inference_steps=5							,       negative_prompt=''''''							,       ).to_tuple()

								__magic_name__		:int   =   torch.Generator(device='''cuda'''  ).manual_seed(0  )
								__magic_name__		:Optional[Any]   =   pipeline(
								    image_embeds=__lowerCAmelCase							,       negative_image_embeds=__lowerCAmelCase							,       generator=__lowerCAmelCase							,       num_inference_steps=1_0_0							,       output_type='''np'''							,       )

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

								assert image.shape == (5_1_2, 5_1_2, 3)

								assert_mean_pixel_difference(__lowerCAmelCase							,       __lowerCAmelCase  )













import unittest

from transformers import LiltConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import (
       LiltForQuestionAnswering,
       LiltForSequenceClassification,
       LiltForTokenClassification,
       LiltModel,
   )
   from transformers.models.lilt.modeling_lilt import LILT_PRETRAINED_MODEL_ARCHIVE_LIST




class 			__UpperCamelCase    :





     def __init__(     self				,							_UpperCamelCase				,							_UpperCamelCase=13				,							_UpperCamelCase=7				,							_UpperCamelCase=True				,							_UpperCamelCase=True				,							_UpperCamelCase=True				,							_UpperCamelCase=True				,							_UpperCamelCase=99				,							_UpperCamelCase=24				,							_UpperCamelCase=2				,							_UpperCamelCase=6				,							_UpperCamelCase=37				,							_UpperCamelCase="gelu"				,							_UpperCamelCase=0.1				,							_UpperCamelCase=0.1				,							_UpperCamelCase=512				,							_UpperCamelCase=16				,							_UpperCamelCase=2				,							_UpperCamelCase=0.02				,							_UpperCamelCase=3				,							_UpperCamelCase=None				,							_UpperCamelCase=1000				,							):
        _UpperCAmelCase           =   parent
        _UpperCAmelCase           =   batch_size
        _UpperCAmelCase           =   seq_length
        _UpperCAmelCase           =   is_training
        _UpperCAmelCase           =   use_input_mask
        _UpperCAmelCase           =   use_token_type_ids
        _UpperCAmelCase           =   use_labels
        _UpperCAmelCase           =   vocab_size
        _UpperCAmelCase           =   hidden_size
        _UpperCAmelCase           =   num_hidden_layers
        _UpperCAmelCase           =   num_attention_heads
        _UpperCAmelCase           =   intermediate_size
        _UpperCAmelCase           =   hidden_act
        _UpperCAmelCase           =   hidden_dropout_prob
        _UpperCAmelCase           =   attention_probs_dropout_prob
        _UpperCAmelCase           =   max_position_embeddings
        _UpperCAmelCase           =   type_vocab_size
        _UpperCAmelCase           =   type_sequence_label_size
        _UpperCAmelCase           =   initializer_range
        _UpperCAmelCase           =   num_labels
        _UpperCAmelCase           =   scope
        _UpperCAmelCase           =   range_bbox





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   ids_tensor([self.batch_size, self.seq_length]				,							self.vocab_size					)

        _UpperCAmelCase           =   ids_tensor([self.batch_size, self.seq_length, 4]				,							self.range_bbox					)
        # Ensure that bbox is legal
        for i in range(bbox.shape[0]					):
           for j in range(bbox.shape[1]					):
              if bbox[i, j, 3] < bbox[i, j, 1]:
                 _UpperCAmelCase           =   bbox[i, j, 3]
                 _UpperCAmelCase           =   bbox[i, j, 1]
                 _UpperCAmelCase           =   t
              if bbox[i, j, 2] < bbox[i, j, 0]:
                 _UpperCAmelCase           =   bbox[i, j, 2]
                 _UpperCAmelCase           =   bbox[i, j, 0]
                 _UpperCAmelCase           =   t

        _UpperCAmelCase           =   None
        if self.use_input_mask:
           _UpperCAmelCase           =   ids_tensor([self.batch_size, self.seq_length]				,							vocab_size=2					)

        _UpperCAmelCase           =   None
        if self.use_token_type_ids:
           _UpperCAmelCase           =   ids_tensor([self.batch_size, self.seq_length]				,							self.type_vocab_size					)

        _UpperCAmelCase           =   None
        _UpperCAmelCase           =   None
        if self.use_labels:
           _UpperCAmelCase           =   ids_tensor([self.batch_size]				,							self.type_sequence_label_size					)
           _UpperCAmelCase           =   ids_tensor([self.batch_size, self.seq_length]				,							self.num_labels					)

        _UpperCAmelCase           =   self.get_config()

        return config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels





     def 		UpperCamelCase(     self					):
        return LiltConfig(
            vocab_size=self.vocab_size				,							hidden_size=self.hidden_size				,							num_hidden_layers=self.num_hidden_layers				,							num_attention_heads=self.num_attention_heads				,							intermediate_size=self.intermediate_size				,							hidden_act=self.hidden_act				,							hidden_dropout_prob=self.hidden_dropout_prob				,							attention_probs_dropout_prob=self.attention_probs_dropout_prob				,							max_position_embeddings=self.max_position_embeddings				,							type_vocab_size=self.type_vocab_size				,							initializer_range=self.initializer_range				,							)





     def 		UpperCamelCase(     self				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							):
        _UpperCAmelCase           =   LiltModel(config=_UpperCamelCase					)
        model.to(_UpperCamelCase					)
        model.eval()
        _UpperCAmelCase           =   model(_UpperCamelCase				,							bbox=_UpperCamelCase				,							attention_mask=_UpperCamelCase				,							token_type_ids=_UpperCamelCase					)
        _UpperCAmelCase           =   model(_UpperCamelCase				,							bbox=_UpperCamelCase				,							token_type_ids=_UpperCamelCase					)
        _UpperCAmelCase           =   model(_UpperCamelCase				,							bbox=_UpperCamelCase					)

        self.parent.assertEqual(result.last_hidden_state.shape				,							(self.batch_size, self.seq_length, self.hidden_size)					)
        self.parent.assertEqual(result.pooler_output.shape				,							(self.batch_size, self.hidden_size)					)





     def 		UpperCamelCase(     self				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							):
        _UpperCAmelCase           =   self.num_labels
        _UpperCAmelCase           =   LiltForTokenClassification(config=_UpperCamelCase					)
        model.to(_UpperCamelCase					)
        model.eval()
        _UpperCAmelCase           =   model(
            _UpperCamelCase				,							bbox=_UpperCamelCase				,							attention_mask=_UpperCamelCase				,							token_type_ids=_UpperCamelCase				,							labels=_UpperCamelCase					)
        self.parent.assertEqual(result.logits.shape				,							(self.batch_size, self.seq_length, self.num_labels)					)





     def 		UpperCamelCase(     self				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							):
        _UpperCAmelCase           =   LiltForQuestionAnswering(config=_UpperCamelCase					)
        model.to(_UpperCamelCase					)
        model.eval()
        _UpperCAmelCase           =   model(
            _UpperCamelCase				,							bbox=_UpperCamelCase				,							attention_mask=_UpperCamelCase				,							token_type_ids=_UpperCamelCase				,							start_positions=_UpperCamelCase				,							end_positions=_UpperCamelCase				,							)
        self.parent.assertEqual(result.start_logits.shape				,							(self.batch_size, self.seq_length)					)
        self.parent.assertEqual(result.end_logits.shape				,							(self.batch_size, self.seq_length)					)





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   self.prepare_config_and_inputs()
        (
            (
            _UpperCAmelCase
        )       ,      (
            _UpperCAmelCase
        )       ,      (
            _UpperCAmelCase
        )       ,      (
            _UpperCAmelCase
        )       ,      (
            _UpperCAmelCase
        )       ,      (
            _UpperCAmelCase
        )       ,      (
            _UpperCAmelCase
        )       ,      
        )           =   config_and_inputs
        _UpperCAmelCase           =   {
            '''input_ids''': input_ids,
            '''bbox''': bbox,
            '''token_type_ids''': token_type_ids,
            '''attention_mask''': input_mask,
        }
        return config, inputs_dict




@require_torch
class 			__UpperCamelCase    (							A__	, A__	, A__	, unittest.TestCase    ):
     __A	:  Dict							  =   (
         (
             LiltModel,
             LiltForSequenceClassification,
             LiltForTokenClassification,
             LiltForQuestionAnswering,
         )
         if is_torch_available()
         else ()
     )
     __A	:  Optional[Any]							  =   (
         {
             """feature-extraction""": LiltModel,
             """question-answering""": LiltForQuestionAnswering,
             """text-classification""": LiltForSequenceClassification,
             """token-classification""": LiltForTokenClassification,
             """zero-shot""": LiltForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     __A	:  List[Any]							  =   False
     __A	:  Optional[int]							  =   False





     def 		UpperCamelCase(     self				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase					):
        return True





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   LiltModelTester(self					)
        _UpperCAmelCase           =   ConfigTester(self				,							config_class=_UpperCamelCase				,							hidden_size=37					)





     def 		UpperCamelCase(     self					):
        self.config_tester.run_common_tests()





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*_UpperCamelCase					)





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
           _UpperCAmelCase           =   type
           self.model_tester.create_and_check_model(*_UpperCamelCase					)





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*_UpperCamelCase					)





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*_UpperCamelCase					)





     @slow
     def 		UpperCamelCase(     self					):
        for model_name in LILT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           _UpperCAmelCase           =   LiltModel.from_pretrained(_UpperCamelCase					)
           self.assertIsNotNone(_UpperCamelCase					)




@require_torch
@slow
class 			__UpperCamelCase    (							unittest.TestCase    ):





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   LiltModel.from_pretrained('''SCUT-DLVCLab/lilt-roberta-en-base'''					).to(_UpperCamelCase					)

        _UpperCAmelCase           =   torch.tensor([[1, 2]]				,							device=_UpperCamelCase					)
        _UpperCAmelCase           =   torch.tensor([[[1, 2, 3, 4], [5, 6, 7, 8]]]				,							device=_UpperCamelCase					)

        # forward pass
        with torch.no_grad():
           _UpperCAmelCase           =   model(input_ids=_UpperCamelCase				,							bbox=_UpperCamelCase					)

        _UpperCAmelCase           =   torch.Size([1, 2, 768]					)
        _UpperCAmelCase           =   torch.tensor(
            [[-0.0653, 0.0950, -0.0061], [-0.0545, 0.0926, -0.0324]]				,							device=_UpperCamelCase				,							)

        self.assertTrue(outputs.last_hidden_state.shape				,							_UpperCamelCase					)
        self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :, :3]				,							_UpperCamelCase				,							atol=1e-3					)					)