






"""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__      (							snake_case__,       snake_case__,       unittest.TestCase       ):
			_UpperCAmelCase		:Union[str, Any]             =			AutoencoderKL
			_UpperCAmelCase		:List[Any]             =			"sample"
			_UpperCAmelCase		:int             =			1e-2


			@property
			def      UpperCAmelCase__						(				self					:  Any ):
										lowerCamelCase_	:							str									=4
										lowerCamelCase_	:							Any									=3
										lowerCamelCase_	:							Union[str, Any]									=(32, 32)

										lowerCamelCase_	:							List[str]									=floats_tensor((batch_size, num_channels) + sizes ).to(snake_case__ )

										return {"sample": image}


			@property
			def      UpperCAmelCase__						(				self					:  Union[str, Any] ):
										return (3, 32, 32)


			@property
			def      UpperCAmelCase__						(				self					:  List[str] ):
										return (3, 32, 32)


			def      UpperCAmelCase__						(				self					:  List[Any] ):
										lowerCamelCase_	:							Optional[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,
										}
										lowerCamelCase_	:							Dict									=self.dummy_input
										return init_dict, inputs_dict


			def      UpperCAmelCase__						(				self					:  Union[str, Any] ):
										pass


			def      UpperCAmelCase__						(				self					:  Dict ):
										pass


			@unittest.skipIf(torch_device == "mps"       ,	"Gradient checkpointing skipped on MPS" )
			def      UpperCAmelCase__						(				self					:  int ):
										# enable deterministic behavior for gradient checkpointing
										lowerCamelCase_							,						lowerCamelCase_	:							Union[str, Any]									=self.prepare_init_args_and_inputs_for_common()
										lowerCamelCase_	:							List[Any]									=self.model_class(**snake_case__ )
										model.to(snake_case__ )

										assert not model.is_gradient_checkpointing and model.training

										lowerCamelCase_	:							Optional[int]									=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()

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

										# re-instantiate the model now enabling gradient checkpointing
										lowerCamelCase_	:							Dict									=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

										lowerCamelCase_	:							List[Any]									=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()
										lowerCamelCase_	:							str									=(out_a - labels).mean()
										loss_a.backward()

										# compare the output and parameters gradients
										self.assertTrue((loss - loss_a).abs() < 1E-5 )
										lowerCamelCase_	:							Optional[Any]									=dict(model.named_parameters() )
										lowerCamelCase_	:							Dict									=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      UpperCAmelCase__						(				self					:  Any ):
										lowerCamelCase_							,						lowerCamelCase_	:							Optional[Any]									=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__ )
										lowerCamelCase_	:							Tuple									=model(**self.dummy_input )

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


			def      UpperCAmelCase__						(				self					:  Optional[int] ):
										lowerCamelCase_	:							Dict									=AutoencoderKL.from_pretrained("fusing/autoencoder-kl-dummy" )
										lowerCamelCase_	:							Any									=model.to(snake_case__ )
										model.eval()

										if torch_device == "mps":
																	lowerCamelCase_	:							List[str]									=torch.manual_seed(0 )
										else:
																	lowerCamelCase_	:							List[Any]									=torch.Generator(device=snake_case__ ).manual_seed(0 )

										lowerCamelCase_	:							List[Any]									=torch.randn(
										    1       ,	model.config.in_channels       ,	model.config.sample_size       ,	model.config.sample_size       ,	generator=torch.manual_seed(0 )       ,	)
										lowerCamelCase_	:							Tuple									=image.to(snake_case__ )
										with torch.no_grad():
																	lowerCamelCase_	:							Optional[Any]									=model(snake_case__       ,	sample_posterior=snake_case__       ,	generator=snake_case__ ).sample

										lowerCamelCase_	:							List[str]									=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":
																	lowerCamelCase_	:							List[Any]									=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":
																	lowerCamelCase_	:							List[str]									=torch.tensor(
																	    [-0.1_352, 0.0_878, 0.0_419, -0.0_818, -0.1_069, 0.0_688, -0.1_458, -0.4_446, -0.0_026] )
										else:
																	lowerCamelCase_	:							str									=torch.tensor(
																	    [-0.2_421, 0.4_642, 0.2_507, -0.0_438, 0.0_682, 0.3_160, -0.2_018, -0.0_727, 0.2_485] )

										self.assertTrue(torch_all_close(snake_case__       ,	snake_case__       ,	rtol=1E-2 ) )







@slow
class        lowercase__      (							unittest.TestCase       ):


			def      UpperCAmelCase__						(				self					:  List[Any]       ,	snake_case__					:  Optional[Any]       ,	snake_case__					:  Dict ):
										return F"""gaussian_noise_s={seed}_shape={'_'.join([str(snake_case__ ) for s in shape] )}.npy"""


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


			def      UpperCAmelCase__						(				self					:  List[Any]       ,	snake_case__					:  int=0       ,	snake_case__					:  Union[str, Any]=(4, 3, 512, 512)       ,	snake_case__					:  Tuple=False ):
										lowerCamelCase_	:							Optional[int]									=torch.floataa if fpaa else torch.floataa
										lowerCamelCase_	:							Tuple									=torch.from_numpy(load_hf_numpy(self.get_file_format(snake_case__       ,	snake_case__ ) ) ).to(snake_case__ ).to(snake_case__ )
										return image


			def      UpperCAmelCase__						(				self					:  Tuple       ,	snake_case__					:  Optional[int]="CompVis/stable-diffusion-v1-4"       ,	snake_case__					:  Optional[int]=False ):
										lowerCamelCase_	:							Tuple									="fp16" if fpaa else None
										lowerCamelCase_	:							int									=torch.floataa if fpaa else torch.floataa

										lowerCamelCase_	:							str									=AutoencoderKL.from_pretrained(
										    snake_case__       ,	subfolder="vae"       ,	torch_dtype=snake_case__       ,	revision=snake_case__       ,	)
										model.to(snake_case__ ).eval()

										return model


			def      UpperCAmelCase__						(				self					:  List[Any]       ,	snake_case__					:  Optional[int]=0 ):
										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.1_603, 0.9_878, -0.0_495, -0.0_790, -0.2_709, 0.8_375, -0.2_060, -0.0_824], [-0.2_395, 0.0_098, 0.0_102, -0.0_709, -0.2_840, -0.0_274, -0.0_718, -0.1_824]],
			        [47, [-0.2_376, 0.1_168, 0.1_332, -0.4_840, -0.2_508, -0.0_791, -0.0_493, -0.4_089], [0.0_350, 0.0_847, 0.0_467, 0.0_344, -0.0_842, -0.0_547, -0.0_633, -0.1_131]],
			        # fmt: on
			    ] )
			def      UpperCAmelCase__						(				self					:  Optional[int]       ,	snake_case__					:  Dict       ,	snake_case__					:  Dict       ,	snake_case__					:  int ):
										lowerCamelCase_	:							List[str]									=self.get_sd_vae_model()
										lowerCamelCase_	:							str									=self.get_sd_image(snake_case__ )
										lowerCamelCase_	:							List[Any]									=self.get_generator(snake_case__ )

										with torch.no_grad():
																	lowerCamelCase_	:							Any									=model(snake_case__       ,	generator=snake_case__       ,	sample_posterior=snake_case__ ).sample

										assert sample.shape == image.shape

										lowerCamelCase_	:							List[str]									=sample[-1, -2:, -2:, :2].flatten().float().cpu()
										lowerCamelCase_	:							List[Any]									=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.0_513, 0.0_289, 1.3_799, 0.2_166, -0.2_573, -0.0_871, 0.5_103, -0.0_999]],
			        [47, [-0.4_128, -0.1_320, -0.3_704, 0.1_965, -0.4_116, -0.2_332, -0.3_340, 0.2_247]],
			        # fmt: on
			    ] )
			@require_torch_gpu
			def      UpperCAmelCase__						(				self					:  int       ,	snake_case__					:  Optional[int]       ,	snake_case__					:  List[str] ):
										lowerCamelCase_	:							Optional[int]									=self.get_sd_vae_model(fpaa=snake_case__ )
										lowerCamelCase_	:							Optional[int]									=self.get_sd_image(snake_case__       ,	fpaa=snake_case__ )
										lowerCamelCase_	:							Optional[Any]									=self.get_generator(snake_case__ )

										with torch.no_grad():
																	lowerCamelCase_	:							Any									=model(snake_case__       ,	generator=snake_case__       ,	sample_posterior=snake_case__ ).sample

										assert sample.shape == image.shape

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

										assert torch_all_close(snake_case__       ,	snake_case__       ,	atol=1E-2 )


			@parameterized.expand(
			    [
			        # fmt: off
			        [33, [-0.1_609, 0.9_866, -0.0_487, -0.0_777, -0.2_716, 0.8_368, -0.2_055, -0.0_814], [-0.2_395, 0.0_098, 0.0_102, -0.0_709, -0.2_840, -0.0_274, -0.0_718, -0.1_824]],
			        [47, [-0.2_377, 0.1_147, 0.1_333, -0.4_841, -0.2_506, -0.0_805, -0.0_491, -0.4_085], [0.0_350, 0.0_847, 0.0_467, 0.0_344, -0.0_842, -0.0_547, -0.0_633, -0.1_131]],
			        # fmt: on
			    ] )
			def      UpperCAmelCase__						(				self					:  int       ,	snake_case__					:  Any       ,	snake_case__					:  Tuple       ,	snake_case__					:  Tuple ):
										lowerCamelCase_	:							List[str]									=self.get_sd_vae_model()
										lowerCamelCase_	:							str									=self.get_sd_image(snake_case__ )

										with torch.no_grad():
																	lowerCamelCase_	:							int									=model(snake_case__ ).sample

										assert sample.shape == image.shape

										lowerCamelCase_	:							Optional[Any]									=sample[-1, -2:, -2:, :2].flatten().float().cpu()
										lowerCamelCase_	:							Union[str, Any]									=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.2_051, -0.1_803, -0.2_311, -0.2_114, -0.3_292, -0.3_574, -0.2_953, -0.3_323]],
			        [37, [-0.2_632, -0.2_625, -0.2_199, -0.2_741, -0.4_539, -0.4_990, -0.3_720, -0.4_925]],
			        # fmt: on
			    ] )
			@require_torch_gpu
			def      UpperCAmelCase__						(				self					:  Optional[Any]       ,	snake_case__					:  List[Any]       ,	snake_case__					:  Any ):
										lowerCamelCase_	:							Optional[int]									=self.get_sd_vae_model()
										lowerCamelCase_	:							Optional[Any]									=self.get_sd_image(snake_case__       ,	shape=(3, 4, 64, 64) )

										with torch.no_grad():
																	lowerCamelCase_	:							Optional[int]									=model.decode(snake_case__ ).sample

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

										lowerCamelCase_	:							Dict									=sample[-1, -2:, :2, -2:].flatten().cpu()
										lowerCamelCase_	:							Tuple									=torch.tensor(snake_case__ )

										assert torch_all_close(snake_case__       ,	snake_case__       ,	atol=1E-3 )


			@parameterized.expand(
			    [
			        # fmt: off
			        [27, [-0.0_369, 0.0_207, -0.0_776, -0.0_682, -0.1_747, -0.1_930, -0.1_465, -0.2_039]],
			        [16, [-0.1_628, -0.2_134, -0.2_747, -0.2_642, -0.3_774, -0.4_404, -0.3_687, -0.4_277]],
			        # fmt: on
			    ] )
			@require_torch_gpu
			def      UpperCAmelCase__						(				self					:  Optional[Any]       ,	snake_case__					:  List[Any]       ,	snake_case__					:  Tuple ):
										lowerCamelCase_	:							List[Any]									=self.get_sd_vae_model(fpaa=snake_case__ )
										lowerCamelCase_	:							Tuple									=self.get_sd_image(snake_case__       ,	shape=(3, 4, 64, 64)       ,	fpaa=snake_case__ )

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

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

										lowerCamelCase_	:							Tuple									=sample[-1, -2:, :2, -2:].flatten().float().cpu()
										lowerCamelCase_	:							Optional[Any]									=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      UpperCAmelCase__						(				self					:  Dict       ,	snake_case__					:  Tuple ):
										lowerCamelCase_	:							Optional[int]									=self.get_sd_vae_model(fpaa=snake_case__ )
										lowerCamelCase_	:							Optional[Any]									=self.get_sd_image(snake_case__       ,	shape=(3, 4, 64, 64)       ,	fpaa=snake_case__ )

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

										model.enable_xformers_memory_efficient_attention()
										with torch.no_grad():
																	lowerCamelCase_	:							Tuple									=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      UpperCAmelCase__						(				self					:  Dict       ,	snake_case__					:  List[Any] ):
										lowerCamelCase_	:							Union[str, Any]									=self.get_sd_vae_model()
										lowerCamelCase_	:							Optional[Any]									=self.get_sd_image(snake_case__       ,	shape=(3, 4, 64, 64) )

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

										model.enable_xformers_memory_efficient_attention()
										with torch.no_grad():
																	lowerCamelCase_	:							str									=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.3_001, 0.0_918, -2.6_984, -3.9_720, -3.2_099, -5.0_353, 1.7_338, -0.2_065, 3.4_267]],
			        [47, [-1.5_030, -4.3_871, -6.0_355, -9.1_157, -1.6_661, -2.7_853, 2.1_607, -5.0_823, 2.5_633]],
			        # fmt: on
			    ] )
			def      UpperCAmelCase__						(				self					:  List[str]       ,	snake_case__					:  Optional[Any]       ,	snake_case__					:  Dict ):
										lowerCamelCase_	:							Union[str, Any]									=self.get_sd_vae_model()
										lowerCamelCase_	:							Union[str, Any]									=self.get_sd_image(snake_case__ )
										lowerCamelCase_	:							Tuple									=self.get_generator(snake_case__ )

										with torch.no_grad():
																	lowerCamelCase_	:							Dict									=model.encode(snake_case__ ).latent_dist
																	lowerCamelCase_	:							Any									=dist.sample(generator=snake_case__ )

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

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

										lowerCamelCase_	:							int									=3E-3 if torch_device != "mps" else 1E-2
										assert torch_all_close(snake_case__       ,	snake_case__       ,	atol=snake_case__ )











"""simple docstring"""



import torch

def 			_snake_case    (					)    ->     Union[str, Any]:
							if torch.cuda.is_available():
														lowerCamelCase_	:							int									=torch.cuda.device_count()
							else:
														lowerCamelCase_	:							List[str]									=0
							print(F"""Successfully ran on {num_gpus} GPUs"""  )


if __name__ == "__main__":
						main()



