



'''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 _UpperCamelCase		(     A      ,       A      ,       unittest.TestCase		):





				'''simple docstring'''





				lowerCAmelCase__					    =    AutoencoderKL
				lowerCAmelCase__					    =    """sample"""
				lowerCAmelCase__					    =    1E-2

				@property
				def 				__lowerCamelCase     (							self  :  int):







											'''simple docstring'''

											__lowercase       	=4
											__lowercase       	=3
											__lowercase       	=(3_2, 3_2)

											__lowercase       	=floats_tensor((batch_size, num_channels) + sizes).to(_lowerCAmelCase)

											return {"sample": image}

				@property
				def 				__lowerCamelCase     (							self  :  Any):







											'''simple docstring'''

											return (3, 3_2, 3_2)

				@property
				def 				__lowerCamelCase     (							self  :  int):







											'''simple docstring'''

											return (3, 3_2, 3_2)

				def 				__lowerCamelCase     (							self  :  Tuple):







											'''simple docstring'''

											__lowercase       	={
											    'block_out_channels': [3_2, 6_4],
											    'in_channels': 3,
											    'out_channels': 3,
											    'down_block_types': ['DownEncoderBlock2D', 'DownEncoderBlock2D'],
											    'up_block_types': ['UpDecoderBlock2D', 'UpDecoderBlock2D'],
											    'latent_channels': 4,
											}
											__lowercase       	=self.dummy_input
											return init_dict, inputs_dict

				def 				__lowerCamelCase     (							self  :  List[str]):







											'''simple docstring'''

											pass

				def 				__lowerCamelCase     (							self  :  List[Any]):







											'''simple docstring'''

											pass

				@unittest.skipIf(torch_device == 'mps'  ,       'Gradient checkpointing skipped on MPS')
				def 				__lowerCamelCase     (							self  :  Any):







											'''simple docstring'''

											__lowercase	,				__lowercase       	=self.prepare_init_args_and_inputs_for_common()
											__lowercase       	=self.model_class(**_lowerCAmelCase)
											model.to(_lowerCAmelCase)

											assert not model.is_gradient_checkpointing and model.training

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

											__lowercase       	=torch.randn_like(_lowerCAmelCase)
											__lowercase       	=(out - labels).mean()
											loss.backward()

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

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

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







											'''simple docstring'''

											__lowercase	,				__lowercase       	=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)
											__lowercase       	=model(**self.dummy_input)

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

				def 				__lowerCamelCase     (							self  :  List[Any]):







											'''simple docstring'''

											__lowercase       	=AutoencoderKL.from_pretrained('fusing/autoencoder-kl-dummy')
											__lowercase       	=model.to(_lowerCAmelCase)
											model.eval()

											if torch_device == "mps":
																		__lowercase       	=torch.manual_seed(0)
											else:
																		__lowercase       	=torch.Generator(device=_lowerCAmelCase).manual_seed(0)

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

											__lowercase       	=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":
																		__lowercase       	=torch.tensor(
																		    [
																		        -4.00_78e-01,
																		        -3.83_23e-04,
																		        -1.26_81e-01,
																		        -1.14_62e-01,
																		        2.00_95e-01,
																		        1.08_93e-01,
																		        -8.82_47e-02,
																		        -3.03_61e-01,
																		        -9.86_44e-03,
																		    ])
											elif torch_device == "cpu":
																		__lowercase       	=torch.tensor(
																		    [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026])
											else:
																		__lowercase       	=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 _UpperCamelCase		(     unittest.TestCase		):





				'''simple docstring'''





				def 				__lowerCamelCase     (							self  :  str  ,       _lowerCAmelCase  :  int  ,       _lowerCAmelCase  :  List[str]):







											'''simple docstring'''

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

				def 				__lowerCamelCase     (							self  :  List[Any]):







											'''simple docstring'''

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

				def 				__lowerCamelCase     (							self  :  Optional[Any]  ,       _lowerCAmelCase  :  Dict=0  ,       _lowerCAmelCase  :  Tuple=(4, 3, 5_1_2, 5_1_2)  ,       _lowerCAmelCase  :  Optional[Any]=False):







											'''simple docstring'''

											__lowercase       	=torch.floataa if fpaa else torch.floataa
											__lowercase       	=torch.from_numpy(load_hf_numpy(self.get_file_format(_lowerCAmelCase  ,       _lowerCAmelCase))).to(_lowerCAmelCase).to(_lowerCAmelCase)
											return image

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







											'''simple docstring'''

											__lowercase       	='fp16' if fpaa else None
											__lowercase       	=torch.floataa if fpaa else torch.floataa

											__lowercase       	=AutoencoderKL.from_pretrained(
											    _lowerCAmelCase  ,       subfolder='vae'  ,       torch_dtype=_lowerCAmelCase  ,       revision=_lowerCAmelCase  ,       )
											model.to(_lowerCAmelCase).eval()

											return model

				def 				__lowerCamelCase     (							self  :  List[Any]  ,       _lowerCAmelCase  :  Dict=0):







											'''simple docstring'''

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

				@parameterized.expand(
				    [
				        # fmt: off
				        [3_3, [-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]],
				        [4_7, [-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 				__lowerCamelCase     (							self  :  List[str]  ,       _lowerCAmelCase  :  Optional[Any]  ,       _lowerCAmelCase  :  str  ,       _lowerCAmelCase  :  Optional[int]):







											'''simple docstring'''

											__lowercase       	=self.get_sd_vae_model()
											__lowercase       	=self.get_sd_image(_lowerCAmelCase)
											__lowercase       	=self.get_generator(_lowerCAmelCase)

											with torch.no_grad():
																		__lowercase       	=model(_lowerCAmelCase  ,       generator=_lowerCAmelCase  ,       sample_posterior=_lowerCAmelCase).sample

											assert sample.shape == image.shape

											__lowercase       	=sample[-1, -2:, -2:, :2].flatten().float().cpu()
											__lowercase       	=torch.tensor(expected_slice_mps if torch_device == 'mps' else expected_slice)

											assert torch_all_close(_lowerCAmelCase  ,       _lowerCAmelCase  ,       atol=3e-3)

				@parameterized.expand(
				    [
				        # fmt: off
				        [3_3, [-0.0513, 0.0289, 1.3799, 0.2166, -0.2573, -0.0871, 0.5103, -0.0999]],
				        [4_7, [-0.4128, -0.1320, -0.3704, 0.1965, -0.4116, -0.2332, -0.3340, 0.2247]],
				        # fmt: on
				    ])
				@require_torch_gpu
				def 				__lowerCamelCase     (							self  :  Optional[Any]  ,       _lowerCAmelCase  :  str  ,       _lowerCAmelCase  :  Tuple):







											'''simple docstring'''

											__lowercase       	=self.get_sd_vae_model(fpaa=_lowerCAmelCase)
											__lowercase       	=self.get_sd_image(_lowerCAmelCase  ,       fpaa=_lowerCAmelCase)
											__lowercase       	=self.get_generator(_lowerCAmelCase)

											with torch.no_grad():
																		__lowercase       	=model(_lowerCAmelCase  ,       generator=_lowerCAmelCase  ,       sample_posterior=_lowerCAmelCase).sample

											assert sample.shape == image.shape

											__lowercase       	=sample[-1, -2:, :2, -2:].flatten().float().cpu()
											__lowercase       	=torch.tensor(_lowerCAmelCase)

											assert torch_all_close(_lowerCAmelCase  ,       _lowerCAmelCase  ,       atol=1e-2)

				@parameterized.expand(
				    [
				        # fmt: off
				        [3_3, [-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]],
				        [4_7, [-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 				__lowerCamelCase     (							self  :  Dict  ,       _lowerCAmelCase  :  int  ,       _lowerCAmelCase  :  Dict  ,       _lowerCAmelCase  :  int):







											'''simple docstring'''

											__lowercase       	=self.get_sd_vae_model()
											__lowercase       	=self.get_sd_image(_lowerCAmelCase)

											with torch.no_grad():
																		__lowercase       	=model(_lowerCAmelCase).sample

											assert sample.shape == image.shape

											__lowercase       	=sample[-1, -2:, -2:, :2].flatten().float().cpu()
											__lowercase       	=torch.tensor(expected_slice_mps if torch_device == 'mps' else expected_slice)

											assert torch_all_close(_lowerCAmelCase  ,       _lowerCAmelCase  ,       atol=3e-3)

				@parameterized.expand(
				    [
				        # fmt: off
				        [1_3, [-0.2051, -0.1803, -0.2311, -0.2114, -0.3292, -0.3574, -0.2953, -0.3323]],
				        [3_7, [-0.2632, -0.2625, -0.2199, -0.2741, -0.4539, -0.4990, -0.3720, -0.4925]],
				        # fmt: on
				    ])
				@require_torch_gpu
				def 				__lowerCamelCase     (							self  :  Dict  ,       _lowerCAmelCase  :  Union[str, Any]  ,       _lowerCAmelCase  :  Optional[Any]):







											'''simple docstring'''

											__lowercase       	=self.get_sd_vae_model()
											__lowercase       	=self.get_sd_image(_lowerCAmelCase  ,       shape=(3, 4, 6_4, 6_4))

											with torch.no_grad():
																		__lowercase       	=model.decode(_lowerCAmelCase).sample

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

											__lowercase       	=sample[-1, -2:, :2, -2:].flatten().cpu()
											__lowercase       	=torch.tensor(_lowerCAmelCase)

											assert torch_all_close(_lowerCAmelCase  ,       _lowerCAmelCase  ,       atol=1e-3)

				@parameterized.expand(
				    [
				        # fmt: off
				        [2_7, [-0.0369, 0.0207, -0.0776, -0.0682, -0.1747, -0.1930, -0.1465, -0.2039]],
				        [1_6, [-0.1628, -0.2134, -0.2747, -0.2642, -0.3774, -0.4404, -0.3687, -0.4277]],
				        # fmt: on
				    ])
				@require_torch_gpu
				def 				__lowerCamelCase     (							self  :  str  ,       _lowerCAmelCase  :  List[str]  ,       _lowerCAmelCase  :  int):







											'''simple docstring'''

											__lowercase       	=self.get_sd_vae_model(fpaa=_lowerCAmelCase)
											__lowercase       	=self.get_sd_image(_lowerCAmelCase  ,       shape=(3, 4, 6_4, 6_4)  ,       fpaa=_lowerCAmelCase)

											with torch.no_grad():
																		__lowercase       	=model.decode(_lowerCAmelCase).sample

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

											__lowercase       	=sample[-1, -2:, :2, -2:].flatten().float().cpu()
											__lowercase       	=torch.tensor(_lowerCAmelCase)

											assert torch_all_close(_lowerCAmelCase  ,       _lowerCAmelCase  ,       atol=5e-3)

				@parameterized.expand([(1_3,), (1_6,), (2_7,)])
				@require_torch_gpu
				@unittest.skipIf(not is_xformers_available()  ,       reason='xformers is not required when using PyTorch 2.0.')
				def 				__lowerCamelCase     (							self  :  List[str]  ,       _lowerCAmelCase  :  Dict):







											'''simple docstring'''

											__lowercase       	=self.get_sd_vae_model(fpaa=_lowerCAmelCase)
											__lowercase       	=self.get_sd_image(_lowerCAmelCase  ,       shape=(3, 4, 6_4, 6_4)  ,       fpaa=_lowerCAmelCase)

											with torch.no_grad():
																		__lowercase       	=model.decode(_lowerCAmelCase).sample

											model.enable_xformers_memory_efficient_attention()
											with torch.no_grad():
																		__lowercase       	=model.decode(_lowerCAmelCase).sample

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

											assert torch_all_close(_lowerCAmelCase  ,       _lowerCAmelCase  ,       atol=1e-1)

				@parameterized.expand([(1_3,), (1_6,), (3_7,)])
				@require_torch_gpu
				@unittest.skipIf(not is_xformers_available()  ,       reason='xformers is not required when using PyTorch 2.0.')
				def 				__lowerCamelCase     (							self  :  Tuple  ,       _lowerCAmelCase  :  Tuple):







											'''simple docstring'''

											__lowercase       	=self.get_sd_vae_model()
											__lowercase       	=self.get_sd_image(_lowerCAmelCase  ,       shape=(3, 4, 6_4, 6_4))

											with torch.no_grad():
																		__lowercase       	=model.decode(_lowerCAmelCase).sample

											model.enable_xformers_memory_efficient_attention()
											with torch.no_grad():
																		__lowercase       	=model.decode(_lowerCAmelCase).sample

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

											assert torch_all_close(_lowerCAmelCase  ,       _lowerCAmelCase  ,       atol=1e-2)

				@parameterized.expand(
				    [
				        # fmt: off
				        [3_3, [-0.3001, 0.0918, -2.6984, -3.9720, -3.2099, -5.0353, 1.7338, -0.2065, 3.4267]],
				        [4_7, [-1.5030, -4.3871, -6.0355, -9.1157, -1.6661, -2.7853, 2.1607, -5.0823, 2.5633]],
				        # fmt: on
				    ])
				def 				__lowerCamelCase     (							self  :  Any  ,       _lowerCAmelCase  :  int  ,       _lowerCAmelCase  :  int):







											'''simple docstring'''

											__lowercase       	=self.get_sd_vae_model()
											__lowercase       	=self.get_sd_image(_lowerCAmelCase)
											__lowercase       	=self.get_generator(_lowerCAmelCase)

											with torch.no_grad():
																		__lowercase       	=model.encode(_lowerCAmelCase).latent_dist
																		__lowercase       	=dist.sample(generator=_lowerCAmelCase)

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

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

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








'''simple docstring'''



import warnings
from typing import List, Optional, Union

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType





class _UpperCamelCase		(     A		):





				'''simple docstring'''





				lowerCAmelCase__					    =    ["""image_processor""", """tokenizer"""]
				lowerCAmelCase__					    =    """ViltImageProcessor"""
				lowerCAmelCase__					    =    ("""BertTokenizer""", """BertTokenizerFast""")

				def __init__(							self  :  str  ,       _lowerCAmelCase  :  Optional[Any]=None  ,       _lowerCAmelCase  :  int=None  ,       **_lowerCAmelCase  :  Optional[Any]):







											'''simple docstring'''

											__lowercase       	=None
											if "feature_extractor" in kwargs:
																		warnings.warn(
																		    'The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'
																		    ' instead.'  ,       _lowerCAmelCase  ,       )
																		__lowercase       	=kwargs.pop('feature_extractor')

											__lowercase       	=image_processor if image_processor is not None else feature_extractor
											if image_processor is None:
																		raise ValueError('You need to specify an `image_processor`.')
											if tokenizer is None:
																		raise ValueError('You need to specify a `tokenizer`.')

											super().__init__(_lowerCAmelCase  ,       _lowerCAmelCase)
											__lowercase       	=self.image_processor

				def __call__(							self  :  List[Any]  ,       _lowerCAmelCase  :  Tuple  ,       _lowerCAmelCase  :  Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None  ,       _lowerCAmelCase  :  bool = True  ,       _lowerCAmelCase  :  Union[bool, str, PaddingStrategy] = False  ,       _lowerCAmelCase  :  Union[bool, str, TruncationStrategy] = None  ,       _lowerCAmelCase  :  Optional[int] = None  ,       _lowerCAmelCase  :  int = 0  ,       _lowerCAmelCase  :  Optional[int] = None  ,       _lowerCAmelCase  :  Optional[bool] = None  ,       _lowerCAmelCase  :  Optional[bool] = None  ,       _lowerCAmelCase  :  bool = False  ,       _lowerCAmelCase  :  bool = False  ,       _lowerCAmelCase  :  bool = False  ,       _lowerCAmelCase  :  bool = False  ,       _lowerCAmelCase  :  bool = True  ,       _lowerCAmelCase  :  Optional[Union[str, TensorType]] = None  ,       **_lowerCAmelCase  :  Any  ,       ):







											'''simple docstring'''

											__lowercase       	=self.tokenizer(
											    text=_lowerCAmelCase  ,       add_special_tokens=_lowerCAmelCase  ,       padding=_lowerCAmelCase  ,       truncation=_lowerCAmelCase  ,       max_length=_lowerCAmelCase  ,       stride=_lowerCAmelCase  ,       pad_to_multiple_of=_lowerCAmelCase  ,       return_token_type_ids=_lowerCAmelCase  ,       return_attention_mask=_lowerCAmelCase  ,       return_overflowing_tokens=_lowerCAmelCase  ,       return_special_tokens_mask=_lowerCAmelCase  ,       return_offsets_mapping=_lowerCAmelCase  ,       return_length=_lowerCAmelCase  ,       verbose=_lowerCAmelCase  ,       return_tensors=_lowerCAmelCase  ,       **_lowerCAmelCase  ,       )
											# add pixel_values + pixel_mask
											__lowercase       	=self.image_processor(_lowerCAmelCase  ,       return_tensors=_lowerCAmelCase)
											encoding.update(_lowerCAmelCase)

											return encoding

				def 				__lowerCamelCase     (							self  :  List[Any]  ,       *_lowerCAmelCase  :  str  ,       **_lowerCAmelCase  :  Union[str, Any]):







											'''simple docstring'''

											return self.tokenizer.batch_decode(*_lowerCAmelCase  ,       **_lowerCAmelCase)

				def 				__lowerCamelCase     (							self  :  Union[str, Any]  ,       *_lowerCAmelCase  :  str  ,       **_lowerCAmelCase  :  Optional[Any]):







											'''simple docstring'''

											return self.tokenizer.decode(*_lowerCAmelCase  ,       **_lowerCAmelCase)

				@property
				def 				__lowerCamelCase     (							self  :  Optional[int]):







											'''simple docstring'''

											__lowercase       	=self.tokenizer.model_input_names
											__lowercase       	=self.image_processor.model_input_names
											return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))

				@property
				def 				__lowerCamelCase     (							self  :  Union[str, Any]):







											'''simple docstring'''

											warnings.warn(
											    '`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.'  ,       _lowerCAmelCase  ,       )
											return self.image_processor_class

				@property
				def 				__lowerCamelCase     (							self  :  Tuple):







											'''simple docstring'''

											warnings.warn(
											    '`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.'  ,       _lowerCAmelCase  ,       )
											return self.image_processor



