

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



    '''simple docstring'''




    _lowerCamelCase					=   AutoencoderKL
    _lowerCamelCase					=   '''sample'''
    _lowerCamelCase					=   1E-2




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

          A				  =		floats_tensor((batch_size, num_channels) + sizes    ).to(lowerCamelCase_    )

          return {"sample": image}




    @property
    def UpperCamelCase__		( self    )    ->      List[Any]:
          return (3, 3_2, 3_2)




    @property
    def UpperCamelCase__		( self    )    ->      List[Any]:
          return (3, 3_2, 3_2)




    def UpperCamelCase__		( self    )    ->      Any:
          A				  =		{
              """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				  =		self.dummy_input
          return init_dict, inputs_dict




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




    def UpperCamelCase__		( self    )    ->      List[Any]:
          pass




    @unittest.skipIf(torch_device == """mps""" ,"""Gradient checkpointing skipped on MPS"""    )
    def UpperCamelCase__		( self    )    ->      List[str]:
          # enable deterministic behavior for gradient checkpointing
          A  ,			A				  =		self.prepare_init_args_and_inputs_for_common()
          A				  =		self.model_class(**lowerCamelCase_    )
          model.to(lowerCamelCase_    )

          assert not model.is_gradient_checkpointing and model.training

          A				  =		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				  =		torch.randn_like(lowerCamelCase_    )
          A				  =		(out - labels).mean()
          loss.backward()

          # re-instantiate the model now enabling gradient checkpointing
          A				  =		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				  =		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				  =		(out_a - labels).mean()
          loss_a.backward()

          # compare the output and parameters gradients
          self.assertTrue((loss - loss_a).abs() < 1E-5    )
          A				  =		dict(model.named_parameters()    )
          A				  =		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    )    ->      Tuple:
          A  ,			A				  =		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				  =		model(**self.dummy_input    )

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




    def UpperCamelCase__		( self    )    ->      Union[str, Any]:
          A				  =		AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""    )
          A				  =		model.to(lowerCamelCase_    )
          model.eval()

          if torch_device == "mps":
                A				  =		torch.manual_seed(0    )
          else:
                A				  =		torch.Generator(device=lowerCamelCase_    ).manual_seed(0    )

          A				  =		torch.randn(
              1 ,model.config.in_channels ,model.config.sample_size ,model.config.sample_size ,generator=torch.manual_seed(0    ) ,)
          A				  =		image.to(lowerCamelCase_    )
          with torch.no_grad():
                A				  =		model(lowerCamelCase_ ,sample_posterior=lowerCamelCase_ ,generator=lowerCamelCase_    ).sample

          A				  =		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				  =		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				  =		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				  =		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    ):



    '''simple docstring'''




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_    )    ->      List[Any]:
          return f'gaussian_noise_s={seed}_shape={"_".join([str(lowerCamelCase_    ) for s in shape]    )}.npy'




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




    def UpperCamelCase__		( self ,lowerCamelCase_=0 ,lowerCamelCase_=(4, 3, 5_1_2, 5_1_2) ,lowerCamelCase_=False    )    ->      Any:
          A				  =		torch.floataa if fpaa else torch.floataa
          A				  =		torch.from_numpy(load_hf_numpy(self.get_file_format(lowerCamelCase_ ,lowerCamelCase_    )    )    ).to(lowerCamelCase_    ).to(lowerCamelCase_    )
          return image




    def UpperCamelCase__		( self ,lowerCamelCase_="CompVis/stable-diffusion-v1-4" ,lowerCamelCase_=False    )    ->      str:
          A				  =		"""fp16""" if fpaa else None
          A				  =		torch.floataa if fpaa else torch.floataa

          A				  =		AutoencoderKL.from_pretrained(
              lowerCamelCase_ ,subfolder="""vae""" ,torch_dtype=lowerCamelCase_ ,revision=lowerCamelCase_ ,)
          model.to(lowerCamelCase_    ).eval()

          return model




    def UpperCamelCase__		( self ,lowerCamelCase_=0    )    ->      List[str]:
          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 UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )    ->      Any:
          A				  =		self.get_sd_vae_model()
          A				  =		self.get_sd_image(lowerCamelCase_    )
          A				  =		self.get_generator(lowerCamelCase_    )

          with torch.no_grad():
                A				  =		model(lowerCamelCase_ ,generator=lowerCamelCase_ ,sample_posterior=lowerCamelCase_    ).sample

          assert sample.shape == image.shape

          A				  =		sample[-1, -2:, -2:, :2].flatten().float().cpu()
          A				  =		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 UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_    )    ->      str:
          A				  =		self.get_sd_vae_model(fpaa=lowerCamelCase_    )
          A				  =		self.get_sd_image(lowerCamelCase_ ,fpaa=lowerCamelCase_    )
          A				  =		self.get_generator(lowerCamelCase_    )

          with torch.no_grad():
                A				  =		model(lowerCamelCase_ ,generator=lowerCamelCase_ ,sample_posterior=lowerCamelCase_    ).sample

          assert sample.shape == image.shape

          A				  =		sample[-1, -2:, :2, -2:].flatten().float().cpu()
          A				  =		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 UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )    ->      str:
          A				  =		self.get_sd_vae_model()
          A				  =		self.get_sd_image(lowerCamelCase_    )

          with torch.no_grad():
                A				  =		model(lowerCamelCase_    ).sample

          assert sample.shape == image.shape

          A				  =		sample[-1, -2:, -2:, :2].flatten().float().cpu()
          A				  =		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 UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_    )    ->      str:
          A				  =		self.get_sd_vae_model()
          A				  =		self.get_sd_image(lowerCamelCase_ ,shape=(3, 4, 6_4, 6_4)    )

          with torch.no_grad():
                A				  =		model.decode(lowerCamelCase_    ).sample

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

          A				  =		sample[-1, -2:, :2, -2:].flatten().cpu()
          A				  =		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 UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_    )    ->      Optional[int]:
          A				  =		self.get_sd_vae_model(fpaa=lowerCamelCase_    )
          A				  =		self.get_sd_image(lowerCamelCase_ ,shape=(3, 4, 6_4, 6_4) ,fpaa=lowerCamelCase_    )

          with torch.no_grad():
                A				  =		model.decode(lowerCamelCase_    ).sample

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

          A				  =		sample[-1, -2:, :2, -2:].flatten().float().cpu()
          A				  =		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 UpperCamelCase__		( self ,lowerCamelCase_    )    ->      Dict:
          A				  =		self.get_sd_vae_model(fpaa=lowerCamelCase_    )
          A				  =		self.get_sd_image(lowerCamelCase_ ,shape=(3, 4, 6_4, 6_4) ,fpaa=lowerCamelCase_    )

          with torch.no_grad():
                A				  =		model.decode(lowerCamelCase_    ).sample

          model.enable_xformers_memory_efficient_attention()
          with torch.no_grad():
                A				  =		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 UpperCamelCase__		( self ,lowerCamelCase_    )    ->      List[Any]:
          A				  =		self.get_sd_vae_model()
          A				  =		self.get_sd_image(lowerCamelCase_ ,shape=(3, 4, 6_4, 6_4)    )

          with torch.no_grad():
                A				  =		model.decode(lowerCamelCase_    ).sample

          model.enable_xformers_memory_efficient_attention()
          with torch.no_grad():
                A				  =		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 UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_    )    ->      int:
          A				  =		self.get_sd_vae_model()
          A				  =		self.get_sd_image(lowerCamelCase_    )
          A				  =		self.get_generator(lowerCamelCase_    )

          with torch.no_grad():
                A				  =		model.encode(lowerCamelCase_    ).latent_dist
                A				  =		dist.sample(generator=lowerCamelCase_    )

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

          A				  =		sample[0, -1, -3:, -3:].flatten().cpu()
          A				  =		torch.tensor(lowerCamelCase_    )

          A				  =		3E-3 if torch_device != """mps""" else 1E-2
          assert torch_all_close(lowerCamelCase_ ,lowerCamelCase_ ,atol=lowerCamelCase_    )







"""simple docstring"""


from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


UpperCAmelCase													={
    "configuration_mega": ["MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP", "MegaConfig", "MegaOnnxConfig"],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase													=[
      "MEGA_PRETRAINED_MODEL_ARCHIVE_LIST",
      "MegaForCausalLM",
      "MegaForMaskedLM",
      "MegaForMultipleChoice",
      "MegaForQuestionAnswering",
      "MegaForSequenceClassification",
      "MegaForTokenClassification",
      "MegaModel",
      "MegaPreTrainedModel",
  ]

if TYPE_CHECKING:
  from .configuration_mega import MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP, MegaConfig, MegaOnnxConfig

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_mega import (
        MEGA_PRETRAINED_MODEL_ARCHIVE_LIST,
        MegaForCausalLM,
        MegaForMaskedLM,
        MegaForMultipleChoice,
        MegaForQuestionAnswering,
        MegaForSequenceClassification,
        MegaForTokenClassification,
        MegaModel,
        MegaPreTrainedModel,
    )

else:
  import sys

  UpperCAmelCase													=_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)




