






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 							_UpperCamelCase			(a_							,		a_							,		unittest.TestCase			):
      snake_case_  =							AutoencoderKL
      snake_case_  =							"""sample"""
      snake_case_  =							1E-2


      @property
      def  __UpperCAmelCase		(	self	)->	Union[str, Any]:
           __lowerCAmelCase								=					4
           __lowerCAmelCase								=					3
           __lowerCAmelCase								=					(3_2, 3_2)

           __lowerCAmelCase								=					floats_tensor((batch_size, num_channels) + sizes	).to(__UpperCamelCase	)

           return {"sample": image}


      @property
      def  __UpperCAmelCase		(	self	)->	Optional[int]:
           return (3, 3_2, 3_2)


      @property
      def  __UpperCAmelCase		(	self	)->	Any:
           return (3, 3_2, 3_2)


      def  __UpperCAmelCase		(	self	)->	int:
           __lowerCAmelCase								=					{
               "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,
           }
           __lowerCAmelCase								=					self.dummy_input
           return init_dict, inputs_dict


      def  __UpperCAmelCase		(	self	)->	Optional[Any]:
           pass


      def  __UpperCAmelCase		(	self	)->	List[str]:
           pass


      @unittest.skipIf(torch_device == "mps"					,						"Gradient checkpointing skipped on MPS"	)
      def  __UpperCAmelCase		(	self	)->	List[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  __UpperCAmelCase		(	self	)->	Optional[int]:
           __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  __UpperCAmelCase		(	self	)->	Optional[Any]:
           __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.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":
                __lowerCAmelCase								=					torch.tensor(
                    [-0.1_3_5_2, 0.0_8_7_8, 0.0_4_1_9, -0.0_8_1_8, -0.1_0_6_9, 0.0_6_8_8, -0.1_4_5_8, -0.4_4_4_6, -0.0_0_2_6]	)
           else:
                __lowerCAmelCase								=					torch.tensor(
                    [-0.2_4_2_1, 0.4_6_4_2, 0.2_5_0_7, -0.0_4_3_8, 0.0_6_8_2, 0.3_1_6_0, -0.2_0_1_8, -0.0_7_2_7, 0.2_4_8_5]	)

           self.assertTrue(torch_all_close(__UpperCamelCase					,						__UpperCamelCase					,						rtol=1e-2	)	)



@slow
class 							_UpperCamelCase			(unittest.TestCase			):


      def  __UpperCAmelCase		(	self					,						__UpperCamelCase					,						__UpperCamelCase	)->	Dict:
           return F"""gaussian_noise_s={seed}_shape={"_".join([str(__UpperCamelCase	) for s in shape]	)}.npy"""


      def  __UpperCAmelCase		(	self	)->	Union[str, Any]:
           # clean up the VRAM after each test
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()


      def  __UpperCAmelCase		(	self					,						__UpperCamelCase=0					,						__UpperCamelCase=(4, 3, 5_1_2, 5_1_2)					,						__UpperCamelCase=False	)->	Dict:
           __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  __UpperCAmelCase		(	self					,						__UpperCamelCase="CompVis/stable-diffusion-v1-4"					,						__UpperCamelCase=False	)->	List[Any]:
           __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  __UpperCAmelCase		(	self					,						__UpperCamelCase=0	)->	List[str]:
           if torch_device == "mps":
                return torch.manual_seed(__UpperCamelCase	)
           return torch.Generator(device=__UpperCamelCase	).manual_seed(__UpperCamelCase	)


      @parameterized.expand(
          [
              # fmt: off
              [3_3, [-0.1_6_0_3, 0.9_8_7_8, -0.0_4_9_5, -0.0_7_9_0, -0.2_7_0_9, 0.8_3_7_5, -0.2_0_6_0, -0.0_8_2_4], [-0.2_3_9_5, 0.0_0_9_8, 0.0_1_0_2, -0.0_7_0_9, -0.2_8_4_0, -0.0_2_7_4, -0.0_7_1_8, -0.1_8_2_4]],
              [4_7, [-0.2_3_7_6, 0.1_1_6_8, 0.1_3_3_2, -0.4_8_4_0, -0.2_5_0_8, -0.0_7_9_1, -0.0_4_9_3, -0.4_0_8_9], [0.0_3_5_0, 0.0_8_4_7, 0.0_4_6_7, 0.0_3_4_4, -0.0_8_4_2, -0.0_5_4_7, -0.0_6_3_3, -0.1_1_3_1]],
              # fmt: on
          ]	)
      def  __UpperCAmelCase		(	self					,						__UpperCamelCase					,						__UpperCamelCase					,						__UpperCamelCase	)->	List[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
              [3_3, [-0.0_5_1_3, 0.0_2_8_9, 1.3_7_9_9, 0.2_1_6_6, -0.2_5_7_3, -0.0_8_7_1, 0.5_1_0_3, -0.0_9_9_9]],
              [4_7, [-0.4_1_2_8, -0.1_3_2_0, -0.3_7_0_4, 0.1_9_6_5, -0.4_1_1_6, -0.2_3_3_2, -0.3_3_4_0, 0.2_2_4_7]],
              # fmt: on
          ]	)
      @require_torch_gpu
      def  __UpperCAmelCase		(	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
              [3_3, [-0.1_6_0_9, 0.9_8_6_6, -0.0_4_8_7, -0.0_7_7_7, -0.2_7_1_6, 0.8_3_6_8, -0.2_0_5_5, -0.0_8_1_4], [-0.2_3_9_5, 0.0_0_9_8, 0.0_1_0_2, -0.0_7_0_9, -0.2_8_4_0, -0.0_2_7_4, -0.0_7_1_8, -0.1_8_2_4]],
              [4_7, [-0.2_3_7_7, 0.1_1_4_7, 0.1_3_3_3, -0.4_8_4_1, -0.2_5_0_6, -0.0_8_0_5, -0.0_4_9_1, -0.4_0_8_5], [0.0_3_5_0, 0.0_8_4_7, 0.0_4_6_7, 0.0_3_4_4, -0.0_8_4_2, -0.0_5_4_7, -0.0_6_3_3, -0.1_1_3_1]],
              # fmt: on
          ]	)
      def  __UpperCAmelCase		(	self					,						__UpperCamelCase					,						__UpperCamelCase					,						__UpperCamelCase	)->	Optional[int]:
           __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
              [1_3, [-0.2_0_5_1, -0.1_8_0_3, -0.2_3_1_1, -0.2_1_1_4, -0.3_2_9_2, -0.3_5_7_4, -0.2_9_5_3, -0.3_3_2_3]],
              [3_7, [-0.2_6_3_2, -0.2_6_2_5, -0.2_1_9_9, -0.2_7_4_1, -0.4_5_3_9, -0.4_9_9_0, -0.3_7_2_0, -0.4_9_2_5]],
              # fmt: on
          ]	)
      @require_torch_gpu
      def  __UpperCAmelCase		(	self					,						__UpperCamelCase					,						__UpperCamelCase	)->	Tuple:
           __lowerCAmelCase								=					self.get_sd_vae_model()
           __lowerCAmelCase								=					self.get_sd_image(__UpperCamelCase					,						shape=(3, 4, 6_4, 6_4)	)

           with torch.no_grad():
                __lowerCAmelCase								=					model.decode(__UpperCamelCase	).sample

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

           __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
              [2_7, [-0.0_3_6_9, 0.0_2_0_7, -0.0_7_7_6, -0.0_6_8_2, -0.1_7_4_7, -0.1_9_3_0, -0.1_4_6_5, -0.2_0_3_9]],
              [1_6, [-0.1_6_2_8, -0.2_1_3_4, -0.2_7_4_7, -0.2_6_4_2, -0.3_7_7_4, -0.4_4_0_4, -0.3_6_8_7, -0.4_2_7_7]],
              # fmt: on
          ]	)
      @require_torch_gpu
      def  __UpperCAmelCase		(	self					,						__UpperCamelCase					,						__UpperCamelCase	)->	List[Any]:
           __lowerCAmelCase								=					self.get_sd_vae_model(fpaa=__UpperCamelCase	)
           __lowerCAmelCase								=					self.get_sd_image(__UpperCamelCase					,						shape=(3, 4, 6_4, 6_4)					,						fpaa=__UpperCamelCase	)

           with torch.no_grad():
                __lowerCAmelCase								=					model.decode(__UpperCamelCase	).sample

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

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

           assert torch_all_close(__UpperCamelCase					,						__UpperCamelCase					,						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  __UpperCAmelCase		(	self					,						__UpperCamelCase	)->	int:
           __lowerCAmelCase								=					self.get_sd_vae_model(fpaa=__UpperCamelCase	)
           __lowerCAmelCase								=					self.get_sd_image(__UpperCamelCase					,						shape=(3, 4, 6_4, 6_4)					,						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, 5_1_2, 5_1_2]

           assert torch_all_close(__UpperCamelCase					,						__UpperCamelCase					,						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  __UpperCAmelCase		(	self					,						__UpperCamelCase	)->	Dict:
           __lowerCAmelCase								=					self.get_sd_vae_model()
           __lowerCAmelCase								=					self.get_sd_image(__UpperCamelCase					,						shape=(3, 4, 6_4, 6_4)	)

           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, 5_1_2, 5_1_2]

           assert torch_all_close(__UpperCamelCase					,						__UpperCamelCase					,						atol=1e-2	)


      @parameterized.expand(
          [
              # fmt: off
              [3_3, [-0.3_0_0_1, 0.0_9_1_8, -2.6_9_8_4, -3.9_7_2_0, -3.2_0_9_9, -5.0_3_5_3, 1.7_3_3_8, -0.2_0_6_5, 3.4_2_6_7]],
              [4_7, [-1.5_0_3_0, -4.3_8_7_1, -6.0_3_5_5, -9.1_1_5_7, -1.6_6_6_1, -2.7_8_5_3, 2.1_6_0_7, -5.0_8_2_3, 2.5_6_3_3]],
              # fmt: on
          ]	)
      def  __UpperCAmelCase		(	self					,						__UpperCamelCase					,						__UpperCamelCase	)->	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 ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase			:					Any         =	logging.get_logger(__name__)

lowerCamelCase			:					Dict         =	{
    '''microsoft/markuplm-base''': '''https://huggingface.co/microsoft/markuplm-base/resolve/main/config.json''',
    '''microsoft/markuplm-large''': '''https://huggingface.co/microsoft/markuplm-large/resolve/main/config.json''',
}



class 							_UpperCamelCase			(a_			):
      snake_case_  =							"""markuplm"""


      def __init__(	self					,						__UpperCamelCase=3_0_5_2_2					,						__UpperCamelCase=7_6_8					,						__UpperCamelCase=1_2					,						__UpperCamelCase=1_2					,						__UpperCamelCase=3_0_7_2					,						__UpperCamelCase="gelu"					,						__UpperCamelCase=0.1					,						__UpperCamelCase=0.1					,						__UpperCamelCase=5_1_2					,						__UpperCamelCase=2					,						__UpperCamelCase=0.0_2					,						__UpperCamelCase=1e-12					,						__UpperCamelCase=0					,						__UpperCamelCase=0					,						__UpperCamelCase=2					,						__UpperCamelCase=2_5_6					,						__UpperCamelCase=1_0_2_4					,						__UpperCamelCase=2_1_6					,						__UpperCamelCase=1_0_0_1					,						__UpperCamelCase=3_2					,						__UpperCamelCase=5_0					,						__UpperCamelCase="absolute"					,						__UpperCamelCase=True					,						__UpperCamelCase=None					,						**__UpperCamelCase					,						)->	Dict:
           super().__init__(
               pad_token_id=__UpperCamelCase					,						bos_token_id=__UpperCamelCase					,						eos_token_id=__UpperCamelCase					,						**__UpperCamelCase					,						)
           __lowerCAmelCase								=					vocab_size
           __lowerCAmelCase								=					hidden_size
           __lowerCAmelCase								=					num_hidden_layers
           __lowerCAmelCase								=					num_attention_heads
           __lowerCAmelCase								=					hidden_act
           __lowerCAmelCase								=					intermediate_size
           __lowerCAmelCase								=					hidden_dropout_prob
           __lowerCAmelCase								=					attention_probs_dropout_prob
           __lowerCAmelCase								=					max_position_embeddings
           __lowerCAmelCase								=					type_vocab_size
           __lowerCAmelCase								=					initializer_range
           __lowerCAmelCase								=					layer_norm_eps
           __lowerCAmelCase								=					position_embedding_type
           __lowerCAmelCase								=					use_cache
           __lowerCAmelCase								=					classifier_dropout
           # additional properties
           __lowerCAmelCase								=					max_depth
           __lowerCAmelCase								=					max_xpath_tag_unit_embeddings
           __lowerCAmelCase								=					max_xpath_subs_unit_embeddings
           __lowerCAmelCase								=					tag_pad_id
           __lowerCAmelCase								=					subs_pad_id
           __lowerCAmelCase								=					xpath_unit_hidden_size






