




"""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 					__magic_name__							( lowercase_		,      lowercase_		,      unittest.TestCase       ):






				'''simple docstring'''
				__UpperCamelCase    						=		AutoencoderKL
				__UpperCamelCase    						=		'''sample'''
				__UpperCamelCase    						=		1E-2





				@property
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      4
							lowerCamelCase       =      3
							lowerCamelCase       =      (32, 32)

							lowerCamelCase       =      floats_tensor((batch_size, num_channels) + sizes				).to(_a				)

							return {"sample": image}





				@property
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							return (3, 32, 32)





				@property
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							return (3, 32, 32)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							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      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							pass





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							pass





				@unittest.skipIf(torch_device == """mps"""	,				"""Gradient checkpointing skipped on MPS"""				)
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							# enable deterministic behavior for gradient checkpointing
							lowerCamelCase   ,     lowerCamelCase       =      self.prepare_init_args_and_inputs_for_common()
							lowerCamelCase       =      self.model_class(**_a				)
							model.to(_a				)

							assert not model.is_gradient_checkpointing and model.training

							lowerCamelCase       =      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()

							lowerCamelCase       =      torch.randn_like(_a				)
							lowerCamelCase       =      (out - labels).mean()
							loss.backward()

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

							lowerCamelCase       =      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()
							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      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase   ,     lowerCamelCase       =      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				)
							lowerCamelCase       =      model(**self.dummy_input				)

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""				)
							lowerCamelCase       =      model.to(_a				)
							model.eval()

							if torch_device == "mps":
										lowerCamelCase       =      torch.manual_seed(0				)
							else:
										lowerCamelCase       =      torch.Generator(device=_a				).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(_a				)
							with torch.no_grad():
										lowerCamelCase       =      model(_a	,				sample_posterior=_a	,				generator=_a				).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.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       =      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(_a	,				_a	,				rtol=1e-2				)				)






@slow
class 					__magic_name__							( unittest.TestCase       ):






				'''simple docstring'''





				def      _lowerCAmelCase       ( self	,				_a	,				_a				):





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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()





				def      _lowerCAmelCase       ( self	,				_a=0	,				_a=(4, 3, 512, 512)	,				_a=False				):





							"""simple docstring"""
							lowerCamelCase       =      torch.floataa if fpaa else torch.floataa
							lowerCamelCase       =      torch.from_numpy(load_hf_numpy(self.get_file_format(_a	,				_a				)				)				).to(_a				).to(_a				)
							return image





				def      _lowerCAmelCase       ( self	,				_a="CompVis/stable-diffusion-v1-4"	,				_a=False				):





							"""simple docstring"""
							lowerCamelCase       =      """fp16""" if fpaa else None
							lowerCamelCase       =      torch.floataa if fpaa else torch.floataa

							lowerCamelCase       =      AutoencoderKL.from_pretrained(
							    _a	,				subfolder="""vae"""	,				torch_dtype=_a	,				revision=_a	,				)
							model.to(_a				).eval()

							return model





				def      _lowerCAmelCase       ( self	,				_a=0				):





							"""simple docstring"""
							if torch_device == "mps":
										return torch.manual_seed(_a				)
							return torch.Generator(device=_a				).manual_seed(_a				)





				@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      _lowerCAmelCase       ( self	,				_a	,				_a	,				_a				):





							"""simple docstring"""
							lowerCamelCase       =      self.get_sd_vae_model()
							lowerCamelCase       =      self.get_sd_image(_a				)
							lowerCamelCase       =      self.get_generator(_a				)

							with torch.no_grad():
										lowerCamelCase       =      model(_a	,				generator=_a	,				sample_posterior=_a				).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(_a	,				_a	,				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      _lowerCAmelCase       ( self	,				_a	,				_a				):





							"""simple docstring"""
							lowerCamelCase       =      self.get_sd_vae_model(fpaa=_a				)
							lowerCamelCase       =      self.get_sd_image(_a	,				fpaa=_a				)
							lowerCamelCase       =      self.get_generator(_a				)

							with torch.no_grad():
										lowerCamelCase       =      model(_a	,				generator=_a	,				sample_posterior=_a				).sample

							assert sample.shape == image.shape

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

							assert torch_all_close(_a	,				_a	,				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      _lowerCAmelCase       ( self	,				_a	,				_a	,				_a				):





							"""simple docstring"""
							lowerCamelCase       =      self.get_sd_vae_model()
							lowerCamelCase       =      self.get_sd_image(_a				)

							with torch.no_grad():
										lowerCamelCase       =      model(_a				).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(_a	,				_a	,				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      _lowerCAmelCase       ( self	,				_a	,				_a				):





							"""simple docstring"""
							lowerCamelCase       =      self.get_sd_vae_model()
							lowerCamelCase       =      self.get_sd_image(_a	,				shape=(3, 4, 64, 64)				)

							with torch.no_grad():
										lowerCamelCase       =      model.decode(_a				).sample

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

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

							assert torch_all_close(_a	,				_a	,				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      _lowerCAmelCase       ( self	,				_a	,				_a				):





							"""simple docstring"""
							lowerCamelCase       =      self.get_sd_vae_model(fpaa=_a				)
							lowerCamelCase       =      self.get_sd_image(_a	,				shape=(3, 4, 64, 64)	,				fpaa=_a				)

							with torch.no_grad():
										lowerCamelCase       =      model.decode(_a				).sample

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

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





							"""simple docstring"""
							lowerCamelCase       =      self.get_sd_vae_model(fpaa=_a				)
							lowerCamelCase       =      self.get_sd_image(_a	,				shape=(3, 4, 64, 64)	,				fpaa=_a				)

							with torch.no_grad():
										lowerCamelCase       =      model.decode(_a				).sample

							model.enable_xformers_memory_efficient_attention()
							with torch.no_grad():
										lowerCamelCase       =      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	,				_a				):





							"""simple docstring"""
							lowerCamelCase       =      self.get_sd_vae_model()
							lowerCamelCase       =      self.get_sd_image(_a	,				shape=(3, 4, 64, 64)				)

							with torch.no_grad():
										lowerCamelCase       =      model.decode(_a				).sample

							model.enable_xformers_memory_efficient_attention()
							with torch.no_grad():
										lowerCamelCase       =      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.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      _lowerCAmelCase       ( self	,				_a	,				_a				):





							"""simple docstring"""
							lowerCamelCase       =      self.get_sd_vae_model()
							lowerCamelCase       =      self.get_sd_image(_a				)
							lowerCamelCase       =      self.get_generator(_a				)

							with torch.no_grad():
										lowerCamelCase       =      model.encode(_a				).latent_dist
										lowerCamelCase       =      dist.sample(generator=_a				)

							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(_a				)

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









"""simple docstring"""




import unittest

import numpy as np
import torch

from diffusers import ScoreSdeVePipeline, ScoreSdeVeScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()






class 					__magic_name__							( unittest.TestCase       ):






				'''simple docstring'''





				@property
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							torch.manual_seed(0				)
							lowerCamelCase       =      UNetaDModel(
							    block_out_channels=(32, 64)	,				layers_per_block=2	,				sample_size=32	,				in_channels=3	,				out_channels=3	,				down_block_types=("""DownBlock2D""", """AttnDownBlock2D""")	,				up_block_types=("""AttnUpBlock2D""", """UpBlock2D""")	,				)
							return model





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.dummy_uncond_unet
							lowerCamelCase       =      ScoreSdeVeScheduler()

							lowerCamelCase       =      ScoreSdeVePipeline(unet=_a	,				scheduler=_a				)
							sde_ve.to(_a				)
							sde_ve.set_progress_bar_config(disable=_a				)

							lowerCamelCase       =      torch.manual_seed(0				)
							lowerCamelCase       =      sde_ve(num_inference_steps=2	,				output_type="""numpy"""	,				generator=_a				).images

							lowerCamelCase       =      torch.manual_seed(0				)
							lowerCamelCase       =      sde_ve(num_inference_steps=2	,				output_type="""numpy"""	,				generator=_a	,				return_dict=_a				)[
							    0
							]

							lowerCamelCase       =      image[0, -3:, -3:, -1]
							lowerCamelCase       =      image_from_tuple[0, -3:, -3:, -1]

							assert image.shape == (1, 32, 32, 3)
							lowerCamelCase       =      np.array([0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]				)

							assert np.abs(image_slice.flatten() - expected_slice				).max() < 1e-2
							assert np.abs(image_from_tuple_slice.flatten() - expected_slice				).max() < 1e-2






@slow
@require_torch
class 					__magic_name__							( unittest.TestCase       ):






				'''simple docstring'''





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      """google/ncsnpp-church-256"""
							lowerCamelCase       =      UNetaDModel.from_pretrained(_a				)

							lowerCamelCase       =      ScoreSdeVeScheduler.from_pretrained(_a				)

							lowerCamelCase       =      ScoreSdeVePipeline(unet=_a	,				scheduler=_a				)
							sde_ve.to(_a				)
							sde_ve.set_progress_bar_config(disable=_a				)

							lowerCamelCase       =      torch.manual_seed(0				)
							lowerCamelCase       =      sde_ve(num_inference_steps=10	,				output_type="""numpy"""	,				generator=_a				).images

							lowerCamelCase       =      image[0, -3:, -3:, -1]

							assert image.shape == (1, 256, 256, 3)

							lowerCamelCase       =      np.array([0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0]				)

							assert np.abs(image_slice.flatten() - expected_slice				).max() < 1e-2



