
import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    MobileViTConfig,
    MobileViTForImageClassification,
    MobileViTForSemanticSegmentation,
    MobileViTImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
A     :	int    = logging.get_logger(__name__)


def     __lowerCAmelCase			(			a__  )       ->  List[Any]:
				__a										=    MobileViTConfig()

				# size of the architecture
				if "mobilevit_s" in mobilevit_name:
								__a										=    [144, 192, 240]
								__a										=    [16, 32, 64, 96, 128, 160, 640]
				elif "mobilevit_xs" in mobilevit_name:
								__a										=    [96, 120, 144]
								__a										=    [16, 32, 48, 64, 80, 96, 384]
				elif "mobilevit_xxs" in mobilevit_name:
								__a										=    [64, 80, 96]
								__a										=    [16, 16, 24, 48, 64, 80, 320]
								__a										=    0.05
								__a										=    2.0

				if mobilevit_name.startswith('''deeplabv3_'''  ):
								__a										=    512
								__a										=    16
								__a										=    21
								__a										=    '''pascal-voc-id2label.json'''
				else:
								__a										=    1000
								__a										=    '''imagenet-1k-id2label.json'''

				__a										=    '''huggingface/label-files'''
				__a										=    json.load(open(hf_hub_download(a__						,    a__						,    repo_type='''dataset'''  )						,    '''r'''  )  )
				__a										=    {int(a__  ): v for k, v in idalabel.items()}
				__a										=    idalabel
				__a										=    {v: k for k, v in idalabel.items()}

				return config


def     __lowerCAmelCase			(			a__						,    a__=False  )       ->  List[Any]:
				for i in range(1						,    6  ):
								if F"""layer_{i}.""" in name:
												__a										=    name.replace(F"""layer_{i}."""						,    F"""encoder.layer.{i - 1}."""  )

				if "conv_1." in name:
								__a										=    name.replace('''conv_1.'''						,    '''conv_stem.'''  )
				if ".block." in name:
								__a										=    name.replace('''.block.'''						,    '''.'''  )

				if "exp_1x1" in name:
								__a										=    name.replace('''exp_1x1'''						,    '''expand_1x1'''  )
				if "red_1x1" in name:
								__a										=    name.replace('''red_1x1'''						,    '''reduce_1x1'''  )
				if ".local_rep.conv_3x3." in name:
								__a										=    name.replace('''.local_rep.conv_3x3.'''						,    '''.conv_kxk.'''  )
				if ".local_rep.conv_1x1." in name:
								__a										=    name.replace('''.local_rep.conv_1x1.'''						,    '''.conv_1x1.'''  )
				if ".norm." in name:
								__a										=    name.replace('''.norm.'''						,    '''.normalization.'''  )
				if ".conv." in name:
								__a										=    name.replace('''.conv.'''						,    '''.convolution.'''  )
				if ".conv_proj." in name:
								__a										=    name.replace('''.conv_proj.'''						,    '''.conv_projection.'''  )

				for i in range(0						,    2  ):
								for j in range(0						,    4  ):
												if F""".{i}.{j}.""" in name:
																__a										=    name.replace(F""".{i}.{j}."""						,    F""".{i}.layer.{j}."""  )

				for i in range(2						,    6  ):
								for j in range(0						,    4  ):
												if F""".{i}.{j}.""" in name:
																__a										=    name.replace(F""".{i}.{j}."""						,    F""".{i}."""  )
																if "expand_1x1" in name:
																				__a										=    name.replace('''expand_1x1'''						,    '''downsampling_layer.expand_1x1'''  )
																if "conv_3x3" in name:
																				__a										=    name.replace('''conv_3x3'''						,    '''downsampling_layer.conv_3x3'''  )
																if "reduce_1x1" in name:
																				__a										=    name.replace('''reduce_1x1'''						,    '''downsampling_layer.reduce_1x1'''  )

				for i in range(2						,    5  ):
								if F""".global_rep.{i}.weight""" in name:
												__a										=    name.replace(F""".global_rep.{i}.weight"""						,    '''.layernorm.weight'''  )
								if F""".global_rep.{i}.bias""" in name:
												__a										=    name.replace(F""".global_rep.{i}.bias"""						,    '''.layernorm.bias'''  )

				if ".global_rep." in name:
								__a										=    name.replace('''.global_rep.'''						,    '''.transformer.'''  )
				if ".pre_norm_mha.0." in name:
								__a										=    name.replace('''.pre_norm_mha.0.'''						,    '''.layernorm_before.'''  )
				if ".pre_norm_mha.1.out_proj." in name:
								__a										=    name.replace('''.pre_norm_mha.1.out_proj.'''						,    '''.attention.output.dense.'''  )
				if ".pre_norm_ffn.0." in name:
								__a										=    name.replace('''.pre_norm_ffn.0.'''						,    '''.layernorm_after.'''  )
				if ".pre_norm_ffn.1." in name:
								__a										=    name.replace('''.pre_norm_ffn.1.'''						,    '''.intermediate.dense.'''  )
				if ".pre_norm_ffn.4." in name:
								__a										=    name.replace('''.pre_norm_ffn.4.'''						,    '''.output.dense.'''  )
				if ".transformer." in name:
								__a										=    name.replace('''.transformer.'''						,    '''.transformer.layer.'''  )

				if ".aspp_layer." in name:
								__a										=    name.replace('''.aspp_layer.'''						,    '''.'''  )
				if ".aspp_pool." in name:
								__a										=    name.replace('''.aspp_pool.'''						,    '''.'''  )
				if "seg_head." in name:
								__a										=    name.replace('''seg_head.'''						,    '''segmentation_head.'''  )
				if "segmentation_head.classifier.classifier." in name:
								__a										=    name.replace('''segmentation_head.classifier.classifier.'''						,    '''segmentation_head.classifier.'''  )

				if "classifier.fc." in name:
								__a										=    name.replace('''classifier.fc.'''						,    '''classifier.'''  )
				elif (not base_model) and ("segmentation_head." not in name):
								__a										=    '''mobilevit.''' + name

				return name


def     __lowerCAmelCase			(			a__						,    a__						,    a__=False  )       ->  Union[str, Any]:
				if base_model:
								__a										=    ''''''
				else:
								__a										=    '''mobilevit.'''

				for key in orig_state_dict.copy().keys():
								__a										=    orig_state_dict.pop(a__  )

								if key[:8] == "encoder.":
												__a										=    key[8:]

								if "qkv" in key:
												__a										=    key.split('''.'''  )
												__a										=    int(key_split[0][6:]  ) - 1
												__a										=    int(key_split[3]  )
												__a										=    model.get_submodule(F"""{model_prefix}encoder.layer.{layer_num}"""  )
												__a										=    layer.transformer.layer[transformer_num].attention.attention.all_head_size
												__a										=    (
												    F"""{model_prefix}encoder.layer.{layer_num}.transformer.layer.{transformer_num}.attention.attention."""
												)
												if "weight" in key:
																__a										=    val[:dim, :]
																__a										=    val[dim : dim * 2, :]
																__a										=    val[-dim:, :]
												else:
																__a										=    val[:dim]
																__a										=    val[dim : dim * 2]
																__a										=    val[-dim:]
								else:
												__a										=    val

				return orig_state_dict


def     __lowerCAmelCase			(			)       ->  List[Any]:
				__a										=    '''http://images.cocodataset.org/val2017/000000039769.jpg'''
				__a										=    Image.open(requests.get(a__						,    stream=a__  ).raw  )
				return im


@torch.no_grad()
def     __lowerCAmelCase			(			a__						,    a__						,    a__						,    a__=False  )       ->  Dict:
				__a										=    get_mobilevit_config(a__  )

				# load original state_dict
				__a										=    torch.load(a__						,    map_location='''cpu'''  )

				# load 🤗 model
				if mobilevit_name.startswith('''deeplabv3_'''  ):
								__a										=    MobileViTForSemanticSegmentation(a__  ).eval()
				else:
								__a										=    MobileViTForImageClassification(a__  ).eval()

				__a										=    convert_state_dict(a__						,    a__  )
				model.load_state_dict(a__  )

				# Check outputs on an image, prepared by MobileViTImageProcessor
				__a										=    MobileViTImageProcessor(crop_size=config.image_size						,    size=config.image_size + 32  )
				__a										=    image_processor(images=prepare_img()						,    return_tensors='''pt'''  )
				__a										=    model(**a__  )
				__a										=    outputs.logits

				if mobilevit_name.startswith('''deeplabv3_'''  ):
								assert logits.shape == (1, 21, 32, 32)

								if mobilevit_name == "deeplabv3_mobilevit_s":
												__a										=    torch.tensor(
												    [
												        [[6.2_065, 6.1_292, 6.2_070], [6.1_079, 6.1_254, 6.1_747], [6.0_042, 6.1_071, 6.1_034]],
												        [[-6.9_253, -6.8_653, -7.0_398], [-7.3_218, -7.3_983, -7.3_670], [-7.1_961, -7.2_482, -7.1_569]],
												        [[-4.4_723, -4.4_348, -4.3_769], [-5.3_629, -5.4_632, -5.4_598], [-5.1_587, -5.3_402, -5.5_059]],
												    ]  )
								elif mobilevit_name == "deeplabv3_mobilevit_xs":
												__a										=    torch.tensor(
												    [
												        [[5.4_449, 5.5_733, 5.6_314], [5.1_815, 5.3_930, 5.5_963], [5.1_656, 5.4_333, 5.4_853]],
												        [[-9.4_423, -9.7_766, -9.6_714], [-9.1_581, -9.5_720, -9.5_519], [-9.1_006, -9.6_458, -9.5_703]],
												        [[-7.7_721, -7.3_716, -7.1_583], [-8.4_599, -8.0_624, -7.7_944], [-8.4_172, -7.8_366, -7.5_025]],
												    ]  )
								elif mobilevit_name == "deeplabv3_mobilevit_xxs":
												__a										=    torch.tensor(
												    [
												        [[6.9_811, 6.9_743, 7.3_123], [7.1_777, 7.1_931, 7.3_938], [7.5_633, 7.8_050, 7.8_901]],
												        [[-10.5_536, -10.2_332, -10.2_924], [-10.2_336, -9.8_624, -9.5_964], [-10.8_840, -10.8_158, -10.6_659]],
												        [[-3.4_938, -3.0_631, -2.8_620], [-3.4_205, -2.8_135, -2.6_875], [-3.4_179, -2.7_945, -2.8_750]],
												    ]  )
								else:
												raise ValueError(F"""Unknown mobilevit_name: {mobilevit_name}"""  )

								assert torch.allclose(logits[0, :3, :3, :3]						,    a__						,    atol=1e-4  )
				else:
								assert logits.shape == (1, 1000)

								if mobilevit_name == "mobilevit_s":
												__a										=    torch.tensor([-0.9_866, 0.2_392, -1.1_241]  )
								elif mobilevit_name == "mobilevit_xs":
												__a										=    torch.tensor([-2.4_761, -0.9_399, -1.9_587]  )
								elif mobilevit_name == "mobilevit_xxs":
												__a										=    torch.tensor([-1.9_364, -1.2_327, -0.4_653]  )
								else:
												raise ValueError(F"""Unknown mobilevit_name: {mobilevit_name}"""  )

								assert torch.allclose(logits[0, :3]						,    a__						,    atol=1e-4  )

				Path(a__  ).mkdir(exist_ok=a__  )
				print(F"""Saving model {mobilevit_name} to {pytorch_dump_folder_path}"""  )
				model.save_pretrained(a__  )
				print(F"""Saving image processor to {pytorch_dump_folder_path}"""  )
				image_processor.save_pretrained(a__  )

				if push_to_hub:
								__a										=    {
								    '''mobilevit_s''': '''mobilevit-small''',
								    '''mobilevit_xs''': '''mobilevit-x-small''',
								    '''mobilevit_xxs''': '''mobilevit-xx-small''',
								    '''deeplabv3_mobilevit_s''': '''deeplabv3-mobilevit-small''',
								    '''deeplabv3_mobilevit_xs''': '''deeplabv3-mobilevit-x-small''',
								    '''deeplabv3_mobilevit_xxs''': '''deeplabv3-mobilevit-xx-small''',
								}

								print('''Pushing to the hub...'''  )
								__a										=    model_mapping[mobilevit_name]
								image_processor.push_to_hub(a__						,    organization='''apple'''  )
								model.push_to_hub(a__						,    organization='''apple'''  )


if __name__ == "__main__":
					A     :	Dict    = argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    '--mobilevit_name',
					    default='mobilevit_s',
					    type=str,
					    help=(
					        'Name of the MobileViT model you\'d like to convert. Should be one of \'mobilevit_s\', \'mobilevit_xs\','
					        ' \'mobilevit_xxs\', \'deeplabv3_mobilevit_s\', \'deeplabv3_mobilevit_xs\', \'deeplabv3_mobilevit_xxs\'.'
					    ),
					)
					parser.add_argument(
					    '--checkpoint_path', required=True, type=str, help='Path to the original state dict (.pt file).'
					)
					parser.add_argument(
					    '--pytorch_dump_folder_path', required=True, type=str, help='Path to the output PyTorch model directory.'
					)
					parser.add_argument(
					    '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
					)

					A     :	Optional[int]    = parser.parse_args()
					convert_movilevit_checkpoint(
					    args.mobilevit_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub
					)

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 			__A( a  ,   a  ,   unittest.TestCase			):
				snake_case_						   =    AutoencoderKL
				snake_case_						   =    '''sample'''
				snake_case_						   =    1E-2




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


								'''simple docstring'''




								__a										=    4
								__a										=    3
								__a										=    (32, 32)

								__a										=    floats_tensor((batch_size, num_channels) + sizes				).to(_snake_case				)

								return {"sample": image}




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


								'''simple docstring'''




								return (3, 32, 32)




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


								'''simple docstring'''




								return (3, 32, 32)




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Dict:


								'''simple docstring'''




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




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Tuple:


								'''simple docstring'''




								pass




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


								'''simple docstring'''




								pass




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


								'''simple docstring'''




								__a							,     __a										=    self.prepare_init_args_and_inputs_for_common()
								__a										=    self.model_class(**_snake_case				)
								model.to(_snake_case				)

								assert not model.is_gradient_checkpointing and model.training

								__a										=    model(**_snake_case				).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()

								__a										=    torch.randn_like(_snake_case				)
								__a										=    (out - labels).mean()
								loss.backward()

								# re-instantiate the model now enabling gradient checkpointing
								__a										=    self.model_class(**_snake_case				)
								# clone model
								model_a.load_state_dict(model.state_dict()				)
								model_a.to(_snake_case				)
								model_a.enable_gradient_checkpointing()

								assert model_a.is_gradient_checkpointing and model_a.training

								__a										=    model_a(**_snake_case				).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()
								__a										=    (out_a - labels).mean()
								loss_a.backward()

								# compare the output and parameters gradients
								self.assertTrue((loss - loss_a).abs() < 1E-5				)
								__a										=    dict(model.named_parameters()				)
								__a										=    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 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Dict:


								'''simple docstring'''




								__a							,     __a										=    AutoencoderKL.from_pretrained('''fusing/autoencoder-kl-dummy'''		,     output_loading_info=_snake_case				)
								self.assertIsNotNone(_snake_case				)
								self.assertEqual(len(loading_info['''missing_keys''']				)		,     0				)

								model.to(_snake_case				)
								__a										=    model(**self.dummy_input				)

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




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


								'''simple docstring'''




								__a										=    AutoencoderKL.from_pretrained('''fusing/autoencoder-kl-dummy'''				)
								__a										=    model.to(_snake_case				)
								model.eval()

								if torch_device == "mps":
												__a										=    torch.manual_seed(0				)
								else:
												__a										=    torch.Generator(device=_snake_case				).manual_seed(0				)

								__a										=    torch.randn(
								    1		,     model.config.in_channels		,     model.config.sample_size		,     model.config.sample_size		,     generator=torch.manual_seed(0				)		,     )
								__a										=    image.to(_snake_case				)
								with torch.no_grad():
												__a										=    model(_snake_case		,     sample_posterior=_snake_case		,     generator=_snake_case				).sample

								__a										=    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":
												__a										=    torch.tensor(
												    [
												        -4.0_078E-01,
												        -3.8_323E-04,
												        -1.2_681E-01,
												        -1.1_462E-01,
												        2.0_095E-01,
												        1.0_893E-01,
												        -8.8_247E-02,
												        -3.0_361E-01,
												        -9.8_644E-03,
												    ]				)
								elif torch_device == "cpu":
												__a										=    torch.tensor(
												    [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026]				)
								else:
												__a										=    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(_snake_case		,     _snake_case		,     rtol=1E-2				)				)







@slow
class 			__A( unittest.TestCase			):




				def 				SCREAMING_SNAKE_CASE_		(     self		,     _snake_case		,     _snake_case				)    ->      Optional[Any]:


								'''simple docstring'''




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




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Dict:


								'''simple docstring'''




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




				def 				SCREAMING_SNAKE_CASE_		(     self		,     _snake_case=0		,     _snake_case=(4, 3, 512, 512)		,     _snake_case=False				)    ->      Any:


								'''simple docstring'''




								__a										=    torch.floataa if fpaa else torch.floataa
								__a										=    torch.from_numpy(load_hf_numpy(self.get_file_format(_snake_case		,     _snake_case				)				)				).to(_snake_case				).to(_snake_case				)
								return image




				def 				SCREAMING_SNAKE_CASE_		(     self		,     _snake_case="CompVis/stable-diffusion-v1-4"		,     _snake_case=False				)    ->      Optional[Any]:


								'''simple docstring'''




								__a										=    '''fp16''' if fpaa else None
								__a										=    torch.floataa if fpaa else torch.floataa

								__a										=    AutoencoderKL.from_pretrained(
								    _snake_case		,     subfolder='''vae'''		,     torch_dtype=_snake_case		,     revision=_snake_case		,     )
								model.to(_snake_case				).eval()

								return model




				def 				SCREAMING_SNAKE_CASE_		(     self		,     _snake_case=0				)    ->      Tuple:


								'''simple docstring'''




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




				@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 				SCREAMING_SNAKE_CASE_		(     self		,     _snake_case		,     _snake_case		,     _snake_case				)    ->      List[Any]:


								'''simple docstring'''




								__a										=    self.get_sd_vae_model()
								__a										=    self.get_sd_image(_snake_case				)
								__a										=    self.get_generator(_snake_case				)

								with torch.no_grad():
												__a										=    model(_snake_case		,     generator=_snake_case		,     sample_posterior=_snake_case				).sample

								assert sample.shape == image.shape

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

								assert torch_all_close(_snake_case		,     _snake_case		,     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 				SCREAMING_SNAKE_CASE_		(     self		,     _snake_case		,     _snake_case				)    ->      Tuple:


								'''simple docstring'''




								__a										=    self.get_sd_vae_model(fpaa=_snake_case				)
								__a										=    self.get_sd_image(_snake_case		,     fpaa=_snake_case				)
								__a										=    self.get_generator(_snake_case				)

								with torch.no_grad():
												__a										=    model(_snake_case		,     generator=_snake_case		,     sample_posterior=_snake_case				).sample

								assert sample.shape == image.shape

								__a										=    sample[-1, -2:, :2, -2:].flatten().float().cpu()
								__a										=    torch.tensor(_snake_case				)

								assert torch_all_close(_snake_case		,     _snake_case		,     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 				SCREAMING_SNAKE_CASE_		(     self		,     _snake_case		,     _snake_case		,     _snake_case				)    ->      Optional[int]:


								'''simple docstring'''




								__a										=    self.get_sd_vae_model()
								__a										=    self.get_sd_image(_snake_case				)

								with torch.no_grad():
												__a										=    model(_snake_case				).sample

								assert sample.shape == image.shape

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

								assert torch_all_close(_snake_case		,     _snake_case		,     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 				SCREAMING_SNAKE_CASE_		(     self		,     _snake_case		,     _snake_case				)    ->      Optional[int]:


								'''simple docstring'''




								__a										=    self.get_sd_vae_model()
								__a										=    self.get_sd_image(_snake_case		,     shape=(3, 4, 64, 64)				)

								with torch.no_grad():
												__a										=    model.decode(_snake_case				).sample

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

								__a										=    sample[-1, -2:, :2, -2:].flatten().cpu()
								__a										=    torch.tensor(_snake_case				)

								assert torch_all_close(_snake_case		,     _snake_case		,     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 				SCREAMING_SNAKE_CASE_		(     self		,     _snake_case		,     _snake_case				)    ->      Optional[Any]:


								'''simple docstring'''




								__a										=    self.get_sd_vae_model(fpaa=_snake_case				)
								__a										=    self.get_sd_image(_snake_case		,     shape=(3, 4, 64, 64)		,     fpaa=_snake_case				)

								with torch.no_grad():
												__a										=    model.decode(_snake_case				).sample

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

								__a										=    sample[-1, -2:, :2, -2:].flatten().float().cpu()
								__a										=    torch.tensor(_snake_case				)

								assert torch_all_close(_snake_case		,     _snake_case		,     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 				SCREAMING_SNAKE_CASE_		(     self		,     _snake_case				)    ->      Union[str, Any]:


								'''simple docstring'''




								__a										=    self.get_sd_vae_model(fpaa=_snake_case				)
								__a										=    self.get_sd_image(_snake_case		,     shape=(3, 4, 64, 64)		,     fpaa=_snake_case				)

								with torch.no_grad():
												__a										=    model.decode(_snake_case				).sample

								model.enable_xformers_memory_efficient_attention()
								with torch.no_grad():
												__a										=    model.decode(_snake_case				).sample

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

								assert torch_all_close(_snake_case		,     _snake_case		,     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 				SCREAMING_SNAKE_CASE_		(     self		,     _snake_case				)    ->      List[str]:


								'''simple docstring'''




								__a										=    self.get_sd_vae_model()
								__a										=    self.get_sd_image(_snake_case		,     shape=(3, 4, 64, 64)				)

								with torch.no_grad():
												__a										=    model.decode(_snake_case				).sample

								model.enable_xformers_memory_efficient_attention()
								with torch.no_grad():
												__a										=    model.decode(_snake_case				).sample

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

								assert torch_all_close(_snake_case		,     _snake_case		,     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 				SCREAMING_SNAKE_CASE_		(     self		,     _snake_case		,     _snake_case				)    ->      Optional[int]:


								'''simple docstring'''




								__a										=    self.get_sd_vae_model()
								__a										=    self.get_sd_image(_snake_case				)
								__a										=    self.get_generator(_snake_case				)

								with torch.no_grad():
												__a										=    model.encode(_snake_case				).latent_dist
												__a										=    dist.sample(generator=_snake_case				)

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

								__a										=    sample[0, -1, -3:, -3:].flatten().cpu()
								__a										=    torch.tensor(_snake_case				)

								__a										=    3E-3 if torch_device != '''mps''' else 1E-2
								assert torch_all_close(_snake_case		,     _snake_case		,     atol=_snake_case				)