


'''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        lowerCAmelCase_  :






       def __init__( self			,     _UpperCamelCase			,     _UpperCamelCase=13			,     _UpperCamelCase=30			,     _UpperCamelCase=2			,     _UpperCamelCase=3			,     _UpperCamelCase=True			,     _UpperCamelCase=True			,     _UpperCamelCase=32			,     _UpperCamelCase=5			,     _UpperCamelCase=4			,     _UpperCamelCase=37			,     _UpperCamelCase="gelu"			,     _UpperCamelCase=0.1			,     _UpperCamelCase=0.1			,     _UpperCamelCase=10			,     _UpperCamelCase=0.02			,     _UpperCamelCase=3			,     _UpperCamelCase=0.6			,     _UpperCamelCase=None			,     )->      Any:
            _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				=							mask_ratio
            _A				=							scope

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






       def   UpperCamelCase						( self   )->      List[Any]:
            _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   UpperCamelCase						( self   )->      Optional[Any]:
            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=_UpperCamelCase			,     initializer_range=self.initializer_range			,     mask_ratio=self.mask_ratio			,     )






       def   UpperCamelCase						( self			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase   )->      Optional[Any]:
            _A				=							ViTMAEModel(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   UpperCamelCase						( self			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase   )->      Dict:
            _A				=							ViTMAEForPreTraining(_UpperCamelCase   )
            model.to(_UpperCamelCase   )
            model.eval()
            _A				=							model(_UpperCamelCase   )
            _A				=							(self.image_size // self.patch_size) ** 2
            _A				=							self.patch_size**2 * self.num_channels
            self.parent.assertEqual(result.logits.shape			,     (self.batch_size, num_patches, expected_num_channels)   )

            # test greyscale images
            _A				=							1
            _A				=							ViTMAEForPreTraining(_UpperCamelCase   )
            model.to(_UpperCamelCase   )
            model.eval()
            _A				=							floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
            _A				=							model(_UpperCamelCase   )
            _A				=							self.patch_size**2
            self.parent.assertEqual(result.logits.shape			,     (self.batch_size, num_patches, expected_num_channels)   )






       def   UpperCamelCase						( self   )->      List[str]:
            _A				=							self.prepare_config_and_inputs()
            _A		,				_A		,				_A				=							config_and_inputs
            _A				=							{'pixel_values': pixel_values}
            return config, inputs_dict






@require_torch
class        lowerCAmelCase_  (     UpperCAmelCase      ,							UpperCAmelCase      ,							unittest.TestCase	):
       __UpperCAmelCase					=(ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
       __UpperCAmelCase					={'feature-extraction': ViTMAEModel} if is_torch_available() else {}

       __UpperCAmelCase					=False
       __UpperCAmelCase					=False
       __UpperCAmelCase					=False
       __UpperCAmelCase					=False






       def   UpperCamelCase						( self   )->      Union[str, Any]:
            _A				=							ViTMAEModelTester(self   )
            _A				=							ConfigTester(self			,     config_class=_UpperCamelCase			,     has_text_modality=_UpperCamelCase			,     hidden_size=37   )






       def   UpperCamelCase						( self   )->      Tuple:
            self.config_tester.run_common_tests()






       @unittest.skip(reason='ViTMAE does not use inputs_embeds'   )
       def   UpperCamelCase						( self   )->      int:
            pass






       def   UpperCamelCase						( self   )->      Dict:
            _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   UpperCamelCase						( self   )->      List[str]:
            _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   UpperCamelCase						( self   )->      int:
            _A				=							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*_UpperCamelCase   )






       def   UpperCamelCase						( self   )->      Optional[int]:
            _A				=							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_pretraining(*_UpperCamelCase   )






       def   UpperCamelCase						( self			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase   )->      List[Any]:
            # make masks reproducible
            np.random.seed(2   )

            _A				=							int((pt_model.config.image_size // pt_model.config.patch_size) ** 2   )
            _A				=							np.random.uniform(size=(self.model_tester.batch_size, num_patches)   )
            _A				=							torch.from_numpy(_UpperCamelCase   )

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

            super().check_pt_tf_models(_UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase   )






       def   UpperCamelCase						( self   )->      List[str]:
            _A		,				_A				=							self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                 _A				=							model_class(_UpperCamelCase   )
                 model.to(_UpperCamelCase   )
                 model.eval()
                 # make random mask reproducible
                 torch.manual_seed(2   )
                 with torch.no_grad():
                      _A				=							model(**self._prepare_for_class(_UpperCamelCase			,     _UpperCamelCase   )   )

                 _A				=							outputs[0].cpu().numpy()
                 _A				=							0

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

                      # Make sure we don't have nans
                      _A				=							after_outputs[0].cpu().numpy()
                      _A				=							0
                      _A				=							np.amax(np.abs(out_a - out_a   )   )
                      self.assertLessEqual(_UpperCamelCase			,     1e-5   )






       @unittest.skip(
           reason='ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results.'   )
       def   UpperCamelCase						( self   )->      str:
            pass






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






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






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






       @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'   )
       def   UpperCamelCase						( self   )->      str:
            pass






       @slow
       def   UpperCamelCase						( self   )->      Tuple:
            for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 _A				=							ViTMAEModel.from_pretrained(_UpperCamelCase   )
                 self.assertIsNotNone(_UpperCamelCase   )






def 					lowerCamelCase_ ( )					->	Optional[int]:







     """simple docstring"""
     _A				=							Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'		)
     return image






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






       @cached_property
       def   UpperCamelCase						( self   )->      str:
            return ViTImageProcessor.from_pretrained('facebook/vit-mae-base'   ) if is_vision_available() else None






       @slow
       def   UpperCamelCase						( self   )->      Union[str, Any]:
            # make random mask reproducible across the PT and TF model
            np.random.seed(2   )

            _A				=							ViTMAEForPreTraining.from_pretrained('facebook/vit-mae-base'   ).to(_UpperCamelCase   )

            _A				=							self.default_image_processor
            _A				=							prepare_img()
            _A				=							image_processor(images=_UpperCamelCase			,     return_tensors='pt'   ).to(_UpperCamelCase   )

            # 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)
            _A				=							ViTMAEConfig()
            _A				=							int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2   )
            _A				=							np.random.uniform(size=(1, num_patches)   )

            # forward pass
            with torch.no_grad():
                 _A				=							model(**_UpperCamelCase			,     noise=torch.from_numpy(_UpperCamelCase   ).to(device=_UpperCamelCase   )   )

            # verify the logits
            _A				=							torch.Size((1, 196, 768)   )
            self.assertEqual(outputs.logits.shape			,     _UpperCamelCase   )

            _A				=							torch.tensor(
                [[-0.0548, -1.7023, -0.9325], [0.3721, -0.5670, -0.2233], [0.8235, -1.3878, -0.3524]]   )

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









'''simple docstring'''
import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()






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






       def   UpperCamelCase						( self   )->      Optional[int]:
            # clean up the VRAM after each test
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()






       def   UpperCamelCase						( self   )->      Union[str, Any]:
            _A				=							StableDiffusionKDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4'   )
            _A				=							sd_pipe.to(_UpperCamelCase   )
            sd_pipe.set_progress_bar_config(disable=_UpperCamelCase   )

            sd_pipe.set_scheduler('sample_euler'   )

            _A				=							'A painting of a squirrel eating a burger'
            _A				=							torch.manual_seed(0   )
            _A				=							sd_pipe([prompt]			,     generator=_UpperCamelCase			,     guidance_scale=9.0			,     num_inference_steps=20			,     output_type='np'   )

            _A				=							output.images

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

            assert image.shape == (1, 512, 512, 3)
            _A				=							np.array([0.0447, 0.0492, 0.0468, 0.0408, 0.0383, 0.0408, 0.0354, 0.0380, 0.0339]   )

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






       def   UpperCamelCase						( self   )->      Optional[Any]:
            _A				=							StableDiffusionKDiffusionPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'   )
            _A				=							sd_pipe.to(_UpperCamelCase   )
            sd_pipe.set_progress_bar_config(disable=_UpperCamelCase   )

            sd_pipe.set_scheduler('sample_euler'   )

            _A				=							'A painting of a squirrel eating a burger'
            _A				=							torch.manual_seed(0   )
            _A				=							sd_pipe([prompt]			,     generator=_UpperCamelCase			,     guidance_scale=9.0			,     num_inference_steps=20			,     output_type='np'   )

            _A				=							output.images

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

            assert image.shape == (1, 512, 512, 3)
            _A				=							np.array([0.1237, 0.1320, 0.1438, 0.1359, 0.1390, 0.1132, 0.1277, 0.1175, 0.1112]   )

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






       def   UpperCamelCase						( self   )->      Optional[int]:
            _A				=							StableDiffusionKDiffusionPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'   )
            _A				=							sd_pipe.to(_UpperCamelCase   )
            sd_pipe.set_progress_bar_config(disable=_UpperCamelCase   )

            sd_pipe.set_scheduler('sample_dpmpp_2m'   )

            _A				=							'A painting of a squirrel eating a burger'
            _A				=							torch.manual_seed(0   )
            _A				=							sd_pipe(
                [prompt]			,     generator=_UpperCamelCase			,     guidance_scale=7.5			,     num_inference_steps=15			,     output_type='np'			,     use_karras_sigmas=_UpperCamelCase			,     )

            _A				=							output.images

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

            assert image.shape == (1, 512, 512, 3)
            _A				=							np.array(
                [0.1138_1689, 0.1211_2921, 0.138_9457, 0.1254_9606, 0.124_4964, 0.1083_1517, 0.1156_2866, 0.1086_7816, 0.1049_9048]   )

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





