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




				@property
				def 	__a		(   self	)							-> Tuple:
							lowerCAmelCase_         =      4
							lowerCAmelCase_         =      3
							lowerCAmelCase_         =      (32, 32)

							lowerCAmelCase_         =      floats_tensor((batch_size, num_channels) + sizes	).to(_UpperCamelCase	)

							return {"sample": image}




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




				@property
				def 	__a		(   self	)							-> Dict:
							return (3, 32, 32)




				def 	__a		(   self	)							-> str:
							lowerCAmelCase_         =      {
							    "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_         =      self.dummy_input
							return init_dict, inputs_dict




				def 	__a		(   self	)							-> Tuple:
							pass




				def 	__a		(   self	)							-> Any:
							pass




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

							assert not model.is_gradient_checkpointing and model.training

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

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

							assert model_a.is_gradient_checkpointing and model_a.training

							lowerCAmelCase_         =      model_a(**_UpperCamelCase	).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_         =      (out_a - labels).mean()
							loss_a.backward()

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

							model.to(_UpperCamelCase	)
							lowerCAmelCase_         =      model(**self.dummy_input	)

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




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

							if torch_device == "mps":
										lowerCAmelCase_         =      torch.manual_seed(0	)
							else:
										lowerCAmelCase_         =      torch.Generator(device=_UpperCamelCase	).manual_seed(0	)

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

							lowerCAmelCase_         =      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_         =      torch.tensor(
										    [
										        -4.0_0_7_8e-0_1,
										        -3.8_3_2_3e-0_4,
										        -1.2_6_8_1e-0_1,
										        -1.1_4_6_2e-0_1,
										        2.0_0_9_5e-0_1,
										        1.0_8_9_3e-0_1,
										        -8.8_2_4_7e-0_2,
										        -3.0_3_6_1e-0_1,
										        -9.8_6_4_4e-0_3,
										    ]	)
							elif torch_device == "cpu":
										lowerCAmelCase_         =      torch.tensor(
										    [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026]	)
							else:
										lowerCAmelCase_         =      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(_UpperCamelCase							,							_UpperCamelCase							,							rtol=1e-2	)	)





@slow
class       _lowerCAmelCase   (      unittest.TestCase			):




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




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




				def 	__a		(   self							,							_UpperCamelCase=0							,							_UpperCamelCase=(4, 3, 512, 512)							,							_UpperCamelCase=False	)							-> int:
							lowerCAmelCase_         =      torch.floataa if fpaa else torch.floataa
							lowerCAmelCase_         =      torch.from_numpy(load_hf_numpy(self.get_file_format(_UpperCamelCase							,							_UpperCamelCase	)	)	).to(_UpperCamelCase	).to(_UpperCamelCase	)
							return image




				def 	__a		(   self							,							_UpperCamelCase="CompVis/stable-diffusion-v1-4"							,							_UpperCamelCase=False	)							-> List[str]:
							lowerCAmelCase_         =      "fp16" if fpaa else None
							lowerCAmelCase_         =      torch.floataa if fpaa else torch.floataa

							lowerCAmelCase_         =      AutoencoderKL.from_pretrained(
							    _UpperCamelCase							,							subfolder="vae"							,							torch_dtype=_UpperCamelCase							,							revision=_UpperCamelCase							,							)
							model.to(_UpperCamelCase	).eval()

							return model




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




				@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 	__a		(   self							,							_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase	)							-> Optional[Any]:
							lowerCAmelCase_         =      self.get_sd_vae_model()
							lowerCAmelCase_         =      self.get_sd_image(_UpperCamelCase	)
							lowerCAmelCase_         =      self.get_generator(_UpperCamelCase	)

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

							assert sample.shape == image.shape

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

							assert torch_all_close(_UpperCamelCase							,							_UpperCamelCase							,							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 	__a		(   self							,							_UpperCamelCase							,							_UpperCamelCase	)							-> Tuple:
							lowerCAmelCase_         =      self.get_sd_vae_model(fpaa=_UpperCamelCase	)
							lowerCAmelCase_         =      self.get_sd_image(_UpperCamelCase							,							fpaa=_UpperCamelCase	)
							lowerCAmelCase_         =      self.get_generator(_UpperCamelCase	)

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

							assert sample.shape == image.shape

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

							assert torch_all_close(_UpperCamelCase							,							_UpperCamelCase							,							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 	__a		(   self							,							_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase	)							-> Tuple:
							lowerCAmelCase_         =      self.get_sd_vae_model()
							lowerCAmelCase_         =      self.get_sd_image(_UpperCamelCase	)

							with torch.no_grad():
										lowerCAmelCase_         =      model(_UpperCamelCase	).sample

							assert sample.shape == image.shape

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

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

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

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

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

							assert torch_all_close(_UpperCamelCase							,							_UpperCamelCase							,							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 	__a		(   self							,							_UpperCamelCase							,							_UpperCamelCase	)							-> List[str]:
							lowerCAmelCase_         =      self.get_sd_vae_model(fpaa=_UpperCamelCase	)
							lowerCAmelCase_         =      self.get_sd_image(_UpperCamelCase							,							shape=(3, 4, 64, 64)							,							fpaa=_UpperCamelCase	)

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

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

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

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

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

							model.enable_xformers_memory_efficient_attention()
							with torch.no_grad():
										lowerCAmelCase_         =      model.decode(_UpperCamelCase	).sample

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

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

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

							model.enable_xformers_memory_efficient_attention()
							with torch.no_grad():
										lowerCAmelCase_         =      model.decode(_UpperCamelCase	).sample

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

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

							with torch.no_grad():
										lowerCAmelCase_         =      model.encode(_UpperCamelCase	).latent_dist
										lowerCAmelCase_         =      dist.sample(generator=_UpperCamelCase	)

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

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

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





from dataclasses import dataclass
from enum import Enum
from typing import List, Optional, Union

import numpy as np
import PIL
from PIL import Image

from ...utils import BaseOutput, is_torch_available, is_transformers_available





@dataclass
class       _lowerCAmelCase   (      __a			):
				_lowercase     						=42
				_lowercase     						=42


if is_transformers_available() and is_torch_available():
			from .pipeline_semantic_stable_diffusion import SemanticStableDiffusionPipeline



