




'''simple docstring'''

import math
import os
import re
import sys
import unittest
from pathlib import Path
from typing import Tuple
from unittest.mock import patch

from parameterized import parameterized

from transformers.testing_utils import (
    CaptureStderr,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    get_torch_dist_unique_port,
    require_apex,
    require_bitsandbytes,
    require_fairscale,
    require_torch,
    require_torch_gpu,
    require_torch_multi_gpu,
    require_torch_non_multi_gpu,
    slow,
)
from transformers.trainer_callback import TrainerState
from transformers.trainer_utils import set_seed


a  :				Union[str, Any]									=  os.path.abspath(os.path.dirname(__file__))
with ExtendSysPath(f'''{bindir}/../../examples/pytorch/translation'''):
							from run_translation import main  # noqa


set_seed(42)
a  :				List[str]									=  'sshleifer/student_marian_en_ro_6_1'
a  :				Dict									=  'sshleifer/tiny-mbart'






@require_torch
class 						a				(   _lowerCamelCase     ):


				def   A_ (		self    :					int       ,				lowercase_    :					Any=False       ,				lowercase_    :					int=None       ,				lowercase_    :					List[str]=True       ,				lowercase_    :					Union[str, Any]=True       ,				lowercase_    :					Tuple=True       ,				lowercase_    :					str=True       ,				):
									snake_case_									=				self.run_trainer(
									    eval_steps=1       ,				max_len=12       ,				model_name=lowercase_       ,				num_train_epochs=1       ,				distributed=lowercase_       ,				extra_args_str=lowercase_       ,				predict_with_generate=lowercase_       ,				do_train=lowercase_       ,				do_eval=lowercase_       ,				do_predict=lowercase_       ,				)
									snake_case_									=				TrainerState.load_from_json(os.path.join(lowercase_       ,				'''trainer_state.json'''					)					).log_history

									if not do_eval:
														return

									snake_case_									=				[log for log in logs if '''eval_loss''' in log.keys()]

									snake_case_									=				eval_metrics[0]
									if predict_with_generate:
														assert "eval_bleu" in first_step_stats

														snake_case_									=				eval_metrics[-1]
														assert isinstance(last_step_stats['''eval_bleu''']       ,				lowercase_					)
														assert not math.isnan(float(last_step_stats['''eval_loss''']					)					), "eval_loss must not be `nan`"


				@require_torch_non_multi_gpu
				def   A_ (		self    :					Any					):
									self.run_seqaseq_quick()


				@require_torch_multi_gpu
				def   A_ (		self    :					Optional[int]					):
									self.run_seqaseq_quick(distributed=lowercase_					)


				@require_torch_multi_gpu
				def   A_ (		self    :					str					):
									self.run_seqaseq_quick(distributed=lowercase_					)


				@unittest.skip('''Requires an update of the env running those tests'''					)
				@require_torch_multi_gpu
				@require_fairscale
				def   A_ (		self    :					int					):
									self.run_seqaseq_quick(distributed=lowercase_       ,				extra_args_str='''--sharded_ddp simple'''					)


				@unittest.skip('''Requires an update of the env running those tests'''					)
				@require_torch_multi_gpu
				@require_fairscale
				def   A_ (		self    :					str					):
									self.run_seqaseq_quick(distributed=lowercase_       ,				extra_args_str='''--sharded_ddp simple --fp16'''					)


				@unittest.skip('''Requires an update of the env running those tests'''					)
				@require_torch_multi_gpu
				@require_fairscale
				def   A_ (		self    :					Tuple					):
									self.run_seqaseq_quick(distributed=lowercase_       ,				extra_args_str='''--sharded_ddp zero_dp_2'''       ,				predict_with_generate=lowercase_					)


				@unittest.skip('''Requires an update of the env running those tests'''					)
				@require_torch_multi_gpu
				@require_fairscale
				def   A_ (		self    :					int					):
									self.run_seqaseq_quick(
									    distributed=lowercase_       ,				extra_args_str='''--sharded_ddp zero_dp_2 --fp16'''       ,				predict_with_generate=lowercase_					)


				@require_apex
				@require_torch_gpu
				def   A_ (		self    :					int					):
									# XXX: apex breaks the trainer if it's run twice e.g. run_seq2seq.main() from the same
									# program and it breaks other tests that run from the same pytest worker, therefore until this is
									# sorted out it must be run only in an external program, that is distributed=True in this
									# test and only under one or more gpus - if we want cpu will need to make a special test
									#
									# specifically to the problem traced it to self.optimizer.step() - if it's run 2nd time via
									# 2nd main() call it botches the future eval.
									#
									self.run_seqaseq_quick(distributed=lowercase_       ,				extra_args_str='''--fp16 --fp16_backend=apex'''					)
									# test 2nd time - was getting eval_loss': nan'
									# to reproduce the problem set distributed=False
									self.run_seqaseq_quick(distributed=lowercase_       ,				extra_args_str='''--fp16 --fp16_backend=apex'''					)


				@parameterized.expand(['''base''', '''low''', '''high''', '''mixed''']					)
				@require_torch_multi_gpu
				def   A_ (		self    :					Optional[int]       ,				lowercase_    :					Dict					):
									# as each sub-test is slow-ish split into multiple sub-tests to avoid CI timeout
									snake_case_									=				{
									    # test with the default log_level - should be info and thus log info once
									    '''base''': {'''extra_args_str''': '''''', '''n_matches''': 1},
									    # test with low log_level and log_level_replica - should be noisy on all processes
									    # now the info string should appear twice on 2 processes
									    '''low''': {'''extra_args_str''': '''--log_level debug --log_level_replica debug''', '''n_matches''': 2},
									    # test with high log_level and low log_level_replica
									    # now the info string should appear once only on the replica
									    '''high''': {'''extra_args_str''': '''--log_level error --log_level_replica debug''', '''n_matches''': 1},
									    # test with high log_level and log_level_replica - should be quiet on all processes
									    '''mixed''': {'''extra_args_str''': '''--log_level error --log_level_replica error''', '''n_matches''': 0},
									}

									snake_case_									=				experiments[experiment_id]
									snake_case_									=				{'''distributed''': True, '''predict_with_generate''': False, '''do_eval''': False, '''do_predict''': False}
									snake_case_									=				'''Running training'''
									with CaptureStderr() as cl:
														self.run_seqaseq_quick(**lowercase_       ,				extra_args_str=data['''extra_args_str''']					)
									snake_case_									=				len(re.findall(lowercase_       ,				cl.err					)					)
									self.assertEqual(lowercase_       ,				data['''n_matches''']					)


				@slow
				def   A_ (		self    :					Optional[Any]					):
									snake_case_									=				self.run_trainer(
									    eval_steps=2       ,				max_len=128       ,				model_name=lowercase_       ,				learning_rate=3e-4       ,				num_train_epochs=10       ,				distributed=lowercase_       ,				)

									# Check metrics
									snake_case_									=				TrainerState.load_from_json(os.path.join(lowercase_       ,				'''trainer_state.json'''					)					).log_history
									snake_case_									=				[log for log in logs if '''eval_loss''' in log.keys()]
									snake_case_									=				eval_metrics[0]
									snake_case_									=				eval_metrics[-1]

									assert first_step_stats["eval_loss"] > last_step_stats["eval_loss"], "model learned nothing"
									assert isinstance(last_step_stats['''eval_bleu''']       ,				lowercase_					)

									# test if do_predict saves generations and metrics
									snake_case_									=				os.listdir(lowercase_					)
									snake_case_									=				{os.path.basename(lowercase_					) for p in contents}
									assert "generated_predictions.txt" in contents
									assert "predict_results.json" in contents


				@slow
				@require_bitsandbytes
				def   A_ (		self    :					Tuple					):
									from transformers.training_args import OptimizerNames

									def train_and_return_metrics(lowercase_    :					str					) -> Tuple[int, float]:
														snake_case_									=				'''--skip_memory_metrics 0'''

														snake_case_									=				self.run_trainer(
														    max_len=128       ,				model_name=lowercase_       ,				learning_rate=3e-4       ,				num_train_epochs=1       ,				optim=lowercase_       ,				distributed=lowercase_       ,				extra_args_str=lowercase_       ,				do_eval=lowercase_       ,				do_predict=lowercase_       ,				n_gpus_to_use=1       ,				)

														# Check metrics
														snake_case_									=				TrainerState.load_from_json(Path(lowercase_       ,				'''trainer_state.json'''					)					).log_history
														snake_case_									=				int(logs[0]['''train_mem_gpu_peaked_delta'''] / 2**20					)
														snake_case_									=				int(logs[0]['''train_mem_gpu_alloc_delta'''] / 2**20					)

														snake_case_									=				logs[0]['''train_loss''']
														return gpu_peak_mem_mb, gpu_alloc_mem_mb, loss

									snake_case_			,snake_case_			,snake_case_									=				train_and_return_metrics(OptimizerNames.ADAMW_TORCH.value					)
									snake_case_			,snake_case_			,snake_case_									=				train_and_return_metrics(OptimizerNames.ADAMW_BNB.value					)

									snake_case_									=				gpu_alloc_mem_orig - gpu_alloc_mem_bnb

									snake_case_									=				gpu_peak_mem_orig + gpu_alloc_mem_orig
									snake_case_									=				gpu_peak_mem_bnb + gpu_alloc_mem_bnb
									snake_case_									=				gpu_total_mem_orig - gpu_total_mem_bnb

									# sshleifer/student_marian_en_ro_6_1 has 54M parameter, 29M of which is `nn.Embedding` which
									# doesn't get quantized and remains in fp32. Therefore we only have 25M parameters quantized
									# in 2 bytes and the diff in optim memory usage is derived as so:
									#
									# - normal 25*8=~200MB (8 bytes per param)
									# - bnb    25*2= ~50MB (2 bytes per param)
									#
									# Thus we should expect ~150MB total memory saved.
									#
									# Peak memory should be the same - the total should be different by about that same margin
									#
									# After leaving a small margin to accommodate for differences between gpus let's check
									# that we have at least 120MB in savings
									snake_case_									=				120

									# uncomment the following if this test starts failing - requires py38 for a new print feature
									# gpu_peak_mem_diff = gpu_peak_mem_orig - gpu_peak_mem_bnb
									# print(f"{gpu_alloc_mem_orig=}MB {gpu_peak_mem_orig=}MB {gpu_alloc_mem_orig+gpu_peak_mem_orig=}MB")
									# print(f" {gpu_alloc_mem_bnb=}MB  {gpu_peak_mem_bnb=}MB  {gpu_alloc_mem_bnb+gpu_peak_mem_bnb=}MB")
									# print(f"{gpu_alloc_mem_diff=}MB")
									# print(f"{gpu_peak_mem_diff=}MB")
									# print(f"{gpu_total_mem_orig=}MB, {gpu_total_mem_bnb=}MB")
									# print(f"{gpu_total_mem_diff=}MB, {gpu_total_mem_diff=}MB")

									self.assertGreater(
									    lowercase_       ,				lowercase_       ,				'''should use ~150MB less alloc gpu memory with BNB, compared to without it for this model but got'''
									    F" a difference of {gpu_alloc_mem_diff}MB, with gpu_alloc_mem_orig={gpu_alloc_mem_orig}MB and"
									    F" gpu_alloc_mem_bnb={gpu_alloc_mem_bnb}MB"       ,				)

									self.assertGreater(
									    lowercase_       ,				lowercase_       ,				'''should use ~150MB less total gpu memory with BNB, compared to without it for this model but got'''
									    F" a difference of {gpu_total_mem_diff}MB, with gpu_total_mem_orig={gpu_total_mem_orig}MB and"
									    F" gpu_total_mem_bnb={gpu_total_mem_bnb}MB"       ,				)

									self.assertEqual(
									    lowercase_       ,				lowercase_       ,				F"loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}"					)


				def   A_ (		self    :					int       ,				lowercase_    :					int       ,				lowercase_    :					str       ,				lowercase_    :					int       ,				lowercase_    :					float = 3e-3       ,				lowercase_    :					str = "adafactor"       ,				lowercase_    :					bool = False       ,				lowercase_    :					str = None       ,				lowercase_    :					int = 0       ,				lowercase_    :					bool = True       ,				lowercase_    :					bool = True       ,				lowercase_    :					bool = True       ,				lowercase_    :					bool = True       ,				lowercase_    :					int = None       ,				):
									snake_case_									=				self.test_file_dir / '''../fixtures/tests_samples/wmt_en_ro'''
									snake_case_									=				self.get_auto_remove_tmp_dir()
									snake_case_									=				F"\n            --model_name_or_path {model_name}\n            --train_file {data_dir}/train.json\n            --validation_file {data_dir}/val.json\n            --test_file {data_dir}/test.json\n            --output_dir {output_dir}\n            --overwrite_output_dir\n            --max_train_samples 8\n            --max_source_length {max_len}\n            --max_target_length {max_len}\n            --do_train\n            --num_train_epochs {str(lowercase_					)}\n            --per_device_train_batch_size 4\n            --learning_rate {learning_rate}\n            --warmup_steps 8\n            --logging_steps 0\n            --logging_strategy no\n            --save_steps {str(lowercase_					)}\n            --group_by_length\n            --label_smoothing_factor 0.1\n            --target_lang ro_RO\n            --source_lang en_XX\n        ".split()

									snake_case_									=				F"\n            --do_eval\n            --per_device_eval_batch_size 4\n            --max_eval_samples 8\n            --val_max_target_length {max_len}\n            --evaluation_strategy steps\n            --eval_steps {str(lowercase_					)}\n        ".split()

									snake_case_									=				'''
            --do_predict
        '''.split()

									snake_case_									=				[]
									if do_train:
														args += args_train

									if do_eval:
														args += args_eval

									if do_predict:
														args += args_predict

									if predict_with_generate:
														args += "--predict_with_generate".split()

									if do_train:
														if optim == "adafactor":
																			args += "--adafactor".split()
														else:
																			args += F"--optim {optim}".split()

									if extra_args_str is not None:
														args += extra_args_str.split()

									if distributed:
														if n_gpus_to_use is None:
																			snake_case_									=				get_gpu_count()
														snake_case_									=				get_torch_dist_unique_port()
														snake_case_									=				F"\n                -m torch.distributed.run\n                --nproc_per_node={n_gpus_to_use}\n                --master_port={master_port}\n                {self.examples_dir_str}/pytorch/translation/run_translation.py\n            ".split()
														snake_case_									=				[sys.executable] + distributed_args + args
														# keep for quick debug
														# print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
														execute_subprocess_async(lowercase_       ,				env=self.get_env()					)
									else:
														snake_case_									=				['''run_translation.py'''] + args
														with patch.object(lowercase_       ,				'''argv'''       ,				lowercase_					):
																			main()

									return output_dir







'''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 						a				(   _lowerCamelCase   ,     _lowerCamelCase   ,     unittest.TestCase     ):
				snake_case_      =		AutoencoderKL
				snake_case_      =		"sample"
				snake_case_      =		1e-2


				@property
				def   A_ (		self    :					Dict					):
									snake_case_									=				4
									snake_case_									=				3
									snake_case_									=				(32, 32)

									snake_case_									=				floats_tensor((batch_size, num_channels) + sizes					).to(lowercase_					)

									return {"sample": image}


				@property
				def   A_ (		self    :					List[Any]					):
									return (3, 32, 32)


				@property
				def   A_ (		self    :					Dict					):
									return (3, 32, 32)


				def   A_ (		self    :					Union[str, Any]					):
									snake_case_									=				{
									    '''block_out_channels''': [32, 64],
									    '''in_channels''': 3,
									    '''out_channels''': 3,
									    '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
									    '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
									    '''latent_channels''': 4,
									}
									snake_case_									=				self.dummy_input
									return init_dict, inputs_dict


				def   A_ (		self    :					Any					):
									pass


				def   A_ (		self    :					str					):
									pass


				@unittest.skipIf(torch_device == '''mps'''       ,				'''Gradient checkpointing skipped on MPS'''					)
				def   A_ (		self    :					Dict					):
									# enable deterministic behavior for gradient checkpointing
									snake_case_			,snake_case_									=				self.prepare_init_args_and_inputs_for_common()
									snake_case_									=				self.model_class(**lowercase_					)
									model.to(lowercase_					)

									assert not model.is_gradient_checkpointing and model.training

									snake_case_									=				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()

									snake_case_									=				torch.randn_like(lowercase_					)
									snake_case_									=				(out - labels).mean()
									loss.backward()

									# re-instantiate the model now enabling gradient checkpointing
									snake_case_									=				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

									snake_case_									=				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()
									snake_case_									=				(out_a - labels).mean()
									loss_a.backward()

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

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


				def   A_ (		self    :					Tuple					):
									snake_case_									=				AutoencoderKL.from_pretrained('''fusing/autoencoder-kl-dummy'''					)
									snake_case_									=				model.to(lowercase_					)
									model.eval()

									if torch_device == "mps":
														snake_case_									=				torch.manual_seed(0					)
									else:
														snake_case_									=				torch.Generator(device=lowercase_					).manual_seed(0					)

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

									snake_case_									=				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":
														snake_case_									=				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":
														snake_case_									=				torch.tensor(
														    [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026]					)
									else:
														snake_case_									=				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 						a				(   unittest.TestCase     ):


				def   A_ (		self    :					Union[str, Any]       ,				lowercase_    :					Optional[Any]       ,				lowercase_    :					Optional[int]					):
									return F"gaussian_noise_s={seed}_shape={'_'.join([str(lowercase_					) for s in shape]					)}.npy"


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


				def   A_ (		self    :					Dict       ,				lowercase_    :					List[Any]=0       ,				lowercase_    :					Union[str, Any]=(4, 3, 512, 512)       ,				lowercase_    :					Optional[Any]=False					):
									snake_case_									=				torch.floataa if fpaa else torch.floataa
									snake_case_									=				torch.from_numpy(load_hf_numpy(self.get_file_format(lowercase_       ,				lowercase_					)					)					).to(lowercase_					).to(lowercase_					)
									return image


				def   A_ (		self    :					Any       ,				lowercase_    :					Dict="CompVis/stable-diffusion-v1-4"       ,				lowercase_    :					List[str]=False					):
									snake_case_									=				'''fp16''' if fpaa else None
									snake_case_									=				torch.floataa if fpaa else torch.floataa

									snake_case_									=				AutoencoderKL.from_pretrained(
									    lowercase_       ,				subfolder='''vae'''       ,				torch_dtype=lowercase_       ,				revision=lowercase_       ,				)
									model.to(lowercase_					).eval()

									return model


				def   A_ (		self    :					Any       ,				lowercase_    :					int=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   A_ (		self    :					Union[str, Any]       ,				lowercase_    :					Tuple       ,				lowercase_    :					Tuple       ,				lowercase_    :					Tuple					):
									snake_case_									=				self.get_sd_vae_model()
									snake_case_									=				self.get_sd_image(lowercase_					)
									snake_case_									=				self.get_generator(lowercase_					)

									with torch.no_grad():
														snake_case_									=				model(lowercase_       ,				generator=lowercase_       ,				sample_posterior=lowercase_					).sample

									assert sample.shape == image.shape

									snake_case_									=				sample[-1, -2:, -2:, :2].flatten().float().cpu()
									snake_case_									=				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   A_ (		self    :					Optional[int]       ,				lowercase_    :					Union[str, Any]       ,				lowercase_    :					Dict					):
									snake_case_									=				self.get_sd_vae_model(fpaa=lowercase_					)
									snake_case_									=				self.get_sd_image(lowercase_       ,				fpaa=lowercase_					)
									snake_case_									=				self.get_generator(lowercase_					)

									with torch.no_grad():
														snake_case_									=				model(lowercase_       ,				generator=lowercase_       ,				sample_posterior=lowercase_					).sample

									assert sample.shape == image.shape

									snake_case_									=				sample[-1, -2:, :2, -2:].flatten().float().cpu()
									snake_case_									=				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   A_ (		self    :					Tuple       ,				lowercase_    :					Dict       ,				lowercase_    :					str       ,				lowercase_    :					Optional[int]					):
									snake_case_									=				self.get_sd_vae_model()
									snake_case_									=				self.get_sd_image(lowercase_					)

									with torch.no_grad():
														snake_case_									=				model(lowercase_					).sample

									assert sample.shape == image.shape

									snake_case_									=				sample[-1, -2:, -2:, :2].flatten().float().cpu()
									snake_case_									=				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   A_ (		self    :					Dict       ,				lowercase_    :					Tuple       ,				lowercase_    :					Optional[int]					):
									snake_case_									=				self.get_sd_vae_model()
									snake_case_									=				self.get_sd_image(lowercase_       ,				shape=(3, 4, 64, 64)					)

									with torch.no_grad():
														snake_case_									=				model.decode(lowercase_					).sample

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

									snake_case_									=				sample[-1, -2:, :2, -2:].flatten().cpu()
									snake_case_									=				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   A_ (		self    :					Union[str, Any]       ,				lowercase_    :					Any       ,				lowercase_    :					Optional[Any]					):
									snake_case_									=				self.get_sd_vae_model(fpaa=lowercase_					)
									snake_case_									=				self.get_sd_image(lowercase_       ,				shape=(3, 4, 64, 64)       ,				fpaa=lowercase_					)

									with torch.no_grad():
														snake_case_									=				model.decode(lowercase_					).sample

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

									snake_case_									=				sample[-1, -2:, :2, -2:].flatten().float().cpu()
									snake_case_									=				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   A_ (		self    :					Optional[Any]       ,				lowercase_    :					List[str]					):
									snake_case_									=				self.get_sd_vae_model(fpaa=lowercase_					)
									snake_case_									=				self.get_sd_image(lowercase_       ,				shape=(3, 4, 64, 64)       ,				fpaa=lowercase_					)

									with torch.no_grad():
														snake_case_									=				model.decode(lowercase_					).sample

									model.enable_xformers_memory_efficient_attention()
									with torch.no_grad():
														snake_case_									=				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   A_ (		self    :					Optional[Any]       ,				lowercase_    :					Any					):
									snake_case_									=				self.get_sd_vae_model()
									snake_case_									=				self.get_sd_image(lowercase_       ,				shape=(3, 4, 64, 64)					)

									with torch.no_grad():
														snake_case_									=				model.decode(lowercase_					).sample

									model.enable_xformers_memory_efficient_attention()
									with torch.no_grad():
														snake_case_									=				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   A_ (		self    :					str       ,				lowercase_    :					Optional[int]       ,				lowercase_    :					Tuple					):
									snake_case_									=				self.get_sd_vae_model()
									snake_case_									=				self.get_sd_image(lowercase_					)
									snake_case_									=				self.get_generator(lowercase_					)

									with torch.no_grad():
														snake_case_									=				model.encode(lowercase_					).latent_dist
														snake_case_									=				dist.sample(generator=lowercase_					)

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

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

									snake_case_									=				3e-3 if torch_device != '''mps''' else 1e-2
									assert torch_all_close(lowercase_       ,				lowercase_       ,				atol=lowercase_					)

