





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  ):
  __UpperCamelCase	     =  AutoencoderKL
  __UpperCamelCase	     =  """sample"""
  __UpperCamelCase	     =  1E-2







  @property
  def 	UpperCAmelCase__		(	self						):
       lowerCAmelCase_					:       List[Any]    =		4
       lowerCAmelCase_					:       Any    =		3
       lowerCAmelCase_					:       int    =		(32, 32)

       lowerCAmelCase_					:       Tuple    =		floats_tensor((batch_size, num_channels) + sizes						).to(_lowercase						)

       return {"sample": image}







  @property
  def 	UpperCAmelCase__		(	self						):
       return (3, 32, 32)







  @property
  def 	UpperCAmelCase__		(	self						):
       return (3, 32, 32)







  def 	UpperCAmelCase__		(	self						):
       lowerCAmelCase_					:       int    =		{
           """block_out_channels""": [32, 64],
           """in_channels""": 3,
           """out_channels""": 3,
           """down_block_types""": ["""DownEncoderBlock2D""", """DownEncoderBlock2D"""],
           """up_block_types""": ["""UpDecoderBlock2D""", """UpDecoderBlock2D"""],
           """latent_channels""": 4,
       }
       lowerCAmelCase_					:       Dict    =		self.dummy_input
       return init_dict, inputs_dict







  def 	UpperCAmelCase__		(	self						):
       pass







  def 	UpperCAmelCase__		(	self						):
       pass







  @unittest.skipIf(torch_device == """mps"""    ,						"""Gradient checkpointing skipped on MPS"""						)
  def 	UpperCAmelCase__		(	self						):
       # enable deterministic behavior for gradient checkpointing
       lowerCAmelCase_   ,  lowerCAmelCase_					:       Optional[Any]    =		self.prepare_init_args_and_inputs_for_common()
       lowerCAmelCase_					:       List[str]    =		self.model_class(**_lowercase						)
       model.to(_lowercase						)

       assert not model.is_gradient_checkpointing and model.training

       lowerCAmelCase_					:       Union[str, Any]    =		model(**_lowercase						).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_					:       Tuple    =		torch.randn_like(_lowercase						)
       lowerCAmelCase_					:       Optional[Any]    =		(out - labels).mean()
       loss.backward()

       # re-instantiate the model now enabling gradient checkpointing
       lowerCAmelCase_					:       List[str]    =		self.model_class(**_lowercase						)
       # clone model
       model_a.load_state_dict(model.state_dict()						)
       model_a.to(_lowercase						)
       model_a.enable_gradient_checkpointing()

       assert model_a.is_gradient_checkpointing and model_a.training

       lowerCAmelCase_					:       str    =		model_a(**_lowercase						).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_					:       Any    =		(out_a - labels).mean()
       loss_a.backward()

       # compare the output and parameters gradients
       self.assertTrue((loss - loss_a).abs() < 1e-5						)
       lowerCAmelCase_					:       Dict    =		dict(model.named_parameters()						)
       lowerCAmelCase_					:       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						):
       lowerCAmelCase_   ,  lowerCAmelCase_					:       str    =		AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""    ,						output_loading_info=_lowercase						)
       self.assertIsNotNone(_lowercase						)
       self.assertEqual(len(loading_info["""missing_keys"""]						)    ,						0						)

       model.to(_lowercase						)
       lowerCAmelCase_					:       Union[str, Any]    =		model(**self.dummy_input						)

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







  def 	UpperCAmelCase__		(	self						):
       lowerCAmelCase_					:       List[Any]    =		AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""						)
       lowerCAmelCase_					:       Optional[Any]    =		model.to(_lowercase						)
       model.eval()

       if torch_device == "mps":
            lowerCAmelCase_					:       int    =		torch.manual_seed(0						)
       else:
            lowerCAmelCase_					:       List[Any]    =		torch.Generator(device=_lowercase						).manual_seed(0						)

       lowerCAmelCase_					:       Tuple    =		torch.randn(
           1    ,						model.config.in_channels    ,						model.config.sample_size    ,						model.config.sample_size    ,						generator=torch.manual_seed(0						)    ,						)
       lowerCAmelCase_					:       Tuple    =		image.to(_lowercase						)
       with torch.no_grad():
            lowerCAmelCase_					:       int    =		model(_lowercase    ,						sample_posterior=_lowercase    ,						generator=_lowercase						).sample

       lowerCAmelCase_					:       Union[str, 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":
            lowerCAmelCase_					:       Tuple    =		torch.tensor(
                [
                    -4.0_078e-01,
                    -3.8_323e-04,
                    -1.2_681e-01,
                    -1.1_462e-01,
                    2.0_095e-01,
                    1.0_893e-01,
                    -8.8_247e-02,
                    -3.0_361e-01,
                    -9.8_644e-03,
                ]						)
       elif torch_device == "cpu":
            lowerCAmelCase_					:       Optional[Any]    =		torch.tensor(
                [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026]						)
       else:
            lowerCAmelCase_					:       Tuple    =		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(_lowercase    ,						_lowercase    ,						rtol=1e-2						)						)





@slow
class      lowercase__    (							unittest.TestCase  ):







  def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase						):
       return F'gaussian_noise_s={seed}_shape={"_".join([str(_lowercase						) for s in shape]						)}.npy'







  def 	UpperCAmelCase__		(	self						):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()







  def 	UpperCAmelCase__		(	self    ,						_lowercase=0    ,						_lowercase=(4, 3, 512, 512)    ,						_lowercase=False						):
       lowerCAmelCase_					:       List[Any]    =		torch.floataa if fpaa else torch.floataa
       lowerCAmelCase_					:       int    =		torch.from_numpy(load_hf_numpy(self.get_file_format(_lowercase    ,						_lowercase						)						)						).to(_lowercase						).to(_lowercase						)
       return image







  def 	UpperCAmelCase__		(	self    ,						_lowercase="CompVis/stable-diffusion-v1-4"    ,						_lowercase=False						):
       lowerCAmelCase_					:       Any    =		"""fp16""" if fpaa else None
       lowerCAmelCase_					:       List[str]    =		torch.floataa if fpaa else torch.floataa

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

       return model







  def 	UpperCAmelCase__		(	self    ,						_lowercase=0						):
       if torch_device == "mps":
            return torch.manual_seed(_lowercase						)
       return torch.Generator(device=_lowercase						).manual_seed(_lowercase						)







  @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 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase    ,						_lowercase						):
       lowerCAmelCase_					:       int    =		self.get_sd_vae_model()
       lowerCAmelCase_					:       Any    =		self.get_sd_image(_lowercase						)
       lowerCAmelCase_					:       str    =		self.get_generator(_lowercase						)

       with torch.no_grad():
            lowerCAmelCase_					:       List[str]    =		model(_lowercase    ,						generator=_lowercase    ,						sample_posterior=_lowercase						).sample

       assert sample.shape == image.shape

       lowerCAmelCase_					:       int    =		sample[-1, -2:, -2:, :2].flatten().float().cpu()
       lowerCAmelCase_					:       List[str]    =		torch.tensor(expected_slice_mps if torch_device == """mps""" else expected_slice						)

       assert torch_all_close(_lowercase    ,						_lowercase    ,						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 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase						):
       lowerCAmelCase_					:       Optional[Any]    =		self.get_sd_vae_model(fpaa=_lowercase						)
       lowerCAmelCase_					:       Any    =		self.get_sd_image(_lowercase    ,						fpaa=_lowercase						)
       lowerCAmelCase_					:       Dict    =		self.get_generator(_lowercase						)

       with torch.no_grad():
            lowerCAmelCase_					:       Optional[Any]    =		model(_lowercase    ,						generator=_lowercase    ,						sample_posterior=_lowercase						).sample

       assert sample.shape == image.shape

       lowerCAmelCase_					:       List[str]    =		sample[-1, -2:, :2, -2:].flatten().float().cpu()
       lowerCAmelCase_					:       str    =		torch.tensor(_lowercase						)

       assert torch_all_close(_lowercase    ,						_lowercase    ,						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 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase    ,						_lowercase						):
       lowerCAmelCase_					:       List[str]    =		self.get_sd_vae_model()
       lowerCAmelCase_					:       Optional[int]    =		self.get_sd_image(_lowercase						)

       with torch.no_grad():
            lowerCAmelCase_					:       Dict    =		model(_lowercase						).sample

       assert sample.shape == image.shape

       lowerCAmelCase_					:       Optional[Any]    =		sample[-1, -2:, -2:, :2].flatten().float().cpu()
       lowerCAmelCase_					:       Optional[Any]    =		torch.tensor(expected_slice_mps if torch_device == """mps""" else expected_slice						)

       assert torch_all_close(_lowercase    ,						_lowercase    ,						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 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase						):
       lowerCAmelCase_					:       Dict    =		self.get_sd_vae_model()
       lowerCAmelCase_					:       Optional[int]    =		self.get_sd_image(_lowercase    ,						shape=(3, 4, 64, 64)						)

       with torch.no_grad():
            lowerCAmelCase_					:       Any    =		model.decode(_lowercase						).sample

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

       lowerCAmelCase_					:       Tuple    =		sample[-1, -2:, :2, -2:].flatten().cpu()
       lowerCAmelCase_					:       str    =		torch.tensor(_lowercase						)

       assert torch_all_close(_lowercase    ,						_lowercase    ,						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 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase						):
       lowerCAmelCase_					:       Union[str, Any]    =		self.get_sd_vae_model(fpaa=_lowercase						)
       lowerCAmelCase_					:       Union[str, Any]    =		self.get_sd_image(_lowercase    ,						shape=(3, 4, 64, 64)    ,						fpaa=_lowercase						)

       with torch.no_grad():
            lowerCAmelCase_					:       List[str]    =		model.decode(_lowercase						).sample

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

       lowerCAmelCase_					:       int    =		sample[-1, -2:, :2, -2:].flatten().float().cpu()
       lowerCAmelCase_					:       Union[str, Any]    =		torch.tensor(_lowercase						)

       assert torch_all_close(_lowercase    ,						_lowercase    ,						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    ,						_lowercase						):
       lowerCAmelCase_					:       Optional[int]    =		self.get_sd_vae_model(fpaa=_lowercase						)
       lowerCAmelCase_					:       List[Any]    =		self.get_sd_image(_lowercase    ,						shape=(3, 4, 64, 64)    ,						fpaa=_lowercase						)

       with torch.no_grad():
            lowerCAmelCase_					:       List[Any]    =		model.decode(_lowercase						).sample

       model.enable_xformers_memory_efficient_attention()
       with torch.no_grad():
            lowerCAmelCase_					:       int    =		model.decode(_lowercase						).sample

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

       assert torch_all_close(_lowercase    ,						_lowercase    ,						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    ,						_lowercase						):
       lowerCAmelCase_					:       int    =		self.get_sd_vae_model()
       lowerCAmelCase_					:       Any    =		self.get_sd_image(_lowercase    ,						shape=(3, 4, 64, 64)						)

       with torch.no_grad():
            lowerCAmelCase_					:       Tuple    =		model.decode(_lowercase						).sample

       model.enable_xformers_memory_efficient_attention()
       with torch.no_grad():
            lowerCAmelCase_					:       List[str]    =		model.decode(_lowercase						).sample

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

       assert torch_all_close(_lowercase    ,						_lowercase    ,						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 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase						):
       lowerCAmelCase_					:       Dict    =		self.get_sd_vae_model()
       lowerCAmelCase_					:       Any    =		self.get_sd_image(_lowercase						)
       lowerCAmelCase_					:       int    =		self.get_generator(_lowercase						)

       with torch.no_grad():
            lowerCAmelCase_					:       List[Any]    =		model.encode(_lowercase						).latent_dist
            lowerCAmelCase_					:       List[str]    =		dist.sample(generator=_lowercase						)

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

       lowerCAmelCase_					:       Optional[int]    =		sample[0, -1, -3:, -3:].flatten().cpu()
       lowerCAmelCase_					:       Dict    =		torch.tensor(_lowercase						)

       lowerCAmelCase_					:       List[str]    =		3e-3 if torch_device != """mps""" else 1e-2
       assert torch_all_close(_lowercase    ,						_lowercase    ,						atol=_lowercase						)







# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


UpperCAmelCase_					:   Tuple 							=						{
    """configuration_vivit""": ["""VIVIT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """VivitConfig"""],
}
try:
      if not is_vision_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      UpperCAmelCase_					:   Tuple 							=						["""VivitImageProcessor"""]


try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      UpperCAmelCase_					:   Any 							=						[
          """VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """VivitModel""",
          """VivitPreTrainedModel""",
          """VivitForVideoClassification""",
      ]


if TYPE_CHECKING:
      from .configuration_vivit import VIVIT_PRETRAINED_CONFIG_ARCHIVE_MAP, VivitConfig

      try:
            if not is_vision_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .image_processing_vivit import VivitImageProcessor

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_vivit import (
                VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
                VivitForVideoClassification,
                VivitModel,
                VivitPreTrainedModel,
            )


else:
      import sys

      UpperCAmelCase_					:   Dict 							=						_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
