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

 """simple docstring"""
 SCREAMING_SNAKE_CASE_				:					List[str]       						=AutoencoderKL
 SCREAMING_SNAKE_CASE_				:					Optional[int]       						="sample"
 SCREAMING_SNAKE_CASE_				:					Optional[Any]       						=1E-2

 @property
 def      __lowerCAmelCase  (						self			:			Optional[Any]   ):






      """simple docstring"""



      UpperCamelCase						    =					4
      UpperCamelCase						    =					3
      UpperCamelCase						    =					(32, 32)

      UpperCamelCase						    =					floats_tensor((batch_size, num_channels) + sizes   ).to(SCREAMING_SNAKE_CASE__   )

      return {"sample": image}

 @property
 def      __lowerCAmelCase  (						self			:			Optional[int]   ):






      """simple docstring"""



      return (3, 32, 32)

 @property
 def      __lowerCAmelCase  (						self			:			Optional[Any]   ):






      """simple docstring"""



      return (3, 32, 32)

 def      __lowerCAmelCase  (						self			:			List[str]   ):






      """simple docstring"""



      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			:			Optional[int]   ):






      """simple docstring"""



      pass

 def      __lowerCAmelCase  (						self			:			int   ):






      """simple docstring"""



      pass

 @unittest.skipIf(torch_device == 'mps'					, 'Gradient checkpointing skipped on MPS'   )
 def      __lowerCAmelCase  (						self			:			Tuple   ):






      """simple docstring"""



      UpperCamelCase    ,		UpperCamelCase						    =					self.prepare_init_args_and_inputs_for_common()
      UpperCamelCase						    =					self.model_class(**SCREAMING_SNAKE_CASE__   )
      model.to(SCREAMING_SNAKE_CASE__   )

      assert not model.is_gradient_checkpointing and model.training

      UpperCamelCase						    =					model(**SCREAMING_SNAKE_CASE__   ).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(SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					(out - labels).mean()
      loss.backward()

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

      assert model_a.is_gradient_checkpointing and model_a.training

      UpperCamelCase						    =					model_a(**SCREAMING_SNAKE_CASE__   ).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			:			Any   ):






      """simple docstring"""



      UpperCamelCase    ,		UpperCamelCase						    =					AutoencoderKL.from_pretrained('fusing/autoencoder-kl-dummy'					, output_loading_info=SCREAMING_SNAKE_CASE__   )
      self.assertIsNotNone(SCREAMING_SNAKE_CASE__   )
      self.assertEqual(len(loading_info['missing_keys']   )					, 0   )

      model.to(SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					model(**self.dummy_input   )

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






 def      __lowerCAmelCase  (						self			:			int   ):






      """simple docstring"""



      UpperCamelCase						    =					AutoencoderKL.from_pretrained('fusing/autoencoder-kl-dummy'   )
      UpperCamelCase						    =					model.to(SCREAMING_SNAKE_CASE__   )
      model.eval()

      if torch_device == "mps":
           UpperCamelCase						    =					torch.manual_seed(0   )
      else:
           UpperCamelCase						    =					torch.Generator(device=SCREAMING_SNAKE_CASE__   ).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(SCREAMING_SNAKE_CASE__   )
      with torch.no_grad():
           UpperCamelCase						    =					model(SCREAMING_SNAKE_CASE__					, sample_posterior=SCREAMING_SNAKE_CASE__					, generator=SCREAMING_SNAKE_CASE__   ).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.0078e-01,
                   -3.8323e-04,
                   -1.2681e-01,
                   -1.1462e-01,
                   2.0095e-01,
                   1.0893e-01,
                   -8.8247e-02,
                   -3.0361e-01,
                   -9.8644e-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(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, rtol=1e-2   )   )




@slow
class       _lowerCAmelCase  (		unittest.TestCase							):

 """simple docstring"""

 def      __lowerCAmelCase  (						self			:			int					, SCREAMING_SNAKE_CASE__			:			Optional[Any]					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]   ):






      """simple docstring"""



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

 def      __lowerCAmelCase  (						self			:			Tuple   ):






      """simple docstring"""



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

 def      __lowerCAmelCase  (						self			:			List[Any]					, SCREAMING_SNAKE_CASE__			:			int=0					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]=(4, 3, 5_12, 5_12)					, SCREAMING_SNAKE_CASE__			:			Optional[int]=False   ):






      """simple docstring"""



      UpperCamelCase						    =					torch.floataa if fpaa else torch.floataa
      UpperCamelCase						    =					torch.from_numpy(load_hf_numpy(self.get_file_format(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )   )   ).to(SCREAMING_SNAKE_CASE__   ).to(SCREAMING_SNAKE_CASE__   )
      return image

 def      __lowerCAmelCase  (						self			:			Dict					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]="CompVis/stable-diffusion-v1-4"					, SCREAMING_SNAKE_CASE__			:			Optional[Any]=False   ):






      """simple docstring"""



      UpperCamelCase						    =					'fp16' if fpaa else None
      UpperCamelCase						    =					torch.floataa if fpaa else torch.floataa

      UpperCamelCase						    =					AutoencoderKL.from_pretrained(
          SCREAMING_SNAKE_CASE__					, subfolder='vae'					, torch_dtype=SCREAMING_SNAKE_CASE__					, revision=SCREAMING_SNAKE_CASE__					, )
      model.to(SCREAMING_SNAKE_CASE__   ).eval()

      return model

 def      __lowerCAmelCase  (						self			:			Union[str, Any]					, SCREAMING_SNAKE_CASE__			:			str=0   ):






      """simple docstring"""



      if torch_device == "mps":
           return torch.manual_seed(SCREAMING_SNAKE_CASE__   )
      return torch.Generator(device=SCREAMING_SNAKE_CASE__   ).manual_seed(SCREAMING_SNAKE_CASE__   )

 @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			:			Dict					, SCREAMING_SNAKE_CASE__			:			Optional[Any]					, SCREAMING_SNAKE_CASE__			:			Tuple					, SCREAMING_SNAKE_CASE__			:			Tuple   ):






      """simple docstring"""



      UpperCamelCase						    =					self.get_sd_vae_model()
      UpperCamelCase						    =					self.get_sd_image(SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					self.get_generator(SCREAMING_SNAKE_CASE__   )

      with torch.no_grad():
           UpperCamelCase						    =					model(SCREAMING_SNAKE_CASE__					, generator=SCREAMING_SNAKE_CASE__					, sample_posterior=SCREAMING_SNAKE_CASE__   ).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(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, 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			:			Tuple					, SCREAMING_SNAKE_CASE__			:			Dict					, SCREAMING_SNAKE_CASE__			:			List[str]   ):






      """simple docstring"""



      UpperCamelCase						    =					self.get_sd_vae_model(fpaa=SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					self.get_sd_image(SCREAMING_SNAKE_CASE__					, fpaa=SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					self.get_generator(SCREAMING_SNAKE_CASE__   )

      with torch.no_grad():
           UpperCamelCase						    =					model(SCREAMING_SNAKE_CASE__					, generator=SCREAMING_SNAKE_CASE__					, sample_posterior=SCREAMING_SNAKE_CASE__   ).sample

      assert sample.shape == image.shape

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

      assert torch_all_close(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, 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			:			List[str]					, SCREAMING_SNAKE_CASE__			:			List[Any]					, SCREAMING_SNAKE_CASE__			:			int					, SCREAMING_SNAKE_CASE__			:			Any   ):






      """simple docstring"""



      UpperCamelCase						    =					self.get_sd_vae_model()
      UpperCamelCase						    =					self.get_sd_image(SCREAMING_SNAKE_CASE__   )

      with torch.no_grad():
           UpperCamelCase						    =					model(SCREAMING_SNAKE_CASE__   ).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(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, 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			:			str					, SCREAMING_SNAKE_CASE__			:			List[str]					, SCREAMING_SNAKE_CASE__			:			Dict   ):






      """simple docstring"""



      UpperCamelCase						    =					self.get_sd_vae_model()
      UpperCamelCase						    =					self.get_sd_image(SCREAMING_SNAKE_CASE__					, shape=(3, 4, 64, 64)   )

      with torch.no_grad():
           UpperCamelCase						    =					model.decode(SCREAMING_SNAKE_CASE__   ).sample

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

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

      assert torch_all_close(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, 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			:			Tuple					, SCREAMING_SNAKE_CASE__			:			Optional[Any]					, SCREAMING_SNAKE_CASE__			:			List[str]   ):






      """simple docstring"""



      UpperCamelCase						    =					self.get_sd_vae_model(fpaa=SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					self.get_sd_image(SCREAMING_SNAKE_CASE__					, shape=(3, 4, 64, 64)					, fpaa=SCREAMING_SNAKE_CASE__   )

      with torch.no_grad():
           UpperCamelCase						    =					model.decode(SCREAMING_SNAKE_CASE__   ).sample

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

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

      assert torch_all_close(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, 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			:			Optional[Any]					, SCREAMING_SNAKE_CASE__			:			Optional[int]   ):






      """simple docstring"""



      UpperCamelCase						    =					self.get_sd_vae_model(fpaa=SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					self.get_sd_image(SCREAMING_SNAKE_CASE__					, shape=(3, 4, 64, 64)					, fpaa=SCREAMING_SNAKE_CASE__   )

      with torch.no_grad():
           UpperCamelCase						    =					model.decode(SCREAMING_SNAKE_CASE__   ).sample

      model.enable_xformers_memory_efficient_attention()
      with torch.no_grad():
           UpperCamelCase						    =					model.decode(SCREAMING_SNAKE_CASE__   ).sample

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

      assert torch_all_close(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, 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			:			Union[str, Any]					, SCREAMING_SNAKE_CASE__			:			Optional[int]   ):






      """simple docstring"""



      UpperCamelCase						    =					self.get_sd_vae_model()
      UpperCamelCase						    =					self.get_sd_image(SCREAMING_SNAKE_CASE__					, shape=(3, 4, 64, 64)   )

      with torch.no_grad():
           UpperCamelCase						    =					model.decode(SCREAMING_SNAKE_CASE__   ).sample

      model.enable_xformers_memory_efficient_attention()
      with torch.no_grad():
           UpperCamelCase						    =					model.decode(SCREAMING_SNAKE_CASE__   ).sample

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

      assert torch_all_close(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, 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			:			Union[str, Any]					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]					, SCREAMING_SNAKE_CASE__			:			Optional[Any]   ):






      """simple docstring"""



      UpperCamelCase						    =					self.get_sd_vae_model()
      UpperCamelCase						    =					self.get_sd_image(SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					self.get_generator(SCREAMING_SNAKE_CASE__   )

      with torch.no_grad():
           UpperCamelCase						    =					model.encode(SCREAMING_SNAKE_CASE__   ).latent_dist
           UpperCamelCase						    =					dist.sample(generator=SCREAMING_SNAKE_CASE__   )

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

      UpperCamelCase						    =					3e-3 if torch_device != 'mps' else 1e-2
      assert torch_all_close(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, atol=SCREAMING_SNAKE_CASE__   )




from functools import lru_cache

@lru_cache
def  __lowerCamelCase					( _lowercase			)       ->      int:
     if num < 0:
          raise ValueError('Number should not be negative.'			)

     return 1 if num in (0, 1) else num * factorial(num - 1			)


if __name__ == "__main__":
 import doctest

 doctest.testmod()


