


'''simple docstring'''






# This model implementation is heavily inspired by https://github.com/haofanwang/ControlNet-for-Diffusers/

import gc
import random
import tempfile
import unittest

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

from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DDIMScheduler,
    StableDiffusionControlNetImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel
from diffusers.utils import floats_tensor, load_image, load_numpy, randn_tensor, slow, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
)


enable_full_determinism()






class        __snake_case	(       _SCREAMING_SNAKE_CASE     ,_SCREAMING_SNAKE_CASE     ,_SCREAMING_SNAKE_CASE     ,unittest.TestCase):







    """simple docstring"""



    lowercase       			= StableDiffusionControlNetImgaImgPipeline
    lowercase       			= TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width'}
    lowercase       			= TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    lowercase       			= IMAGE_TO_IMAGE_IMAGE_PARAMS.union({'control_image'})
    lowercase       			= IMAGE_TO_IMAGE_IMAGE_PARAMS

    def __lowercase       (       self   :  Union[str, Any]		)       ->  List[str]:
      torch.manual_seed(0		)
      lowerCAmelCase_      :				Optional[int]      =       UNetaDConditionModel(
          block_out_channels=(32, 64)   ,						layers_per_block=2   ,						sample_size=32   ,						in_channels=4   ,						out_channels=4   ,						down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")   ,						up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")   ,						cross_attention_dim=32   ,						)
      torch.manual_seed(0		)
      lowerCAmelCase_      :				str      =       ControlNetModel(
          block_out_channels=(32, 64)   ,						layers_per_block=2   ,						in_channels=4   ,						down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")   ,						cross_attention_dim=32   ,						conditioning_embedding_out_channels=(16, 32)   ,						)
      torch.manual_seed(0		)
      lowerCAmelCase_      :				int      =       DDIMScheduler(
          beta_start=0.00_085   ,						beta_end=0.012   ,						beta_schedule="""scaled_linear"""   ,						clip_sample=lowerCamelCase   ,						set_alpha_to_one=lowerCamelCase   ,						)
      torch.manual_seed(0		)
      lowerCAmelCase_      :				Union[str, Any]      =       AutoencoderKL(
          block_out_channels=[32, 64]   ,						in_channels=3   ,						out_channels=3   ,						down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]   ,						up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]   ,						latent_channels=4   ,						)
      torch.manual_seed(0		)
      lowerCAmelCase_      :				Any      =       CLIPTextConfig(
          bos_token_id=0   ,						eos_token_id=2   ,						hidden_size=32   ,						intermediate_size=37   ,						layer_norm_eps=1E-05   ,						num_attention_heads=4   ,						num_hidden_layers=5   ,						pad_token_id=1   ,						vocab_size=10_00   ,						)
      lowerCAmelCase_      :				Optional[Any]      =       CLIPTextModel(lowerCamelCase		)
      lowerCAmelCase_      :				int      =       CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""		)

      lowerCAmelCase_      :				Optional[Any]      =       {
          """unet""": unet,
          """controlnet""": controlnet,
          """scheduler""": scheduler,
          """vae""": vae,
          """text_encoder""": text_encoder,
          """tokenizer""": tokenizer,
          """safety_checker""": None,
          """feature_extractor""": None,
      }
      return components

    def __lowercase       (       self   :  str   ,						lowerCamelCase   :  Dict   ,						lowerCamelCase   :  int=0		)       ->  List[Any]:
      if str(lowerCamelCase		).startswith("""mps"""		):
        lowerCAmelCase_      :				str      =       torch.manual_seed(lowerCamelCase		)
      else:
        lowerCAmelCase_      :				str      =       torch.Generator(device=lowerCamelCase		).manual_seed(lowerCamelCase		)

      lowerCAmelCase_      :				str      =       2
      lowerCAmelCase_      :				Any      =       randn_tensor(
          (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)   ,						generator=lowerCamelCase   ,						device=torch.device(lowerCamelCase		)   ,						)
      lowerCAmelCase_      :				Union[str, Any]      =       floats_tensor(control_image.shape   ,						rng=random.Random(lowerCamelCase		)		).to(lowerCamelCase		)
      lowerCAmelCase_      :				Dict      =       image.cpu().permute(0   ,						2   ,						3   ,						1		)[0]
      lowerCAmelCase_      :				Union[str, Any]      =       Image.fromarray(np.uinta(lowerCamelCase		)		).convert("""RGB"""		).resize((64, 64)		)
      lowerCAmelCase_      :				Optional[int]      =       {
          """prompt""": """A painting of a squirrel eating a burger""",
          """generator""": generator,
          """num_inference_steps""": 2,
          """guidance_scale""": 6.0,
          """output_type""": """numpy""",
          """image""": image,
          """control_image""": control_image,
      }

      return inputs

    def __lowercase       (       self   :  str		)       ->  Tuple:
      return self._test_attention_slicing_forward_pass(expected_max_diff=2E-3		)

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





    def __lowercase       (       self   :  List[Any]		)       ->  Any:
      self._test_inference_batch_single_identical(expected_max_diff=2E-3		)







class        __snake_case	(       _SCREAMING_SNAKE_CASE     ,_SCREAMING_SNAKE_CASE     ,unittest.TestCase):







    """simple docstring"""



    lowercase       			= StableDiffusionControlNetImgaImgPipeline
    lowercase       			= TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width'}
    lowercase       			= TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    lowercase       			= frozenset([])  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess

    def __lowercase       (       self   :  List[Any]		)       ->  Tuple:
      torch.manual_seed(0		)
      lowerCAmelCase_      :				Optional[Any]      =       UNetaDConditionModel(
          block_out_channels=(32, 64)   ,						layers_per_block=2   ,						sample_size=32   ,						in_channels=4   ,						out_channels=4   ,						down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")   ,						up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")   ,						cross_attention_dim=32   ,						)
      torch.manual_seed(0		)

      def init_weights(lowerCamelCase   :  Optional[Any]		):
        if isinstance(lowerCamelCase   ,						torch.nn.Convad		):
          torch.nn.init.normal(m.weight		)
          m.bias.data.fill_(1.0		)

      lowerCAmelCase_      :				List[Any]      =       ControlNetModel(
          block_out_channels=(32, 64)   ,						layers_per_block=2   ,						in_channels=4   ,						down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")   ,						cross_attention_dim=32   ,						conditioning_embedding_out_channels=(16, 32)   ,						)
      controlneta.controlnet_down_blocks.apply(lowerCamelCase		)

      torch.manual_seed(0		)
      lowerCAmelCase_      :				str      =       ControlNetModel(
          block_out_channels=(32, 64)   ,						layers_per_block=2   ,						in_channels=4   ,						down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")   ,						cross_attention_dim=32   ,						conditioning_embedding_out_channels=(16, 32)   ,						)
      controlneta.controlnet_down_blocks.apply(lowerCamelCase		)

      torch.manual_seed(0		)
      lowerCAmelCase_      :				Tuple      =       DDIMScheduler(
          beta_start=0.00_085   ,						beta_end=0.012   ,						beta_schedule="""scaled_linear"""   ,						clip_sample=lowerCamelCase   ,						set_alpha_to_one=lowerCamelCase   ,						)
      torch.manual_seed(0		)
      lowerCAmelCase_      :				Optional[int]      =       AutoencoderKL(
          block_out_channels=[32, 64]   ,						in_channels=3   ,						out_channels=3   ,						down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]   ,						up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]   ,						latent_channels=4   ,						)
      torch.manual_seed(0		)
      lowerCAmelCase_      :				Tuple      =       CLIPTextConfig(
          bos_token_id=0   ,						eos_token_id=2   ,						hidden_size=32   ,						intermediate_size=37   ,						layer_norm_eps=1E-05   ,						num_attention_heads=4   ,						num_hidden_layers=5   ,						pad_token_id=1   ,						vocab_size=10_00   ,						)
      lowerCAmelCase_      :				Optional[Any]      =       CLIPTextModel(lowerCamelCase		)
      lowerCAmelCase_      :				List[Any]      =       CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""		)

      lowerCAmelCase_      :				str      =       MultiControlNetModel([controlneta, controlneta]		)

      lowerCAmelCase_      :				Any      =       {
          """unet""": unet,
          """controlnet""": controlnet,
          """scheduler""": scheduler,
          """vae""": vae,
          """text_encoder""": text_encoder,
          """tokenizer""": tokenizer,
          """safety_checker""": None,
          """feature_extractor""": None,
      }
      return components

    def __lowercase       (       self   :  int   ,						lowerCamelCase   :  Dict   ,						lowerCamelCase   :  int=0		)       ->  List[Any]:
      if str(lowerCamelCase		).startswith("""mps"""		):
        lowerCAmelCase_      :				Optional[Any]      =       torch.manual_seed(lowerCamelCase		)
      else:
        lowerCAmelCase_      :				Union[str, Any]      =       torch.Generator(device=lowerCamelCase		).manual_seed(lowerCamelCase		)

      lowerCAmelCase_      :				Tuple      =       2

      lowerCAmelCase_      :				Tuple      =       [
          randn_tensor(
              (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)   ,						generator=lowerCamelCase   ,						device=torch.device(lowerCamelCase		)   ,						),
          randn_tensor(
              (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)   ,						generator=lowerCamelCase   ,						device=torch.device(lowerCamelCase		)   ,						),
      ]

      lowerCAmelCase_      :				Any      =       floats_tensor(control_image[0].shape   ,						rng=random.Random(lowerCamelCase		)		).to(lowerCamelCase		)
      lowerCAmelCase_      :				List[Any]      =       image.cpu().permute(0   ,						2   ,						3   ,						1		)[0]
      lowerCAmelCase_      :				Union[str, Any]      =       Image.fromarray(np.uinta(lowerCamelCase		)		).convert("""RGB"""		).resize((64, 64)		)
      lowerCAmelCase_      :				Any      =       {
          """prompt""": """A painting of a squirrel eating a burger""",
          """generator""": generator,
          """num_inference_steps""": 2,
          """guidance_scale""": 6.0,
          """output_type""": """numpy""",
          """image""": image,
          """control_image""": control_image,
      }

      return inputs

    def __lowercase       (       self   :  Tuple		)       ->  str:
      lowerCAmelCase_      :				Tuple      =       self.get_dummy_components()
      lowerCAmelCase_      :				Optional[Any]      =       self.pipeline_class(**lowerCamelCase		)
      pipe.to(lowerCamelCase		)

      lowerCAmelCase_      :				Optional[Any]      =       10.0
      lowerCAmelCase_      :				Union[str, Any]      =       4

      lowerCAmelCase_      :				Any      =       self.get_dummy_inputs(lowerCamelCase		)
      lowerCAmelCase_      :				Optional[Any]      =       steps
      lowerCAmelCase_      :				Tuple      =       scale
      lowerCAmelCase_      :				Tuple      =       pipe(**lowerCamelCase		)[0]

      lowerCAmelCase_      :				Optional[Any]      =       self.get_dummy_inputs(lowerCamelCase		)
      lowerCAmelCase_      :				Tuple      =       steps
      lowerCAmelCase_      :				List[Any]      =       scale
      lowerCAmelCase_      :				str      =       pipe(**lowerCamelCase   ,						control_guidance_start=0.1   ,						control_guidance_end=0.2		)[0]

      lowerCAmelCase_      :				Union[str, Any]      =       self.get_dummy_inputs(lowerCamelCase		)
      lowerCAmelCase_      :				Tuple      =       steps
      lowerCAmelCase_      :				Union[str, Any]      =       scale
      lowerCAmelCase_      :				Tuple      =       pipe(**lowerCamelCase   ,						control_guidance_start=[0.1, 0.3]   ,						control_guidance_end=[0.2, 0.7]		)[0]

      lowerCAmelCase_      :				str      =       self.get_dummy_inputs(lowerCamelCase		)
      lowerCAmelCase_      :				Optional[Any]      =       steps
      lowerCAmelCase_      :				Union[str, Any]      =       scale
      lowerCAmelCase_      :				Union[str, Any]      =       pipe(**lowerCamelCase   ,						control_guidance_start=0.4   ,						control_guidance_end=[0.5, 0.8]		)[0]

      # make sure that all outputs are different
      assert np.sum(np.abs(output_a - output_a		)		) > 1E-3
      assert np.sum(np.abs(output_a - output_a		)		) > 1E-3
      assert np.sum(np.abs(output_a - output_a		)		) > 1E-3

    def __lowercase       (       self   :  Any		)       ->  Any:
      return self._test_attention_slicing_forward_pass(expected_max_diff=2E-3		)

    @unittest.skipIf(
        torch_device != """cuda""" or not is_xformers_available()   ,						reason="""XFormers attention is only available with CUDA and `xformers` installed"""   ,						)
    def __lowercase       (       self   :  int		)       ->  Dict:
      self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3		)

    def __lowercase       (       self   :  Optional[int]		)       ->  List[str]:
      self._test_inference_batch_single_identical(expected_max_diff=2E-3		)





    def __lowercase       (       self   :  Optional[Any]		)       ->  Optional[int]:
      lowerCAmelCase_      :				Dict      =       self.get_dummy_components()
      lowerCAmelCase_      :				List[Any]      =       self.pipeline_class(**lowerCamelCase		)
      pipe.to(lowerCamelCase		)
      pipe.set_progress_bar_config(disable=lowerCamelCase		)
      with tempfile.TemporaryDirectory() as tmpdir:
        try:
          # save_pretrained is not implemented for Multi-ControlNet
          pipe.save_pretrained(lowerCamelCase		)
        except NotImplementedError:
          pass







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







    """simple docstring"""



    def __lowercase       (       self   :  Dict		)       ->  Union[str, Any]:
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()





    def __lowercase       (       self   :  List[str]		)       ->  int:
      lowerCAmelCase_      :				Union[str, Any]      =       ControlNetModel.from_pretrained("""lllyasviel/sd-controlnet-canny"""		)

      lowerCAmelCase_      :				List[str]      =       StableDiffusionControlNetImgaImgPipeline.from_pretrained(
          """runwayml/stable-diffusion-v1-5"""   ,						safety_checker=lowerCamelCase   ,						controlnet=lowerCamelCase		)
      pipe.enable_model_cpu_offload()
      pipe.set_progress_bar_config(disable=lowerCamelCase		)

      lowerCAmelCase_      :				Optional[Any]      =       torch.Generator(device="""cpu"""		).manual_seed(0		)
      lowerCAmelCase_      :				str      =       """evil space-punk bird"""
      lowerCAmelCase_      :				str      =       load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png"""		).resize((5_12, 5_12)		)
      lowerCAmelCase_      :				Union[str, Any]      =       load_image(
          """https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png"""		).resize((5_12, 5_12)		)

      lowerCAmelCase_      :				Dict      =       pipe(
          lowerCamelCase   ,						lowerCamelCase   ,						control_image=lowerCamelCase   ,						generator=lowerCamelCase   ,						output_type="""np"""   ,						num_inference_steps=50   ,						strength=0.6   ,						)

      lowerCAmelCase_      :				str      =       output.images[0]

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

      lowerCAmelCase_      :				Optional[Any]      =       load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/img2img.npy"""		)

      assert np.abs(expected_image - image		).max() < 9E-2







'''simple docstring'''






import inspect
import math
import tempfile
import unittest

import numpy as np

from transformers import ViTMAEConfig
from transformers.testing_utils import require_torch, 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 ViTMAEForPreTraining, ViTMAEModel
  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        __snake_case	:







    """simple docstring"""



    def __init__(       self   :  Optional[Any]   ,						lowerCamelCase   :  Tuple   ,						lowerCamelCase   :  Tuple=13   ,						lowerCamelCase   :  Dict=30   ,						lowerCamelCase   :  Dict=2   ,						lowerCamelCase   :  Optional[int]=3   ,						lowerCamelCase   :  List[Any]=True   ,						lowerCamelCase   :  Any=True   ,						lowerCamelCase   :  str=32   ,						lowerCamelCase   :  Any=5   ,						lowerCamelCase   :  int=4   ,						lowerCamelCase   :  List[str]=37   ,						lowerCamelCase   :  Any="gelu"   ,						lowerCamelCase   :  Optional[Any]=0.1   ,						lowerCamelCase   :  List[str]=0.1   ,						lowerCamelCase   :  str=10   ,						lowerCamelCase   :  Optional[Any]=0.02   ,						lowerCamelCase   :  List[str]=3   ,						lowerCamelCase   :  Union[str, Any]=0.6   ,						lowerCamelCase   :  List[Any]=None   ,						)       ->  Optional[int]:
      lowerCAmelCase_      :				Optional[Any]      =       parent
      lowerCAmelCase_      :				Optional[int]      =       batch_size
      lowerCAmelCase_      :				int      =       image_size
      lowerCAmelCase_      :				List[Any]      =       patch_size
      lowerCAmelCase_      :				int      =       num_channels
      lowerCAmelCase_      :				Any      =       is_training
      lowerCAmelCase_      :				Tuple      =       use_labels
      lowerCAmelCase_      :				Optional[Any]      =       hidden_size
      lowerCAmelCase_      :				List[Any]      =       num_hidden_layers
      lowerCAmelCase_      :				Optional[Any]      =       num_attention_heads
      lowerCAmelCase_      :				Dict      =       intermediate_size
      lowerCAmelCase_      :				Union[str, Any]      =       hidden_act
      lowerCAmelCase_      :				Union[str, Any]      =       hidden_dropout_prob
      lowerCAmelCase_      :				Any      =       attention_probs_dropout_prob
      lowerCAmelCase_      :				List[Any]      =       type_sequence_label_size
      lowerCAmelCase_      :				Dict      =       initializer_range
      lowerCAmelCase_      :				List[str]      =       mask_ratio
      lowerCAmelCase_      :				Tuple      =       scope

      # in ViTMAE, the expected sequence length = (num_patches + 1) * (1 - config.mask_ratio), rounded above
      # (we add 1 for the [CLS] token)
      lowerCAmelCase_      :				Union[str, Any]      =       (image_size // patch_size) ** 2
      lowerCAmelCase_      :				Any      =       int(math.ceil((1 - mask_ratio) * (num_patches + 1)		)		)

    def __lowercase       (       self   :  Optional[int]		)       ->  str:
      lowerCAmelCase_      :				str      =       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

      lowerCAmelCase_      :				Optional[int]      =       None
      if self.use_labels:
        lowerCAmelCase_      :				int      =       ids_tensor([self.batch_size]   ,						self.type_sequence_label_size		)

      lowerCAmelCase_      :				str      =       self.get_config()

      return config, pixel_values, labels

    def __lowercase       (       self   :  Optional[int]		)       ->  Optional[int]:
      return ViTMAEConfig(
          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   ,						mask_ratio=self.mask_ratio   ,						)

    def __lowercase       (       self   :  Any   ,						lowerCamelCase   :  Any   ,						lowerCamelCase   :  Union[str, Any]   ,						lowerCamelCase   :  Dict		)       ->  Tuple:
      lowerCAmelCase_      :				Tuple      =       ViTMAEModel(config=lowerCamelCase		)
      model.to(lowerCamelCase		)
      model.eval()
      lowerCAmelCase_      :				Dict      =       model(lowerCamelCase		)
      self.parent.assertEqual(result.last_hidden_state.shape   ,						(self.batch_size, self.seq_length, self.hidden_size)		)

    def __lowercase       (       self   :  List[str]   ,						lowerCamelCase   :  Union[str, Any]   ,						lowerCamelCase   :  List[str]   ,						lowerCamelCase   :  Union[str, Any]		)       ->  Dict:
      lowerCAmelCase_      :				Tuple      =       ViTMAEForPreTraining(lowerCamelCase		)
      model.to(lowerCamelCase		)
      model.eval()
      lowerCAmelCase_      :				List[str]      =       model(lowerCamelCase		)
      lowerCAmelCase_      :				int      =       (self.image_size // self.patch_size) ** 2
      lowerCAmelCase_      :				int      =       self.patch_size**2 * self.num_channels
      self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, num_patches, expected_num_channels)		)

      # test greyscale images
      lowerCAmelCase_      :				List[Any]      =       1
      lowerCAmelCase_      :				List[str]      =       ViTMAEForPreTraining(lowerCamelCase		)
      model.to(lowerCamelCase		)
      model.eval()
      lowerCAmelCase_      :				Optional[Any]      =       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
      lowerCAmelCase_      :				Tuple      =       model(lowerCamelCase		)
      lowerCAmelCase_      :				List[Any]      =       self.patch_size**2
      self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, num_patches, expected_num_channels)		)





    def __lowercase       (       self   :  Optional[int]		)       ->  str:
      lowerCAmelCase_      :				Any      =       self.prepare_config_and_inputs()
      lowerCAmelCase_,							lowerCAmelCase_,							lowerCAmelCase_      :				List[Any]      =       config_and_inputs
      lowerCAmelCase_      :				Any      =       {"""pixel_values""": pixel_values}
      return config, inputs_dict







@require_torch
class        __snake_case	(       _SCREAMING_SNAKE_CASE     ,_SCREAMING_SNAKE_CASE     ,unittest.TestCase):







    """simple docstring"""



    lowercase       			= (ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
    lowercase       			= {'feature-extraction': ViTMAEModel} if is_torch_available() else {}

    lowercase       			= False
    lowercase       			= False
    lowercase       			= False
    lowercase       			= False

    def __lowercase       (       self   :  Optional[Any]		)       ->  List[Any]:
      lowerCAmelCase_      :				Optional[int]      =       ViTMAEModelTester(self		)
      lowerCAmelCase_      :				Optional[int]      =       ConfigTester(self   ,						config_class=lowerCamelCase   ,						has_text_modality=lowerCamelCase   ,						hidden_size=37		)

    def __lowercase       (       self   :  Dict		)       ->  Tuple:
      self.config_tester.run_common_tests()

    @unittest.skip(reason="""ViTMAE does not use inputs_embeds"""		)
    def __lowercase       (       self   :  Optional[int]		)       ->  Optional[int]:
      pass

    def __lowercase       (       self   :  List[str]		)       ->  Tuple:
      lowerCAmelCase_,							lowerCAmelCase_      :				Optional[Any]      =       self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        lowerCAmelCase_      :				List[str]      =       model_class(lowerCamelCase		)
        self.assertIsInstance(model.get_input_embeddings()   ,						(nn.Module)		)
        lowerCAmelCase_      :				Tuple      =       model.get_output_embeddings()
        self.assertTrue(x is None or isinstance(lowerCamelCase   ,						nn.Linear		)		)

    def __lowercase       (       self   :  Optional[Any]		)       ->  Any:
      lowerCAmelCase_,							lowerCAmelCase_      :				List[Any]      =       self.model_tester.prepare_config_and_inputs_for_common()

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

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

    def __lowercase       (       self   :  Tuple		)       ->  str:
      lowerCAmelCase_      :				Any      =       self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*lowerCamelCase		)

    def __lowercase       (       self   :  Optional[int]		)       ->  str:
      lowerCAmelCase_      :				Tuple      =       self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_pretraining(*lowerCamelCase		)

    def __lowercase       (       self   :  Optional[int]   ,						lowerCamelCase   :  Dict   ,						lowerCamelCase   :  List[Any]   ,						lowerCamelCase   :  Optional[Any]		)       ->  str:
      # make masks reproducible
      np.random.seed(2		)

      lowerCAmelCase_      :				Tuple      =       int((pt_model.config.image_size // pt_model.config.patch_size) ** 2		)
      lowerCAmelCase_      :				Optional[int]      =       np.random.uniform(size=(self.model_tester.batch_size, num_patches)		)
      lowerCAmelCase_      :				Optional[Any]      =       torch.from_numpy(lowerCamelCase		)

      # Add `noise` argument.
      # PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
      lowerCAmelCase_      :				int      =       pt_noise

      super().check_pt_tf_models(lowerCamelCase   ,						lowerCamelCase   ,						lowerCamelCase		)

    def __lowercase       (       self   :  int		)       ->  Dict:
      lowerCAmelCase_,							lowerCAmelCase_      :				Optional[Any]      =       self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        lowerCAmelCase_      :				Optional[int]      =       model_class(lowerCamelCase		)
        model.to(lowerCamelCase		)
        model.eval()
        # make random mask reproducible
        torch.manual_seed(2		)
        with torch.no_grad():
          lowerCAmelCase_      :				Any      =       model(**self._prepare_for_class(lowerCamelCase   ,						lowerCamelCase		)		)

        lowerCAmelCase_      :				Any      =       outputs[0].cpu().numpy()
        lowerCAmelCase_      :				List[str]      =       0

        with tempfile.TemporaryDirectory() as tmpdirname:
          model.save_pretrained(lowerCamelCase		)
          lowerCAmelCase_      :				int      =       model_class.from_pretrained(lowerCamelCase		)
          model.to(lowerCamelCase		)
          # make random mask reproducible
          torch.manual_seed(2		)
          with torch.no_grad():
            lowerCAmelCase_      :				str      =       model(**self._prepare_for_class(lowerCamelCase   ,						lowerCamelCase		)		)

          # Make sure we don't have nans
          lowerCAmelCase_      :				Optional[Any]      =       after_outputs[0].cpu().numpy()
          lowerCAmelCase_      :				str      =       0
          lowerCAmelCase_      :				List[Any]      =       np.amax(np.abs(out_a - out_a		)		)
          self.assertLessEqual(lowerCamelCase   ,						1E-5		)

    @unittest.skip(
        reason="""ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load
    to get deterministic results."""		)
    def __lowercase       (       self   :  Optional[int]		)       ->  List[Any]:
      pass

    @unittest.skip(
        reason="""ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load
    to get deterministic results."""		)
    def __lowercase       (       self   :  Union[str, Any]		)       ->  str:
      pass

    @unittest.skip(
        reason="""ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load
    to get deterministic results."""		)
    def __lowercase       (       self   :  Optional[Any]		)       ->  Union[str, Any]:
      pass

    @unittest.skip(reason="""ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load"""		)
    def __lowercase       (       self   :  Tuple		)       ->  Optional[Any]:
      pass

    @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""		)
    def __lowercase       (       self   :  List[Any]		)       ->  str:
      pass





    @slow
    def __lowercase       (       self   :  List[str]		)       ->  List[Any]:
      for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        lowerCAmelCase_      :				List[Any]      =       ViTMAEModel.from_pretrained(lowerCamelCase		)
        self.assertIsNotNone(lowerCamelCase		)







def     UpperCamelCase_      (   ):
  '''simple docstring'''





  lowerCAmelCase_      :				List[str]      =       Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png""" )
  return image






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







    """simple docstring"""



    @cached_property
    def __lowercase       (       self   :  Union[str, Any]		)       ->  Union[str, Any]:
      return ViTImageProcessor.from_pretrained("""facebook/vit-mae-base"""		) if is_vision_available() else None





    @slow
    def __lowercase       (       self   :  int		)       ->  List[Any]:
      # make random mask reproducible across the PT and TF model
      np.random.seed(2		)

      lowerCAmelCase_      :				Dict      =       ViTMAEForPreTraining.from_pretrained("""facebook/vit-mae-base"""		).to(lowerCamelCase		)

      lowerCAmelCase_      :				Union[str, Any]      =       self.default_image_processor
      lowerCAmelCase_      :				Union[str, Any]      =       prepare_img()
      lowerCAmelCase_      :				Dict      =       image_processor(images=lowerCamelCase   ,						return_tensors="""pt"""		).to(lowerCamelCase		)

      # prepare a noise vector that will be also used for testing the TF model
      # (this way we can ensure that the PT and TF models operate on the same inputs)
      lowerCAmelCase_      :				Optional[int]      =       ViTMAEConfig()
      lowerCAmelCase_      :				Optional[Any]      =       int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2		)
      lowerCAmelCase_      :				Optional[int]      =       np.random.uniform(size=(1, num_patches)		)

      # forward pass
      with torch.no_grad():
        lowerCAmelCase_      :				str      =       model(**lowerCamelCase   ,						noise=torch.from_numpy(lowerCamelCase		).to(device=lowerCamelCase		)		)

      # verify the logits
      lowerCAmelCase_      :				str      =       torch.Size((1, 1_96, 7_68)		)
      self.assertEqual(outputs.logits.shape   ,						lowerCamelCase		)

      lowerCAmelCase_      :				str      =       torch.tensor(
          [[-0.0_548, -1.7_023, -0.9_325], [0.3_721, -0.5_670, -0.2_233], [0.8_235, -1.3_878, -0.3_524]]		)

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



