






"""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 			UpperCamelCase_       (__A				,						__A				,						unittest.TestCase     ):
	__magic_name__	      =				AutoencoderKL
	__magic_name__	      =				'''sample'''
	__magic_name__	      =				1e-2





	@property
	def      _SCREAMING_SNAKE_CASE      (      self					: Any						)						->			Union[str, Any]:
				UpperCAmelCase_     : List[str]			     =					4
				UpperCAmelCase_     : str			     =					3
				UpperCAmelCase_     : List[Any]			     =					(32, 32)

				UpperCAmelCase_     : Dict			     =					floats_tensor((batch_size, num_channels) + sizes						).to(lowerCAmelCase_						)

				return {"sample": image}





	@property
	def      _SCREAMING_SNAKE_CASE      (      self					: Tuple						)						->			int:
				return (3, 32, 32)





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





	def      _SCREAMING_SNAKE_CASE      (      self					: Any						)						->			int:
				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_     : Union[str, Any]			     =					self.dummy_input
				return init_dict, inputs_dict





	def      _SCREAMING_SNAKE_CASE      (      self					: List[str]						)						->			Tuple:
				pass





	def      _SCREAMING_SNAKE_CASE      (      self					: List[Any]						)						->			int:
				pass





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

				assert not model.is_gradient_checkpointing and model.training

				UpperCAmelCase_     : 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()

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

				# re-instantiate the model now enabling gradient checkpointing
				UpperCAmelCase_     : Union[str, Any]			     =					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_     : Tuple			     =					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_     : Optional[int]			     =					(out_a - labels).mean()
				loss_a.backward()

				# compare the output and parameters gradients
				self.assertTrue((loss - loss_a).abs() < 1e-5						)
				UpperCAmelCase_     : Tuple			     =					dict(model.named_parameters()						)
				UpperCAmelCase_     : 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      _SCREAMING_SNAKE_CASE      (      self					: Tuple						)						->			Any:
				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_     : Union[str, Any]			     =					model(**self.dummy_input						)

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





	def      _SCREAMING_SNAKE_CASE      (      self					: Optional[Any]						)						->			str:
				UpperCAmelCase_     : Union[str, Any]			     =					AutoencoderKL.from_pretrained("fusing/autoencoder-kl-dummy"						)
				UpperCAmelCase_     : Dict			     =					model.to(lowerCAmelCase_						)
				model.eval()

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

				UpperCAmelCase_     : Dict			     =					torch.randn(
				    1	,	model.config.in_channels	,	model.config.sample_size	,	model.config.sample_size	,	generator=torch.manual_seed(0						)	,	)
				UpperCAmelCase_     : Union[str, Any]			     =					image.to(lowerCAmelCase_						)
				with torch.no_grad():
							UpperCAmelCase_     : List[str]			     =					model(lowerCAmelCase_	,	sample_posterior=lowerCAmelCase_	,	generator=lowerCAmelCase_						).sample

				UpperCAmelCase_     : 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":
							UpperCAmelCase_     : Any			     =					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":
							UpperCAmelCase_     : List[str]			     =					torch.tensor(
							    [-0.1_3_5_2, 0.0_8_7_8, 0.0_4_1_9, -0.0_8_1_8, -0.1_0_6_9, 0.0_6_8_8, -0.1_4_5_8, -0.4_4_4_6, -0.0_0_2_6]						)
				else:
							UpperCAmelCase_     : Dict			     =					torch.tensor(
							    [-0.2_4_2_1, 0.4_6_4_2, 0.2_5_0_7, -0.0_4_3_8, 0.0_6_8_2, 0.3_1_6_0, -0.2_0_1_8, -0.0_7_2_7, 0.2_4_8_5]						)

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







@slow
class 			UpperCamelCase_       (unittest.TestCase     ):





	def      _SCREAMING_SNAKE_CASE      (      self					: Optional[int]	,	lowerCAmelCase_					: Tuple	,	lowerCAmelCase_					: Union[str, Any]						)						->			Dict:
				return f"""gaussian_noise_s={seed}_shape={'_'.join([str(lowerCAmelCase_						) for s in shape]						)}.npy"""





	def      _SCREAMING_SNAKE_CASE      (      self					: Any						)						->			List[str]:
				# clean up the VRAM after each test
				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()





	def      _SCREAMING_SNAKE_CASE      (      self					: Dict	,	lowerCAmelCase_					: Tuple=0	,	lowerCAmelCase_					: List[Any]=(4, 3, 512, 512)	,	lowerCAmelCase_					: Union[str, Any]=False						)						->			Optional[int]:
				UpperCAmelCase_     : List[str]			     =					torch.floataa if fpaa else torch.floataa
				UpperCAmelCase_     : Optional[Any]			     =					torch.from_numpy(load_hf_numpy(self.get_file_format(lowerCAmelCase_	,	lowerCAmelCase_						)						)						).to(lowerCAmelCase_						).to(lowerCAmelCase_						)
				return image





	def      _SCREAMING_SNAKE_CASE      (      self					: Any	,	lowerCAmelCase_					: int="CompVis/stable-diffusion-v1-4"	,	lowerCAmelCase_					: Tuple=False						)						->			Optional[Any]:
				UpperCAmelCase_     : str			     =					"fp16" if fpaa else None
				UpperCAmelCase_     : List[Any]			     =					torch.floataa if fpaa else torch.floataa

				UpperCAmelCase_     : Union[str, Any]			     =					AutoencoderKL.from_pretrained(
				    lowerCAmelCase_	,	subfolder="vae"	,	torch_dtype=lowerCAmelCase_	,	revision=lowerCAmelCase_	,	)
				model.to(lowerCAmelCase_						).eval()

				return model





	def      _SCREAMING_SNAKE_CASE      (      self					: Any	,	lowerCAmelCase_					: Optional[int]=0						)						->			str:
				if torch_device == "mps":
							return torch.manual_seed(lowerCAmelCase_						)
				return torch.Generator(device=lowerCAmelCase_						).manual_seed(lowerCAmelCase_						)





	@parameterized.expand(
	    [
	        # fmt: off
	        [33, [-0.1_6_0_3, 0.9_8_7_8, -0.0_4_9_5, -0.0_7_9_0, -0.2_7_0_9, 0.8_3_7_5, -0.2_0_6_0, -0.0_8_2_4], [-0.2_3_9_5, 0.0_0_9_8, 0.0_1_0_2, -0.0_7_0_9, -0.2_8_4_0, -0.0_2_7_4, -0.0_7_1_8, -0.1_8_2_4]],
	        [47, [-0.2_3_7_6, 0.1_1_6_8, 0.1_3_3_2, -0.4_8_4_0, -0.2_5_0_8, -0.0_7_9_1, -0.0_4_9_3, -0.4_0_8_9], [0.0_3_5_0, 0.0_8_4_7, 0.0_4_6_7, 0.0_3_4_4, -0.0_8_4_2, -0.0_5_4_7, -0.0_6_3_3, -0.1_1_3_1]],
	        # fmt: on
	    ]						)
	def      _SCREAMING_SNAKE_CASE      (      self					: List[Any]	,	lowerCAmelCase_					: Tuple	,	lowerCAmelCase_					: Union[str, Any]	,	lowerCAmelCase_					: int						)						->			int:
				UpperCAmelCase_     : str			     =					self.get_sd_vae_model()
				UpperCAmelCase_     : str			     =					self.get_sd_image(lowerCAmelCase_						)
				UpperCAmelCase_     : Any			     =					self.get_generator(lowerCAmelCase_						)

				with torch.no_grad():
							UpperCAmelCase_     : Any			     =					model(lowerCAmelCase_	,	generator=lowerCAmelCase_	,	sample_posterior=lowerCAmelCase_						).sample

				assert sample.shape == image.shape

				UpperCAmelCase_     : Optional[int]			     =					sample[-1, -2:, -2:, :2].flatten().float().cpu()
				UpperCAmelCase_     : int			     =					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.0_5_1_3, 0.0_2_8_9, 1.3_7_9_9, 0.2_1_6_6, -0.2_5_7_3, -0.0_8_7_1, 0.5_1_0_3, -0.0_9_9_9]],
	        [47, [-0.4_1_2_8, -0.1_3_2_0, -0.3_7_0_4, 0.1_9_6_5, -0.4_1_1_6, -0.2_3_3_2, -0.3_3_4_0, 0.2_2_4_7]],
	        # fmt: on
	    ]						)
	@require_torch_gpu
	def      _SCREAMING_SNAKE_CASE      (      self					: Union[str, Any]	,	lowerCAmelCase_					: List[str]	,	lowerCAmelCase_					: List[Any]						)						->			Tuple:
				UpperCAmelCase_     : int			     =					self.get_sd_vae_model(fpaa=lowerCAmelCase_						)
				UpperCAmelCase_     : List[Any]			     =					self.get_sd_image(lowerCAmelCase_	,	fpaa=lowerCAmelCase_						)
				UpperCAmelCase_     : int			     =					self.get_generator(lowerCAmelCase_						)

				with torch.no_grad():
							UpperCAmelCase_     : Union[str, Any]			     =					model(lowerCAmelCase_	,	generator=lowerCAmelCase_	,	sample_posterior=lowerCAmelCase_						).sample

				assert sample.shape == image.shape

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

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





	@parameterized.expand(
	    [
	        # fmt: off
	        [33, [-0.1_6_0_9, 0.9_8_6_6, -0.0_4_8_7, -0.0_7_7_7, -0.2_7_1_6, 0.8_3_6_8, -0.2_0_5_5, -0.0_8_1_4], [-0.2_3_9_5, 0.0_0_9_8, 0.0_1_0_2, -0.0_7_0_9, -0.2_8_4_0, -0.0_2_7_4, -0.0_7_1_8, -0.1_8_2_4]],
	        [47, [-0.2_3_7_7, 0.1_1_4_7, 0.1_3_3_3, -0.4_8_4_1, -0.2_5_0_6, -0.0_8_0_5, -0.0_4_9_1, -0.4_0_8_5], [0.0_3_5_0, 0.0_8_4_7, 0.0_4_6_7, 0.0_3_4_4, -0.0_8_4_2, -0.0_5_4_7, -0.0_6_3_3, -0.1_1_3_1]],
	        # fmt: on
	    ]						)
	def      _SCREAMING_SNAKE_CASE      (      self					: Dict	,	lowerCAmelCase_					: List[Any]	,	lowerCAmelCase_					: Tuple	,	lowerCAmelCase_					: Optional[Any]						)						->			Union[str, Any]:
				UpperCAmelCase_     : Optional[Any]			     =					self.get_sd_vae_model()
				UpperCAmelCase_     : Tuple			     =					self.get_sd_image(lowerCAmelCase_						)

				with torch.no_grad():
							UpperCAmelCase_     : Union[str, Any]			     =					model(lowerCAmelCase_						).sample

				assert sample.shape == image.shape

				UpperCAmelCase_     : Tuple			     =					sample[-1, -2:, -2:, :2].flatten().float().cpu()
				UpperCAmelCase_     : 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.2_0_5_1, -0.1_8_0_3, -0.2_3_1_1, -0.2_1_1_4, -0.3_2_9_2, -0.3_5_7_4, -0.2_9_5_3, -0.3_3_2_3]],
	        [37, [-0.2_6_3_2, -0.2_6_2_5, -0.2_1_9_9, -0.2_7_4_1, -0.4_5_3_9, -0.4_9_9_0, -0.3_7_2_0, -0.4_9_2_5]],
	        # fmt: on
	    ]						)
	@require_torch_gpu
	def      _SCREAMING_SNAKE_CASE      (      self					: Optional[int]	,	lowerCAmelCase_					: Any	,	lowerCAmelCase_					: Optional[Any]						)						->			int:
				UpperCAmelCase_     : Any			     =					self.get_sd_vae_model()
				UpperCAmelCase_     : List[Any]			     =					self.get_sd_image(lowerCAmelCase_	,	shape=(3, 4, 64, 64)						)

				with torch.no_grad():
							UpperCAmelCase_     : List[Any]			     =					model.decode(lowerCAmelCase_						).sample

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

				UpperCAmelCase_     : Dict			     =					sample[-1, -2:, :2, -2:].flatten().cpu()
				UpperCAmelCase_     : Any			     =					torch.tensor(lowerCAmelCase_						)

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





	@parameterized.expand(
	    [
	        # fmt: off
	        [27, [-0.0_3_6_9, 0.0_2_0_7, -0.0_7_7_6, -0.0_6_8_2, -0.1_7_4_7, -0.1_9_3_0, -0.1_4_6_5, -0.2_0_3_9]],
	        [16, [-0.1_6_2_8, -0.2_1_3_4, -0.2_7_4_7, -0.2_6_4_2, -0.3_7_7_4, -0.4_4_0_4, -0.3_6_8_7, -0.4_2_7_7]],
	        # fmt: on
	    ]						)
	@require_torch_gpu
	def      _SCREAMING_SNAKE_CASE      (      self					: Tuple	,	lowerCAmelCase_					: Optional[int]	,	lowerCAmelCase_					: Optional[Any]						)						->			List[str]:
				UpperCAmelCase_     : str			     =					self.get_sd_vae_model(fpaa=lowerCAmelCase_						)
				UpperCAmelCase_     : Tuple			     =					self.get_sd_image(lowerCAmelCase_	,	shape=(3, 4, 64, 64)	,	fpaa=lowerCAmelCase_						)

				with torch.no_grad():
							UpperCAmelCase_     : int			     =					model.decode(lowerCAmelCase_						).sample

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

				UpperCAmelCase_     : str			     =					sample[-1, -2:, :2, -2:].flatten().float().cpu()
				UpperCAmelCase_     : List[str]			     =					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      _SCREAMING_SNAKE_CASE      (      self					: Dict	,	lowerCAmelCase_					: List[Any]						)						->			int:
				UpperCAmelCase_     : List[Any]			     =					self.get_sd_vae_model(fpaa=lowerCAmelCase_						)
				UpperCAmelCase_     : List[str]			     =					self.get_sd_image(lowerCAmelCase_	,	shape=(3, 4, 64, 64)	,	fpaa=lowerCAmelCase_						)

				with torch.no_grad():
							UpperCAmelCase_     : List[Any]			     =					model.decode(lowerCAmelCase_						).sample

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

				with torch.no_grad():
							UpperCAmelCase_     : int			     =					model.decode(lowerCAmelCase_						).sample

				model.enable_xformers_memory_efficient_attention()
				with torch.no_grad():
							UpperCAmelCase_     : 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.3_0_0_1, 0.0_9_1_8, -2.6_9_8_4, -3.9_7_2_0, -3.2_0_9_9, -5.0_3_5_3, 1.7_3_3_8, -0.2_0_6_5, 3.4_2_6_7]],
	        [47, [-1.5_0_3_0, -4.3_8_7_1, -6.0_3_5_5, -9.1_1_5_7, -1.6_6_6_1, -2.7_8_5_3, 2.1_6_0_7, -5.0_8_2_3, 2.5_6_3_3]],
	        # fmt: on
	    ]						)
	def      _SCREAMING_SNAKE_CASE      (      self					: str	,	lowerCAmelCase_					: int	,	lowerCAmelCase_					: Dict						)						->			Dict:
				UpperCAmelCase_     : Dict			     =					self.get_sd_vae_model()
				UpperCAmelCase_     : int			     =					self.get_sd_image(lowerCAmelCase_						)
				UpperCAmelCase_     : Tuple			     =					self.get_generator(lowerCAmelCase_						)

				with torch.no_grad():
							UpperCAmelCase_     : List[str]			     =					model.encode(lowerCAmelCase_						).latent_dist
							UpperCAmelCase_     : List[str]			     =					dist.sample(generator=lowerCAmelCase_						)

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

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

				UpperCAmelCase_     : Dict			     =					3e-3 if torch_device != "mps" else 1e-2
				assert torch_all_close(lowerCAmelCase_	,	lowerCAmelCase_	,	atol=lowerCAmelCase_						)











"""simple docstring"""
import warnings

from diffusers import StableDiffusionInpaintPipeline as StableDiffusionInpaintPipeline  # noqa F401


warnings.warn(
    '''The `inpainting.py` script is outdated. Please use directly `from diffusers import'''
    ''' StableDiffusionInpaintPipeline` instead.'''
)



