'''simple docstring'''

import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyInpaintPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()


class 					_lowercase						(   lowerCAmelCase		,unittest.TestCase ):
      '''simple docstring'''


      UpperCAmelCase_   :   Any   				= KandinskyInpaintPipeline
      UpperCAmelCase_   :   str   				= ['''prompt''', '''image_embeds''', '''negative_image_embeds''', '''image''', '''mask_image''']
      UpperCAmelCase_   :   Union[str, Any]   				= [
          '''prompt''',
          '''negative_prompt''',
          '''image_embeds''',
          '''negative_image_embeds''',
          '''image''',
          '''mask_image''',
      ]
      UpperCAmelCase_   :   Union[str, Any]   				= [
          '''generator''',
          '''height''',
          '''width''',
          '''latents''',
          '''guidance_scale''',
          '''negative_prompt''',
          '''num_inference_steps''',
          '''return_dict''',
          '''guidance_scale''',
          '''num_images_per_prompt''',
          '''output_type''',
          '''return_dict''',
      ]
      UpperCAmelCase_   :   Optional[int]   				= False




      @property
      def 		lowerCAmelCase__       (					self				)      ->      Dict:



         '''simple docstring'''



         return 32




      @property
      def 		lowerCAmelCase__       (					self				)      ->      Optional[Any]:



         '''simple docstring'''



         return 32




      @property
      def 		lowerCAmelCase__       (					self				)      ->      Dict:



         '''simple docstring'''



         return self.time_input_dim




      @property
      def 		lowerCAmelCase__       (					self				)      ->      Optional[int]:



         '''simple docstring'''



         return self.time_input_dim * 4




      @property
      def 		lowerCAmelCase__       (					self				)      ->      Dict:



         '''simple docstring'''



         return 100




      @property
      def 		lowerCAmelCase__       (					self				)      ->      int:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[Any] 					=	XLMRobertaTokenizerFast.from_pretrained('''YiYiXu/tiny-random-mclip-base'''				)
         return tokenizer




      @property
      def 		lowerCAmelCase__       (					self				)      ->      int:



         '''simple docstring'''



         torch.manual_seed(0				)
         UpperCAmelCase__				:					Any 					=	MCLIPConfig(
             numDims=self.cross_attention_dim       ,transformerDimensions=self.text_embedder_hidden_size       ,hidden_size=self.text_embedder_hidden_size       ,intermediate_size=37       ,num_attention_heads=4       ,num_hidden_layers=5       ,vocab_size=1005       ,)

         UpperCAmelCase__				:					Any 					=	MultilingualCLIP(lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[Any] 					=	text_encoder.eval()

         return text_encoder




      @property
      def 		lowerCAmelCase__       (					self				)      ->      Union[str, Any]:



         '''simple docstring'''



         torch.manual_seed(0				)

         UpperCAmelCase__				:					Tuple 					=	{
             '''in_channels''': 9,
             # Out channels is double in channels because predicts mean and variance
             '''out_channels''': 8,
             '''addition_embed_type''': '''text_image''',
             '''down_block_types''': ('''ResnetDownsampleBlock2D''', '''SimpleCrossAttnDownBlock2D'''),
             '''up_block_types''': ('''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D'''),
             '''mid_block_type''': '''UNetMidBlock2DSimpleCrossAttn''',
             '''block_out_channels''': (self.block_out_channels_a, self.block_out_channels_a * 2),
             '''layers_per_block''': 1,
             '''encoder_hid_dim''': self.text_embedder_hidden_size,
             '''encoder_hid_dim_type''': '''text_image_proj''',
             '''cross_attention_dim''': self.cross_attention_dim,
             '''attention_head_dim''': 4,
             '''resnet_time_scale_shift''': '''scale_shift''',
             '''class_embed_type''': None,
         }

         UpperCAmelCase__				:					List[str] 					=	UNetaDConditionModel(**lowerCamelCase_				)
         return model




      @property
      def 		lowerCAmelCase__       (					self				)      ->      List[str]:



         '''simple docstring'''



         return {
             "block_out_channels": [32, 64],
             "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
             "in_channels": 3,
             "latent_channels": 4,
             "layers_per_block": 1,
             "norm_num_groups": 8,
             "norm_type": "spatial",
             "num_vq_embeddings": 12,
             "out_channels": 3,
             "up_block_types": [
                 "AttnUpDecoderBlock2D",
                 "UpDecoderBlock2D",
             ],
             "vq_embed_dim": 4,
         }




      @property
      def 		lowerCAmelCase__       (					self				)      ->      Tuple:



         '''simple docstring'''



         torch.manual_seed(0				)
         UpperCAmelCase__				:					Union[str, Any] 					=	VQModel(**self.dummy_movq_kwargs				)
         return model




      def 		lowerCAmelCase__       (					self				)      ->      Any:



         '''simple docstring'''



         UpperCAmelCase__				:					List[Any] 					=	self.dummy_text_encoder
         UpperCAmelCase__				:					Optional[int] 					=	self.dummy_tokenizer
         UpperCAmelCase__				:					Optional[Any] 					=	self.dummy_unet
         UpperCAmelCase__				:					Any 					=	self.dummy_movq

         UpperCAmelCase__				:					Tuple 					=	DDIMScheduler(
             num_train_timesteps=1000       ,beta_schedule='''linear'''       ,beta_start=0.0_0085       ,beta_end=0.012       ,clip_sample=lowerCamelCase_       ,set_alpha_to_one=lowerCamelCase_       ,steps_offset=1       ,prediction_type='''epsilon'''       ,thresholding=lowerCamelCase_       ,)

         UpperCAmelCase__				:					Dict 					=	{
             '''text_encoder''': text_encoder,
             '''tokenizer''': tokenizer,
             '''unet''': unet,
             '''scheduler''': scheduler,
             '''movq''': movq,
         }

         return components




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_=0				)      ->      Tuple:



         '''simple docstring'''



         UpperCAmelCase__				:					Dict 					=	floats_tensor((1, self.cross_attention_dim)       ,rng=random.Random(lowerCamelCase_				)				).to(lowerCamelCase_				)
         UpperCAmelCase__				:					str 					=	floats_tensor((1, self.cross_attention_dim)       ,rng=random.Random(seed + 1				)				).to(lowerCamelCase_				)
         # create init_image
         UpperCAmelCase__				:					str 					=	floats_tensor((1, 3, 64, 64)       ,rng=random.Random(lowerCamelCase_				)				).to(lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[Any] 					=	image.cpu().permute(0       ,2       ,3       ,1				)[0]
         UpperCAmelCase__				:					Union[str, Any] 					=	Image.fromarray(np.uinta(lowerCamelCase_				)				).convert('''RGB'''				).resize((256, 256)				)
         # create mask
         UpperCAmelCase__				:					List[str] 					=	np.ones((64, 64)       ,dtype=np.floataa				)
         UpperCAmelCase__				:					Optional[Any] 					=	0

         if str(lowerCamelCase_				).startswith('''mps'''				):
            UpperCAmelCase__				:					Any 					=	torch.manual_seed(lowerCamelCase_				)
         else:
            UpperCAmelCase__				:					Dict 					=	torch.Generator(device=lowerCamelCase_				).manual_seed(lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[Any] 					=	{
             '''prompt''': '''horse''',
             '''image''': init_image,
             '''mask_image''': mask,
             '''image_embeds''': image_embeds,
             '''negative_image_embeds''': negative_image_embeds,
             '''generator''': generator,
             '''height''': 64,
             '''width''': 64,
             '''num_inference_steps''': 2,
             '''guidance_scale''': 4.0,
             '''output_type''': '''np''',
         }
         return inputs




      def 		lowerCAmelCase__       (					self				)      ->      Dict:



         '''simple docstring'''



         UpperCAmelCase__				:					Tuple 					=	'''cpu'''

         UpperCAmelCase__				:					List[str] 					=	self.get_dummy_components()

         UpperCAmelCase__				:					Tuple 					=	self.pipeline_class(**lowerCamelCase_				)
         UpperCAmelCase__				:					int 					=	pipe.to(lowerCamelCase_				)

         pipe.set_progress_bar_config(disable=lowerCamelCase_				)

         UpperCAmelCase__				:					Tuple 					=	pipe(**self.get_dummy_inputs(lowerCamelCase_				)				)
         UpperCAmelCase__				:					List[Any] 					=	output.images

         UpperCAmelCase__				:					Union[str, Any] 					=	pipe(
             **self.get_dummy_inputs(lowerCamelCase_				)       ,return_dict=lowerCamelCase_       ,)[0]

         UpperCAmelCase__				:					Tuple 					=	image[0, -3:, -3:, -1]
         UpperCAmelCase__				:					int 					=	image_from_tuple[0, -3:, -3:, -1]

         print(f'''image.shape {image.shape}'''				)

         assert image.shape == (1, 64, 64, 3)

         UpperCAmelCase__				:					List[Any] 					=	np.array(
             [0.832_6919, 0.7379_0467, 0.2091_8581, 0.930_9612, 0.551_1791, 0.4371_3328, 0.551_3321, 0.4992_2934, 0.5949_7786]				)

         assert (
             np.abs(image_slice.flatten() - expected_slice				).max() < 1e-2
         ), f''' expected_slice {expected_slice}, but got {image_slice.flatten()}'''
         assert (
             np.abs(image_from_tuple_slice.flatten() - expected_slice				).max() < 1e-2
         ), f''' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'''




      def 		lowerCAmelCase__       (					self				)      ->      Dict:



         '''simple docstring'''



         super().test_inference_batch_single_identical(expected_max_diff=3e-3				)




@slow
@require_torch_gpu
class 					_lowercase						(   unittest.TestCase ):
      '''simple docstring'''


      def 		lowerCAmelCase__       (					self				)      ->      List[str]:



         '''simple docstring'''



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




      def 		lowerCAmelCase__       (					self				)      ->      str:



         '''simple docstring'''



         UpperCAmelCase__				:					Dict 					=	load_numpy(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/kandinsky/kandinsky_inpaint_cat_with_hat_fp16.npy'''				)

         UpperCAmelCase__				:					str 					=	load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/kandinsky/cat.png'''				)
         UpperCAmelCase__				:					Union[str, Any] 					=	np.ones((768, 768)       ,dtype=np.floataa				)
         UpperCAmelCase__				:					int 					=	0

         UpperCAmelCase__				:					List[Any] 					=	'''a hat'''

         UpperCAmelCase__				:					str 					=	KandinskyPriorPipeline.from_pretrained(
             '''kandinsky-community/kandinsky-2-1-prior'''       ,torch_dtype=torch.floataa				)
         pipe_prior.to(lowerCamelCase_				)

         UpperCAmelCase__				:					List[str] 					=	KandinskyInpaintPipeline.from_pretrained(
             '''kandinsky-community/kandinsky-2-1-inpaint'''       ,torch_dtype=torch.floataa				)
         UpperCAmelCase__				:					Tuple 					=	pipeline.to(lowerCamelCase_				)
         pipeline.set_progress_bar_config(disable=lowerCamelCase_				)

         UpperCAmelCase__				:					List[Any] 					=	torch.Generator(device='''cpu'''				).manual_seed(0				)
         UpperCAmelCase__					,   UpperCAmelCase__				:					int 					=	pipe_prior(
             lowerCamelCase_       ,generator=lowerCamelCase_       ,num_inference_steps=5       ,negative_prompt=''''''       ,).to_tuple()

         UpperCAmelCase__				:					Optional[Any] 					=	pipeline(
             lowerCamelCase_       ,image=lowerCamelCase_       ,mask_image=lowerCamelCase_       ,image_embeds=lowerCamelCase_       ,negative_image_embeds=lowerCamelCase_       ,generator=lowerCamelCase_       ,num_inference_steps=100       ,height=768       ,width=768       ,output_type='''np'''       ,)

         UpperCAmelCase__				:					str 					=	output.images[0]

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

         assert_mean_pixel_difference(lowerCamelCase_       ,lowerCamelCase_				)





'''simple docstring'''

import gc
import unittest

import torch
from parameterized import parameterized

from diffusers import AutoencoderKL
from diffusers.utils import floats_tensor, load_hf_numpy, require_torch_gpu, slow, torch_all_close, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()


class 					_lowercase						(   lowerCAmelCase		,lowerCAmelCase		,unittest.TestCase ):
      '''simple docstring'''


      UpperCAmelCase_   :   int   				= AutoencoderKL
      UpperCAmelCase_   :   Any   				= '''sample'''
      UpperCAmelCase_   :   List[str]   				= 1E-2




      @property
      def 		lowerCAmelCase__       (					self				)      ->      int:



         '''simple docstring'''



         UpperCAmelCase__				:					Dict 					=	4
         UpperCAmelCase__				:					Optional[int] 					=	3
         UpperCAmelCase__				:					Tuple 					=	(32, 32)

         UpperCAmelCase__				:					Optional[int] 					=	floats_tensor((batch_size, num_channels) + sizes				).to(lowerCamelCase_				)

         return {"sample": image}




      @property
      def 		lowerCAmelCase__       (					self				)      ->      List[str]:



         '''simple docstring'''



         return (3, 32, 32)




      @property
      def 		lowerCAmelCase__       (					self				)      ->      Any:



         '''simple docstring'''



         return (3, 32, 32)




      def 		lowerCAmelCase__       (					self				)      ->      Dict:



         '''simple docstring'''



         UpperCAmelCase__				:					List[Any] 					=	{
             '''block_out_channels''': [32, 64],
             '''in_channels''': 3,
             '''out_channels''': 3,
             '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
             '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
             '''latent_channels''': 4,
         }
         UpperCAmelCase__				:					Optional[Any] 					=	self.dummy_input
         return init_dict, inputs_dict




      def 		lowerCAmelCase__       (					self				)      ->      Union[str, Any]:



         '''simple docstring'''



         pass




      def 		lowerCAmelCase__       (					self				)      ->      Union[str, Any]:



         '''simple docstring'''



         pass




      @unittest.skipIf(torch_device == '''mps'''       ,'''Gradient checkpointing skipped on MPS'''				)
      def 		lowerCAmelCase__       (					self				)      ->      Tuple:



         '''simple docstring'''



         UpperCAmelCase__					,   UpperCAmelCase__				:					Optional[Any] 					=	self.prepare_init_args_and_inputs_for_common()
         UpperCAmelCase__				:					Union[str, Any] 					=	self.model_class(**lowerCamelCase_				)
         model.to(lowerCamelCase_				)

         assert not model.is_gradient_checkpointing and model.training

         UpperCAmelCase__				:					Any 					=	model(**lowerCamelCase_				).sample
         # run the backwards pass on the model. For backwards pass, for simplicity purpose,
         # we won't calculate the loss and rather backprop on out.sum()
         model.zero_grad()

         UpperCAmelCase__				:					Optional[int] 					=	torch.randn_like(lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[Any] 					=	(out - labels).mean()
         loss.backward()

         # re-instantiate the model now enabling gradient checkpointing
         UpperCAmelCase__				:					Optional[int] 					=	self.model_class(**lowerCamelCase_				)
         # clone model
         model_a.load_state_dict(model.state_dict()				)
         model_a.to(lowerCamelCase_				)
         model_a.enable_gradient_checkpointing()

         assert model_a.is_gradient_checkpointing and model_a.training

         UpperCAmelCase__				:					int 					=	model_a(**lowerCamelCase_				).sample
         # run the backwards pass on the model. For backwards pass, for simplicity purpose,
         # we won't calculate the loss and rather backprop on out.sum()
         model_a.zero_grad()
         UpperCAmelCase__				:					Union[str, Any] 					=	(out_a - labels).mean()
         loss_a.backward()

         # compare the output and parameters gradients
         self.assertTrue((loss - loss_a).abs() < 1e-5				)
         UpperCAmelCase__				:					Union[str, Any] 					=	dict(model.named_parameters()				)
         UpperCAmelCase__				:					int 					=	dict(model_a.named_parameters()				)
         for name, param in named_params.items():
            self.assertTrue(torch_all_close(param.grad.data       ,named_params_a[name].grad.data       ,atol=5e-5				)				)




      def 		lowerCAmelCase__       (					self				)      ->      Tuple:



         '''simple docstring'''



         UpperCAmelCase__					,   UpperCAmelCase__				:					Union[str, Any] 					=	AutoencoderKL.from_pretrained('''fusing/autoencoder-kl-dummy'''       ,output_loading_info=lowerCamelCase_				)
         self.assertIsNotNone(lowerCamelCase_				)
         self.assertEqual(len(loading_info['''missing_keys''']				)       ,0				)

         model.to(lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[int] 					=	model(**self.dummy_input				)

         assert image is not None, "Make sure output is not None"




      def 		lowerCAmelCase__       (					self				)      ->      str:



         '''simple docstring'''



         UpperCAmelCase__				:					Tuple 					=	AutoencoderKL.from_pretrained('''fusing/autoencoder-kl-dummy'''				)
         UpperCAmelCase__				:					Optional[Any] 					=	model.to(lowerCamelCase_				)
         model.eval()

         if torch_device == "mps":
            UpperCAmelCase__				:					Optional[Any] 					=	torch.manual_seed(0				)
         else:
            UpperCAmelCase__				:					Tuple 					=	torch.Generator(device=lowerCamelCase_				).manual_seed(0				)

         UpperCAmelCase__				:					str 					=	torch.randn(
             1       ,model.config.in_channels       ,model.config.sample_size       ,model.config.sample_size       ,generator=torch.manual_seed(0				)       ,)
         UpperCAmelCase__				:					int 					=	image.to(lowerCamelCase_				)
         with torch.no_grad():
            UpperCAmelCase__				:					str 					=	model(lowerCamelCase_       ,sample_posterior=lowerCamelCase_       ,generator=lowerCamelCase_				).sample

         UpperCAmelCase__				:					int 					=	output[0, -1, -3:, -3:].flatten().cpu()

         # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
         # the expected output slices are not the same for CPU and GPU.
         if torch_device == "mps":
            UpperCAmelCase__				:					int 					=	torch.tensor(
                [
                    -4.0078e-01,
                    -3.8323e-04,
                    -1.2681e-01,
                    -1.1462e-01,
                    2.0095e-01,
                    1.0893e-01,
                    -8.8247e-02,
                    -3.0361e-01,
                    -9.8644e-03,
                ]				)
         elif torch_device == "cpu":
            UpperCAmelCase__				:					List[Any] 					=	torch.tensor(
                [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026]				)
         else:
            UpperCAmelCase__				:					Optional[Any] 					=	torch.tensor(
                [-0.2421, 0.4642, 0.2507, -0.0438, 0.0682, 0.3160, -0.2018, -0.0727, 0.2485]				)

         self.assertTrue(torch_all_close(lowerCamelCase_       ,lowerCamelCase_       ,rtol=1e-2				)				)




@slow
class 					_lowercase						(   unittest.TestCase ):
      '''simple docstring'''


      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      int:



         '''simple docstring'''



         return f'''gaussian_noise_s={seed}_shape={'_'.join([str(lowerCamelCase_				) for s in shape]				)}.npy'''




      def 		lowerCAmelCase__       (					self				)      ->      Optional[int]:



         '''simple docstring'''



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




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_=0       ,lowerCamelCase_=(4, 3, 512, 512)       ,lowerCamelCase_=False				)      ->      int:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[Any] 					=	torch.floataa if fpaa else torch.floataa
         UpperCAmelCase__				:					Dict 					=	torch.from_numpy(load_hf_numpy(self.get_file_format(lowerCamelCase_       ,lowerCamelCase_				)				)				).to(lowerCamelCase_				).to(lowerCamelCase_				)
         return image




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_="CompVis/stable-diffusion-v1-4"       ,lowerCamelCase_=False				)      ->      Dict:



         '''simple docstring'''



         UpperCAmelCase__				:					Union[str, Any] 					=	'''fp16''' if fpaa else None
         UpperCAmelCase__				:					str 					=	torch.floataa if fpaa else torch.floataa

         UpperCAmelCase__				:					List[Any] 					=	AutoencoderKL.from_pretrained(
             lowerCamelCase_       ,subfolder='''vae'''       ,torch_dtype=lowerCamelCase_       ,revision=lowerCamelCase_       ,)
         model.to(lowerCamelCase_				).eval()

         return model




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_=0				)      ->      Optional[Any]:



         '''simple docstring'''



         if torch_device == "mps":
            return torch.manual_seed(lowerCamelCase_				)
         return torch.Generator(device=lowerCamelCase_				).manual_seed(lowerCamelCase_				)




      @parameterized.expand(
          [
              # fmt: off
              [33, [-0.1603, 0.9878, -0.0495, -0.0790, -0.2709, 0.8375, -0.2060, -0.0824], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
              [47, [-0.2376, 0.1168, 0.1332, -0.4840, -0.2508, -0.0791, -0.0493, -0.4089], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
              # fmt: on
          ]				)
      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      int:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[Any] 					=	self.get_sd_vae_model()
         UpperCAmelCase__				:					str 					=	self.get_sd_image(lowerCamelCase_				)
         UpperCAmelCase__				:					Dict 					=	self.get_generator(lowerCamelCase_				)

         with torch.no_grad():
            UpperCAmelCase__				:					Tuple 					=	model(lowerCamelCase_       ,generator=lowerCamelCase_       ,sample_posterior=lowerCamelCase_				).sample

         assert sample.shape == image.shape

         UpperCAmelCase__				:					List[Any] 					=	sample[-1, -2:, -2:, :2].flatten().float().cpu()
         UpperCAmelCase__				:					Tuple 					=	torch.tensor(expected_slice_mps if torch_device == '''mps''' else expected_slice				)

         assert torch_all_close(lowerCamelCase_       ,lowerCamelCase_       ,atol=3e-3				)




      @parameterized.expand(
          [
              # fmt: off
              [33, [-0.0513, 0.0289, 1.3799, 0.2166, -0.2573, -0.0871, 0.5103, -0.0999]],
              [47, [-0.4128, -0.1320, -0.3704, 0.1965, -0.4116, -0.2332, -0.3340, 0.2247]],
              # fmt: on
          ]				)
      @require_torch_gpu
      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      List[Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					Tuple 					=	self.get_sd_vae_model(fpaa=lowerCamelCase_				)
         UpperCAmelCase__				:					Tuple 					=	self.get_sd_image(lowerCamelCase_       ,fpaa=lowerCamelCase_				)
         UpperCAmelCase__				:					Union[str, Any] 					=	self.get_generator(lowerCamelCase_				)

         with torch.no_grad():
            UpperCAmelCase__				:					List[str] 					=	model(lowerCamelCase_       ,generator=lowerCamelCase_       ,sample_posterior=lowerCamelCase_				).sample

         assert sample.shape == image.shape

         UpperCAmelCase__				:					Optional[Any] 					=	sample[-1, -2:, :2, -2:].flatten().float().cpu()
         UpperCAmelCase__				:					List[str] 					=	torch.tensor(lowerCamelCase_				)

         assert torch_all_close(lowerCamelCase_       ,lowerCamelCase_       ,atol=1e-2				)




      @parameterized.expand(
          [
              # fmt: off
              [33, [-0.1609, 0.9866, -0.0487, -0.0777, -0.2716, 0.8368, -0.2055, -0.0814], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
              [47, [-0.2377, 0.1147, 0.1333, -0.4841, -0.2506, -0.0805, -0.0491, -0.4085], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
              # fmt: on
          ]				)
      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      str:



         '''simple docstring'''



         UpperCAmelCase__				:					Any 					=	self.get_sd_vae_model()
         UpperCAmelCase__				:					Optional[int] 					=	self.get_sd_image(lowerCamelCase_				)

         with torch.no_grad():
            UpperCAmelCase__				:					List[str] 					=	model(lowerCamelCase_				).sample

         assert sample.shape == image.shape

         UpperCAmelCase__				:					str 					=	sample[-1, -2:, -2:, :2].flatten().float().cpu()
         UpperCAmelCase__				:					Optional[Any] 					=	torch.tensor(expected_slice_mps if torch_device == '''mps''' else expected_slice				)

         assert torch_all_close(lowerCamelCase_       ,lowerCamelCase_       ,atol=3e-3				)




      @parameterized.expand(
          [
              # fmt: off
              [13, [-0.2051, -0.1803, -0.2311, -0.2114, -0.3292, -0.3574, -0.2953, -0.3323]],
              [37, [-0.2632, -0.2625, -0.2199, -0.2741, -0.4539, -0.4990, -0.3720, -0.4925]],
              # fmt: on
          ]				)
      @require_torch_gpu
      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      Dict:



         '''simple docstring'''



         UpperCAmelCase__				:					Dict 					=	self.get_sd_vae_model()
         UpperCAmelCase__				:					Dict 					=	self.get_sd_image(lowerCamelCase_       ,shape=(3, 4, 64, 64)				)

         with torch.no_grad():
            UpperCAmelCase__				:					Optional[Any] 					=	model.decode(lowerCamelCase_				).sample

         assert list(sample.shape				) == [3, 3, 512, 512]

         UpperCAmelCase__				:					Optional[int] 					=	sample[-1, -2:, :2, -2:].flatten().cpu()
         UpperCAmelCase__				:					str 					=	torch.tensor(lowerCamelCase_				)

         assert torch_all_close(lowerCamelCase_       ,lowerCamelCase_       ,atol=1e-3				)




      @parameterized.expand(
          [
              # fmt: off
              [27, [-0.0369, 0.0207, -0.0776, -0.0682, -0.1747, -0.1930, -0.1465, -0.2039]],
              [16, [-0.1628, -0.2134, -0.2747, -0.2642, -0.3774, -0.4404, -0.3687, -0.4277]],
              # fmt: on
          ]				)
      @require_torch_gpu
      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      Any:



         '''simple docstring'''



         UpperCAmelCase__				:					int 					=	self.get_sd_vae_model(fpaa=lowerCamelCase_				)
         UpperCAmelCase__				:					str 					=	self.get_sd_image(lowerCamelCase_       ,shape=(3, 4, 64, 64)       ,fpaa=lowerCamelCase_				)

         with torch.no_grad():
            UpperCAmelCase__				:					Dict 					=	model.decode(lowerCamelCase_				).sample

         assert list(sample.shape				) == [3, 3, 512, 512]

         UpperCAmelCase__				:					List[str] 					=	sample[-1, -2:, :2, -2:].flatten().float().cpu()
         UpperCAmelCase__				:					Any 					=	torch.tensor(lowerCamelCase_				)

         assert torch_all_close(lowerCamelCase_       ,lowerCamelCase_       ,atol=5e-3				)




      @parameterized.expand([(13,), (16,), (27,)]				)
      @require_torch_gpu
      @unittest.skipIf(not is_xformers_available()       ,reason='''xformers is not required when using PyTorch 2.0.'''				)
      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_				)      ->      Tuple:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[Any] 					=	self.get_sd_vae_model(fpaa=lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[int] 					=	self.get_sd_image(lowerCamelCase_       ,shape=(3, 4, 64, 64)       ,fpaa=lowerCamelCase_				)

         with torch.no_grad():
            UpperCAmelCase__				:					Any 					=	model.decode(lowerCamelCase_				).sample

         model.enable_xformers_memory_efficient_attention()
         with torch.no_grad():
            UpperCAmelCase__				:					int 					=	model.decode(lowerCamelCase_				).sample

         assert list(sample.shape				) == [3, 3, 512, 512]

         assert torch_all_close(lowerCamelCase_       ,lowerCamelCase_       ,atol=1e-1				)




      @parameterized.expand([(13,), (16,), (37,)]				)
      @require_torch_gpu
      @unittest.skipIf(not is_xformers_available()       ,reason='''xformers is not required when using PyTorch 2.0.'''				)
      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_				)      ->      List[str]:



         '''simple docstring'''



         UpperCAmelCase__				:					List[Any] 					=	self.get_sd_vae_model()
         UpperCAmelCase__				:					List[Any] 					=	self.get_sd_image(lowerCamelCase_       ,shape=(3, 4, 64, 64)				)

         with torch.no_grad():
            UpperCAmelCase__				:					Tuple 					=	model.decode(lowerCamelCase_				).sample

         model.enable_xformers_memory_efficient_attention()
         with torch.no_grad():
            UpperCAmelCase__				:					List[str] 					=	model.decode(lowerCamelCase_				).sample

         assert list(sample.shape				) == [3, 3, 512, 512]

         assert torch_all_close(lowerCamelCase_       ,lowerCamelCase_       ,atol=1e-2				)




      @parameterized.expand(
          [
              # fmt: off
              [33, [-0.3001, 0.0918, -2.6984, -3.9720, -3.2099, -5.0353, 1.7338, -0.2065, 3.4267]],
              [47, [-1.5030, -4.3871, -6.0355, -9.1157, -1.6661, -2.7853, 2.1607, -5.0823, 2.5633]],
              # fmt: on
          ]				)
      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      List[str]:



         '''simple docstring'''



         UpperCAmelCase__				:					Any 					=	self.get_sd_vae_model()
         UpperCAmelCase__				:					Optional[Any] 					=	self.get_sd_image(lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[int] 					=	self.get_generator(lowerCamelCase_				)

         with torch.no_grad():
            UpperCAmelCase__				:					Optional[int] 					=	model.encode(lowerCamelCase_				).latent_dist
            UpperCAmelCase__				:					Tuple 					=	dist.sample(generator=lowerCamelCase_				)

         assert list(sample.shape				) == [image.shape[0], 4] + [i // 8 for i in image.shape[2:]]

         UpperCAmelCase__				:					List[str] 					=	sample[0, -1, -3:, -3:].flatten().cpu()
         UpperCAmelCase__				:					Optional[Any] 					=	torch.tensor(lowerCamelCase_				)

         UpperCAmelCase__				:					List[Any] 					=	3e-3 if torch_device != '''mps''' else 1e-2
         assert torch_all_close(lowerCamelCase_       ,lowerCamelCase_       ,atol=lowerCamelCase_				)




