

from sklearn.metrics import mean_squared_error

import datasets


__UpperCAmelCase								=			"""\
@article{scikit-learn,
 title={Scikit-learn: Machine Learning in {P}ython},
 author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.
         and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.
         and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and
         Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},
 journal={Journal of Machine Learning Research},
 volume={12},
 pages={2825--2830},
 year={2011}
}
"""

__UpperCAmelCase								=			"""\
Mean Squared Error(MSE) is the average of the square of difference between the predicted
and actual values.
"""


__UpperCAmelCase								=			"""
Args:
    predictions: array-like of shape (n_samples,) or (n_samples, n_outputs)
        Estimated target values.
    references: array-like of shape (n_samples,) or (n_samples, n_outputs)
        Ground truth (correct) target values.
    sample_weight: array-like of shape (n_samples,), default=None
        Sample weights.
    multioutput: {\"raw_values\", \"uniform_average\"} or array-like of shape (n_outputs,), default=\"uniform_average\"
        Defines aggregating of multiple output values. Array-like value defines weights used to average errors.

        \"raw_values\" : Returns a full set of errors in case of multioutput input.

        \"uniform_average\" : Errors of all outputs are averaged with uniform weight.

    squared : bool, default=True
        If True returns MSE value, if False returns RMSE (Root Mean Squared Error) value.

Returns:
    mse : mean squared error.
Examples:

    >>> mse_metric = datasets.load_metric(\"mse\")
    >>> predictions = [2.5, 0.0, 2, 8]
    >>> references = [3, -0.5, 2, 7]
    >>> results = mse_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {'mse': 0.375}
    >>> rmse_result = mse_metric.compute(predictions=predictions, references=references, squared=False)
    >>> print(rmse_result)
    {'mse': 0.6123724356957945}

    If you're using multi-dimensional lists, then set the config as follows :

    >>> mse_metric = datasets.load_metric(\"mse\", \"multilist\")
    >>> predictions = [[0.5, 1], [-1, 1], [7, -6]]
    >>> references = [[0, 2], [-1, 2], [8, -5]]
    >>> results = mse_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {'mse': 0.7083333333333334}
    >>> results = mse_metric.compute(predictions=predictions, references=references, multioutput='raw_values')
    >>> print(results) # doctest: +NORMALIZE_WHITESPACE
    {'mse': array([0.41666667, 1. ])}
"""



@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION       ,							_KWARGS_DESCRIPTION )
class SCREAMING_SNAKE_CASE     (  datasets.Metric ):
  """simple docstring"""






  def 	SCREAMING_SNAKE_CASE   (      self			:   Optional[Any]   )							->     Optional[int]:






     """simple docstring"""







     return datasets.MetricInfo(
         description=_DESCRIPTION	,							citation=_CITATION	,							inputs_description=_KWARGS_DESCRIPTION	,							features=datasets.Features(self._get_feature_types()   )	,							reference_urls=[
             """https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_error.html"""
         ]	,							)







  def 	SCREAMING_SNAKE_CASE   (      self			:   str   )							->     Tuple:






     """simple docstring"""







     if self.config_name == "multilist":
        return {
            "predictions": datasets.Sequence(datasets.Value("""float"""   )   ),
            "references": datasets.Sequence(datasets.Value("""float"""   )   ),
        }
     else:
        return {
            "predictions": datasets.Value("""float"""   ),
            "references": datasets.Value("""float"""   ),
        }







  def 	SCREAMING_SNAKE_CASE   (      self			:   Optional[Any]	,							lowerCAmelCase			:   Optional[int]	,							lowerCAmelCase			:   Optional[int]	,							lowerCAmelCase			:   List[Any]=None	,							lowerCAmelCase			:   List[Any]="uniform_average"	,							lowerCAmelCase			:   Optional[Any]=True   )							->     int:






     """simple docstring"""







     __lowerCAmelCase  :      Union[str, Any]    =						mean_squared_error(
         lowerCAmelCase	,							lowerCAmelCase	,							sample_weight=lowerCAmelCase	,							multioutput=lowerCAmelCase	,							squared=lowerCAmelCase   )

     return {"mse": mse}





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






  lowerCamelCase   :      Tuple										=AutoencoderKL
  lowerCamelCase   :      Tuple										="sample"
  lowerCamelCase   :      Dict										=1e-2







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






     """simple docstring"""







     __lowerCAmelCase  :      str    =						4
     __lowerCAmelCase  :      Dict    =						3
     __lowerCAmelCase  :      Optional[Any]    =						(32, 32)

     __lowerCAmelCase  :      Union[str, Any]    =						floats_tensor((batch_size, num_channels) + sizes   ).to(lowerCAmelCase   )

     return {"sample": image}







  @property
  def 	SCREAMING_SNAKE_CASE   (      self			:   Any   )							->     int:






     """simple docstring"""







     return (3, 32, 32)







  @property
  def 	SCREAMING_SNAKE_CASE   (      self			:   List[str]   )							->     str:






     """simple docstring"""







     return (3, 32, 32)







  def 	SCREAMING_SNAKE_CASE   (      self			:   Tuple   )							->     List[str]:






     """simple docstring"""







     __lowerCAmelCase  :      List[Any]    =						{
         """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  :      Optional[int]    =						self.dummy_input
     return init_dict, inputs_dict







  def 	SCREAMING_SNAKE_CASE   (      self			:   int   )							->     str:






     """simple docstring"""







     pass







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






     """simple docstring"""







     pass







  @unittest.skipIf(torch_device == """mps"""	,							"""Gradient checkpointing skipped on MPS"""   )
  def 	SCREAMING_SNAKE_CASE   (      self			:   Dict   )							->     List[Any]:






     """simple docstring"""







     __lowerCAmelCase   ,__lowerCAmelCase  :      str    =						self.prepare_init_args_and_inputs_for_common()
     __lowerCAmelCase  :      Dict    =						self.model_class(**lowerCAmelCase   )
     model.to(lowerCAmelCase   )

     assert not model.is_gradient_checkpointing and model.training

     __lowerCAmelCase  :      str    =						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()

     __lowerCAmelCase  :      Any    =						torch.randn_like(lowerCAmelCase   )
     __lowerCAmelCase  :      str    =						(out - labels).mean()
     loss.backward()

     # re-instantiate the model now enabling gradient checkpointing
     __lowerCAmelCase  :      List[str]    =						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

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

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






     """simple docstring"""







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

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







  def 	SCREAMING_SNAKE_CASE   (      self			:   Any   )							->     Any:






     """simple docstring"""







     __lowerCAmelCase  :      Dict    =						AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""   )
     __lowerCAmelCase  :      Optional[Any]    =						model.to(lowerCAmelCase   )
     model.eval()

     if torch_device == "mps":
        __lowerCAmelCase  :      List[Any]    =						torch.manual_seed(0   )
     else:
        __lowerCAmelCase  :      Any    =						torch.Generator(device=lowerCAmelCase   ).manual_seed(0   )

     __lowerCAmelCase  :      Optional[int]    =						torch.randn(
         1	,							model.config.in_channels	,							model.config.sample_size	,							model.config.sample_size	,							generator=torch.manual_seed(0   )	,							)
     __lowerCAmelCase  :      Optional[int]    =						image.to(lowerCAmelCase   )
     with torch.no_grad():
        __lowerCAmelCase  :      Union[str, Any]    =						model(lowerCAmelCase	,							sample_posterior=lowerCAmelCase	,							generator=lowerCAmelCase   ).sample

     __lowerCAmelCase  :      Dict    =						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  :      List[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":
        __lowerCAmelCase  :      Union[str, 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(lowerCAmelCase	,							lowerCAmelCase	,							rtol=1e-2   )   )






@slow
class SCREAMING_SNAKE_CASE     (  unittest.TestCase			):
  """simple docstring"""






  def 	SCREAMING_SNAKE_CASE   (      self			:   Union[str, Any]	,							lowerCAmelCase			:   Union[str, Any]	,							lowerCAmelCase			:   str   )							->     int:






     """simple docstring"""







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







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






     """simple docstring"""







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







  def 	SCREAMING_SNAKE_CASE   (      self			:   Any	,							lowerCAmelCase			:   Union[str, Any]=0	,							lowerCAmelCase			:   Any=(4, 3, 5_12, 5_12)	,							lowerCAmelCase			:   Any=False   )							->     Union[str, Any]:






     """simple docstring"""







     __lowerCAmelCase  :      Union[str, Any]    =						torch.floataa if fpaa else torch.floataa
     __lowerCAmelCase  :      Optional[int]    =						torch.from_numpy(load_hf_numpy(self.get_file_format(lowerCAmelCase	,							lowerCAmelCase   )   )   ).to(lowerCAmelCase   ).to(lowerCAmelCase   )
     return image







  def 	SCREAMING_SNAKE_CASE   (      self			:   Union[str, Any]	,							lowerCAmelCase			:   Optional[Any]="CompVis/stable-diffusion-v1-4"	,							lowerCAmelCase			:   int=False   )							->     Union[str, Any]:






     """simple docstring"""







     __lowerCAmelCase  :      Optional[Any]    =						"""fp16""" if fpaa else None
     __lowerCAmelCase  :      List[str]    =						torch.floataa if fpaa else torch.floataa

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

     return model







  def 	SCREAMING_SNAKE_CASE   (      self			:   Tuple	,							lowerCAmelCase			:   Tuple=0   )							->     Tuple:






     """simple docstring"""







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







  @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 	SCREAMING_SNAKE_CASE   (      self			:   Optional[int]	,							lowerCAmelCase			:   Union[str, Any]	,							lowerCAmelCase			:   Optional[int]	,							lowerCAmelCase			:   Dict   )							->     List[Any]:






     """simple docstring"""







     __lowerCAmelCase  :      Dict    =						self.get_sd_vae_model()
     __lowerCAmelCase  :      Optional[int]    =						self.get_sd_image(lowerCAmelCase   )
     __lowerCAmelCase  :      List[str]    =						self.get_generator(lowerCAmelCase   )

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

     assert sample.shape == image.shape

     __lowerCAmelCase  :      Any    =						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(lowerCAmelCase	,							lowerCAmelCase	,							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 	SCREAMING_SNAKE_CASE   (      self			:   Union[str, Any]	,							lowerCAmelCase			:   Optional[Any]	,							lowerCAmelCase			:   Optional[int]   )							->     Any:






     """simple docstring"""







     __lowerCAmelCase  :      List[Any]    =						self.get_sd_vae_model(fpaa=lowerCAmelCase   )
     __lowerCAmelCase  :      Tuple    =						self.get_sd_image(lowerCAmelCase	,							fpaa=lowerCAmelCase   )
     __lowerCAmelCase  :      Optional[int]    =						self.get_generator(lowerCAmelCase   )

     with torch.no_grad():
        __lowerCAmelCase  :      Dict    =						model(lowerCAmelCase	,							generator=lowerCAmelCase	,							sample_posterior=lowerCAmelCase   ).sample

     assert sample.shape == image.shape

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

     assert torch_all_close(lowerCAmelCase	,							lowerCAmelCase	,							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 	SCREAMING_SNAKE_CASE   (      self			:   Any	,							lowerCAmelCase			:   Any	,							lowerCAmelCase			:   List[str]	,							lowerCAmelCase			:   Any   )							->     str:






     """simple docstring"""







     __lowerCAmelCase  :      Union[str, Any]    =						self.get_sd_vae_model()
     __lowerCAmelCase  :      Optional[int]    =						self.get_sd_image(lowerCAmelCase   )

     with torch.no_grad():
        __lowerCAmelCase  :      List[Any]    =						model(lowerCAmelCase   ).sample

     assert sample.shape == image.shape

     __lowerCAmelCase  :      Union[str, Any]    =						sample[-1, -2:, -2:, :2].flatten().float().cpu()
     __lowerCAmelCase  :      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
          [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 	SCREAMING_SNAKE_CASE   (      self			:   Any	,							lowerCAmelCase			:   Optional[Any]	,							lowerCAmelCase			:   int   )							->     str:






     """simple docstring"""







     __lowerCAmelCase  :      Dict    =						self.get_sd_vae_model()
     __lowerCAmelCase  :      Optional[Any]    =						self.get_sd_image(lowerCAmelCase	,							shape=(3, 4, 64, 64)   )

     with torch.no_grad():
        __lowerCAmelCase  :      Optional[Any]    =						model.decode(lowerCAmelCase   ).sample

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

     __lowerCAmelCase  :      Union[str, Any]    =						sample[-1, -2:, :2, -2:].flatten().cpu()
     __lowerCAmelCase  :      Tuple    =						torch.tensor(lowerCAmelCase   )

     assert torch_all_close(lowerCAmelCase	,							lowerCAmelCase	,							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 	SCREAMING_SNAKE_CASE   (      self			:   List[Any]	,							lowerCAmelCase			:   Optional[Any]	,							lowerCAmelCase			:   Any   )							->     List[Any]:






     """simple docstring"""







     __lowerCAmelCase  :      Tuple    =						self.get_sd_vae_model(fpaa=lowerCAmelCase   )
     __lowerCAmelCase  :      str    =						self.get_sd_image(lowerCAmelCase	,							shape=(3, 4, 64, 64)	,							fpaa=lowerCAmelCase   )

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

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

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

     assert torch_all_close(lowerCAmelCase	,							lowerCAmelCase	,							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 	SCREAMING_SNAKE_CASE   (      self			:   Optional[Any]	,							lowerCAmelCase			:   Any   )							->     Any:






     """simple docstring"""







     __lowerCAmelCase  :      List[Any]    =						self.get_sd_vae_model(fpaa=lowerCAmelCase   )
     __lowerCAmelCase  :      Union[str, Any]    =						self.get_sd_image(lowerCAmelCase	,							shape=(3, 4, 64, 64)	,							fpaa=lowerCAmelCase   )

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

     model.enable_xformers_memory_efficient_attention()
     with torch.no_grad():
        __lowerCAmelCase  :      int    =						model.decode(lowerCAmelCase   ).sample

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

     assert torch_all_close(lowerCAmelCase	,							lowerCAmelCase	,							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 	SCREAMING_SNAKE_CASE   (      self			:   Dict	,							lowerCAmelCase			:   List[Any]   )							->     Dict:






     """simple docstring"""







     __lowerCAmelCase  :      Optional[int]    =						self.get_sd_vae_model()
     __lowerCAmelCase  :      Optional[Any]    =						self.get_sd_image(lowerCAmelCase	,							shape=(3, 4, 64, 64)   )

     with torch.no_grad():
        __lowerCAmelCase  :      Optional[Any]    =						model.decode(lowerCAmelCase   ).sample

     model.enable_xformers_memory_efficient_attention()
     with torch.no_grad():
        __lowerCAmelCase  :      Tuple    =						model.decode(lowerCAmelCase   ).sample

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

     assert torch_all_close(lowerCAmelCase	,							lowerCAmelCase	,							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 	SCREAMING_SNAKE_CASE   (      self			:   Any	,							lowerCAmelCase			:   int	,							lowerCAmelCase			:   List[Any]   )							->     Any:






     """simple docstring"""







     __lowerCAmelCase  :      Optional[Any]    =						self.get_sd_vae_model()
     __lowerCAmelCase  :      List[str]    =						self.get_sd_image(lowerCAmelCase   )
     __lowerCAmelCase  :      Any    =						self.get_generator(lowerCAmelCase   )

     with torch.no_grad():
        __lowerCAmelCase  :      Optional[int]    =						model.encode(lowerCAmelCase   ).latent_dist
        __lowerCAmelCase  :      Union[str, Any]    =						dist.sample(generator=lowerCAmelCase   )

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

     __lowerCAmelCase  :      Any    =						sample[0, -1, -3:, -3:].flatten().cpu()
     __lowerCAmelCase  :      int    =						torch.tensor(lowerCAmelCase   )

     __lowerCAmelCase  :      str    =						3e-3 if torch_device != """mps""" else 1e-2
     assert torch_all_close(lowerCAmelCase	,							lowerCAmelCase	,							atol=lowerCAmelCase   )


