



"""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__			(      A,					A,					unittest.TestCase    ):
							'''simple docstring'''







							_UpperCAmelCase      				=					AutoencoderKL
							_UpperCAmelCase      				=					'''sample'''
							_UpperCAmelCase      				=					1E-2






							@property
							def        lowerCamelCase_   (  self						)			->							Dict:
									_UpperCAmelCase							 =		4
									_UpperCAmelCase							 =		3
									_UpperCAmelCase							 =		(32, 32)

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

									return {"sample": image}






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






							@property
							def        lowerCamelCase_   (  self						)			->							Optional[int]:
									return (3, 32, 32)






							def        lowerCamelCase_   (  self						)			->							Any:
									_UpperCAmelCase							 =		{
									    '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							 =		self.dummy_input
									return init_dict, inputs_dict






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






							def        lowerCamelCase_   (  self						)			->							str:
									pass






							@unittest.skipIf(torch_device == 'mps'	,       'Gradient checkpointing skipped on MPS'						)
							def        lowerCamelCase_   (  self						)			->							List[Any]:
									_UpperCAmelCase ,					_UpperCAmelCase							 =		self.prepare_init_args_and_inputs_for_common()
									_UpperCAmelCase							 =		self.model_class(**__a						)
									model.to(__a						)

									assert not model.is_gradient_checkpointing and model.training

									_UpperCAmelCase							 =		model(**__a						).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							 =		torch.randn_like(__a						)
									_UpperCAmelCase							 =		(out - labels).mean()
									loss.backward()

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

									assert model_a.is_gradient_checkpointing and model_a.training

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

									# compare the output and parameters gradients
									self.assertTrue((loss - loss_a).abs() < 1E-5						)
									_UpperCAmelCase							 =		dict(model.named_parameters()						)
									_UpperCAmelCase							 =		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						)			->							Union[str, Any]:
									_UpperCAmelCase ,					_UpperCAmelCase							 =		AutoencoderKL.from_pretrained('fusing/autoencoder-kl-dummy'	,       output_loading_info=__a						)
									self.assertIsNotNone(__a						)
									self.assertEqual(len(loading_info['missing_keys']						)	,       0						)

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

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






							def        lowerCamelCase_   (  self						)			->							Optional[int]:
									_UpperCAmelCase							 =		AutoencoderKL.from_pretrained('fusing/autoencoder-kl-dummy'						)
									_UpperCAmelCase							 =		model.to(__a						)
									model.eval()

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

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

									_UpperCAmelCase							 =		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							 =		torch.tensor(
											    [
											        -4.00_78E-01,
											        -3.83_23E-04,
											        -1.26_81E-01,
											        -1.14_62E-01,
											        2.00_95E-01,
											        1.08_93E-01,
											        -8.82_47E-02,
											        -3.03_61E-01,
											        -9.86_44E-03,
											    ]						)
									elif torch_device == "cpu":
											_UpperCAmelCase							 =		torch.tensor(
											    [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026]						)
									else:
											_UpperCAmelCase							 =		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(__a	,       __a	,       rtol=1E-2						)						)







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







							def        lowerCamelCase_   (  self	,       snake_case	,       snake_case						)			->							Optional[Any]:
									return f'gaussian_noise_s={seed}_shape={"_".join([str(__a						) for s in shape]						)}.npy'






							def        lowerCamelCase_   (  self						)			->							Optional[int]:
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()






							def        lowerCamelCase_   (  self	,       snake_case=0	,       snake_case=(4, 3, 512, 512)	,       snake_case=False						)			->							str:
									_UpperCAmelCase							 =		torch.floataa if fpaa else torch.floataa
									_UpperCAmelCase							 =		torch.from_numpy(load_hf_numpy(self.get_file_format(__a	,       __a						)						)						).to(__a						).to(__a						)
									return image






							def        lowerCamelCase_   (  self	,       snake_case="CompVis/stable-diffusion-v1-4"	,       snake_case=False						)			->							Union[str, Any]:
									_UpperCAmelCase							 =		'fp16' if fpaa else None
									_UpperCAmelCase							 =		torch.floataa if fpaa else torch.floataa

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

									return model






							def        lowerCamelCase_   (  self	,       snake_case=0						)			->							Tuple:
									if torch_device == "mps":
											return torch.manual_seed(__a						)
									return torch.Generator(device=__a						).manual_seed(__a						)






							@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	,       snake_case	,       snake_case	,       snake_case						)			->							Any:
									_UpperCAmelCase							 =		self.get_sd_vae_model()
									_UpperCAmelCase							 =		self.get_sd_image(__a						)
									_UpperCAmelCase							 =		self.get_generator(__a						)

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

									assert sample.shape == image.shape

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

									assert torch_all_close(__a	,       __a	,       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	,       snake_case	,       snake_case						)			->							List[Any]:
									_UpperCAmelCase							 =		self.get_sd_vae_model(fpaa=__a						)
									_UpperCAmelCase							 =		self.get_sd_image(__a	,       fpaa=__a						)
									_UpperCAmelCase							 =		self.get_generator(__a						)

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

									assert sample.shape == image.shape

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

									assert torch_all_close(__a	,       __a	,       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	,       snake_case	,       snake_case	,       snake_case						)			->							Union[str, Any]:
									_UpperCAmelCase							 =		self.get_sd_vae_model()
									_UpperCAmelCase							 =		self.get_sd_image(__a						)

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

									assert sample.shape == image.shape

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

									assert torch_all_close(__a	,       __a	,       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	,       snake_case	,       snake_case						)			->							Union[str, Any]:
									_UpperCAmelCase							 =		self.get_sd_vae_model()
									_UpperCAmelCase							 =		self.get_sd_image(__a	,       shape=(3, 4, 64, 64)						)

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

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

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

									assert torch_all_close(__a	,       __a	,       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	,       snake_case	,       snake_case						)			->							Optional[Any]:
									_UpperCAmelCase							 =		self.get_sd_vae_model(fpaa=__a						)
									_UpperCAmelCase							 =		self.get_sd_image(__a	,       shape=(3, 4, 64, 64)	,       fpaa=__a						)

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

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

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

									assert torch_all_close(__a	,       __a	,       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	,       snake_case						)			->							str:
									_UpperCAmelCase							 =		self.get_sd_vae_model(fpaa=__a						)
									_UpperCAmelCase							 =		self.get_sd_image(__a	,       shape=(3, 4, 64, 64)	,       fpaa=__a						)

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

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

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

									assert torch_all_close(__a	,       __a	,       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	,       snake_case						)			->							int:
									_UpperCAmelCase							 =		self.get_sd_vae_model()
									_UpperCAmelCase							 =		self.get_sd_image(__a	,       shape=(3, 4, 64, 64)						)

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

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

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

									assert torch_all_close(__a	,       __a	,       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	,       snake_case	,       snake_case						)			->							List[Any]:
									_UpperCAmelCase							 =		self.get_sd_vae_model()
									_UpperCAmelCase							 =		self.get_sd_image(__a						)
									_UpperCAmelCase							 =		self.get_generator(__a						)

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

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

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

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







"""simple docstring"""





def 	UpperCAmelCase     (		A	:						int       ):

		'''simple docstring'''


		_UpperCAmelCase							 =		abs(A       )
		_UpperCAmelCase							 =		0
		while n > 0:
				res += n % 10
				n //= 10
		return res




def 	UpperCAmelCase     (		A	:						int       ):

		'''simple docstring'''


		_UpperCAmelCase							 =		abs(A       )
		return n if n < 10 else n % 10 + sum_of_digits(n // 10       )




def 	UpperCAmelCase     (		A	:						int       ):

		'''simple docstring'''


		return sum(int(A       ) for c in str(abs(A       )       )       )




def 	UpperCAmelCase     (		):

		'''simple docstring'''


		from collections.abc import Callable
		from timeit import timeit

		def benchmark_a_function(A	:						Callable     ,				A	:						int       ) -> None:
				_UpperCAmelCase							 =		f'{func.__name__}({value})'
				_UpperCAmelCase							 =		timeit(f'__main__.{call}'     ,				setup='import __main__'       )
				print(f'{call:56} = {func(A       )} -- {timing:.4f} seconds'       )

		for value in (26_2144, 1125_8999_0684_2624, 126_7650_6002_2822_9401_4967_0320_5376):
				for func in (sum_of_digits, sum_of_digits_recursion, sum_of_digits_compact):
						benchmark_a_function(A     ,				A       )
				print()


if __name__ == "__main__":
					import doctest

					doctest.testmod()
					benchmark()


