


import unittest

import numpy as np
import torch

from diffusers import DDIMPipeline, DDIMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow, torch_device

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()






class   __lowercase				(  UpperCamelCase__ ,					unittest.TestCase  ):
      UpperCamelCase								=       DDIMPipeline
      UpperCamelCase								=       UNCONDITIONAL_IMAGE_GENERATION_PARAMS
      UpperCamelCase								=       PipelineTesterMixin.required_optional_params - {
          '''num_images_per_prompt''',
          '''latents''',
          '''callback''',
          '''callback_steps''',
      }
      UpperCamelCase								=       UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS
      UpperCamelCase								=       False





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






         """simple docstring"""







         torch.manual_seed(0       )
         UpperCAmelCase							=     UNetaDModel(
             block_out_channels=(3_2, 6_4)					,       layers_per_block=2					,       sample_size=3_2					,       in_channels=3					,       out_channels=3					,       down_block_types=("""DownBlock2D""", """AttnDownBlock2D""")					,       up_block_types=("""AttnUpBlock2D""", """UpBlock2D""")					,       )
         UpperCAmelCase							=     DDIMScheduler()
         UpperCAmelCase							=     {"""unet""": unet, """scheduler""": scheduler}
         return components





      def        _lowercase				(		self						:					Tuple					,       __lowerCamelCase						:					Optional[Any]					,       __lowerCamelCase						:					Any=0       )    ->   List[str]:






         """simple docstring"""







         if str(_a       ).startswith("""mps"""       ):
            UpperCAmelCase							=     torch.manual_seed(_a       )
         else:
            UpperCAmelCase							=     torch.Generator(device=_a       ).manual_seed(_a       )
         UpperCAmelCase							=     {
             """batch_size""": 1,
             """generator""": generator,
             """num_inference_steps""": 2,
             """output_type""": """numpy""",
         }
         return inputs





      def        _lowercase				(		self						:					List[Any]       )    ->   int:






         """simple docstring"""







         UpperCAmelCase							=     """cpu"""

         UpperCAmelCase							=     self.get_dummy_components()
         UpperCAmelCase							=     self.pipeline_class(**_a       )
         pipe.to(_a       )
         pipe.set_progress_bar_config(disable=_a       )

         UpperCAmelCase							=     self.get_dummy_inputs(_a       )
         UpperCAmelCase							=     pipe(**_a       ).images
         UpperCAmelCase							=     image[0, -3:, -3:, -1]

         self.assertEqual(image.shape					,       (1, 3_2, 3_2, 3)       )
         UpperCAmelCase							=     np.array(
             [1.0_0_0e0_0, 5.7_1_7e-0_1, 4.7_1_7e-0_1, 1.0_0_0e0_0, 0.0_0_0e0_0, 1.0_0_0e0_0, 3.0_0_0e-0_4, 0.0_0_0e0_0, 9.0_0_0e-0_4]       )
         UpperCAmelCase							=     np.abs(image_slice.flatten() - expected_slice       ).max()
         self.assertLessEqual(_a					,       1e-3       )





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






         """simple docstring"""







         super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3       )





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






         """simple docstring"""







         super().test_save_load_local(expected_max_difference=3e-3       )





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






         """simple docstring"""







         super().test_save_load_optional_components(expected_max_difference=3e-3       )





      def        _lowercase				(		self						:					str       )    ->   str:






         """simple docstring"""







         super().test_inference_batch_single_identical(expected_max_diff=3e-3       )






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





      def        _lowercase				(		self						:					List[str]       )    ->   int:






         """simple docstring"""







         UpperCAmelCase							=     """google/ddpm-cifar10-32"""

         UpperCAmelCase							=     UNetaDModel.from_pretrained(_a       )
         UpperCAmelCase							=     DDIMScheduler()

         UpperCAmelCase							=     DDIMPipeline(unet=_a					,       scheduler=_a       )
         ddim.to(_a       )
         ddim.set_progress_bar_config(disable=_a       )

         UpperCAmelCase							=     torch.manual_seed(0       )
         UpperCAmelCase							=     ddim(generator=_a					,       eta=0.0					,       output_type="""numpy"""       ).images

         UpperCAmelCase							=     image[0, -3:, -3:, -1]

         assert image.shape == (1, 3_2, 3_2, 3)
         UpperCAmelCase							=     np.array([0.1_723, 0.1_617, 0.1_600, 0.1_626, 0.1_497, 0.1_513, 0.1_505, 0.1_442, 0.1_453]       )

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





      def        _lowercase				(		self						:					int       )    ->   List[Any]:






         """simple docstring"""







         UpperCAmelCase							=     """google/ddpm-ema-bedroom-256"""

         UpperCAmelCase							=     UNetaDModel.from_pretrained(_a       )
         UpperCAmelCase							=     DDIMScheduler.from_pretrained(_a       )

         UpperCAmelCase							=     DDIMPipeline(unet=_a					,       scheduler=_a       )
         ddpm.to(_a       )
         ddpm.set_progress_bar_config(disable=_a       )

         UpperCAmelCase							=     torch.manual_seed(0       )
         UpperCAmelCase							=     ddpm(generator=_a					,       output_type="""numpy"""       ).images

         UpperCAmelCase							=     image[0, -3:, -3:, -1]

         assert image.shape == (1, 2_5_6, 2_5_6, 3)
         UpperCAmelCase							=     np.array([0.0_060, 0.0_201, 0.0_344, 0.0_024, 0.0_018, 0.0_002, 0.0_022, 0.0_000, 0.0_069]       )

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







from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import DeiTConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
 import tensorflow as tf

 from transformers import (
     TFDeiTForImageClassification,
     TFDeiTForImageClassificationWithTeacher,
     TFDeiTForMaskedImageModeling,
     TFDeiTModel,
 )
 from transformers.models.deit.modeling_tf_deit import TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import DeiTImageProcessor






class   __lowercase				:





      def __init__(		self						:					str					,       __lowerCamelCase						:					Optional[Any]					,       __lowerCamelCase						:					Optional[Any]=1_3					,       __lowerCamelCase						:					Optional[Any]=3_0					,       __lowerCamelCase						:					Any=2					,       __lowerCamelCase						:					List[str]=3					,       __lowerCamelCase						:					str=True					,       __lowerCamelCase						:					str=True					,       __lowerCamelCase						:					Union[str, Any]=3_2					,       __lowerCamelCase						:					Optional[int]=2					,       __lowerCamelCase						:					Optional[int]=4					,       __lowerCamelCase						:					Union[str, Any]=3_7					,       __lowerCamelCase						:					str="gelu"					,       __lowerCamelCase						:					str=0.1					,       __lowerCamelCase						:					List[str]=0.1					,       __lowerCamelCase						:					str=1_0					,       __lowerCamelCase						:					Union[str, Any]=0.02					,       __lowerCamelCase						:					List[str]=3					,       __lowerCamelCase						:					Optional[Any]=None					,       __lowerCamelCase						:					Any=2					,       )    ->   Tuple:






         """simple docstring"""







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

         # in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
         UpperCAmelCase							=     (image_size // patch_size) ** 2
         UpperCAmelCase							=     num_patches + 2





      def        _lowercase				(		self						:					Tuple       )    ->   int:






         """simple docstring"""







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

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

         UpperCAmelCase							=     self.get_config()

         return config, pixel_values, labels





      def        _lowercase				(		self						:					Tuple       )    ->   str:






         """simple docstring"""







         return DeiTConfig(
             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=__lowerCamelCase					,       initializer_range=self.initializer_range					,       encoder_stride=self.encoder_stride					,       )





      def        _lowercase				(		self						:					Any					,       __lowerCamelCase						:					Optional[int]					,       __lowerCamelCase						:					Any					,       __lowerCamelCase						:					int       )    ->   Any:






         """simple docstring"""







         UpperCAmelCase							=     TFDeiTModel(config=__lowerCamelCase       )
         UpperCAmelCase							=     model(__lowerCamelCase       )
         self.parent.assertEqual(result.last_hidden_state.shape					,       (self.batch_size, self.seq_length, self.hidden_size)       )





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






         """simple docstring"""







         UpperCAmelCase							=     TFDeiTForMaskedImageModeling(config=__lowerCamelCase       )
         UpperCAmelCase							=     model(__lowerCamelCase       )
         self.parent.assertEqual(
             result.reconstruction.shape					,       (self.batch_size, self.num_channels, self.image_size, self.image_size)       )

         # test greyscale images
         UpperCAmelCase							=     1
         UpperCAmelCase							=     TFDeiTForMaskedImageModeling(__lowerCamelCase       )

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





      def        _lowercase				(		self						:					Union[str, Any]					,       __lowerCamelCase						:					List[Any]					,       __lowerCamelCase						:					Dict					,       __lowerCamelCase						:					int       )    ->   Any:






         """simple docstring"""







         UpperCAmelCase							=     self.type_sequence_label_size
         UpperCAmelCase							=     TFDeiTForImageClassification(__lowerCamelCase       )
         UpperCAmelCase							=     model(__lowerCamelCase					,       labels=__lowerCamelCase       )
         self.parent.assertEqual(result.logits.shape					,       (self.batch_size, self.type_sequence_label_size)       )

         # test greyscale images
         UpperCAmelCase							=     1
         UpperCAmelCase							=     TFDeiTForImageClassification(__lowerCamelCase       )

         UpperCAmelCase							=     floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
         UpperCAmelCase							=     model(__lowerCamelCase					,       labels=__lowerCamelCase       )
         self.parent.assertEqual(result.logits.shape					,       (self.batch_size, self.type_sequence_label_size)       )





      def        _lowercase				(		self						:					Tuple       )    ->   List[str]:






         """simple docstring"""







         UpperCAmelCase							=     self.prepare_config_and_inputs()
         UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase							=     config_and_inputs
         UpperCAmelCase							=     {"""pixel_values""": pixel_values}
         return config, inputs_dict






@require_tf
class   __lowercase				(  __snake_case ,					__snake_case ,					unittest.TestCase  ):
      UpperCamelCase								=       (
          (
              TFDeiTModel,
              TFDeiTForImageClassification,
              TFDeiTForImageClassificationWithTeacher,
              TFDeiTForMaskedImageModeling,
          )
          if is_tf_available()
          else ()
      )
      UpperCamelCase								=       (
          {
              '''feature-extraction''': TFDeiTModel,
              '''image-classification''': (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
          }
          if is_tf_available()
          else {}
      )

      UpperCamelCase								=       False
      UpperCamelCase								=       False
      UpperCamelCase								=       False
      UpperCamelCase								=       False





      def        _lowercase				(		self						:					str       )    ->   str:






         """simple docstring"""







         UpperCAmelCase							=     TFDeiTModelTester(self       )
         UpperCAmelCase							=     ConfigTester(self					,       config_class=__lowerCamelCase					,       has_text_modality=__lowerCamelCase					,       hidden_size=3_7       )





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






         """simple docstring"""







         self.config_tester.run_common_tests()





      @unittest.skip(reason="""DeiT does not use inputs_embeds"""       )
      def        _lowercase				(		self						:					List[Any]       )    ->   Dict:






         """simple docstring"""







         pass





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






         """simple docstring"""







         UpperCAmelCase				,  UpperCAmelCase							=     self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
            UpperCAmelCase							=     model_class(__lowerCamelCase       )
            self.assertIsInstance(model.get_input_embeddings()					,       (tf.keras.layers.Layer)       )
            UpperCAmelCase							=     model.get_output_embeddings()
            self.assertTrue(x is None or isinstance(__lowerCamelCase					,       tf.keras.layers.Dense       )       )





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






         """simple docstring"""







         UpperCAmelCase				,  UpperCAmelCase							=     self.model_tester.prepare_config_and_inputs_for_common()

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

            UpperCAmelCase							=     ["""pixel_values"""]
            self.assertListEqual(arg_names[:1]					,       __lowerCamelCase       )





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






         """simple docstring"""







         UpperCAmelCase							=     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__lowerCamelCase       )





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






         """simple docstring"""







         UpperCAmelCase							=     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_image_modeling(*__lowerCamelCase       )





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






         """simple docstring"""







         UpperCAmelCase							=     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*__lowerCamelCase       )





      def        _lowercase				(		self						:					Optional[int]					,       __lowerCamelCase						:					List[Any]					,       __lowerCamelCase						:					List[Any]					,       __lowerCamelCase						:					Any=False       )    ->   int:






         """simple docstring"""







         UpperCAmelCase							=     super()._prepare_for_class(__lowerCamelCase					,       __lowerCamelCase					,       return_labels=__lowerCamelCase       )

         if return_labels:
            if "labels" in inputs_dict and "labels" not in inspect.signature(model_class.call       ).parameters:
               del inputs_dict["labels"]

         return inputs_dict





      @slow
      def        _lowercase				(		self						:					Optional[int]       )    ->   str:






         """simple docstring"""







         for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            UpperCAmelCase							=     TFDeiTModel.from_pretrained(__lowerCamelCase       )
            self.assertIsNotNone(__lowerCamelCase       )

def   _UpperCamelCase				(						)						->Tuple:
   UpperCAmelCase							=     Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""   )
   return image







@require_tf
@require_vision
class   __lowercase				(  unittest.TestCase  ):





      @cached_property
      def        _lowercase				(		self						:					List[str]       )    ->   Union[str, Any]:






         """simple docstring"""







         return (
             DeiTImageProcessor.from_pretrained("""facebook/deit-base-distilled-patch16-224"""       )
             if is_vision_available()
             else None
         )





      @slow
      def        _lowercase				(		self						:					Optional[int]       )    ->   List[str]:






         """simple docstring"""







         UpperCAmelCase							=     TFDeiTForImageClassificationWithTeacher.from_pretrained("""facebook/deit-base-distilled-patch16-224"""       )

         UpperCAmelCase							=     self.default_image_processor
         UpperCAmelCase							=     prepare_img()
         UpperCAmelCase							=     image_processor(images=__lowerCamelCase					,       return_tensors="""tf"""       )

         # forward pass
         UpperCAmelCase							=     model(**__lowerCamelCase       )

         # verify the logits
         UpperCAmelCase							=     tf.TensorShape((1, 1_0_0_0)       )
         self.assertEqual(outputs.logits.shape					,       __lowerCamelCase       )

         UpperCAmelCase							=     tf.constant([-1.0_266, 0.1_912, -1.2_861]       )

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



