





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					,			snake_case					,			snake_case=13					,			snake_case=30					,			snake_case=2					,			snake_case=3					,			snake_case=True					,			snake_case=True					,			snake_case=32					,			snake_case=5					,			snake_case=4					,			snake_case=37					,			snake_case="gelu"					,			snake_case=0.1					,			snake_case=0.1					,			snake_case=10					,			snake_case=0.02					,			snake_case=3					,			snake_case=0.6					,			snake_case=None					,			)				->  Union[str, Any]:
      """simple docstring"""





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

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





   def 	_snake_case ( self      )				->  int:
      """simple docstring"""





      a__	:						Union[str, Any]      		=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

      a__	:						Tuple      		=  None
      if self.use_labels:
         a__	:						Union[str, Any]      		=  ids_tensor([self.batch_size]					,			self.type_sequence_label_size      )

      a__	:						List[Any]      		=  self.get_config()

      return config, pixel_values, labels





   def 	_snake_case ( self      )				->  Dict:
      """simple docstring"""





      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=snake_case					,			initializer_range=self.initializer_range					,			mask_ratio=self.mask_ratio					,			)





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case      )				->  Tuple:
      """simple docstring"""





      a__	:						List[Any]      		=  ViTMAEModel(config=snake_case      )
      model.to(snake_case      )
      model.eval()
      a__	:						List[Any]      		=  model(snake_case      )
      self.parent.assertEqual(result.last_hidden_state.shape					,			(self.batch_size, self.seq_length, self.hidden_size)      )





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case      )				->  Any:
      """simple docstring"""





      a__	:						int      		=  ViTMAEForPreTraining(snake_case      )
      model.to(snake_case      )
      model.eval()
      a__	:						Any      		=  model(snake_case      )
      a__	:						Any      		=  (self.image_size // self.patch_size) ** 2
      a__	:						List[str]      		=  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__	:						Tuple      		=  1
      a__	:						Optional[Any]      		=  ViTMAEForPreTraining(snake_case      )
      model.to(snake_case      )
      model.eval()
      a__	:						int      		=  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]      )
      a__	:						Tuple      		=  model(snake_case      )
      a__	:						str      		=  self.patch_size**2
      self.parent.assertEqual(result.logits.shape					,			(self.batch_size, num_patches, expected_num_channels)      )





   def 	_snake_case ( self      )				->  Optional[Any]:
      """simple docstring"""





      a__	:						List[str]      		=  self.prepare_config_and_inputs()
      a__		,     a__		,     a__	:						int      		=  config_and_inputs
      a__	:						Optional[Any]      		=  {"pixel_values": pixel_values}
      return config, inputs_dict







@require_torch
class   __lowerCAmelCase			(					_UpperCamelCase		,_UpperCamelCase		,unittest.TestCase  ):
   _UpperCamelCase :						List[Any]     		=				(ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
   _UpperCamelCase :						Optional[int]     		=				{"""feature-extraction""": ViTMAEModel} if is_torch_available() else {}

   _UpperCamelCase :						Dict     		=				False
   _UpperCamelCase :						str     		=				False
   _UpperCamelCase :						str     		=				False
   _UpperCamelCase :						Union[str, Any]     		=				False





   def 	_snake_case ( self      )				->  Tuple:
      """simple docstring"""





      a__	:						Optional[Any]      		=  ViTMAEModelTester(self      )
      a__	:						Any      		=  ConfigTester(self					,			config_class=snake_case					,			has_text_modality=snake_case					,			hidden_size=37      )





   def 	_snake_case ( self      )				->  List[Any]:
      """simple docstring"""





      self.config_tester.run_common_tests()





   @unittest.skip(reason="ViTMAE does not use inputs_embeds"      )
   def 	_snake_case ( self      )				->  List[Any]:
      """simple docstring"""





      pass





   def 	_snake_case ( self      )				->  Dict:
      """simple docstring"""





      a__		,     a__	:						str      		=  self.model_tester.prepare_config_and_inputs_for_common()

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





   def 	_snake_case ( self      )				->  Dict:
      """simple docstring"""





      a__		,     a__	:						Dict      		=  self.model_tester.prepare_config_and_inputs_for_common()

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

         a__	:						Tuple      		=  ["pixel_values"]
         self.assertListEqual(arg_names[:1]					,			snake_case      )





   def 	_snake_case ( self      )				->  int:
      """simple docstring"""





      a__	:						Tuple      		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*snake_case      )





   def 	_snake_case ( self      )				->  Optional[int]:
      """simple docstring"""





      a__	:						Union[str, Any]      		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_pretraining(*snake_case      )





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case      )				->  Tuple:
      """simple docstring"""





      np.random.seed(2      )

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

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

      super().check_pt_tf_models(snake_case					,			snake_case					,			snake_case      )





   def 	_snake_case ( self      )				->  Optional[Any]:
      """simple docstring"""





      a__		,     a__	:						Union[str, Any]      		=  self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
         a__	:						Any      		=  model_class(snake_case      )
         model.to(snake_case      )
         model.eval()
         # make random mask reproducible
         torch.manual_seed(2      )
         with torch.no_grad():
            a__	:						List[str]      		=  model(**self._prepare_for_class(snake_case					,			snake_case      )      )

         a__	:						Union[str, Any]      		=  outputs[0].cpu().numpy()
         a__	:						List[Any]      		=  0

         with tempfile.TemporaryDirectory() as tmpdirname:
            model.save_pretrained(snake_case      )
            a__	:						Union[str, Any]      		=  model_class.from_pretrained(snake_case      )
            model.to(snake_case      )
            # make random mask reproducible
            torch.manual_seed(2      )
            with torch.no_grad():
               a__	:						int      		=  model(**self._prepare_for_class(snake_case					,			snake_case      )      )

            # Make sure we don't have nans
            a__	:						str      		=  after_outputs[0].cpu().numpy()
            a__	:						str      		=  0
            a__	:						Optional[int]      		=  np.amax(np.abs(out_a - out_a      )      )
            self.assertLessEqual(snake_case					,			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 	_snake_case ( self      )				->  Tuple:
      """simple docstring"""





      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 	_snake_case ( self      )				->  Dict:
      """simple docstring"""





      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 	_snake_case ( self      )				->  Optional[Any]:
      """simple docstring"""





      pass





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





      pass





   @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."      )
   def 	_snake_case ( self      )				->  Any:
      """simple docstring"""





      pass





   @slow
   def 	_snake_case ( self      )				->  Dict:
      """simple docstring"""





      for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         a__	:						Union[str, Any]      		=  ViTMAEModel.from_pretrained(snake_case      )
         self.assertIsNotNone(snake_case      )







def       _A  (      ):
   a__	:						int      		=  Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"   )
   return image







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





   @cached_property
   def 	_snake_case ( self      )				->  List[str]:
      """simple docstring"""





      return ViTImageProcessor.from_pretrained("facebook/vit-mae-base"      ) if is_vision_available() else None





   @slow
   def 	_snake_case ( self      )				->  Optional[Any]:
      """simple docstring"""





      np.random.seed(2      )

      a__	:						int      		=  ViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base"      ).to(snake_case      )

      a__	:						int      		=  self.default_image_processor
      a__	:						Optional[Any]      		=  prepare_img()
      a__	:						Optional[int]      		=  image_processor(images=snake_case					,			return_tensors="pt"      ).to(snake_case      )

      # 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__	:						Optional[int]      		=  ViTMAEConfig()
      a__	:						str      		=  int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2      )
      a__	:						Any      		=  np.random.uniform(size=(1, num_patches)      )

      # forward pass
      with torch.no_grad():
         a__	:						Tuple      		=  model(**snake_case					,			noise=torch.from_numpy(snake_case      ).to(device=snake_case      )      )

      # verify the logits
      a__	:						List[Any]      		=  torch.Size((1, 196, 768)      )
      self.assertEqual(outputs.logits.shape					,			snake_case      )

      a__	:						Tuple      		=  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(snake_case      )					,			atol=1E-4      )      )








import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()







class   __lowerCAmelCase			(					_UpperCamelCase		,_UpperCamelCase		,_UpperCamelCase		,unittest.TestCase  ):
   _UpperCamelCase :						Optional[int]     		=				StableUnCLIPPipeline
   _UpperCamelCase :						Optional[int]     		=				TEXT_TO_IMAGE_PARAMS
   _UpperCamelCase :						str     		=				TEXT_TO_IMAGE_BATCH_PARAMS
   _UpperCamelCase :						List[Any]     		=				TEXT_TO_IMAGE_IMAGE_PARAMS
   _UpperCamelCase :						str     		=				TEXT_TO_IMAGE_IMAGE_PARAMS

   # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
   _UpperCamelCase :						Any     		=				False





   def 	_snake_case ( self      )				->  List[str]:
      """simple docstring"""





      a__	:						Any      		=  32
      a__	:						int      		=  embedder_hidden_size

      # prior components

      torch.manual_seed(0      )
      a__	:						Any      		=  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"      )

      torch.manual_seed(0      )
      a__	:						Optional[Any]      		=  CLIPTextModelWithProjection(
          CLIPTextConfig(
              bos_token_id=0					,			eos_token_id=2					,			hidden_size=snake_case					,			projection_dim=snake_case					,			intermediate_size=37					,			layer_norm_eps=1E-05					,			num_attention_heads=4					,			num_hidden_layers=5					,			pad_token_id=1					,			vocab_size=1_000					,			)      )

      torch.manual_seed(0      )
      a__	:						int      		=  PriorTransformer(
          num_attention_heads=2					,			attention_head_dim=12					,			embedding_dim=snake_case					,			num_layers=1					,			)

      torch.manual_seed(0      )
      a__	:						str      		=  DDPMScheduler(
          variance_type="fixed_small_log"					,			prediction_type="sample"					,			num_train_timesteps=1_000					,			clip_sample=snake_case					,			clip_sample_range=5.0					,			beta_schedule="squaredcos_cap_v2"					,			)

      # regular denoising components

      torch.manual_seed(0      )
      a__	:						Any      		=  StableUnCLIPImageNormalizer(embedding_dim=snake_case      )
      a__	:						int      		=  DDPMScheduler(beta_schedule="squaredcos_cap_v2"      )

      torch.manual_seed(0      )
      a__	:						Optional[Any]      		=  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"      )

      torch.manual_seed(0      )
      a__	:						Union[str, Any]      		=  CLIPTextModel(
          CLIPTextConfig(
              bos_token_id=0					,			eos_token_id=2					,			hidden_size=snake_case					,			projection_dim=32					,			intermediate_size=37					,			layer_norm_eps=1E-05					,			num_attention_heads=4					,			num_hidden_layers=5					,			pad_token_id=1					,			vocab_size=1_000					,			)      )

      torch.manual_seed(0      )
      a__	:						Any      		=  UNetaDConditionModel(
          sample_size=32					,			in_channels=4					,			out_channels=4					,			down_block_types=("CrossAttnDownBlock2D", "DownBlock2D")					,			up_block_types=("UpBlock2D", "CrossAttnUpBlock2D")					,			block_out_channels=(32, 64)					,			attention_head_dim=(2, 4)					,			class_embed_type="projection"					,			projection_class_embeddings_input_dim=embedder_projection_dim * 2					,			cross_attention_dim=snake_case					,			layers_per_block=1					,			upcast_attention=snake_case					,			use_linear_projection=snake_case					,			)

      torch.manual_seed(0      )
      a__	:						Tuple      		=  DDIMScheduler(
          beta_schedule="scaled_linear"					,			beta_start=0.00_085					,			beta_end=0.012					,			prediction_type="v_prediction"					,			set_alpha_to_one=snake_case					,			steps_offset=1					,			)

      torch.manual_seed(0      )
      a__	:						Optional[int]      		=  AutoencoderKL()

      a__	:						Any      		=  {
          # prior components
          "prior_tokenizer": prior_tokenizer,
          "prior_text_encoder": prior_text_encoder,
          "prior": prior,
          "prior_scheduler": prior_scheduler,
          # image noising components
          "image_normalizer": image_normalizer,
          "image_noising_scheduler": image_noising_scheduler,
          # regular denoising components
          "tokenizer": tokenizer,
          "text_encoder": text_encoder,
          "unet": unet,
          "scheduler": scheduler,
          "vae": vae,
      }

      return components





   def 	_snake_case ( self					,			snake_case					,			snake_case=0      )				->  Dict:
      """simple docstring"""





      if str(snake_case      ).startswith("mps"      ):
         a__	:						Union[str, Any]      		=  torch.manual_seed(snake_case      )
      else:
         a__	:						List[str]      		=  torch.Generator(device=snake_case      ).manual_seed(snake_case      )
      a__	:						Any      		=  {
          "prompt": "A painting of a squirrel eating a burger",
          "generator": generator,
          "num_inference_steps": 2,
          "prior_num_inference_steps": 2,
          "output_type": "numpy",
      }
      return inputs





   def 	_snake_case ( self      )				->  List[str]:
      """simple docstring"""





      a__	:						Dict      		=  torch_device == "cpu"

      self._test_attention_slicing_forward_pass(test_max_difference=snake_case      )





   def 	_snake_case ( self      )				->  int:
      """simple docstring"""





      a__	:						int      		=  torch_device in ["cpu", "mps"]

      self._test_inference_batch_single_identical(test_max_difference=snake_case      )







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





   def 	_snake_case ( self      )				->  List[str]:
      """simple docstring"""





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





   def 	_snake_case ( self      )				->  Tuple:
      """simple docstring"""





      a__	:						int      		=  load_numpy(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy"      )

      a__	:						int      		=  StableUnCLIPPipeline.from_pretrained("fusing/stable-unclip-2-1-l"					,			torch_dtype=torch.floataa      )
      pipe.to(snake_case      )
      pipe.set_progress_bar_config(disable=snake_case      )
      # stable unclip will oom when integration tests are run on a V100,
      # so turn on memory savings
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      a__	:						Dict      		=  torch.Generator(device="cpu"      ).manual_seed(0      )
      a__	:						Dict      		=  pipe("anime turle"					,			generator=snake_case					,			output_type="np"      )

      a__	:						List[str]      		=  output.images[0]

      assert image.shape == (768, 768, 3)

      assert_mean_pixel_difference(snake_case					,			snake_case      )





   def 	_snake_case ( self      )				->  Tuple:
      """simple docstring"""





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

      a__	:						str      		=  StableUnCLIPPipeline.from_pretrained("fusing/stable-unclip-2-1-l"					,			torch_dtype=torch.floataa      )
      a__	:						Union[str, Any]      		=  pipe.to(snake_case      )
      pipe.set_progress_bar_config(disable=snake_case      )
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      a__	:						Union[str, Any]      		=  pipe(
          "anime turtle"					,			prior_num_inference_steps=2					,			num_inference_steps=2					,			output_type="np"					,			)

      a__	:						List[Any]      		=  torch.cuda.max_memory_allocated()
      # make sure that less than 7 GB is allocated
      assert mem_bytes < 7 * 10**9

