



class  lowerCamelCase__    :




		'''simple docstring'''



		def __init__(self	,__lowerCamelCase				)  ->    List[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  set_counts
						lowerCAmelCase__  :       Dict				=  max(lowercase__				)
						lowerCAmelCase__  :       int				=  len(lowercase__				)
						lowerCAmelCase__  :       Optional[int]				=  [1] * num_sets
						lowerCAmelCase__  :       Any				=  list(range(lowercase__				)				)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__  :       int				=  self.get_parent(lowercase__				)
						lowerCAmelCase__  :       Optional[int]				=  self.get_parent(lowercase__				)

						if src_parent == dst_parent:
										return False

						if self.ranks[dst_parent] >= self.ranks[src_parent]:
										self.set_counts[dst_parent] += self.set_counts[src_parent]
										lowerCAmelCase__  :       int				=  0
										lowerCAmelCase__  :       int				=  dst_parent
										if self.ranks[dst_parent] == self.ranks[src_parent]:
														self.ranks[dst_parent] += 1
										lowerCAmelCase__  :       List[Any]				=  self.set_counts[dst_parent]
						else:
										self.set_counts[src_parent] += self.set_counts[dst_parent]
										lowerCAmelCase__  :       Optional[int]				=  0
										lowerCAmelCase__  :       Dict				=  src_parent
										lowerCAmelCase__  :       int				=  self.set_counts[src_parent]

						lowerCAmelCase__  :       List[str]				=  max(self.max_set	,lowercase__				)
						return True



		def 	lowerCAmelCase__   (self	,__lowerCamelCase				)  ->    Any:

						"""simple docstring"""






						if self.parents[disj_set] == disj_set:
										return disj_set
						lowerCAmelCase__  :       Any				=  self.get_parent(self.parents[disj_set]				)
						return self.parents[disj_set]






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




		'''simple docstring'''



		snake_case_			=AutoencoderKL
		snake_case_			="""sample"""
		snake_case_			=1e-2



		@property
		def 	lowerCAmelCase__   (self				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__  :       str				=  4
						lowerCAmelCase__  :       int				=  3
						lowerCAmelCase__  :       List[Any]				=  (32, 32)

						lowerCAmelCase__  :       Tuple				=  floats_tensor((batch_size, num_channels) + sizes				).to(__lowerCamelCase				)

						return {"sample": image}



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

						"""simple docstring"""






						return (3, 32, 32)



		@property
		def 	lowerCAmelCase__   (self				)  ->    str:

						"""simple docstring"""






						return (3, 32, 32)



		def 	lowerCAmelCase__   (self				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__  :       int				=  {
						    '''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__  :       str				=  self.dummy_input
						return init_dict, inputs_dict



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

						"""simple docstring"""






						pass



		def 	lowerCAmelCase__   (self				)  ->    Any:

						"""simple docstring"""






						pass



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

						"""simple docstring"""






						lowerCAmelCase__		,	lowerCAmelCase__  :       int				=  self.prepare_init_args_and_inputs_for_common()
						lowerCAmelCase__  :       Tuple				=  self.model_class(**__lowerCamelCase				)
						model.to(__lowerCamelCase				)

						assert not model.is_gradient_checkpointing and model.training

						lowerCAmelCase__  :       Tuple				=  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()

						lowerCAmelCase__  :       Optional[Any]				=  torch.randn_like(__lowerCamelCase				)
						lowerCAmelCase__  :       Tuple				=  (out - labels).mean()
						loss.backward()

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

						lowerCAmelCase__  :       Dict				=  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()
						lowerCAmelCase__  :       Any				=  (out_a - labels).mean()
						loss_a.backward()

						# compare the output and parameters gradients
						self.assertTrue((loss - loss_a).abs() < 1e-5				)
						lowerCAmelCase__  :       List[str]				=  dict(model.named_parameters()				)
						lowerCAmelCase__  :       Optional[Any]				=  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				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__		,	lowerCAmelCase__  :       str				=  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				)
						lowerCAmelCase__  :       Optional[int]				=  model(**self.dummy_input				)

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



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

						"""simple docstring"""






						lowerCAmelCase__  :       str				=  AutoencoderKL.from_pretrained('''fusing/autoencoder-kl-dummy'''				)
						lowerCAmelCase__  :       Tuple				=  model.to(__lowerCamelCase				)
						model.eval()

						if torch_device == "mps":
										lowerCAmelCase__  :       Union[str, Any]				=  torch.manual_seed(0				)
						else:
										lowerCAmelCase__  :       Any				=  torch.Generator(device=__lowerCamelCase				).manual_seed(0				)

						lowerCAmelCase__  :       Optional[int]				=  torch.randn(
						    1	,model.config.in_channels	,model.config.sample_size	,model.config.sample_size	,generator=torch.manual_seed(0				)	,)
						lowerCAmelCase__  :       Dict				=  image.to(__lowerCamelCase				)
						with torch.no_grad():
										lowerCAmelCase__  :       List[Any]				=  model(__lowerCamelCase	,sample_posterior=__lowerCamelCase	,generator=__lowerCamelCase				).sample

						lowerCAmelCase__  :       Optional[Any]				=  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[str]				=  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__  :       str				=  torch.tensor(
										    [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026]				)
						else:
										lowerCAmelCase__  :       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  lowerCamelCase__    (							unittest.TestCase):




		'''simple docstring'''



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase				)  ->    List[Any]:

						"""simple docstring"""






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



		def 	lowerCAmelCase__   (self				)  ->    Dict:

						"""simple docstring"""






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



		def 	lowerCAmelCase__   (self	,__lowerCamelCase=0	,__lowerCamelCase=(4, 3, 5_12, 5_12)	,__lowerCamelCase=False				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  torch.floataa if fpaa else torch.floataa
						lowerCAmelCase__  :       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				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__  :       Dict				=  '''fp16''' if fpaa else None
						lowerCAmelCase__  :       Optional[Any]				=  torch.floataa if fpaa else torch.floataa

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

						return model



		def 	lowerCAmelCase__   (self	,__lowerCamelCase=0				)  ->    str:

						"""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				)  ->    Union[str, Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  self.get_sd_vae_model()
						lowerCAmelCase__  :       Any				=  self.get_sd_image(__lowerCamelCase				)
						lowerCAmelCase__  :       List[str]				=  self.get_generator(__lowerCamelCase				)

						with torch.no_grad():
										lowerCAmelCase__  :       Dict				=  model(__lowerCamelCase	,generator=__lowerCamelCase	,sample_posterior=__lowerCamelCase				).sample

						assert sample.shape == image.shape

						lowerCAmelCase__  :       Optional[Any]				=  sample[-1, -2:, -2:, :2].flatten().float().cpu()
						lowerCAmelCase__  :       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
		        [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				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__  :       Union[str, Any]				=  self.get_sd_vae_model(fpaa=__lowerCamelCase				)
						lowerCAmelCase__  :       List[str]				=  self.get_sd_image(__lowerCamelCase	,fpaa=__lowerCamelCase				)
						lowerCAmelCase__  :       List[Any]				=  self.get_generator(__lowerCamelCase				)

						with torch.no_grad():
										lowerCAmelCase__  :       Optional[Any]				=  model(__lowerCamelCase	,generator=__lowerCamelCase	,sample_posterior=__lowerCamelCase				).sample

						assert sample.shape == image.shape

						lowerCAmelCase__  :       Optional[Any]				=  sample[-1, -2:, :2, -2:].flatten().float().cpu()
						lowerCAmelCase__  :       Any				=  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				)  ->    Optional[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       int				=  self.get_sd_vae_model()
						lowerCAmelCase__  :       Dict				=  self.get_sd_image(__lowerCamelCase				)

						with torch.no_grad():
										lowerCAmelCase__  :       Tuple				=  model(__lowerCamelCase				).sample

						assert sample.shape == image.shape

						lowerCAmelCase__  :       Dict				=  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(__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				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__  :       Tuple				=  self.get_sd_vae_model()
						lowerCAmelCase__  :       Union[str, Any]				=  self.get_sd_image(__lowerCamelCase	,shape=(3, 4, 64, 64)				)

						with torch.no_grad():
										lowerCAmelCase__  :       str				=  model.decode(__lowerCamelCase				).sample

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

						lowerCAmelCase__  :       Dict				=  sample[-1, -2:, :2, -2:].flatten().cpu()
						lowerCAmelCase__  :       Any				=  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				)  ->    Optional[int]:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  self.get_sd_vae_model(fpaa=__lowerCamelCase				)
						lowerCAmelCase__  :       Dict				=  self.get_sd_image(__lowerCamelCase	,shape=(3, 4, 64, 64)	,fpaa=__lowerCamelCase				)

						with torch.no_grad():
										lowerCAmelCase__  :       Optional[Any]				=  model.decode(__lowerCamelCase				).sample

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

						lowerCAmelCase__  :       Dict				=  sample[-1, -2:, :2, -2:].flatten().float().cpu()
						lowerCAmelCase__  :       Union[str, 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				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  self.get_sd_vae_model(fpaa=__lowerCamelCase				)
						lowerCAmelCase__  :       str				=  self.get_sd_image(__lowerCamelCase	,shape=(3, 4, 64, 64)	,fpaa=__lowerCamelCase				)

						with torch.no_grad():
										lowerCAmelCase__  :       Any				=  model.decode(__lowerCamelCase				).sample

						model.enable_xformers_memory_efficient_attention()
						with torch.no_grad():
										lowerCAmelCase__  :       Optional[Any]				=  model.decode(__lowerCamelCase				).sample

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

						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				)  ->    Optional[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  self.get_sd_vae_model()
						lowerCAmelCase__  :       Dict				=  self.get_sd_image(__lowerCamelCase	,shape=(3, 4, 64, 64)				)

						with torch.no_grad():
										lowerCAmelCase__  :       Tuple				=  model.decode(__lowerCamelCase				).sample

						model.enable_xformers_memory_efficient_attention()
						with torch.no_grad():
										lowerCAmelCase__  :       Optional[Any]				=  model.decode(__lowerCamelCase				).sample

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

						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"""






						lowerCAmelCase__  :       Tuple				=  self.get_sd_vae_model()
						lowerCAmelCase__  :       Tuple				=  self.get_sd_image(__lowerCamelCase				)
						lowerCAmelCase__  :       Optional[Any]				=  self.get_generator(__lowerCamelCase				)

						with torch.no_grad():
										lowerCAmelCase__  :       List[Any]				=  model.encode(__lowerCamelCase				).latent_dist
										lowerCAmelCase__  :       int				=  dist.sample(generator=__lowerCamelCase				)

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

						lowerCAmelCase__  :       Optional[int]				=  sample[0, -1, -3:, -3:].flatten().cpu()
						lowerCAmelCase__  :       int				=  torch.tensor(__lowerCamelCase				)

						lowerCAmelCase__  :       Any				=  3e-3 if torch_device != '''mps''' else 1e-2
						assert torch_all_close(__lowerCamelCase	,__lowerCamelCase	,atol=__lowerCamelCase				)

