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



						'''simple docstring'''


						_lowerCamelCase:     Optional[Any]   =     AutoencoderKL
						_lowerCamelCase:     str   =     'sample'
						_lowerCamelCase:     Union[str, Any]   =     1E-2


						@property
						def     _SCREAMING_SNAKE_CASE   (			self  :   Union[str, Any]   )  -> Optional[int]:
							A     = 4
							A     = 3
							A     = (32, 32)

							A     = floats_tensor((batch_size, num_channels) + sizes   ).to(lowerCAmelCase__   )

							return {"sample": image}


						@property
						def     _SCREAMING_SNAKE_CASE   (			self  :   str   )  -> Optional[Any]:
							return (3, 32, 32)


						@property
						def     _SCREAMING_SNAKE_CASE   (			self  :   int   )  -> Tuple:
							return (3, 32, 32)


						def     _SCREAMING_SNAKE_CASE   (			self  :   Optional[Any]   )  -> List[str]:
							A     = {
							    'block_out_channels': [32, 64],
							    'in_channels': 3,
							    'out_channels': 3,
							    'down_block_types': ['DownEncoderBlock2D', 'DownEncoderBlock2D'],
							    'up_block_types': ['UpDecoderBlock2D', 'UpDecoderBlock2D'],
							    'latent_channels': 4,
							}
							A     = self.dummy_input
							return init_dict, inputs_dict


						def     _SCREAMING_SNAKE_CASE   (			self  :   str   )  -> Optional[Any]:
							pass


						def     _SCREAMING_SNAKE_CASE   (			self  :   Tuple   )  -> Optional[Any]:
							pass


						@unittest.skipIf(torch_device == 'mps'	,'Gradient checkpointing skipped on MPS'   )
						def     _SCREAMING_SNAKE_CASE   (			self  :   Any   )  -> Optional[int]:
							A ,   A     = self.prepare_init_args_and_inputs_for_common()
							A     = self.model_class(**lowerCAmelCase__   )
							model.to(lowerCAmelCase__   )

							assert not model.is_gradient_checkpointing and model.training

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

							A     = torch.randn_like(lowerCAmelCase__   )
							A     = (out - labels).mean()
							loss.backward()

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

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

							# compare the output and parameters gradients
							self.assertTrue((loss - loss_a).abs() < 1e-5   )
							A     = dict(model.named_parameters()   )
							A     = 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  :   Dict   )  -> Optional[Any]:
							A ,   A     = 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__   )
							A     = model(**self.dummy_input   )

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




						def     _SCREAMING_SNAKE_CASE   (			self  :   Union[str, Any]   )  -> List[Any]:
							A     = AutoencoderKL.from_pretrained('fusing/autoencoder-kl-dummy'   )
							A     = model.to(lowerCAmelCase__   )
							model.eval()

							if torch_device == "mps":
								A     = torch.manual_seed(0   )
							else:
								A     = torch.Generator(device=lowerCAmelCase__   ).manual_seed(0   )

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

							A     = 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":
								A     = 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":
								A     = torch.tensor(
								    [-0.13_52, 0.08_78, 0.04_19, -0.08_18, -0.10_69, 0.06_88, -0.14_58, -0.44_46, -0.00_26]   )
							else:
								A     = torch.tensor(
								    [-0.24_21, 0.46_42, 0.25_07, -0.04_38, 0.06_82, 0.31_60, -0.20_18, -0.07_27, 0.24_85]   )

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







@slow
class 					lowerCAmelCase_					(   unittest.TestCase     ):



						'''simple docstring'''


						def     _SCREAMING_SNAKE_CASE   (			self  :   str	,A_  :   Tuple	,A_  :   Union[str, Any]   )  -> str:
							return F'gaussian_noise_s={seed}_shape={"_".join([str(lowerCAmelCase__   ) for s in shape]   )}.npy'


						def     _SCREAMING_SNAKE_CASE   (			self  :   List[str]   )  -> Tuple:
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()


						def     _SCREAMING_SNAKE_CASE   (			self  :   List[Any]	,A_  :   Any=0	,A_  :   int=(4, 3, 512, 512)	,A_  :   Tuple=False   )  -> Tuple:
							A     = torch.floataa if fpaa else torch.floataa
							A     = torch.from_numpy(load_hf_numpy(self.get_file_format(lowerCAmelCase__	,lowerCAmelCase__   )   )   ).to(lowerCAmelCase__   ).to(lowerCAmelCase__   )
							return image


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

							A     = AutoencoderKL.from_pretrained(
							    lowerCAmelCase__	,subfolder='vae'	,torch_dtype=lowerCAmelCase__	,revision=lowerCAmelCase__	,)
							model.to(lowerCAmelCase__   ).eval()

							return model


						def     _SCREAMING_SNAKE_CASE   (			self  :   int	,A_  :   Tuple=0   )  -> List[Any]:




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


						@parameterized.expand(
						    [
						        # fmt: off
						        [33, [-0.16_03, 0.98_78, -0.04_95, -0.07_90, -0.27_09, 0.83_75, -0.20_60, -0.08_24], [-0.23_95, 0.00_98, 0.01_02, -0.07_09, -0.28_40, -0.02_74, -0.07_18, -0.18_24]],
						        [47, [-0.23_76, 0.11_68, 0.13_32, -0.48_40, -0.25_08, -0.07_91, -0.04_93, -0.40_89], [0.03_50, 0.08_47, 0.04_67, 0.03_44, -0.08_42, -0.05_47, -0.06_33, -0.11_31]],
						        # fmt: on
						    ]   )
						def     _SCREAMING_SNAKE_CASE   (			self  :   Tuple	,A_  :   Dict	,A_  :   Optional[Any]	,A_  :   int   )  -> Tuple:
							A     = self.get_sd_vae_model()
							A     = self.get_sd_image(lowerCAmelCase__   )
							A     = self.get_generator(lowerCAmelCase__   )

							with torch.no_grad():
								A     = model(lowerCAmelCase__	,generator=lowerCAmelCase__	,sample_posterior=lowerCAmelCase__   ).sample

							assert sample.shape == image.shape

							A     = sample[-1, -2:, -2:, :2].flatten().float().cpu()
							A     = 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.05_13, 0.02_89, 1.37_99, 0.21_66, -0.25_73, -0.08_71, 0.51_03, -0.09_99]],
						        [47, [-0.41_28, -0.13_20, -0.37_04, 0.19_65, -0.41_16, -0.23_32, -0.33_40, 0.22_47]],
						        # fmt: on
						    ]   )
						@require_torch_gpu
						def     _SCREAMING_SNAKE_CASE   (			self  :   List[str]	,A_  :   Union[str, Any]	,A_  :   Union[str, Any]   )  -> List[str]:
							A     = self.get_sd_vae_model(fpaa=lowerCAmelCase__   )
							A     = self.get_sd_image(lowerCAmelCase__	,fpaa=lowerCAmelCase__   )
							A     = self.get_generator(lowerCAmelCase__   )

							with torch.no_grad():
								A     = model(lowerCAmelCase__	,generator=lowerCAmelCase__	,sample_posterior=lowerCAmelCase__   ).sample

							assert sample.shape == image.shape

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

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


						@parameterized.expand(
						    [
						        # fmt: off
						        [33, [-0.16_09, 0.98_66, -0.04_87, -0.07_77, -0.27_16, 0.83_68, -0.20_55, -0.08_14], [-0.23_95, 0.00_98, 0.01_02, -0.07_09, -0.28_40, -0.02_74, -0.07_18, -0.18_24]],
						        [47, [-0.23_77, 0.11_47, 0.13_33, -0.48_41, -0.25_06, -0.08_05, -0.04_91, -0.40_85], [0.03_50, 0.08_47, 0.04_67, 0.03_44, -0.08_42, -0.05_47, -0.06_33, -0.11_31]],
						        # fmt: on
						    ]   )
						def     _SCREAMING_SNAKE_CASE   (			self  :   Optional[int]	,A_  :   Union[str, Any]	,A_  :   int	,A_  :   str   )  -> Any:
							A     = self.get_sd_vae_model()
							A     = self.get_sd_image(lowerCAmelCase__   )

							with torch.no_grad():
								A     = model(lowerCAmelCase__   ).sample

							assert sample.shape == image.shape

							A     = sample[-1, -2:, -2:, :2].flatten().float().cpu()
							A     = 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.20_51, -0.18_03, -0.23_11, -0.21_14, -0.32_92, -0.35_74, -0.29_53, -0.33_23]],
						        [37, [-0.26_32, -0.26_25, -0.21_99, -0.27_41, -0.45_39, -0.49_90, -0.37_20, -0.49_25]],
						        # fmt: on
						    ]   )
						@require_torch_gpu
						def     _SCREAMING_SNAKE_CASE   (			self  :   Any	,A_  :   Optional[Any]	,A_  :   int   )  -> Tuple:
							A     = self.get_sd_vae_model()
							A     = self.get_sd_image(lowerCAmelCase__	,shape=(3, 4, 64, 64)   )

							with torch.no_grad():
								A     = model.decode(lowerCAmelCase__   ).sample

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

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

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


						@parameterized.expand(
						    [
						        # fmt: off
						        [27, [-0.03_69, 0.02_07, -0.07_76, -0.06_82, -0.17_47, -0.19_30, -0.14_65, -0.20_39]],
						        [16, [-0.16_28, -0.21_34, -0.27_47, -0.26_42, -0.37_74, -0.44_04, -0.36_87, -0.42_77]],
						        # fmt: on
						    ]   )
						@require_torch_gpu
						def     _SCREAMING_SNAKE_CASE   (			self  :   Optional[Any]	,A_  :   Any	,A_  :   int   )  -> Optional[int]:
							A     = self.get_sd_vae_model(fpaa=lowerCAmelCase__   )
							A     = self.get_sd_image(lowerCAmelCase__	,shape=(3, 4, 64, 64)	,fpaa=lowerCAmelCase__   )

							with torch.no_grad():
								A     = model.decode(lowerCAmelCase__   ).sample

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

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

							with torch.no_grad():
								A     = model.decode(lowerCAmelCase__   ).sample

							model.enable_xformers_memory_efficient_attention()
							with torch.no_grad():
								A     = model.decode(lowerCAmelCase__   ).sample

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

							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  :   Optional[Any]	,A_  :   Optional[int]   )  -> Optional[int]:
							A     = self.get_sd_vae_model()
							A     = self.get_sd_image(lowerCAmelCase__	,shape=(3, 4, 64, 64)   )

							with torch.no_grad():
								A     = model.decode(lowerCAmelCase__   ).sample

							model.enable_xformers_memory_efficient_attention()
							with torch.no_grad():
								A     = model.decode(lowerCAmelCase__   ).sample

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

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




						@parameterized.expand(
						    [
						        # fmt: off
						        [33, [-0.30_01, 0.09_18, -2.69_84, -3.97_20, -3.20_99, -5.03_53, 1.73_38, -0.20_65, 3.42_67]],
						        [47, [-1.50_30, -4.38_71, -6.03_55, -9.11_57, -1.66_61, -2.78_53, 2.16_07, -5.08_23, 2.56_33]],
						        # fmt: on
						    ]   )
						def     _SCREAMING_SNAKE_CASE   (			self  :   List[Any]	,A_  :   List[str]	,A_  :   List[str]   )  -> Union[str, Any]:
							A     = self.get_sd_vae_model()
							A     = self.get_sd_image(lowerCAmelCase__   )
							A     = self.get_generator(lowerCAmelCase__   )

							with torch.no_grad():
								A     = model.encode(lowerCAmelCase__   ).latent_dist
								A     = dist.sample(generator=lowerCAmelCase__   )

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

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

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






'''simple docstring'''

from typing import List, Optional, Union

from ...image_utils import ImageInput
from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType
class 							__lowerCAmelCase	(       __magic_name__					):






       """simple docstring"""
       _snake_case      :	Dict									=       ['image_processor', 'tokenizer']
       _snake_case      :	int									=       'BlipImageProcessor'
       _snake_case      :	Tuple									=       ('BertTokenizer', 'BertTokenizerFast')





       def __init__(				self  :						Union[str, Any]			,					lowerCAmelCase__  :						List[Any]			,					lowerCAmelCase__  :						List[str]	)       ->	int:

            '''simple docstring'''




            _UpperCamelCase							       =			False
            super().__init__(lowerCAmelCase__			,					lowerCAmelCase__	)
            _UpperCamelCase							       =			self.image_processor





       def __call__(				self  :						int			,					lowerCAmelCase__  :						ImageInput = None			,					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__  :						bool = False			,					lowerCAmelCase__  :						bool = False			,					lowerCAmelCase__  :						bool = False			,					lowerCAmelCase__  :						bool = False			,					lowerCAmelCase__  :						bool = False			,					lowerCAmelCase__  :						bool = True			,					lowerCAmelCase__  :						Optional[Union[str, TensorType]] = None			,					**lowerCAmelCase__  :						Optional[Any]			,					)       ->	BatchEncoding:

            '''simple docstring'''




            if images is None and text is None:
                 raise ValueError('''You have to specify either images or text.'''	)

            # Get only text
            if images is None:
                 _UpperCamelCase							       =			self.tokenizer
                 _UpperCamelCase							       =			self.tokenizer(
                     text=lowerCAmelCase__			,					add_special_tokens=lowerCAmelCase__			,					padding=lowerCAmelCase__			,					truncation=lowerCAmelCase__			,					max_length=lowerCAmelCase__			,					stride=lowerCAmelCase__			,					pad_to_multiple_of=lowerCAmelCase__			,					return_attention_mask=lowerCAmelCase__			,					return_overflowing_tokens=lowerCAmelCase__			,					return_special_tokens_mask=lowerCAmelCase__			,					return_offsets_mapping=lowerCAmelCase__			,					return_token_type_ids=lowerCAmelCase__			,					return_length=lowerCAmelCase__			,					verbose=lowerCAmelCase__			,					return_tensors=lowerCAmelCase__			,					**lowerCAmelCase__			,					)
                 return text_encoding

            # add pixel_values
            _UpperCamelCase							       =			self.image_processor(lowerCAmelCase__			,					return_tensors=lowerCAmelCase__	)

            if text is not None:
                 _UpperCamelCase							       =			self.tokenizer(
                     text=lowerCAmelCase__			,					add_special_tokens=lowerCAmelCase__			,					padding=lowerCAmelCase__			,					truncation=lowerCAmelCase__			,					max_length=lowerCAmelCase__			,					stride=lowerCAmelCase__			,					pad_to_multiple_of=lowerCAmelCase__			,					return_attention_mask=lowerCAmelCase__			,					return_overflowing_tokens=lowerCAmelCase__			,					return_special_tokens_mask=lowerCAmelCase__			,					return_offsets_mapping=lowerCAmelCase__			,					return_token_type_ids=lowerCAmelCase__			,					return_length=lowerCAmelCase__			,					verbose=lowerCAmelCase__			,					return_tensors=lowerCAmelCase__			,					**lowerCAmelCase__			,					)
            else:
                 _UpperCamelCase							       =			None

            if text_encoding is not None:
                 encoding_image_processor.update(lowerCAmelCase__	)

            return encoding_image_processor





       def snake_case__		(				self  :						Any			,					*lowerCAmelCase__  :						Union[str, Any]			,					**lowerCAmelCase__  :						Dict	)       ->	List[str]:

            '''simple docstring'''




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





       def snake_case__		(				self  :						Tuple			,					*lowerCAmelCase__  :						int			,					**lowerCAmelCase__  :						Optional[int]	)       ->	Optional[int]:

            '''simple docstring'''




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







       @property
       def snake_case__		(				self  :						Any	)       ->	str:

            '''simple docstring'''




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



