





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       (    UpperCamelCase__						,				UpperCamelCase__						,				unittest.TestCase):
      _lowercase				:					Dict						=			AutoencoderKL
      _lowercase				:					Tuple						=			"""sample"""
      _lowercase				:					Optional[int]						=			1E-2







      @property
      def    _lowercase  (    self							)   ->						Optional[Any]:



        '''simple docstring'''




        a__	:					Union[str, Any]										=4
        a__	:					List[str]										=3
        a__	:					Union[str, Any]										=(3_2, 3_2)

        a__	:					int										=floats_tensor((batch_size, num_channels) + sizes							).to(lowerCAmelCase__							)

        return {"sample": image}







      @property
      def    _lowercase  (    self							)   ->						Optional[Any]:



        '''simple docstring'''




        return (3, 3_2, 3_2)







      @property
      def    _lowercase  (    self							)   ->						List[Any]:



        '''simple docstring'''




        return (3, 3_2, 3_2)







      def    _lowercase  (    self							)   ->						Union[str, Any]:



        '''simple docstring'''




        a__	:					List[Any]										={
            "block_out_channels": [3_2, 6_4],
            "in_channels": 3,
            "out_channels": 3,
            "down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D"],
            "up_block_types": ["UpDecoderBlock2D", "UpDecoderBlock2D"],
            "latent_channels": 4,
        }
        a__	:					List[str]										=self.dummy_input
        return init_dict, inputs_dict







      def    _lowercase  (    self							)   ->						Optional[Any]:



        '''simple docstring'''




        pass







      def    _lowercase  (    self							)   ->						List[Any]:



        '''simple docstring'''




        pass







      @unittest.skipIf(torch_device == "mps"		,      "Gradient checkpointing skipped on MPS"							)
      def    _lowercase  (    self							)   ->						Optional[int]:



        '''simple docstring'''




        a__					,							a__	:					Union[str, Any]										=self.prepare_init_args_and_inputs_for_common()
        a__	:					Tuple										=self.model_class(**lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)

        assert not model.is_gradient_checkpointing and model.training

        a__	:					Optional[Any]										=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()

        a__	:					str										=torch.randn_like(lowerCAmelCase__							)
        a__	:					Tuple										=(out - labels).mean()
        loss.backward()

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

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

        # compare the output and parameters gradients
        self.assertTrue((loss - loss_a).abs() < 1E-5							)
        a__	:					Union[str, Any]										=dict(model.named_parameters()							)
        a__	:					int										=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    _lowercase  (    self							)   ->						Optional[Any]:



        '''simple docstring'''




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

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







      def    _lowercase  (    self							)   ->						Optional[int]:



        '''simple docstring'''




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

        if torch_device == "mps":
          a__	:					Any										=torch.manual_seed(0							)
        else:
          a__	:					Union[str, Any]										=torch.Generator(device=lowerCAmelCase__							).manual_seed(0							)

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

        a__	:					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":
          a__	:					str										=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":
          a__	:					Dict										=torch.tensor(
              [-0.13_52, 0.08_78, 0.04_19, -0.08_18, -0.10_69, 0.06_88, -0.14_58, -0.44_46, -0.00_26]							)
        else:
          a__	:					Optional[int]										=torch.tensor(
              [-0.24_21, 0.46_42, 0.25_07, -0.04_38, 0.06_82, 0.31_60, -0.20_18, -0.07_27, 0.24_85]							)

        self.assertTrue(torch_all_close(lowerCAmelCase__		,      lowerCAmelCase__		,      rtol=1E-2							)							)







@slow
class 				__lowerCAmelCase       (    unittest.TestCase):







      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						Optional[Any]:



        '''simple docstring'''




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







      def    _lowercase  (    self							)   ->						List[str]:



        '''simple docstring'''




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







      def    _lowercase  (    self		,      lowerCAmelCase__=0		,      lowerCAmelCase__=(4, 3, 5_1_2, 5_1_2)		,      lowerCAmelCase__=False							)   ->						Optional[Any]:



        '''simple docstring'''




        a__	:					List[str]										=torch.floataa if fpaa else torch.floataa
        a__	:					Tuple										=torch.from_numpy(load_hf_numpy(self.get_file_format(lowerCAmelCase__		,      lowerCAmelCase__							)							)							).to(lowerCAmelCase__							).to(lowerCAmelCase__							)
        return image







      def    _lowercase  (    self		,      lowerCAmelCase__="CompVis/stable-diffusion-v1-4"		,      lowerCAmelCase__=False							)   ->						Any:



        '''simple docstring'''




        a__	:					str										="fp16" if fpaa else None
        a__	:					str										=torch.floataa if fpaa else torch.floataa

        a__	:					Dict										=AutoencoderKL.from_pretrained(
            lowerCAmelCase__		,      subfolder="vae"		,      torch_dtype=lowerCAmelCase__		,      revision=lowerCAmelCase__		,      )
        model.to(lowerCAmelCase__							).eval()

        return model







      def    _lowercase  (    self		,      lowerCAmelCase__=0							)   ->						Optional[int]:



        '''simple docstring'''




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







      @parameterized.expand(
          [
              # fmt: off
              [3_3, [-0.16_03, 0.98_78, -0.04_95, -0.07_90, -0.27_09, 0.83_75, -0.20_60, -0.08_24], [-0.23_95, 0.00_98, 0.01_02, -0.07_09, -0.28_40, -0.02_74, -0.07_18, -0.18_24]],
              [4_7, [-0.23_76, 0.11_68, 0.13_32, -0.48_40, -0.25_08, -0.07_91, -0.04_93, -0.40_89], [0.03_50, 0.08_47, 0.04_67, 0.03_44, -0.08_42, -0.05_47, -0.06_33, -0.11_31]],
              # fmt: on
          ]							)
      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						List[str]:



        '''simple docstring'''




        a__	:					List[Any]										=self.get_sd_vae_model()
        a__	:					Tuple										=self.get_sd_image(lowerCAmelCase__							)
        a__	:					Tuple										=self.get_generator(lowerCAmelCase__							)

        with torch.no_grad():
          a__	:					Tuple										=model(lowerCAmelCase__		,      generator=lowerCAmelCase__		,      sample_posterior=lowerCAmelCase__							).sample

        assert sample.shape == image.shape

        a__	:					Tuple										=sample[-1, -2:, -2:, :2].flatten().float().cpu()
        a__	:					Tuple										=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
              [3_3, [-0.05_13, 0.02_89, 1.37_99, 0.21_66, -0.25_73, -0.08_71, 0.51_03, -0.09_99]],
              [4_7, [-0.41_28, -0.13_20, -0.37_04, 0.19_65, -0.41_16, -0.23_32, -0.33_40, 0.22_47]],
              # fmt: on
          ]							)
      @require_torch_gpu
      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						int:



        '''simple docstring'''




        a__	:					Optional[int]										=self.get_sd_vae_model(fpaa=lowerCAmelCase__							)
        a__	:					Any										=self.get_sd_image(lowerCAmelCase__		,      fpaa=lowerCAmelCase__							)
        a__	:					Dict										=self.get_generator(lowerCAmelCase__							)

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

        assert sample.shape == image.shape

        a__	:					int										=sample[-1, -2:, :2, -2:].flatten().float().cpu()
        a__	:					List[str]										=torch.tensor(lowerCAmelCase__							)

        assert torch_all_close(lowerCAmelCase__		,      lowerCAmelCase__		,      atol=1E-2							)







      @parameterized.expand(
          [
              # fmt: off
              [3_3, [-0.16_09, 0.98_66, -0.04_87, -0.07_77, -0.27_16, 0.83_68, -0.20_55, -0.08_14], [-0.23_95, 0.00_98, 0.01_02, -0.07_09, -0.28_40, -0.02_74, -0.07_18, -0.18_24]],
              [4_7, [-0.23_77, 0.11_47, 0.13_33, -0.48_41, -0.25_06, -0.08_05, -0.04_91, -0.40_85], [0.03_50, 0.08_47, 0.04_67, 0.03_44, -0.08_42, -0.05_47, -0.06_33, -0.11_31]],
              # fmt: on
          ]							)
      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						List[str]:



        '''simple docstring'''




        a__	:					Any										=self.get_sd_vae_model()
        a__	:					Optional[int]										=self.get_sd_image(lowerCAmelCase__							)

        with torch.no_grad():
          a__	:					Dict										=model(lowerCAmelCase__							).sample

        assert sample.shape == image.shape

        a__	:					List[str]										=sample[-1, -2:, -2:, :2].flatten().float().cpu()
        a__	:					List[str]										=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
              [1_3, [-0.20_51, -0.18_03, -0.23_11, -0.21_14, -0.32_92, -0.35_74, -0.29_53, -0.33_23]],
              [3_7, [-0.26_32, -0.26_25, -0.21_99, -0.27_41, -0.45_39, -0.49_90, -0.37_20, -0.49_25]],
              # fmt: on
          ]							)
      @require_torch_gpu
      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						Optional[int]:



        '''simple docstring'''




        a__	:					List[Any]										=self.get_sd_vae_model()
        a__	:					Optional[Any]										=self.get_sd_image(lowerCAmelCase__		,      shape=(3, 4, 6_4, 6_4)							)

        with torch.no_grad():
          a__	:					Optional[Any]										=model.decode(lowerCAmelCase__							).sample

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

        a__	:					List[Any]										=sample[-1, -2:, :2, -2:].flatten().cpu()
        a__	:					Dict										=torch.tensor(lowerCAmelCase__							)

        assert torch_all_close(lowerCAmelCase__		,      lowerCAmelCase__		,      atol=1E-3							)







      @parameterized.expand(
          [
              # fmt: off
              [2_7, [-0.03_69, 0.02_07, -0.07_76, -0.06_82, -0.17_47, -0.19_30, -0.14_65, -0.20_39]],
              [1_6, [-0.16_28, -0.21_34, -0.27_47, -0.26_42, -0.37_74, -0.44_04, -0.36_87, -0.42_77]],
              # fmt: on
          ]							)
      @require_torch_gpu
      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						int:



        '''simple docstring'''




        a__	:					Optional[Any]										=self.get_sd_vae_model(fpaa=lowerCAmelCase__							)
        a__	:					Tuple										=self.get_sd_image(lowerCAmelCase__		,      shape=(3, 4, 6_4, 6_4)		,      fpaa=lowerCAmelCase__							)

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

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

        a__	:					Tuple										=sample[-1, -2:, :2, -2:].flatten().float().cpu()
        a__	:					Dict										=torch.tensor(lowerCAmelCase__							)

        assert torch_all_close(lowerCAmelCase__		,      lowerCAmelCase__		,      atol=5E-3							)







      @parameterized.expand([(1_3,), (1_6,), (2_7,)]							)
      @require_torch_gpu
      @unittest.skipIf(not is_xformers_available()		,      reason="xformers is not required when using PyTorch 2.0."							)
      def    _lowercase  (    self		,      lowerCAmelCase__							)   ->						Dict:



        '''simple docstring'''




        a__	:					int										=self.get_sd_vae_model(fpaa=lowerCAmelCase__							)
        a__	:					List[Any]										=self.get_sd_image(lowerCAmelCase__		,      shape=(3, 4, 6_4, 6_4)		,      fpaa=lowerCAmelCase__							)

        with torch.no_grad():
          a__	:					Optional[int]										=model.decode(lowerCAmelCase__							).sample

        model.enable_xformers_memory_efficient_attention()
        with torch.no_grad():
          a__	:					Tuple										=model.decode(lowerCAmelCase__							).sample

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

        assert torch_all_close(lowerCAmelCase__		,      lowerCAmelCase__		,      atol=1E-1							)







      @parameterized.expand([(1_3,), (1_6,), (3_7,)]							)
      @require_torch_gpu
      @unittest.skipIf(not is_xformers_available()		,      reason="xformers is not required when using PyTorch 2.0."							)
      def    _lowercase  (    self		,      lowerCAmelCase__							)   ->						List[Any]:



        '''simple docstring'''




        a__	:					Dict										=self.get_sd_vae_model()
        a__	:					str										=self.get_sd_image(lowerCAmelCase__		,      shape=(3, 4, 6_4, 6_4)							)

        with torch.no_grad():
          a__	:					List[str]										=model.decode(lowerCAmelCase__							).sample

        model.enable_xformers_memory_efficient_attention()
        with torch.no_grad():
          a__	:					Optional[int]										=model.decode(lowerCAmelCase__							).sample

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

        assert torch_all_close(lowerCAmelCase__		,      lowerCAmelCase__		,      atol=1E-2							)







      @parameterized.expand(
          [
              # fmt: off
              [3_3, [-0.30_01, 0.09_18, -2.69_84, -3.97_20, -3.20_99, -5.03_53, 1.73_38, -0.20_65, 3.42_67]],
              [4_7, [-1.50_30, -4.38_71, -6.03_55, -9.11_57, -1.66_61, -2.78_53, 2.16_07, -5.08_23, 2.56_33]],
              # fmt: on
          ]							)
      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						Optional[Any]:



        '''simple docstring'''




        a__	:					Tuple										=self.get_sd_vae_model()
        a__	:					int										=self.get_sd_image(lowerCAmelCase__							)
        a__	:					Any										=self.get_generator(lowerCAmelCase__							)

        with torch.no_grad():
          a__	:					str										=model.encode(lowerCAmelCase__							).latent_dist
          a__	:					Union[str, Any]										=dist.sample(generator=lowerCAmelCase__							)

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

        a__	:					Optional[Any]										=sample[0, -1, -3:, -3:].flatten().cpu()
        a__	:					List[str]										=torch.tensor(lowerCAmelCase__							)

        a__	:					Union[str, Any]										=3E-3 if torch_device != "mps" else 1E-2
        assert torch_all_close(lowerCAmelCase__		,      lowerCAmelCase__		,      atol=lowerCAmelCase__							)








import unittest

import numpy as np
import torch

from .utils_summarization import build_mask, compute_token_type_ids, process_story, truncate_or_pad
class  lowercase__(				unittest.TestCase							):



	"""simple docstring"""




	def 						_lowercase  (       self  :			List[str]       )       ->     List[Any]:
					lowercase_       =					1_0
	def 						_lowercase  (       self  :			int       )       ->     List[str]:
					lowercase_       =					[1, 2, 3, 4]
					lowercase_       =					[1, 2, 3, 4, 0, 0, 0, 0, 0, 0]
					self.assertEqual(truncate_or_pad(SCREAMING_SNAKE_CASE_       ,							self.block_size       ,							0       )       ,							SCREAMING_SNAKE_CASE_       )
	def 						_lowercase  (       self  :			int       )       ->     Optional[Any]:
					lowercase_       =					[1, 2, 3, 4, 5, 6, 7, 8, 9, 1_0]
					lowercase_       =					[1, 2, 3, 4, 5, 6, 7, 8, 9, 1_0]
					self.assertEqual(truncate_or_pad(SCREAMING_SNAKE_CASE_       ,							self.block_size       ,							0       )       ,							SCREAMING_SNAKE_CASE_       )
	def 						_lowercase  (       self  :			Union[str, Any]       )       ->     Optional[int]:
					lowercase_       =					[1, 2, 3, 4, 5, 6, 7, 8, 9, 1_0, 1_1, 1_2, 1_3]
					lowercase_       =					[1, 2, 3, 4, 5, 6, 7, 8, 9, 1_0]
					self.assertEqual(truncate_or_pad(SCREAMING_SNAKE_CASE_       ,							self.block_size       ,							0       )       ,							SCREAMING_SNAKE_CASE_       )
	def 						_lowercase  (       self  :			Any       )       ->     List[Any]:
					lowercase_       =					'''It was the year of Our Lord one thousand seven hundred and
        seventy-five.\n\nSpiritual revelations were conceded to England at that
        favoured period, as at this.'''
					lowercase_  ,	lowercase_       =					process_story(SCREAMING_SNAKE_CASE_       )
					self.assertEqual(SCREAMING_SNAKE_CASE_       ,							[]       )
	def 						_lowercase  (       self  :			List[str]       )       ->     List[str]:
					lowercase_       =					''''''
					lowercase_  ,	lowercase_       =					process_story(SCREAMING_SNAKE_CASE_       )
					self.assertEqual(SCREAMING_SNAKE_CASE_       ,							[]       )
					self.assertEqual(SCREAMING_SNAKE_CASE_       ,							[]       )
	def 						_lowercase  (       self  :			Union[str, Any]       )       ->     Union[str, Any]:
					lowercase_       =					(
					    '''It was the year of Our Lord one thousand seven hundred and '''
					    '''seventy-five\n\nSpiritual revelations were conceded to England '''
					    '''at that favoured period, as at this.\n@highlight\n\nIt was the best of times'''
					)
					lowercase_  ,	lowercase_       =					process_story(SCREAMING_SNAKE_CASE_       )

					lowercase_       =					[
					    '''It was the year of Our Lord one thousand seven hundred and seventy-five.''',
					    '''Spiritual revelations were conceded to England at that favoured period, as at this.''',
					]
					self.assertEqual(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       )

					lowercase_       =					['''It was the best of times.''']
					self.assertEqual(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       )
	def 						_lowercase  (       self  :			Union[str, Any]       )       ->     Optional[Any]:
					lowercase_       =					torch.tensor([1, 2, 3, 4]       )
					lowercase_       =					torch.tensor([1, 1, 1, 1]       )
					np.testing.assert_array_equal(build_mask(SCREAMING_SNAKE_CASE_       ,							0       ).numpy()       ,							expected.numpy()       )
	def 						_lowercase  (       self  :			List[Any]       )       ->     Tuple:
					lowercase_       =					torch.tensor([1, 2, 3, 4, 2_3, 2_3, 2_3]       )
					lowercase_       =					torch.tensor([1, 1, 1, 1, 0, 0, 0]       )
					np.testing.assert_array_equal(build_mask(SCREAMING_SNAKE_CASE_       ,							2_3       ).numpy()       ,							expected.numpy()       )
	def 						_lowercase  (       self  :			int       )       ->     Dict:
					lowercase_       =					torch.tensor([8, 2, 3, 4, 1, 1, 1]       )
					lowercase_       =					torch.tensor([1, 1, 1, 1, 0, 0, 0]       )
					np.testing.assert_array_equal(build_mask(SCREAMING_SNAKE_CASE_       ,							1       ).numpy()       ,							expected.numpy()       )




	def 						_lowercase  (       self  :			List[str]       )       ->     Tuple:
					lowercase_       =					1_0_1
					lowercase_       =					torch.tensor([[1, 2, 3, 4, 5, 6], [1, 2, 3, 1_0_1, 5, 6], [1, 1_0_1, 3, 4, 1_0_1, 6]]       )
					lowercase_       =					torch.tensor([[1, 1, 1, 1, 1, 1], [1, 1, 1, 0, 0, 0], [1, 0, 0, 0, 1, 1]]       )

					lowercase_       =					compute_token_type_ids(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       )
					np.testing.assert_array_equal(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       )




