



"""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    __lowercase     (					_UpperCAmelCase					,    _UpperCAmelCase					,    unittest.TestCase):







						"""simple docstring"""







						_A		:							Union[str, Any]							       =		AutoencoderKL
						_A		:							Optional[int]							       =		"""sample"""
						_A		:							Optional[int]							       =		1e-2







						@property
						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Union[str, Any] 					=					4
											snake_case_					:		Tuple 					=					3
											snake_case_					:		Any 					=					(32, 32)

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

											return {"sample": image}







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







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







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		List[str] 					=					{
											    """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_					:		Optional[int] 					=					self.dummy_input
											return init_dict, inputs_dict







						def 						__UpperCamelCase							(self       ):
											pass







						def 						__UpperCamelCase							(self       ):
											pass







						@unittest.skipIf(torch_device == """mps"""   ,		"""Gradient checkpointing skipped on MPS"""       )
						def 						__UpperCamelCase							(self       ):
											# enable deterministic behavior for gradient checkpointing
											snake_case_					,			snake_case_					:		Optional[int] 					=					self.prepare_init_args_and_inputs_for_common()
											snake_case_					:		Optional[Any] 					=					self.model_class(**lowercase__       )
											model.to(lowercase__       )

											assert not model.is_gradient_checkpointing and model.training

											snake_case_					:		Optional[Any] 					=					model(**lowercase__       ).sample
											# run the backwards pass on the model. For backwards pass, for simplicity purpose,
											# we won't calculate the loss and rather backprop on out.sum()
											model.zero_grad()

											snake_case_					:		Any 					=					torch.randn_like(lowercase__       )
											snake_case_					:		Dict 					=					(out - labels).mean()
											loss.backward()

											# re-instantiate the model now enabling gradient checkpointing
											snake_case_					:		Dict 					=					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_					:		List[Any] 					=					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_					:		Any 					=					(out_a - labels).mean()
											loss_a.backward()

											# compare the output and parameters gradients
											self.assertTrue((loss - loss_a).abs() < 1e-5       )
											snake_case_					:		Any 					=					dict(model.named_parameters()       )
											snake_case_					:		List[Any] 					=					dict(model_a.named_parameters()       )
											for name, param in named_params.items():
																self.assertTrue(torch_all_close(param.grad.data   ,		named_params_a[name].grad.data   ,		atol=5e-5       )       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					,			snake_case_					:		Optional[int] 					=					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_					:		List[Any] 					=					model(**self.dummy_input       )

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







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Tuple 					=					AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""       )
											snake_case_					:		Any 					=					model.to(lowercase__       )
											model.eval()

											if torch_device == "mps":
																snake_case_					:		List[str] 					=					torch.manual_seed(0       )
											else:
																snake_case_					:		List[Any] 					=					torch.Generator(device=lowercase__       ).manual_seed(0       )

											snake_case_					:		List[str] 					=					torch.randn(
											    1   ,		model.config.in_channels   ,		model.config.sample_size   ,		model.config.sample_size   ,		generator=torch.manual_seed(0       )   ,		)
											snake_case_					:		str 					=					image.to(lowercase__       )
											with torch.no_grad():
																snake_case_					:		List[Any] 					=					model(lowercase__   ,		sample_posterior=lowercase__   ,		generator=lowercase__       ).sample

											snake_case_					:		Any 					=					output[0, -1, -3:, -3:].flatten().cpu()

											# Since the VAE Gaussian prior's generator is seeded on the appropriate device,
											# the expected output slices are not the same for CPU and GPU.
											if torch_device == "mps":
																snake_case_					:		str 					=					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_					:		int 					=					torch.tensor(
																    [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026]       )
											else:
																snake_case_					:		Optional[Any] 					=					torch.tensor(
																    [-0.2421, 0.4642, 0.2507, -0.0438, 0.0682, 0.3160, -0.2018, -0.0727, 0.2485]       )

											self.assertTrue(torch_all_close(lowercase__   ,		lowercase__   ,		rtol=1e-2       )       )







@slow
class    __lowercase     (					unittest.TestCase):







						"""simple docstring"""







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







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







						def 						__UpperCamelCase							(self   ,		lowercase__=0   ,		lowercase__=(4, 3, 5_12, 5_12)   ,		lowercase__=False       ):
											snake_case_					:		int 					=					torch.floataa if fpaa else torch.floataa
											snake_case_					:		Dict 					=					torch.from_numpy(load_hf_numpy(self.get_file_format(lowercase__   ,		lowercase__       )       )       ).to(lowercase__       ).to(lowercase__       )
											return image







						def 						__UpperCamelCase							(self   ,		lowercase__="CompVis/stable-diffusion-v1-4"   ,		lowercase__=False       ):
											snake_case_					:		str 					=					"""fp16""" if fpaa else None
											snake_case_					:		Dict 					=					torch.floataa if fpaa else torch.floataa

											snake_case_					:		Optional[Any] 					=					AutoencoderKL.from_pretrained(
											    lowercase__   ,		subfolder="""vae"""   ,		torch_dtype=lowercase__   ,		revision=lowercase__   ,		)
											model.to(lowercase__       ).eval()

											return model







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







						@parameterized.expand(
						    [
						        # fmt: off
						        [33, [-0.1603, 0.9878, -0.0495, -0.0790, -0.2709, 0.8375, -0.2060, -0.0824], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
						        [47, [-0.2376, 0.1168, 0.1332, -0.4840, -0.2508, -0.0791, -0.0493, -0.4089], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
						        # fmt: on
						    ]       )
						def 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		Optional[int] 					=					self.get_sd_vae_model()
											snake_case_					:		int 					=					self.get_sd_image(lowercase__       )
											snake_case_					:		Optional[Any] 					=					self.get_generator(lowercase__       )

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

											assert sample.shape == image.shape

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

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







						@parameterized.expand(
						    [
						        # fmt: off
						        [33, [-0.0513, 0.0289, 1.3799, 0.2166, -0.2573, -0.0871, 0.5103, -0.0999]],
						        [47, [-0.4128, -0.1320, -0.3704, 0.1965, -0.4116, -0.2332, -0.3340, 0.2247]],
						        # fmt: on
						    ]       )
						@require_torch_gpu
						def 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		Dict 					=					self.get_sd_vae_model(fpaa=lowercase__       )
											snake_case_					:		List[str] 					=					self.get_sd_image(lowercase__   ,		fpaa=lowercase__       )
											snake_case_					:		Dict 					=					self.get_generator(lowercase__       )

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

											assert sample.shape == image.shape

											snake_case_					:		List[str] 					=					sample[-1, -2:, :2, -2:].flatten().float().cpu()
											snake_case_					:		Tuple 					=					torch.tensor(lowercase__       )

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







						@parameterized.expand(
						    [
						        # fmt: off
						        [33, [-0.1609, 0.9866, -0.0487, -0.0777, -0.2716, 0.8368, -0.2055, -0.0814], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
						        [47, [-0.2377, 0.1147, 0.1333, -0.4841, -0.2506, -0.0805, -0.0491, -0.4085], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
						        # fmt: on
						    ]       )
						def 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		Optional[int] 					=					self.get_sd_vae_model()
											snake_case_					:		Tuple 					=					self.get_sd_image(lowercase__       )

											with torch.no_grad():
																snake_case_					:		Dict 					=					model(lowercase__       ).sample

											assert sample.shape == image.shape

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

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







						@parameterized.expand(
						    [
						        # fmt: off
						        [13, [-0.2051, -0.1803, -0.2311, -0.2114, -0.3292, -0.3574, -0.2953, -0.3323]],
						        [37, [-0.2632, -0.2625, -0.2199, -0.2741, -0.4539, -0.4990, -0.3720, -0.4925]],
						        # fmt: on
						    ]       )
						@require_torch_gpu
						def 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		List[Any] 					=					self.get_sd_vae_model()
											snake_case_					:		Union[str, Any] 					=					self.get_sd_image(lowercase__   ,		shape=(3, 4, 64, 64)       )

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

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

											snake_case_					:		Any 					=					sample[-1, -2:, :2, -2:].flatten().cpu()
											snake_case_					:		List[Any] 					=					torch.tensor(lowercase__       )

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







						@parameterized.expand(
						    [
						        # fmt: off
						        [27, [-0.0369, 0.0207, -0.0776, -0.0682, -0.1747, -0.1930, -0.1465, -0.2039]],
						        [16, [-0.1628, -0.2134, -0.2747, -0.2642, -0.3774, -0.4404, -0.3687, -0.4277]],
						        # fmt: on
						    ]       )
						@require_torch_gpu
						def 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		Optional[Any] 					=					self.get_sd_vae_model(fpaa=lowercase__       )
											snake_case_					:		Union[str, Any] 					=					self.get_sd_image(lowercase__   ,		shape=(3, 4, 64, 64)   ,		fpaa=lowercase__       )

											with torch.no_grad():
																snake_case_					:		List[Any] 					=					model.decode(lowercase__       ).sample

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

											snake_case_					:		Tuple 					=					sample[-1, -2:, :2, -2:].flatten().float().cpu()
											snake_case_					:		Dict 					=					torch.tensor(lowercase__       )

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







						@parameterized.expand([(13,), (16,), (27,)]       )
						@require_torch_gpu
						@unittest.skipIf(not is_xformers_available()   ,		reason="""xformers is not required when using PyTorch 2.0."""       )
						def 						__UpperCamelCase							(self   ,		lowercase__       ):
											snake_case_					:		Union[str, Any] 					=					self.get_sd_vae_model(fpaa=lowercase__       )
											snake_case_					:		Optional[int] 					=					self.get_sd_image(lowercase__   ,		shape=(3, 4, 64, 64)   ,		fpaa=lowercase__       )

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

											model.enable_xformers_memory_efficient_attention()
											with torch.no_grad():
																snake_case_					:		List[Any] 					=					model.decode(lowercase__       ).sample

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

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







						@parameterized.expand([(13,), (16,), (37,)]       )
						@require_torch_gpu
						@unittest.skipIf(not is_xformers_available()   ,		reason="""xformers is not required when using PyTorch 2.0."""       )
						def 						__UpperCamelCase							(self   ,		lowercase__       ):
											snake_case_					:		int 					=					self.get_sd_vae_model()
											snake_case_					:		Any 					=					self.get_sd_image(lowercase__   ,		shape=(3, 4, 64, 64)       )

											with torch.no_grad():
																snake_case_					:		List[Any] 					=					model.decode(lowercase__       ).sample

											model.enable_xformers_memory_efficient_attention()
											with torch.no_grad():
																snake_case_					:		Optional[Any] 					=					model.decode(lowercase__       ).sample

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

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







						@parameterized.expand(
						    [
						        # fmt: off
						        [33, [-0.3001, 0.0918, -2.6984, -3.9720, -3.2099, -5.0353, 1.7338, -0.2065, 3.4267]],
						        [47, [-1.5030, -4.3871, -6.0355, -9.1157, -1.6661, -2.7853, 2.1607, -5.0823, 2.5633]],
						        # fmt: on
						    ]       )
						def 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		Union[str, Any] 					=					self.get_sd_vae_model()
											snake_case_					:		Union[str, Any] 					=					self.get_sd_image(lowercase__       )
											snake_case_					:		int 					=					self.get_generator(lowercase__       )

											with torch.no_grad():
																snake_case_					:		int 					=					model.encode(lowercase__       ).latent_dist
																snake_case_					:		Dict 					=					dist.sample(generator=lowercase__       )

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

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

											snake_case_					:		Tuple 					=					3e-3 if torch_device != """mps""" else 1e-2
											assert torch_all_close(lowercase__   ,		lowercase__   ,		atol=lowercase__       )









"""simple docstring"""

import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding







class    __lowercase     (					_UpperCAmelCase):







						"""simple docstring"""







						_A		:							Union[str, Any]							       =		["""image_processor""", """tokenizer"""]
						_A		:							str							       =		"""ChineseCLIPImageProcessor"""
						_A		:							Tuple							       =		("""BertTokenizer""", """BertTokenizerFast""")







						def __init__(self   ,		lowercase__=None   ,		lowercase__=None   ,		**lowercase__       ):
											snake_case_					:		Any 					=					None
											if "feature_extractor" in kwargs:
																warnings.warn(
																    """The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"""
																    """ instead."""   ,		lowercase__   ,		)
																snake_case_					:		Optional[Any] 					=					kwargs.pop("""feature_extractor"""       )

											snake_case_					:		str 					=					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__(lowercase__   ,		lowercase__       )
											snake_case_					:		Union[str, Any] 					=					self.image_processor







						def __call__(self   ,		lowercase__=None   ,		lowercase__=None   ,		lowercase__=None   ,		**lowercase__       ):

											if text is None and images is None:
																raise ValueError("""You have to specify either text or images. Both cannot be none."""       )

											if text is not None:
																snake_case_					:		Any 					=					self.tokenizer(lowercase__   ,		return_tensors=lowercase__   ,		**lowercase__       )

											if images is not None:
																snake_case_					:		Tuple 					=					self.image_processor(lowercase__   ,		return_tensors=lowercase__   ,		**lowercase__       )

											if text is not None and images is not None:
																snake_case_					:		List[Any] 					=					image_features.pixel_values
																return encoding
											elif text is not None:
																return encoding
											else:
																return BatchEncoding(data=dict(**lowercase__       )   ,		tensor_type=lowercase__       )







						def 						__UpperCamelCase							(self   ,		*lowercase__   ,		**lowercase__       ):
											return self.tokenizer.batch_decode(*lowercase__   ,		**lowercase__       )







						def 						__UpperCamelCase							(self   ,		*lowercase__   ,		**lowercase__       ):
											return self.tokenizer.decode(*lowercase__   ,		**lowercase__       )







						@property
						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Optional[int] 					=					self.tokenizer.model_input_names
											snake_case_					:		Union[str, Any] 					=					self.image_processor.model_input_names
											return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names       )       )







						@property
						def 						__UpperCamelCase							(self       ):
											warnings.warn(
											    """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."""   ,		lowercase__   ,		)
											return self.image_processor_class




