
"""simple docstring"""

import gc
import random
import unittest

import torch

from diffusers import (
    IFImgaImgPipeline,
    IFImgaImgSuperResolutionPipeline,
    IFInpaintingPipeline,
    IFInpaintingSuperResolutionPipeline,
    IFPipeline,
    IFSuperResolutionPipeline,
)
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import floats_tensor, load_numpy, require_torch_gpu, skip_mps, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
from . import IFPipelineTesterMixin
@skip_mps
class   _UpperCAmelCase				(		SCREAMING_SNAKE_CASE_  ,      SCREAMING_SNAKE_CASE_  ,      unittest.TestCase				):
      __SCREAMING_SNAKE_CASE :				Tuple				   =						IFPipeline
      __SCREAMING_SNAKE_CASE :				Dict				   =						TEXT_TO_IMAGE_PARAMS - {"width", "height", "latents"}
      __SCREAMING_SNAKE_CASE :				Optional[int]				   =						TEXT_TO_IMAGE_BATCH_PARAMS
      __SCREAMING_SNAKE_CASE :				str				   =						PipelineTesterMixin.required_optional_params - {"latents"}




      def      a_   (				self				)    ->       int:
          return self._get_dummy_components()




      def      a_   (				self ,							lowercase_ ,							lowercase_=0				)    ->       Tuple:





          if str(lowercase_				).startswith('mps'				):
              UpperCAmelCase          =       torch.manual_seed(lowercase_				)
          else:
              UpperCAmelCase          =       torch.Generator(device=lowercase_				).manual_seed(lowercase_				)

          UpperCAmelCase          =       {
              "prompt": "A painting of a squirrel eating a burger",
              "generator": generator,
              "num_inference_steps": 2,
              "output_type": "numpy",
          }

          return inputs




      def      a_   (				self				)    ->       Optional[Any]:
          self._test_save_load_optional_components()




      @unittest.skipIf(torch_device != 'cuda' ,							reason='float16 requires CUDA'				)
      def      a_   (				self				)    ->       Any:
          super().test_save_load_floataa(expected_max_diff=1E-1				)




      def      a_   (				self				)    ->       Dict:
          self._test_attention_slicing_forward_pass(expected_max_diff=1E-2				)




      def      a_   (				self				)    ->       Tuple:
          self._test_save_load_local()




      def      a_   (				self				)    ->       List[str]:
          self._test_inference_batch_single_identical(
              expected_max_diff=1E-2 ,							)




      @unittest.skipIf(
          torch_device != 'cuda' or not is_xformers_available() ,							reason='XFormers attention is only available with CUDA and `xformers` installed' ,							)
      def      a_   (				self				)    ->       List[str]:
          self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3				)

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




      def      a_   (				self				)    ->       Tuple:
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()




      def      a_   (				self				)    ->       Dict:
          UpperCAmelCase          =       IFPipeline.from_pretrained('DeepFloyd/IF-I-XL-v1.0' ,							variant='fp16' ,							torch_dtype=torch.floataa				)

          UpperCAmelCase          =       IFSuperResolutionPipeline.from_pretrained(
              'DeepFloyd/IF-II-L-v1.0' ,							variant='fp16' ,							torch_dtype=torch.floataa ,							text_encoder=lowercase_ ,							tokenizer=lowercase_				)

          # pre compute text embeddings and remove T5 to save memory

          pipe_a.text_encoder.to('cuda'				)

          UpperCAmelCase          =       pipe_a.encode_prompt('anime turtle' ,							device='cuda'				)

          del pipe_a.tokenizer
          del pipe_a.text_encoder
          gc.collect()

          UpperCAmelCase          =       None
          UpperCAmelCase          =       None

          pipe_a.enable_model_cpu_offload()
          pipe_a.enable_model_cpu_offload()

          pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()				)
          pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()				)

          self._test_if(lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_				)

          pipe_a.remove_all_hooks()
          pipe_a.remove_all_hooks()

          # img2img

          UpperCAmelCase          =       IFImgaImgPipeline(**pipe_a.components				)
          UpperCAmelCase          =       IFImgaImgSuperResolutionPipeline(**pipe_a.components				)

          pipe_a.enable_model_cpu_offload()
          pipe_a.enable_model_cpu_offload()

          pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()				)
          pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()				)

          self._test_if_imgaimg(lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_				)

          pipe_a.remove_all_hooks()
          pipe_a.remove_all_hooks()

          # inpainting

          UpperCAmelCase          =       IFInpaintingPipeline(**pipe_a.components				)
          UpperCAmelCase          =       IFInpaintingSuperResolutionPipeline(**pipe_a.components				)

          pipe_a.enable_model_cpu_offload()
          pipe_a.enable_model_cpu_offload()

          pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()				)
          pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()				)

          self._test_if_inpainting(lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_				)




      def      a_   (				self ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_				)    ->       str:
          _start_torch_memory_measurement()

          UpperCAmelCase          =       torch.Generator(device='cpu'				).manual_seed(0				)
          UpperCAmelCase          =       pipe_a(
              prompt_embeds=lowercase_ ,							negative_prompt_embeds=lowercase_ ,							num_inference_steps=2 ,							generator=lowercase_ ,							output_type='np' ,							)

          UpperCAmelCase          =       output.images[0]

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

          UpperCAmelCase          =       torch.cuda.max_memory_allocated()
          assert mem_bytes < 1_3 * 1_0**9

          UpperCAmelCase          =       load_numpy(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy'				)
          assert_mean_pixel_difference(lowercase_ ,							lowercase_				)

          # pipeline 2

          _start_torch_memory_measurement()

          UpperCAmelCase          =       torch.Generator(device='cpu'				).manual_seed(0				)

          UpperCAmelCase          =       floats_tensor((1, 3, 6_4, 6_4) ,							rng=random.Random(0				)				).to(lowercase_				)

          UpperCAmelCase          =       pipe_a(
              prompt_embeds=lowercase_ ,							negative_prompt_embeds=lowercase_ ,							image=lowercase_ ,							generator=lowercase_ ,							num_inference_steps=2 ,							output_type='np' ,							)

          UpperCAmelCase          =       output.images[0]

          assert image.shape == (2_5_6, 2_5_6, 3)

          UpperCAmelCase          =       torch.cuda.max_memory_allocated()
          assert mem_bytes < 4 * 1_0**9

          UpperCAmelCase          =       load_numpy(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy'				)
          assert_mean_pixel_difference(lowercase_ ,							lowercase_				)




      def      a_   (				self ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_				)    ->       List[str]:
          _start_torch_memory_measurement()

          UpperCAmelCase          =       floats_tensor((1, 3, 6_4, 6_4) ,							rng=random.Random(0				)				).to(lowercase_				)

          UpperCAmelCase          =       torch.Generator(device='cpu'				).manual_seed(0				)

          UpperCAmelCase          =       pipe_a(
              prompt_embeds=lowercase_ ,							negative_prompt_embeds=lowercase_ ,							image=lowercase_ ,							num_inference_steps=2 ,							generator=lowercase_ ,							output_type='np' ,							)

          UpperCAmelCase          =       output.images[0]

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

          UpperCAmelCase          =       torch.cuda.max_memory_allocated()
          assert mem_bytes < 1_0 * 1_0**9

          UpperCAmelCase          =       load_numpy(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy'				)
          assert_mean_pixel_difference(lowercase_ ,							lowercase_				)

          # pipeline 2

          _start_torch_memory_measurement()

          UpperCAmelCase          =       torch.Generator(device='cpu'				).manual_seed(0				)

          UpperCAmelCase          =       floats_tensor((1, 3, 2_5_6, 2_5_6) ,							rng=random.Random(0				)				).to(lowercase_				)
          UpperCAmelCase          =       floats_tensor((1, 3, 6_4, 6_4) ,							rng=random.Random(0				)				).to(lowercase_				)

          UpperCAmelCase          =       pipe_a(
              prompt_embeds=lowercase_ ,							negative_prompt_embeds=lowercase_ ,							image=lowercase_ ,							original_image=lowercase_ ,							generator=lowercase_ ,							num_inference_steps=2 ,							output_type='np' ,							)

          UpperCAmelCase          =       output.images[0]

          assert image.shape == (2_5_6, 2_5_6, 3)

          UpperCAmelCase          =       torch.cuda.max_memory_allocated()
          assert mem_bytes < 4 * 1_0**9

          UpperCAmelCase          =       load_numpy(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img_superresolution_stage_II.npy'				)
          assert_mean_pixel_difference(lowercase_ ,							lowercase_				)




      def      a_   (				self ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_				)    ->       List[str]:
          _start_torch_memory_measurement()

          UpperCAmelCase          =       floats_tensor((1, 3, 6_4, 6_4) ,							rng=random.Random(0				)				).to(lowercase_				)
          UpperCAmelCase          =       floats_tensor((1, 3, 6_4, 6_4) ,							rng=random.Random(1				)				).to(lowercase_				)

          UpperCAmelCase          =       torch.Generator(device='cpu'				).manual_seed(0				)
          UpperCAmelCase          =       pipe_a(
              prompt_embeds=lowercase_ ,							negative_prompt_embeds=lowercase_ ,							image=lowercase_ ,							mask_image=lowercase_ ,							num_inference_steps=2 ,							generator=lowercase_ ,							output_type='np' ,							)

          UpperCAmelCase          =       output.images[0]

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

          UpperCAmelCase          =       torch.cuda.max_memory_allocated()
          assert mem_bytes < 1_0 * 1_0**9

          UpperCAmelCase          =       load_numpy(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy'				)
          assert_mean_pixel_difference(lowercase_ ,							lowercase_				)

          # pipeline 2

          _start_torch_memory_measurement()

          UpperCAmelCase          =       torch.Generator(device='cpu'				).manual_seed(0				)

          UpperCAmelCase          =       floats_tensor((1, 3, 6_4, 6_4) ,							rng=random.Random(0				)				).to(lowercase_				)
          UpperCAmelCase          =       floats_tensor((1, 3, 2_5_6, 2_5_6) ,							rng=random.Random(0				)				).to(lowercase_				)
          UpperCAmelCase          =       floats_tensor((1, 3, 2_5_6, 2_5_6) ,							rng=random.Random(1				)				).to(lowercase_				)

          UpperCAmelCase          =       pipe_a(
              prompt_embeds=lowercase_ ,							negative_prompt_embeds=lowercase_ ,							image=lowercase_ ,							mask_image=lowercase_ ,							original_image=lowercase_ ,							generator=lowercase_ ,							num_inference_steps=2 ,							output_type='np' ,							)

          UpperCAmelCase          =       output.images[0]

          assert image.shape == (2_5_6, 2_5_6, 3)

          UpperCAmelCase          =       torch.cuda.max_memory_allocated()
          assert mem_bytes < 4 * 1_0**9

          UpperCAmelCase          =       load_numpy(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting_superresolution_stage_II.npy'				)
          assert_mean_pixel_difference(lowercase_ ,							lowercase_				)




def      lowercase__    (       )			->   List[str]:






    """simple docstring"""




    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()







'''simple docstring'''


import inspect
import unittest

import numpy as np

from transformers import BeitConfig
from transformers.testing_utils import require_flax, require_vision, slow
from transformers.utils import cached_property, is_flax_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor


if is_flax_available():
      import jax

      from transformers import FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling, FlaxBeitModel

if is_vision_available():
      from PIL import Image

      from transformers import BeitImageProcessor







class 							__magic_name__	(       unittest.TestCase ):

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase			,     lowerCamelCase=100			,     lowerCamelCase=13			,     lowerCamelCase=30			,     lowerCamelCase=2			,     lowerCamelCase=3			,     lowerCamelCase=True			,     lowerCamelCase=True			,     lowerCamelCase=32			,     lowerCamelCase=5			,     lowerCamelCase=4			,     lowerCamelCase=37			,     lowerCamelCase="gelu"			,     lowerCamelCase=0.1			,     lowerCamelCase=0.1			,     lowerCamelCase=10			,     lowerCamelCase=0.02			,     lowerCamelCase=3			,     ):


          '''simple docstring'''





          __A	:       Tuple      =							parent
          __A	:       Dict      =							vocab_size
          __A	:       Union[str, Any]      =							batch_size
          __A	:       str      =							image_size
          __A	:       Optional[Any]      =							patch_size
          __A	:       Optional[Any]      =							num_channels
          __A	:       Optional[Any]      =							is_training
          __A	:       Any      =							use_labels
          __A	:       List[str]      =							hidden_size
          __A	:       Union[str, Any]      =							num_hidden_layers
          __A	:       Optional[Any]      =							num_attention_heads
          __A	:       Optional[int]      =							intermediate_size
          __A	:       Optional[int]      =							hidden_act
          __A	:       Dict      =							hidden_dropout_prob
          __A	:       str      =							attention_probs_dropout_prob
          __A	:       Optional[int]      =							type_sequence_label_size
          __A	:       Optional[int]      =							initializer_range

          # in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
          __A	:       Optional[int]      =							(image_size // patch_size) ** 2
          __A	:       Union[str, Any]      =							num_patches + 1


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       int      =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

          __A	:       str      =							None
          if self.use_labels:
               __A	:       List[str]      =							ids_tensor([self.batch_size]			,     self.type_sequence_label_size		)

          __A	:       Dict      =							BeitConfig(
              vocab_size=self.vocab_size			,     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			,     )

          return config, pixel_values, labels


     def lowerCAmelCase__  (   self			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase		):


          '''simple docstring'''





          __A	:       Optional[int]      =							FlaxBeitModel(config=lowerCamelCase		)
          __A	:       Optional[Any]      =							model(lowerCamelCase		)
          self.parent.assertEqual(result.last_hidden_state.shape			,     (self.batch_size, self.seq_length, self.hidden_size)		)


     def lowerCAmelCase__  (   self			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase		):


          '''simple docstring'''





          __A	:       Dict      =							FlaxBeitForMaskedImageModeling(config=lowerCamelCase		)
          __A	:       Union[str, Any]      =							model(lowerCamelCase		)
          self.parent.assertEqual(result.logits.shape			,     (self.batch_size, self.seq_length - 1, self.vocab_size)		)


     def lowerCAmelCase__  (   self			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase		):


          '''simple docstring'''





          __A	:       Any      =							self.type_sequence_label_size
          __A	:       List[str]      =							FlaxBeitForImageClassification(config=lowerCamelCase		)
          __A	:       Tuple      =							model(lowerCamelCase		)
          self.parent.assertEqual(result.logits.shape			,     (self.batch_size, self.type_sequence_label_size)		)

          # test greyscale images
          __A	:       Tuple      =							1
          __A	:       Optional[Any]      =							FlaxBeitForImageClassification(lowerCamelCase		)

          __A	:       Union[str, Any]      =							floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
          __A	:       str      =							model(lowerCamelCase		)


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





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







@require_flax
class 							__magic_name__	(       lowerCAmelCase    ,		unittest.TestCase ):

     """simple docstring"""






     lowerCamelCase__       						=				(
         (FlaxBeitModel, FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling) if is_flax_available() else ()
     )


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Optional[int]      =							FlaxBeitModelTester(self		)
          __A	:       List[Any]      =							ConfigTester(self			,     config_class=lowerCamelCase			,     has_text_modality=lowerCamelCase			,     hidden_size=37		)


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          self.config_tester.run_common_tests()


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A  ,__A	:       Optional[int]      =							self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
               __A	:       int      =							model_class(lowerCamelCase		)
               __A	:       Dict      =							inspect.signature(model.__call__		)
               # signature.parameters is an OrderedDict => so arg_names order is deterministic
               __A	:       List[str]      =							[*signature.parameters.keys()]

               __A	:       Optional[Any]      =							["pixel_values"]
               self.assertListEqual(arg_names[:1]			,     lowerCamelCase		)


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A  ,__A	:       Optional[int]      =							self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
               with self.subTest(model_class.__name__		):
                    __A	:       Optional[int]      =							self._prepare_for_class(lowerCamelCase			,     lowerCamelCase		)
                    __A	:       Optional[int]      =							model_class(lowerCamelCase		)

                    @jax.jit
                    def model_jitted(lowerCamelCase			,     **lowerCamelCase		):
                         return model(pixel_values=lowerCamelCase			,     **lowerCamelCase		)

                    with self.subTest("JIT Enabled"		):
                         __A	:       Tuple      =							model_jitted(**lowerCamelCase		).to_tuple()

                    with self.subTest("JIT Disabled"		):
                         with jax.disable_jit():
                              __A	:       Dict      =							model_jitted(**lowerCamelCase		).to_tuple()

                    self.assertEqual(len(lowerCamelCase		)			,     len(lowerCamelCase		)		)
                    for jitted_output, output in zip(lowerCamelCase			,     lowerCamelCase		):
                         self.assertEqual(jitted_output.shape			,     output.shape		)


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Dict      =							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*lowerCamelCase		)


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Optional[Any]      =							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase		)


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       List[str]      =							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_classification(*lowerCamelCase		)


     @slow
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          for model_class_name in self.all_model_classes:
               __A	:       Any      =							model_class_name.from_pretrained("microsoft/beit-base-patch16-224"		)
               __A	:       Optional[int]      =							model(np.ones((1, 3, 224, 224)		)		)
               self.assertIsNotNone(lowerCamelCase		)


def 		_lowercase ():




     '''simple docstring'''





     __A	:       Dict      =							Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"							)
     return image







@require_vision
@require_flax
class 							__magic_name__	(       unittest.TestCase ):

     """simple docstring"""






     @cached_property
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          return BeitImageProcessor.from_pretrained("microsoft/beit-base-patch16-224"		) if is_vision_available() else None


     @slow
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Union[str, Any]      =							FlaxBeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k"		)

          __A	:       List[Any]      =							self.default_image_processor
          __A	:       Dict      =							prepare_img()
          __A	:       str      =							image_processor(images=lowerCamelCase			,     return_tensors="np"		).pixel_values

          # prepare bool_masked_pos
          __A	:       List[Any]      =							np.ones((1, 196)			,     dtype=lowerCamelCase		)

          # forward pass
          __A	:       Union[str, Any]      =							model(pixel_values=lowerCamelCase			,     bool_masked_pos=lowerCamelCase		)
          __A	:       Dict      =							outputs.logits

          # verify the logits
          __A	:       int      =							(1, 196, 8192)
          self.assertEqual(logits.shape			,     lowerCamelCase		)

          __A	:       Optional[int]      =							np.array(
              [[-3.2437, 0.5072, -13.9174], [-3.2456, 0.4948, -13.9401], [-3.2033, 0.5121, -13.8550]]		)

          self.assertTrue(np.allclose(logits[bool_masked_pos][:3, :3]			,     lowerCamelCase			,     atol=1E-2		)		)


     @slow
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Union[str, Any]      =							FlaxBeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224"		)

          __A	:       Any      =							self.default_image_processor
          __A	:       List[Any]      =							prepare_img()
          __A	:       List[Any]      =							image_processor(images=lowerCamelCase			,     return_tensors="np"		)

          # forward pass
          __A	:       int      =							model(**lowerCamelCase		)
          __A	:       Any      =							outputs.logits

          # verify the logits
          __A	:       str      =							(1, 1000)
          self.assertEqual(logits.shape			,     lowerCamelCase		)

          __A	:       int      =							np.array([-1.2385, -1.0987, -1.0108]		)

          self.assertTrue(np.allclose(logits[0, :3]			,     lowerCamelCase			,     atol=1E-4		)		)

          __A	:       Optional[int]      =							281
          self.assertEqual(logits.argmax(-1		).item()			,     lowerCamelCase		)


     @slow
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Any      =							FlaxBeitForImageClassification.from_pretrained("microsoft/beit-large-patch16-224-pt22k-ft22k"		)

          __A	:       Optional[int]      =							self.default_image_processor
          __A	:       Tuple      =							prepare_img()
          __A	:       List[str]      =							image_processor(images=lowerCamelCase			,     return_tensors="np"		)

          # forward pass
          __A	:       List[Any]      =							model(**lowerCamelCase		)
          __A	:       Dict      =							outputs.logits

          # verify the logits
          __A	:       Union[str, Any]      =							(1, 2_1841)
          self.assertEqual(logits.shape			,     lowerCamelCase		)

          __A	:       Tuple      =							np.array([1.6881, -0.2787, 0.5901]		)

          self.assertTrue(np.allclose(logits[0, :3]			,     lowerCamelCase			,     atol=1E-4		)		)

          __A	:       Optional[int]      =							2396
          self.assertEqual(logits.argmax(-1		).item()			,     lowerCamelCase		)

