




import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    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 ViTForImageClassification, ViTForMaskedImageModeling, ViTModel
       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 						__lowercase		:




      """simple docstring"""





      def __init__(				self    ,				A    ,				A=1_3    ,				A=3_0    ,				A=2    ,				A=3    ,				A=True    ,				A=True    ,				A=3_2    ,				A=5    ,				A=4    ,				A=3_7    ,				A="gelu"    ,				A=0.1    ,				A=0.1    ,				A=1_0    ,				A=0.02    ,				A=None    ,				A=2    ,				)		->    Dict:
         snake_case				:    str           =	parent
         snake_case				:    Optional[Any]           =	batch_size
         snake_case				:    Optional[Any]           =	image_size
         snake_case				:    List[Any]           =	patch_size
         snake_case				:    Tuple           =	num_channels
         snake_case				:    Dict           =	is_training
         snake_case				:    Tuple           =	use_labels
         snake_case				:    str           =	hidden_size
         snake_case				:    int           =	num_hidden_layers
         snake_case				:    Optional[int]           =	num_attention_heads
         snake_case				:    Tuple           =	intermediate_size
         snake_case				:    Tuple           =	hidden_act
         snake_case				:    Union[str, Any]           =	hidden_dropout_prob
         snake_case				:    Tuple           =	attention_probs_dropout_prob
         snake_case				:    Union[str, Any]           =	type_sequence_label_size
         snake_case				:    List[Any]           =	initializer_range
         snake_case				:    Optional[int]           =	scope
         snake_case				:    List[Any]           =	encoder_stride

         # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
         snake_case				:    Tuple           =	(image_size // patch_size) ** 2
         snake_case				:    Tuple           =	num_patches + 1

      def      UpperCAmelCase				(				self     )		->    Tuple:
         snake_case				:    str           =	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

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

         snake_case				:    str           =	self.get_config()

         return config, pixel_values, labels

      def      UpperCAmelCase				(				self     )		->    int:
         return ViTConfig(
             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=__lowercase    ,				initializer_range=self.initializer_range    ,				encoder_stride=self.encoder_stride    ,				)

      def      UpperCAmelCase				(				self    ,				A    ,				A    ,				A     )		->    str:
         snake_case				:    Union[str, Any]           =	ViTModel(config=__lowercase     )
         model.to(__lowercase     )
         model.eval()
         snake_case				:    Any           =	model(__lowercase     )
         self.parent.assertEqual(result.last_hidden_state.shape    ,				(self.batch_size, self.seq_length, self.hidden_size)     )

      def      UpperCAmelCase				(				self    ,				A    ,				A    ,				A     )		->    Union[str, Any]:
         snake_case				:    Any           =	ViTForMaskedImageModeling(config=__lowercase     )
         model.to(__lowercase     )
         model.eval()
         snake_case				:    int           =	model(__lowercase     )
         self.parent.assertEqual(
             result.reconstruction.shape    ,				(self.batch_size, self.num_channels, self.image_size, self.image_size)     )

         # test greyscale images
         snake_case				:    List[str]           =	1
         snake_case				:    int           =	ViTForMaskedImageModeling(__lowercase     )
         model.to(__lowercase     )
         model.eval()

         snake_case				:    Optional[Any]           =	floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
         snake_case				:    List[Any]           =	model(__lowercase     )
         self.parent.assertEqual(result.reconstruction.shape    ,				(self.batch_size, 1, self.image_size, self.image_size)     )

      def      UpperCAmelCase				(				self    ,				A    ,				A    ,				A     )		->    List[str]:
         snake_case				:    Any           =	self.type_sequence_label_size
         snake_case				:    Any           =	ViTForImageClassification(__lowercase     )
         model.to(__lowercase     )
         model.eval()
         snake_case				:    Optional[int]           =	model(__lowercase    ,				labels=__lowercase     )
         self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.type_sequence_label_size)     )

         # test greyscale images
         snake_case				:    Any           =	1
         snake_case				:    Optional[int]           =	ViTForImageClassification(__lowercase     )
         model.to(__lowercase     )
         model.eval()

         snake_case				:    List[Any]           =	floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
         snake_case				:    str           =	model(__lowercase     )
         self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.type_sequence_label_size)     )

      def      UpperCAmelCase				(				self     )		->    List[Any]:
         snake_case				:    int           =	self.prepare_config_and_inputs()
         (
             snake_case			
         )				:    Optional[Any]           =	config_and_inputs
         snake_case				:    Union[str, Any]           =	{'''pixel_values''': pixel_values}
         return config, inputs_dict




@require_torch
class 						__lowercase		(__lowerCamelCase					,		__lowerCamelCase					,		unittest.TestCase							):




      """simple docstring"""





      _snake_case												=						(
          (
              ViTModel,
              ViTForImageClassification,
              ViTForMaskedImageModeling,
          )
          if is_torch_available()
          else ()
      )
      _snake_case												=						(
          {"""feature-extraction""": ViTModel, """image-classification""": ViTForImageClassification}
          if is_torch_available()
          else {}
      )
      _snake_case												=						True

      _snake_case												=						False
      _snake_case												=						False
      _snake_case												=						False

      def      UpperCAmelCase				(				self     )		->    Tuple:
         snake_case				:    Optional[int]           =	ViTModelTester(self     )
         snake_case				:    Any           =	ConfigTester(self    ,				config_class=__lowercase    ,				has_text_modality=__lowercase    ,				hidden_size=3_7     )

      def      UpperCAmelCase				(				self     )		->    Union[str, Any]:
         self.config_tester.run_common_tests()

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

      def      UpperCAmelCase				(				self     )		->    Union[str, Any]:
         snake_case				:    Optional[Any]           =	self.model_tester.prepare_config_and_inputs_for_common()

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

      def      UpperCAmelCase				(				self     )		->    int:
         snake_case				:    List[str]           =	self.model_tester.prepare_config_and_inputs_for_common()

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

            snake_case				:    Any           =	['''pixel_values''']
            self.assertListEqual(arg_names[:1]    ,				__lowercase     )

      def      UpperCAmelCase				(				self     )		->    str:
         snake_case				:    List[str]           =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__lowercase     )

      def      UpperCAmelCase				(				self     )		->    Union[str, Any]:
         snake_case				:    int           =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_image_modeling(*__lowercase     )

      def      UpperCAmelCase				(				self     )		->    int:
         snake_case				:    List[str]           =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*__lowercase     )

      @slow
      def      UpperCAmelCase				(				self     )		->    str:
         for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            snake_case				:    Any           =	ViTModel.from_pretrained(__lowercase     )
            self.assertIsNotNone(__lowercase     )




def 						SCREAMING_SNAKE_CASE__   (   )  ->       List[str]:
   snake_case				:    str           =	Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""					)
   return image




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




      """simple docstring"""





      @cached_property
      def      UpperCAmelCase				(				self     )		->    Optional[int]:
         return ViTImageProcessor.from_pretrained("""google/vit-base-patch16-224"""     ) if is_vision_available() else None

      @slow
      def      UpperCAmelCase				(				self     )		->    Optional[int]:
         snake_case				:    List[Any]           =	ViTForImageClassification.from_pretrained("""google/vit-base-patch16-224"""     ).to(__lowercase     )

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

         # forward pass
         with torch.no_grad():
            snake_case				:    int           =	model(**__lowercase     )

         # verify the logits
         snake_case				:    Dict           =	torch.Size((1, 1_0_0_0)     )
         self.assertEqual(outputs.logits.shape    ,				__lowercase     )

         snake_case				:    List[str]           =	torch.tensor([-0.27_44, 0.82_15, -0.08_36]     ).to(__lowercase     )

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

      @slow
      def      UpperCAmelCase				(				self     )		->    str:
         # ViT models have an `interpolate_pos_encoding` argument in their forward method,
         # allowing to interpolate the pre-trained position embeddings in order to use
         # the model on higher resolutions. The DINO model by Facebook AI leverages this
         # to visualize self-attention on higher resolution images.
         snake_case				:    str           =	ViTModel.from_pretrained("""facebook/dino-vits8"""     ).to(__lowercase     )

         snake_case				:    Tuple           =	ViTImageProcessor.from_pretrained("""facebook/dino-vits8"""    ,				size=4_8_0     )
         snake_case				:    Tuple           =	prepare_img()
         snake_case				:    Optional[int]           =	image_processor(images=__lowercase    ,				return_tensors="""pt"""     )
         snake_case				:    Tuple           =	inputs.pixel_values.to(__lowercase     )

         # forward pass
         with torch.no_grad():
            snake_case				:    str           =	model(__lowercase    ,				interpolate_pos_encoding=__lowercase     )

         # verify the logits
         snake_case				:    Optional[Any]           =	torch.Size((1, 3_6_0_1, 3_8_4)     )
         self.assertEqual(outputs.last_hidden_state.shape    ,				__lowercase     )

         snake_case				:    Union[str, Any]           =	torch.tensor(
             [[4.23_40, 4.39_06, -6.66_92], [4.54_63, 1.89_28, -6.72_57], [4.44_29, 0.84_96, -5.85_85]]     ).to(__lowercase     )

         self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]    ,				__lowercase    ,				atol=1e-4     )     )

      @slow
      @require_accelerate
      @require_torch_gpu
      def      UpperCAmelCase				(				self     )		->    List[Any]:
         snake_case				:    Union[str, Any]           =	ViTModel.from_pretrained("""facebook/dino-vits8"""    ,				torch_dtype=torch.floataa    ,				device_map="""auto"""     )
         snake_case				:    Tuple           =	self.default_image_processor

         snake_case				:    List[str]           =	prepare_img()
         snake_case				:    int           =	image_processor(images=__lowercase    ,				return_tensors="""pt"""     )
         snake_case				:    Optional[int]           =	inputs.pixel_values.to(__lowercase     )

         # forward pass to make sure inference works in fp16
         with torch.no_grad():
            snake_case				:    List[Any]           =	model(__lowercase     )











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 						__lowercase		(UpperCamelCase__					,		UpperCamelCase__					,		unittest.TestCase							):




      """simple docstring"""





      _snake_case												=						IFPipeline
      _snake_case												=						TEXT_TO_IMAGE_PARAMS - {"""width""", """height""", """latents"""}
      _snake_case												=						TEXT_TO_IMAGE_BATCH_PARAMS
      _snake_case												=						PipelineTesterMixin.required_optional_params - {"""latents"""}

      def      UpperCAmelCase				(				self     )		->    Any:
         return self._get_dummy_components()

      def      UpperCAmelCase				(				self    ,				A    ,				A=0     )		->    Optional[int]:
         if str(A     ).startswith("""mps"""     ):
            snake_case				:    List[str]           =	torch.manual_seed(A     )
         else:
            snake_case				:    Optional[int]           =	torch.Generator(device=A     ).manual_seed(A     )

         snake_case				:    Optional[int]           =	{
             """prompt""": """A painting of a squirrel eating a burger""",
             """generator""": generator,
             """num_inference_steps""": 2,
             """output_type""": """numpy""",
         }

         return inputs

      def      UpperCAmelCase				(				self     )		->    Any:
         self._test_save_load_optional_components()

      @unittest.skipIf(torch_device != """cuda"""    ,				reason="""float16 requires CUDA"""     )
      def      UpperCAmelCase				(				self     )		->    List[str]:
         # Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
         super().test_save_load_floataa(expected_max_diff=1e-1     )

      def      UpperCAmelCase				(				self     )		->    Dict:
         self._test_attention_slicing_forward_pass(expected_max_diff=1e-2     )

      def      UpperCAmelCase				(				self     )		->    List[str]:
         self._test_save_load_local()

      def      UpperCAmelCase				(				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      UpperCAmelCase				(				self     )		->    List[str]:
         self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3     )




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




      """simple docstring"""





      def      UpperCAmelCase				(				self     )		->    Any:
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()

      def      UpperCAmelCase				(				self     )		->    List[Any]:
         # if

         snake_case				:    Tuple           =	IFPipeline.from_pretrained("""DeepFloyd/IF-I-XL-v1.0"""    ,				variant="""fp16"""    ,				torch_dtype=torch.floataa     )

         snake_case				:    Tuple           =	IFSuperResolutionPipeline.from_pretrained(
             """DeepFloyd/IF-II-L-v1.0"""    ,				variant="""fp16"""    ,				torch_dtype=torch.floataa    ,				text_encoder=A    ,				tokenizer=A     )

         # pre compute text embeddings and remove T5 to save memory

         pipe_a.text_encoder.to("""cuda"""     )

         snake_case     , snake_case				:    Optional[int]           =	pipe_a.encode_prompt("""anime turtle"""    ,				device="""cuda"""     )

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

         snake_case				:    List[str]           =	None
         snake_case				:    List[Any]           =	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(A    ,				A    ,				A    ,				A     )

         pipe_a.remove_all_hooks()
         pipe_a.remove_all_hooks()

         # img2img

         snake_case				:    Any           =	IFImgaImgPipeline(**pipe_a.components     )
         snake_case				:    Dict           =	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(A    ,				A    ,				A    ,				A     )

         pipe_a.remove_all_hooks()
         pipe_a.remove_all_hooks()

         # inpainting

         snake_case				:    Optional[Any]           =	IFInpaintingPipeline(**pipe_a.components     )
         snake_case				:    Any           =	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(A    ,				A    ,				A    ,				A     )

      def      UpperCAmelCase				(				self    ,				A    ,				A    ,				A    ,				A     )		->    str:
         # pipeline 1

         _start_torch_memory_measurement()

         snake_case				:    Tuple           =	torch.Generator(device="""cpu"""     ).manual_seed(0     )
         snake_case				:    Tuple           =	pipe_a(
             prompt_embeds=A    ,				negative_prompt_embeds=A    ,				num_inference_steps=2    ,				generator=A    ,				output_type="""np"""    ,				)

         snake_case				:    Optional[int]           =	output.images[0]

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

         snake_case				:    Tuple           =	torch.cuda.max_memory_allocated()
         assert mem_bytes < 1_3 * 1_0**9

         snake_case				:    Dict           =	load_numpy(
             """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy"""     )
         assert_mean_pixel_difference(A    ,				A     )

         # pipeline 2

         _start_torch_memory_measurement()

         snake_case				:    List[Any]           =	torch.Generator(device="""cpu"""     ).manual_seed(0     )

         snake_case				:    Optional[int]           =	floats_tensor((1, 3, 6_4, 6_4)    ,				rng=random.Random(0     )     ).to(A     )

         snake_case				:    str           =	pipe_a(
             prompt_embeds=A    ,				negative_prompt_embeds=A    ,				image=A    ,				generator=A    ,				num_inference_steps=2    ,				output_type="""np"""    ,				)

         snake_case				:    str           =	output.images[0]

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

         snake_case				:    Any           =	torch.cuda.max_memory_allocated()
         assert mem_bytes < 4 * 1_0**9

         snake_case				:    int           =	load_numpy(
             """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy"""     )
         assert_mean_pixel_difference(A    ,				A     )

      def      UpperCAmelCase				(				self    ,				A    ,				A    ,				A    ,				A     )		->    int:
         # pipeline 1

         _start_torch_memory_measurement()

         snake_case				:    Optional[Any]           =	floats_tensor((1, 3, 6_4, 6_4)    ,				rng=random.Random(0     )     ).to(A     )

         snake_case				:    List[str]           =	torch.Generator(device="""cpu"""     ).manual_seed(0     )

         snake_case				:    Union[str, Any]           =	pipe_a(
             prompt_embeds=A    ,				negative_prompt_embeds=A    ,				image=A    ,				num_inference_steps=2    ,				generator=A    ,				output_type="""np"""    ,				)

         snake_case				:    Optional[int]           =	output.images[0]

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

         snake_case				:    Tuple           =	torch.cuda.max_memory_allocated()
         assert mem_bytes < 1_0 * 1_0**9

         snake_case				:    Union[str, Any]           =	load_numpy(
             """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy"""     )
         assert_mean_pixel_difference(A    ,				A     )

         # pipeline 2

         _start_torch_memory_measurement()

         snake_case				:    int           =	torch.Generator(device="""cpu"""     ).manual_seed(0     )

         snake_case				:    int           =	floats_tensor((1, 3, 2_5_6, 2_5_6)    ,				rng=random.Random(0     )     ).to(A     )
         snake_case				:    Optional[Any]           =	floats_tensor((1, 3, 6_4, 6_4)    ,				rng=random.Random(0     )     ).to(A     )

         snake_case				:    int           =	pipe_a(
             prompt_embeds=A    ,				negative_prompt_embeds=A    ,				image=A    ,				original_image=A    ,				generator=A    ,				num_inference_steps=2    ,				output_type="""np"""    ,				)

         snake_case				:    List[Any]           =	output.images[0]

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

         snake_case				:    Tuple           =	torch.cuda.max_memory_allocated()
         assert mem_bytes < 4 * 1_0**9

         snake_case				:    Any           =	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(A    ,				A     )

      def      UpperCAmelCase				(				self    ,				A    ,				A    ,				A    ,				A     )		->    Any:
         # pipeline 1

         _start_torch_memory_measurement()

         snake_case				:    List[Any]           =	floats_tensor((1, 3, 6_4, 6_4)    ,				rng=random.Random(0     )     ).to(A     )
         snake_case				:    Union[str, Any]           =	floats_tensor((1, 3, 6_4, 6_4)    ,				rng=random.Random(1     )     ).to(A     )

         snake_case				:    Dict           =	torch.Generator(device="""cpu"""     ).manual_seed(0     )
         snake_case				:    Tuple           =	pipe_a(
             prompt_embeds=A    ,				negative_prompt_embeds=A    ,				image=A    ,				mask_image=A    ,				num_inference_steps=2    ,				generator=A    ,				output_type="""np"""    ,				)

         snake_case				:    Tuple           =	output.images[0]

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

         snake_case				:    List[str]           =	torch.cuda.max_memory_allocated()
         assert mem_bytes < 1_0 * 1_0**9

         snake_case				:    Dict           =	load_numpy(
             """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy"""     )
         assert_mean_pixel_difference(A    ,				A     )

         # pipeline 2

         _start_torch_memory_measurement()

         snake_case				:    Optional[int]           =	torch.Generator(device="""cpu"""     ).manual_seed(0     )

         snake_case				:    int           =	floats_tensor((1, 3, 6_4, 6_4)    ,				rng=random.Random(0     )     ).to(A     )
         snake_case				:    Any           =	floats_tensor((1, 3, 2_5_6, 2_5_6)    ,				rng=random.Random(0     )     ).to(A     )
         snake_case				:    str           =	floats_tensor((1, 3, 2_5_6, 2_5_6)    ,				rng=random.Random(1     )     ).to(A     )

         snake_case				:    List[str]           =	pipe_a(
             prompt_embeds=A    ,				negative_prompt_embeds=A    ,				image=A    ,				mask_image=A    ,				original_image=A    ,				generator=A    ,				num_inference_steps=2    ,				output_type="""np"""    ,				)

         snake_case				:    List[Any]           =	output.images[0]

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

         snake_case				:    Union[str, Any]           =	torch.cuda.max_memory_allocated()
         assert mem_bytes < 4 * 1_0**9

         snake_case				:    str           =	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(A    ,				A     )




def 						SCREAMING_SNAKE_CASE__ (   )  ->       str:
   torch.cuda.empty_cache()
   torch.cuda.reset_max_memory_allocated()
   torch.cuda.reset_peak_memory_stats()





