






import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch
     from torch import nn

     from transformers import ViTForImageClassification, ViTForMaskedImageModeling, ViTModel
     from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
     from PIL import Image

     from transformers import ViTImageProcessor






class 						lowercase_   :






      '''simple docstring'''




      def __init__( self :		Optional[int]					,   __UpperCAmelCase :		Tuple					,   __UpperCAmelCase :		int=13					,   __UpperCAmelCase :		Optional[Any]=30					,   __UpperCAmelCase :		Optional[Any]=2					,   __UpperCAmelCase :		Any=3					,   __UpperCAmelCase :		Any=True					,   __UpperCAmelCase :		List[str]=True					,   __UpperCAmelCase :		str=32					,   __UpperCAmelCase :		List[Any]=5					,   __UpperCAmelCase :		List[Any]=4					,   __UpperCAmelCase :		Any=37					,   __UpperCAmelCase :		Dict="gelu"					,   __UpperCAmelCase :		Tuple=0.1					,   __UpperCAmelCase :		List[str]=0.1					,   __UpperCAmelCase :		Optional[int]=10					,   __UpperCAmelCase :		Dict=0.02					,   __UpperCAmelCase :		Union[str, Any]=None					,   __UpperCAmelCase :		str=2					,   )       ->List[str]:




       """simple docstring"""


       a						    = parent
       a						    = batch_size
       a						    = image_size
       a						    = patch_size
       a						    = num_channels
       a						    = is_training
       a						    = use_labels
       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						    = type_sequence_label_size
       a						    = initializer_range
       a						    = scope
       a						    = encoder_stride

       # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
       a						    = (image_size // patch_size) ** 2
       a						    = num_patches + 1







      def  __lowerCAmelCase   ( self :		Any    )       ->Optional[int]:




       """simple docstring"""


       a						    = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

       a						    = None
       if self.use_labels:
        a						    = ids_tensor([self.batch_size]					,   self.type_sequence_label_size    )

       a						    = self.get_config()

       return config, pixel_values, labels







      def  __lowerCAmelCase   ( self :		int    )       ->Any:




       """simple docstring"""


       return ViTConfig(
           image_size=self.image_size					,   patch_size=self.patch_size					,   num_channels=self.num_channels					,   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					,   is_decoder=__UpperCAmelCase					,   initializer_range=self.initializer_range					,   encoder_stride=self.encoder_stride					,   )







      def  __lowerCAmelCase   ( self :		List[str]					,   __UpperCAmelCase :		Optional[int]					,   __UpperCAmelCase :		Tuple					,   __UpperCAmelCase :		Dict    )       ->Tuple:




       """simple docstring"""


       a						    = ViTModel(config=__UpperCAmelCase    )
       model.to(__UpperCAmelCase    )
       model.eval()
       a						    = model(__UpperCAmelCase    )
       self.parent.assertEqual(result.last_hidden_state.shape					,   (self.batch_size, self.seq_length, self.hidden_size)    )







      def  __lowerCAmelCase   ( self :		Tuple					,   __UpperCAmelCase :		List[Any]					,   __UpperCAmelCase :		Optional[int]					,   __UpperCAmelCase :		Any    )       ->Optional[Any]:




       """simple docstring"""


       a						    = ViTForMaskedImageModeling(config=__UpperCAmelCase    )
       model.to(__UpperCAmelCase    )
       model.eval()
       a						    = model(__UpperCAmelCase    )
       self.parent.assertEqual(
           result.reconstruction.shape					,   (self.batch_size, self.num_channels, self.image_size, self.image_size)    )

       # test greyscale images
       a						    = 1
       a						    = ViTForMaskedImageModeling(__UpperCAmelCase    )
       model.to(__UpperCAmelCase    )
       model.eval()

       a						    = floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
       a						    = model(__UpperCAmelCase    )
       self.parent.assertEqual(result.reconstruction.shape					,   (self.batch_size, 1, self.image_size, self.image_size)    )







      def  __lowerCAmelCase   ( self :		Optional[int]					,   __UpperCAmelCase :		List[str]					,   __UpperCAmelCase :		int					,   __UpperCAmelCase :		List[Any]    )       ->int:




       """simple docstring"""


       a						    = self.type_sequence_label_size
       a						    = ViTForImageClassification(__UpperCAmelCase    )
       model.to(__UpperCAmelCase    )
       model.eval()
       a						    = model(__UpperCAmelCase					,   labels=__UpperCAmelCase    )
       self.parent.assertEqual(result.logits.shape					,   (self.batch_size, self.type_sequence_label_size)    )

       # test greyscale images
       a						    = 1
       a						    = ViTForImageClassification(__UpperCAmelCase    )
       model.to(__UpperCAmelCase    )
       model.eval()

       a						    = floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
       a						    = model(__UpperCAmelCase    )
       self.parent.assertEqual(result.logits.shape					,   (self.batch_size, self.type_sequence_label_size)    )







      def  __lowerCAmelCase   ( self :		Tuple    )       ->Any:




       """simple docstring"""


       a						    = self.prepare_config_and_inputs()
       (
           (
           a
       )    ,	(
           a
       )    ,	(
           a
       )    ,	
       )						    = config_and_inputs
       a						    = {'''pixel_values''': pixel_values}
       return config, inputs_dict






@require_torch
class 						lowercase_   (      lowercase					,    lowercase					,    unittest.TestCase  ):






      '''simple docstring'''




      __snake_case	    =    (
          (
              ViTModel,
              ViTForImageClassification,
              ViTForMaskedImageModeling,
          )
          if is_torch_available()
          else ()
      )
      __snake_case	    =    (
          {'''feature-extraction''': ViTModel, '''image-classification''': ViTForImageClassification}
          if is_torch_available()
          else {}
      )
      __snake_case	    =    True

      __snake_case	    =    False
      __snake_case	    =    False
      __snake_case	    =    False







      def  __lowerCAmelCase   ( self :		List[Any]    )       ->str:




       """simple docstring"""


       a						    = ViTModelTester(self    )
       a						    = ConfigTester(self					,   config_class=__UpperCAmelCase					,   has_text_modality=__UpperCAmelCase					,   hidden_size=37    )







      def  __lowerCAmelCase   ( self :		Any    )       ->Tuple:




       """simple docstring"""


       self.config_tester.run_common_tests()







      @unittest.skip(reason='''ViT does not use inputs_embeds'''    )
      def  __lowerCAmelCase   ( self :		Tuple    )       ->Union[str, Any]:




       """simple docstring"""


       pass







      def  __lowerCAmelCase   ( self :		Dict    )       ->str:




       """simple docstring"""


       a    ,	a						    = self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        a						    = model_class(__UpperCAmelCase    )
        self.assertIsInstance(model.get_input_embeddings()					,   (nn.Module)    )
        a						    = model.get_output_embeddings()
        self.assertTrue(x is None or isinstance(__UpperCAmelCase					,   nn.Linear    )    )







      def  __lowerCAmelCase   ( self :		Union[str, Any]    )       ->Dict:




       """simple docstring"""


       a    ,	a						    = self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        a						    = model_class(__UpperCAmelCase    )
        a						    = inspect.signature(model.forward    )
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        a						    = [*signature.parameters.keys()]

        a						    = ['''pixel_values''']
        self.assertListEqual(arg_names[:1]					,   __UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Union[str, Any]    )       ->Optional[int]:




       """simple docstring"""


       a						    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Optional[Any]    )       ->Any:




       """simple docstring"""


       a						    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_image_modeling(*__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Any    )       ->Union[str, Any]:




       """simple docstring"""


       a						    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*__UpperCAmelCase    )







      @slow
      def  __lowerCAmelCase   ( self :		Dict    )       ->int:




       """simple docstring"""


       for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        a						    = ViTModel.from_pretrained(__UpperCAmelCase    )
        self.assertIsNotNone(__UpperCAmelCase    )
def 		_a       ( )				->     List[Any]:
 a						    = Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''					)
 return image






@require_torch
@require_vision
class 						lowercase_   (      unittest.TestCase  ):






      '''simple docstring'''




      @cached_property
      def  __lowerCAmelCase   ( self :		List[str]    )       ->str:




       """simple docstring"""


       return ViTImageProcessor.from_pretrained('''google/vit-base-patch16-224'''    ) if is_vision_available() else None







      @slow
      def  __lowerCAmelCase   ( self :		Optional[int]    )       ->Any:




       """simple docstring"""


       a						    = ViTForImageClassification.from_pretrained('''google/vit-base-patch16-224'''    ).to(__UpperCAmelCase    )

       a						    = self.default_image_processor
       a						    = prepare_img()
       a						    = image_processor(images=__UpperCAmelCase					,   return_tensors='''pt'''    ).to(__UpperCAmelCase    )

       # forward pass
       with torch.no_grad():
        a						    = model(**__UpperCAmelCase    )

       # verify the logits
       a						    = torch.Size((1, 1_000)    )
       self.assertEqual(outputs.logits.shape					,   __UpperCAmelCase    )

       a						    = torch.tensor([-0.2744, 0.8215, -0.0836]    ).to(__UpperCAmelCase    )

       self.assertTrue(torch.allclose(outputs.logits[0, :3]					,   __UpperCAmelCase					,   atol=1e-4    )    )







      @slow
      def  __lowerCAmelCase   ( self :		List[Any]    )       ->int:




       """simple docstring"""


       a						    = ViTModel.from_pretrained('''facebook/dino-vits8'''    ).to(__UpperCAmelCase    )

       a						    = ViTImageProcessor.from_pretrained('''facebook/dino-vits8'''					,   size=480    )
       a						    = prepare_img()
       a						    = image_processor(images=__UpperCAmelCase					,   return_tensors='''pt'''    )
       a						    = inputs.pixel_values.to(__UpperCAmelCase    )

       # forward pass
       with torch.no_grad():
        a						    = model(__UpperCAmelCase					,   interpolate_pos_encoding=__UpperCAmelCase    )

       # verify the logits
       a						    = torch.Size((1, 3_601, 384)    )
       self.assertEqual(outputs.last_hidden_state.shape					,   __UpperCAmelCase    )

       a						    = torch.tensor(
           [[4.2340, 4.3906, -6.6692], [4.5463, 1.8928, -6.7257], [4.4429, 0.8496, -5.8585]]    ).to(__UpperCAmelCase    )

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







      @slow
      @require_accelerate
      @require_torch_gpu
      def  __lowerCAmelCase   ( self :		Optional[Any]    )       ->Optional[int]:




       """simple docstring"""


       a						    = ViTModel.from_pretrained('''facebook/dino-vits8'''					,   torch_dtype=torch.floataa					,   device_map='''auto'''    )
       a						    = self.default_image_processor

       a						    = prepare_img()
       a						    = image_processor(images=__UpperCAmelCase					,   return_tensors='''pt'''    )
       a						    = inputs.pixel_values.to(__UpperCAmelCase    )

       # forward pass to make sure inference works in fp16
       with torch.no_grad():
        a						    = model(__UpperCAmelCase    )










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__							  =						False






class 						lowercase_   (      unittest.TestCase  ):






      '''simple docstring'''




      def  __lowerCAmelCase   ( self :		Dict    )       ->str:




       """simple docstring"""


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()







      @property
      def  __lowerCAmelCase   ( self :		List[str]    )       ->Any:




       """simple docstring"""


       return 12







      @property
      def  __lowerCAmelCase   ( self :		Optional[Any]    )       ->Dict:




       """simple docstring"""


       return 12







      @property
      def  __lowerCAmelCase   ( self :		Any    )       ->Any:




       """simple docstring"""


       return 32







      @property
      def  __lowerCAmelCase   ( self :		int    )       ->List[str]:




       """simple docstring"""


       torch.manual_seed(0    )
       a						    = VQModel(
           block_out_channels=[32, 64]					,   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  __lowerCAmelCase   ( self :		str    )       ->str:




       """simple docstring"""


       a						    = CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''    )
       return tokenizer







      @property
      def  __lowerCAmelCase   ( self :		str    )       ->Union[str, Any]:




       """simple docstring"""


       torch.manual_seed(0    )
       a						    = CLIPTextConfig(
           bos_token_id=0					,   eos_token_id=2					,   hidden_size=self.text_embedder_hidden_size					,   intermediate_size=37					,   layer_norm_eps=1e-0_5					,   num_attention_heads=4					,   num_hidden_layers=5					,   pad_token_id=1					,   vocab_size=1_000					,   )
       return CLIPTextModel(__UpperCAmelCase    )







      @property
      def  __lowerCAmelCase   ( self :		Any    )       ->Tuple:




       """simple docstring"""


       torch.manual_seed(0    )

       a						    = 12
       a						    = 12

       a						    = {
           '''attention_bias''': True,
           '''cross_attention_dim''': 32,
           '''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''': 32,
           '''sample_size''': width,
           '''activation_fn''': '''geglu-approximate''',
       }

       a						    = TransformeraDModel(**__UpperCAmelCase    )
       return model







      def  __lowerCAmelCase   ( self :		Dict    )       ->Any:




       """simple docstring"""


       a						    = '''cpu'''

       a						    = self.dummy_vqvae
       a						    = self.dummy_text_encoder
       a						    = self.dummy_tokenizer
       a						    = self.dummy_transformer
       a						    = VQDiffusionScheduler(self.num_embed    )
       a						    = LearnedClassifierFreeSamplingEmbeddings(learnable=__UpperCAmelCase    )

       a						    = VQDiffusionPipeline(
           vqvae=__UpperCAmelCase					,   text_encoder=__UpperCAmelCase					,   tokenizer=__UpperCAmelCase					,   transformer=__UpperCAmelCase					,   scheduler=__UpperCAmelCase					,   learned_classifier_free_sampling_embeddings=__UpperCAmelCase					,   )
       a						    = pipe.to(__UpperCAmelCase    )
       pipe.set_progress_bar_config(disable=__UpperCAmelCase    )

       a						    = '''teddy bear playing in the pool'''

       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipe([prompt]					,   generator=__UpperCAmelCase					,   num_inference_steps=2					,   output_type='''np'''    )
       a						    = output.images

       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipe(
           [prompt]					,   generator=__UpperCAmelCase					,   output_type='''np'''					,   return_dict=__UpperCAmelCase					,   num_inference_steps=2    )[0]

       a						    = image[0, -3:, -3:, -1]
       a						    = image_from_tuple[0, -3:, -3:, -1]

       assert image.shape == (1, 24, 24, 3)

       a						    = np.array([0.6551, 0.6168, 0.5008, 0.5676, 0.5659, 0.4295, 0.6073, 0.5599, 0.4992]    )

       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  __lowerCAmelCase   ( self :		Tuple    )       ->str:




       """simple docstring"""


       a						    = '''cpu'''

       a						    = self.dummy_vqvae
       a						    = self.dummy_text_encoder
       a						    = self.dummy_tokenizer
       a						    = self.dummy_transformer
       a						    = VQDiffusionScheduler(self.num_embed    )
       a						    = LearnedClassifierFreeSamplingEmbeddings(
           learnable=__UpperCAmelCase					,   hidden_size=self.text_embedder_hidden_size					,   length=tokenizer.model_max_length    )

       a						    = VQDiffusionPipeline(
           vqvae=__UpperCAmelCase					,   text_encoder=__UpperCAmelCase					,   tokenizer=__UpperCAmelCase					,   transformer=__UpperCAmelCase					,   scheduler=__UpperCAmelCase					,   learned_classifier_free_sampling_embeddings=__UpperCAmelCase					,   )
       a						    = pipe.to(__UpperCAmelCase    )
       pipe.set_progress_bar_config(disable=__UpperCAmelCase    )

       a						    = '''teddy bear playing in the pool'''

       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipe([prompt]					,   generator=__UpperCAmelCase					,   num_inference_steps=2					,   output_type='''np'''    )
       a						    = output.images

       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipe(
           [prompt]					,   generator=__UpperCAmelCase					,   output_type='''np'''					,   return_dict=__UpperCAmelCase					,   num_inference_steps=2    )[0]

       a						    = image[0, -3:, -3:, -1]
       a						    = image_from_tuple[0, -3:, -3:, -1]

       assert image.shape == (1, 24, 24, 3)

       a						    = np.array([0.6693, 0.6075, 0.4959, 0.5701, 0.5583, 0.4333, 0.6171, 0.5684, 0.4988]    )

       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 						lowercase_   (      unittest.TestCase  ):






      '''simple docstring'''




      def  __lowerCAmelCase   ( self :		Union[str, Any]    )       ->Union[str, Any]:




       """simple docstring"""


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()







      def  __lowerCAmelCase   ( self :		Dict    )       ->int:




       """simple docstring"""


       a						    = load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy'''    )

       a						    = VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq'''    )
       a						    = pipeline.to(__UpperCAmelCase    )
       pipeline.set_progress_bar_config(disable=__UpperCAmelCase    )

       # requires GPU generator for gumbel softmax
       # don't use GPU generator in tests though
       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipeline(
           '''teddy bear playing in the pool'''					,   num_images_per_prompt=1					,   generator=__UpperCAmelCase					,   output_type='''np'''					,   )

       a						    = output.images[0]

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


