




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





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


 @property
 def  UpperCamelCase_   (				self:		Any		):




   '''simple docstring'''



   _snake_case							:		List[Any]								=  4
   _snake_case							:		Tuple								=  3
   _snake_case							:		Dict								=  (32, 32)

   _snake_case							:		Optional[int]								=  floats_tensor((batch_size, num_channels) + sizes		).to(a_		)

   return {"sample": image}


 @property
 def  UpperCamelCase_   (				self:		str		):




   '''simple docstring'''



   return (3, 32, 32)


 @property
 def  UpperCamelCase_   (				self:		List[str]		):




   '''simple docstring'''



   return (3, 32, 32)


 def  UpperCamelCase_   (				self:		int		):




   '''simple docstring'''



   _snake_case							:		Tuple								=  {
       """block_out_channels""": [32, 64],
       """in_channels""": 3,
       """out_channels""": 3,
       """down_block_types""": ["""DownEncoderBlock2D""", """DownEncoderBlock2D"""],
       """up_block_types""": ["""UpDecoderBlock2D""", """UpDecoderBlock2D"""],
       """latent_channels""": 4,
   }
   _snake_case							:		Dict								=  self.dummy_input
   return init_dict, inputs_dict


 def  UpperCamelCase_   (				self:		Tuple		):




   '''simple docstring'''



   pass


 def  UpperCamelCase_   (				self:		Optional[Any]		):




   '''simple docstring'''



   pass


 @unittest.skipIf(torch_device == """mps""",      """Gradient checkpointing skipped on MPS"""		)
 def  UpperCamelCase_   (				self:		List[Any]		):




   '''simple docstring'''



   _snake_case		,					_snake_case							:		Optional[Any]								=  self.prepare_init_args_and_inputs_for_common()
   _snake_case							:		str								=  self.model_class(**a_		)
   model.to(a_		)

   assert not model.is_gradient_checkpointing and model.training

   _snake_case							:		List[str]								=  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()

   _snake_case							:		str								=  torch.randn_like(a_		)
   _snake_case							:		int								=  (out - labels).mean()
   loss.backward()

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

   _snake_case							:		Tuple								=  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()
   _snake_case							:		int								=  (out_a - labels).mean()
   loss_a.backward()

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




   '''simple docstring'''



   _snake_case		,					_snake_case							:		Optional[int]								=  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_		)
   _snake_case							:		List[Any]								=  model(**self.dummy_input		)

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


 def  UpperCamelCase_   (				self:		Optional[Any]		):




   '''simple docstring'''



   _snake_case							:		str								=  AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""		)
   _snake_case							:		Tuple								=  model.to(a_		)
   model.eval()

   if torch_device == "mps":
     _snake_case							:		int								=  torch.manual_seed(0		)
   else:
     _snake_case							:		Tuple								=  torch.Generator(device=a_		).manual_seed(0		)

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

   _snake_case							:		Tuple								=  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":
     _snake_case							:		Any								=  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":
     _snake_case							:		Dict								=  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:
     _snake_case							:		List[Any]								=  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 					lowercase(     unittest.TestCase ):







 '''simple docstring'''





 def  UpperCamelCase_   (				self:		List[str],      a_:		List[Any],      a_:		List[Any]		):




   '''simple docstring'''



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


 def  UpperCamelCase_   (				self:		str		):




   '''simple docstring'''



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


 def  UpperCamelCase_   (				self:		str,      a_:		List[str]=0,      a_:		Tuple=(4, 3, 512, 512),      a_:		Optional[Any]=False		):




   '''simple docstring'''



   _snake_case							:		str								=  torch.floataa if fpaa else torch.floataa
   _snake_case							:		int								=  torch.from_numpy(load_hf_numpy(self.get_file_format(a_,      a_		)		)		).to(a_		).to(a_		)
   return image


 def  UpperCamelCase_   (				self:		Any,      a_:		Optional[int]="CompVis/stable-diffusion-v1-4",      a_:		str=False		):




   '''simple docstring'''



   _snake_case							:		str								=  """fp16""" if fpaa else None
   _snake_case							:		Optional[int]								=  torch.floataa if fpaa else torch.floataa

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

   return model


 def  UpperCamelCase_   (				self:		Union[str, Any],      a_:		List[str]=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  UpperCamelCase_   (				self:		Dict,      a_:		Any,      a_:		Any,      a_:		int		):




   '''simple docstring'''



   _snake_case							:		str								=  self.get_sd_vae_model()
   _snake_case							:		str								=  self.get_sd_image(a_		)
   _snake_case							:		Dict								=  self.get_generator(a_		)

   with torch.no_grad():
     _snake_case							:		str								=  model(a_,      generator=a_,      sample_posterior=a_		).sample

   assert sample.shape == image.shape

   _snake_case							:		Optional[Any]								=  sample[-1, -2:, -2:, :2].flatten().float().cpu()
   _snake_case							:		Optional[int]								=  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  UpperCamelCase_   (				self:		Tuple,      a_:		int,      a_:		Dict		):




   '''simple docstring'''



   _snake_case							:		Tuple								=  self.get_sd_vae_model(fpaa=a_		)
   _snake_case							:		Tuple								=  self.get_sd_image(a_,      fpaa=a_		)
   _snake_case							:		Tuple								=  self.get_generator(a_		)

   with torch.no_grad():
     _snake_case							:		Optional[int]								=  model(a_,      generator=a_,      sample_posterior=a_		).sample

   assert sample.shape == image.shape

   _snake_case							:		Dict								=  sample[-1, -2:, :2, -2:].flatten().float().cpu()
   _snake_case							:		Tuple								=  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  UpperCamelCase_   (				self:		List[str],      a_:		Tuple,      a_:		List[Any],      a_:		int		):




   '''simple docstring'''



   _snake_case							:		List[Any]								=  self.get_sd_vae_model()
   _snake_case							:		Union[str, Any]								=  self.get_sd_image(a_		)

   with torch.no_grad():
     _snake_case							:		Optional[int]								=  model(a_		).sample

   assert sample.shape == image.shape

   _snake_case							:		int								=  sample[-1, -2:, -2:, :2].flatten().float().cpu()
   _snake_case							:		Optional[Any]								=  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  UpperCamelCase_   (				self:		Union[str, Any],      a_:		Optional[int],      a_:		Any		):




   '''simple docstring'''



   _snake_case							:		Union[str, Any]								=  self.get_sd_vae_model()
   _snake_case							:		List[str]								=  self.get_sd_image(a_,      shape=(3, 4, 64, 64)		)

   with torch.no_grad():
     _snake_case							:		Union[str, Any]								=  model.decode(a_		).sample

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

   _snake_case							:		Union[str, Any]								=  sample[-1, -2:, :2, -2:].flatten().cpu()
   _snake_case							:		List[Any]								=  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  UpperCamelCase_   (				self:		Union[str, Any],      a_:		Optional[Any],      a_:		int		):




   '''simple docstring'''



   _snake_case							:		Tuple								=  self.get_sd_vae_model(fpaa=a_		)
   _snake_case							:		Union[str, Any]								=  self.get_sd_image(a_,      shape=(3, 4, 64, 64),      fpaa=a_		)

   with torch.no_grad():
     _snake_case							:		Any								=  model.decode(a_		).sample

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

   _snake_case							:		int								=  sample[-1, -2:, :2, -2:].flatten().float().cpu()
   _snake_case							:		Tuple								=  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  UpperCamelCase_   (				self:		Any,      a_:		Optional[Any]		):




   '''simple docstring'''



   _snake_case							:		Dict								=  self.get_sd_vae_model(fpaa=a_		)
   _snake_case							:		List[Any]								=  self.get_sd_image(a_,      shape=(3, 4, 64, 64),      fpaa=a_		)

   with torch.no_grad():
     _snake_case							:		Optional[int]								=  model.decode(a_		).sample

   model.enable_xformers_memory_efficient_attention()
   with torch.no_grad():
     _snake_case							:		List[str]								=  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  UpperCamelCase_   (				self:		str,      a_:		Union[str, Any]		):




   '''simple docstring'''



   _snake_case							:		List[Any]								=  self.get_sd_vae_model()
   _snake_case							:		Any								=  self.get_sd_image(a_,      shape=(3, 4, 64, 64)		)

   with torch.no_grad():
     _snake_case							:		int								=  model.decode(a_		).sample

   model.enable_xformers_memory_efficient_attention()
   with torch.no_grad():
     _snake_case							:		Union[str, Any]								=  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  UpperCamelCase_   (				self:		Dict,      a_:		Tuple,      a_:		Tuple		):




   '''simple docstring'''



   _snake_case							:		str								=  self.get_sd_vae_model()
   _snake_case							:		int								=  self.get_sd_image(a_		)
   _snake_case							:		Dict								=  self.get_generator(a_		)

   with torch.no_grad():
     _snake_case							:		Dict								=  model.encode(a_		).latent_dist
     _snake_case							:		Dict								=  dist.sample(generator=a_		)

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

   _snake_case							:		Optional[int]								=  sample[0, -1, -3:, -3:].flatten().cpu()
   _snake_case							:		Tuple								=  torch.tensor(a_		)

   _snake_case							:		List[Any]								=  3E-3 if torch_device != """mps""" else 1E-2
   assert torch_all_close(a_,      a_,      atol=a_		)



'''simple docstring'''




from string import ascii_lowercase, ascii_uppercase






def 				snake_case__							(			_A:					str )    ->	str:

		'''simple docstring'''


		if not sentence:
				return ""
		lowerCAmelCase   					=  dict(zip(_A      ,   _A ) )
		return lower_to_upper.get(sentence[0]      ,   sentence[0] ) + sentence[1:]


if __name__ == "__main__":
						from doctest import testmod

						testmod()




