

from __future__ import annotations

import unittest

from transformers import XGLMConfig, XGLMTokenizer, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import tensorflow as tf

   from transformers.models.xglm.modeling_tf_xglm import (
       TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST,
       TFXGLMForCausalLM,
       TFXGLMModel,
   )



@require_tf
class SCREAMING_SNAKE_CASE     :
  """simple docstring"""






  lowerCamelCase   :      Optional[int]										=XGLMConfig
  lowerCamelCase   :      Optional[Any]										={}
  lowerCamelCase   :      Optional[Any]										="gelu"







  def __init__(      self			:   List[Any]	,							lowerCAmelCase			:   Union[str, Any]	,							lowerCAmelCase			:   List[str]=14	,							lowerCAmelCase			:   Any=7	,							lowerCAmelCase			:   Optional[Any]=True	,							lowerCAmelCase			:   List[Any]=True	,							lowerCAmelCase			:   List[str]=True	,							lowerCAmelCase			:   Optional[int]=99	,							lowerCAmelCase			:   List[str]=32	,							lowerCAmelCase			:   Tuple=2	,							lowerCAmelCase			:   Dict=4	,							lowerCAmelCase			:   Optional[Any]=37	,							lowerCAmelCase			:   str="gelu"	,							lowerCAmelCase			:   Dict=0.1	,							lowerCAmelCase			:   Union[str, Any]=0.1	,							lowerCAmelCase			:   Optional[int]=5_12	,							lowerCAmelCase			:   Optional[int]=0.02	,							)							->     List[str]:






     """simple docstring"""







     __lowerCAmelCase  :      Optional[int]    =						parent
     __lowerCAmelCase  :      str    =						batch_size
     __lowerCAmelCase  :      Any    =						seq_length
     __lowerCAmelCase  :      Optional[int]    =						is_training
     __lowerCAmelCase  :      List[Any]    =						use_input_mask
     __lowerCAmelCase  :      int    =						use_labels
     __lowerCAmelCase  :      Dict    =						vocab_size
     __lowerCAmelCase  :      str    =						d_model
     __lowerCAmelCase  :      Any    =						num_hidden_layers
     __lowerCAmelCase  :      Optional[Any]    =						num_attention_heads
     __lowerCAmelCase  :      Optional[Any]    =						ffn_dim
     __lowerCAmelCase  :      List[str]    =						activation_function
     __lowerCAmelCase  :      Union[str, Any]    =						activation_dropout
     __lowerCAmelCase  :      int    =						attention_dropout
     __lowerCAmelCase  :      List[Any]    =						max_position_embeddings
     __lowerCAmelCase  :      Optional[int]    =						initializer_range
     __lowerCAmelCase  :      List[str]    =						None
     __lowerCAmelCase  :      Any    =						0
     __lowerCAmelCase  :      List[str]    =						2
     __lowerCAmelCase  :      Optional[int]    =						1







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






     """simple docstring"""







     return XGLMConfig.from_pretrained("""facebook/xglm-564M"""   )







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






     """simple docstring"""







     __lowerCAmelCase  :      Tuple    =						tf.clip_by_value(
         ids_tensor([self.batch_size, self.seq_length]	,							self.vocab_size   )	,							clip_value_min=0	,							clip_value_max=3   )

     __lowerCAmelCase  :      Any    =						None
     if self.use_input_mask:
        __lowerCAmelCase  :      Union[str, Any]    =						random_attention_mask([self.batch_size, self.seq_length]   )

     __lowerCAmelCase  :      int    =						self.get_config()

     __lowerCAmelCase  :      List[str]    =						floats_tensor([self.num_hidden_layers, self.num_attention_heads]	,							2   )

     return (
         config,
         input_ids,
         input_mask,
         head_mask,
     )







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






     """simple docstring"""







     return XGLMConfig(
         vocab_size=self.vocab_size	,							d_model=self.hidden_size	,							num_layers=self.num_hidden_layers	,							attention_heads=self.num_attention_heads	,							ffn_dim=self.ffn_dim	,							activation_function=self.activation_function	,							activation_dropout=self.activation_dropout	,							attention_dropout=self.attention_dropout	,							max_position_embeddings=self.max_position_embeddings	,							initializer_range=self.initializer_range	,							use_cache=lowerCAmelCase	,							bos_token_id=self.bos_token_id	,							eos_token_id=self.eos_token_id	,							pad_token_id=self.pad_token_id	,							return_dict=lowerCAmelCase	,							)







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






     """simple docstring"""







     __lowerCAmelCase  :      Tuple    =						self.prepare_config_and_inputs()

     (
         (
         __lowerCAmelCase
     )   ,(
         __lowerCAmelCase
     )   ,(
         __lowerCAmelCase
     )   ,(
         __lowerCAmelCase
     )   ,
     )  :      str    =						config_and_inputs

     __lowerCAmelCase  :      List[str]    =						{
         """input_ids""": input_ids,
         """head_mask""": head_mask,
     }

     return config, inputs_dict






@require_tf
class SCREAMING_SNAKE_CASE     (  a_       ,							a_       ,							unittest.TestCase			):
  """simple docstring"""






  lowerCamelCase   :      Tuple										=(TFXGLMModel, TFXGLMForCausalLM) if is_tf_available() else ()
  lowerCamelCase   :      Dict										=(TFXGLMForCausalLM,) if is_tf_available() else ()
  lowerCamelCase   :      Any										=(
      {"feature-extraction": TFXGLMModel, "text-generation": TFXGLMForCausalLM} if is_tf_available() else {}
  )
  lowerCamelCase   :      Tuple										=False
  lowerCamelCase   :      Optional[int]										=False
  lowerCamelCase   :      Optional[int]										=False







  def 	SCREAMING_SNAKE_CASE   (      self			:   Tuple   )							->     Dict:






     """simple docstring"""







     __lowerCAmelCase  :      List[Any]    =						TFXGLMModelTester(self   )
     __lowerCAmelCase  :      str    =						ConfigTester(self	,							config_class=lowerCAmelCase	,							n_embd=37   )







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






     """simple docstring"""







     self.config_tester.run_common_tests()







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






     """simple docstring"""







     for model_name in TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        __lowerCAmelCase  :      List[str]    =						TFXGLMModel.from_pretrained(lowerCAmelCase   )
        self.assertIsNotNone(lowerCAmelCase   )







  @unittest.skip(reason="""Currently, model embeddings are going to undergo a major refactor."""   )
  def 	SCREAMING_SNAKE_CASE   (      self			:   Optional[Any]   )							->     List[Any]:






     """simple docstring"""







     super().test_resize_token_embeddings()






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






  @slow
  def 	SCREAMING_SNAKE_CASE   (      self			:   int	,							lowerCAmelCase			:   List[Any]=True   )							->     Tuple:






     """simple docstring"""







     __lowerCAmelCase  :      str    =						TFXGLMForCausalLM.from_pretrained("""facebook/xglm-564M"""   )
     __lowerCAmelCase  :      Union[str, Any]    =						tf.convert_to_tensor([[2, 2_68, 98_65]]	,							dtype=tf.intaa   )  # The dog
     # </s> The dog is a very friendly dog. He is very affectionate and loves to play with other
     # fmt: off
     __lowerCAmelCase  :      Dict    =						[2, 2_68, 98_65, 67, 11, 19_88, 5_72_52, 98_65, 5, 9_84, 67, 19_88, 21_38_38, 16_58, 53, 7_04_46, 33, 66_57, 2_78, 15_81]
     # fmt: on
     __lowerCAmelCase  :      Union[str, Any]    =						model.generate(lowerCAmelCase	,							do_sample=lowerCAmelCase	,							num_beams=1   )
     if verify_outputs:
        self.assertListEqual(output_ids[0].numpy().tolist()	,							lowerCAmelCase   )







  @slow
  def 	SCREAMING_SNAKE_CASE   (      self			:   Tuple   )							->     Tuple:






     """simple docstring"""







     __lowerCAmelCase  :      Optional[Any]    =						XGLMTokenizer.from_pretrained("""facebook/xglm-564M"""   )
     __lowerCAmelCase  :      str    =						TFXGLMForCausalLM.from_pretrained("""facebook/xglm-564M"""   )

     tf.random.set_seed(0   )
     __lowerCAmelCase  :      str    =						tokenizer("""Today is a nice day and"""	,							return_tensors="""tf"""   )
     __lowerCAmelCase  :      Union[str, Any]    =						tokenized.input_ids
     # forces the generation to happen on CPU, to avoid GPU-related quirks (and assure same output regardless of the available devices)
     with tf.device(""":/CPU:0"""   ):
        __lowerCAmelCase  :      Optional[int]    =						model.generate(lowerCAmelCase	,							do_sample=lowerCAmelCase	,							seed=[7, 0]   )
     __lowerCAmelCase  :      List[Any]    =						tokenizer.decode(output_ids[0]	,							skip_special_tokens=lowerCAmelCase   )

     __lowerCAmelCase  :      Any    =						(
         """Today is a nice day and warm evening here over Southern Alberta!! Today when they closed schools due"""
     )
     self.assertEqual(lowerCAmelCase	,							lowerCAmelCase   )







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






     """simple docstring"""







     __lowerCAmelCase  :      Optional[Any]    =						TFXGLMForCausalLM.from_pretrained("""facebook/xglm-564M"""   )
     __lowerCAmelCase  :      int    =						XGLMTokenizer.from_pretrained("""facebook/xglm-564M"""   )

     __lowerCAmelCase  :      List[str]    =						"""left"""

     # use different length sentences to test batching
     __lowerCAmelCase  :      Union[str, Any]    =						[
         """This is an extremelly long sentence that only exists to test the ability of the model to cope with """
         """left-padding, such as in batched generation. The output for the sequence below should be the same """
         """regardless of whether left padding is applied or not. When""",
         """Hello, my dog is a little""",
     ]

     __lowerCAmelCase  :      List[str]    =						tokenizer(lowerCAmelCase	,							return_tensors="""tf"""	,							padding=lowerCAmelCase   )
     __lowerCAmelCase  :      str    =						inputs["""input_ids"""]

     __lowerCAmelCase  :      Tuple    =						model.generate(input_ids=lowerCAmelCase	,							attention_mask=inputs["""attention_mask"""]	,							max_new_tokens=12   )

     __lowerCAmelCase  :      str    =						tokenizer(sentences[0]	,							return_tensors="""tf"""   ).input_ids
     __lowerCAmelCase  :      Any    =						model.generate(input_ids=lowerCAmelCase	,							max_new_tokens=12   )

     __lowerCAmelCase  :      Optional[Any]    =						tokenizer(sentences[1]	,							return_tensors="""tf"""   ).input_ids
     __lowerCAmelCase  :      Tuple    =						model.generate(input_ids=lowerCAmelCase	,							max_new_tokens=12   )

     __lowerCAmelCase  :      Dict    =						tokenizer.batch_decode(lowerCAmelCase	,							skip_special_tokens=lowerCAmelCase   )
     __lowerCAmelCase  :      Union[str, Any]    =						tokenizer.decode(output_non_padded[0]	,							skip_special_tokens=lowerCAmelCase   )
     __lowerCAmelCase  :      List[Any]    =						tokenizer.decode(output_padded[0]	,							skip_special_tokens=lowerCAmelCase   )

     __lowerCAmelCase  :      str    =						[
         """This is an extremelly long sentence that only exists to test the ability of the model to cope with """
         """left-padding, such as in batched generation. The output for the sequence below should be the same """
         """regardless of whether left padding is applied or not. When left padding is applied, the sequence will be """
         """a single""",
         """Hello, my dog is a little bit of a shy one, but he is very friendly""",
     ]
     self.assertListEqual(lowerCAmelCase	,							lowerCAmelCase   )
     self.assertListEqual(lowerCAmelCase	,							[non_padded_sentence, padded_sentence]   )





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   )


