


'''simple docstring'''





import unittest

import torch

from diffusers import VQModel
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()







class 				snake_case__(							UpperCAmelCase__,						UpperCAmelCase__,						unittest.TestCase							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__      :  Dict					=							VQModel
						SCREAMING_SNAKE_CASE__      :  Union[str, Any]					=							"""sample"""




						@property
						def 						lowercase_ (				self   ,						__lowercase=(3_2, 3_2)							)					->    int:
										lowerCAmelCase_ :      Tuple    				=				4
										lowerCAmelCase_ :      Optional[Any]    				=				3

										lowerCAmelCase_ :      str    				=				floats_tensor((batch_size, num_channels) + sizes							).to(__lowercase							)

										return {"sample": image}




						@property
						def 						lowercase_ (				self							)					->    Optional[int]:
										return (3, 3_2, 3_2)




						@property
						def 						lowercase_ (				self							)					->    Union[str, Any]:
										return (3, 3_2, 3_2)




						def 						lowercase_ (				self							)					->    int:
										lowerCAmelCase_ :      Union[str, Any]    				=				{
										    '''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''': 3,
										}
										lowerCAmelCase_ :      Optional[Any]    				=				self.dummy_input
										return init_dict, inputs_dict




						def 						lowercase_ (				self							)					->    int:
										pass




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




						def 						lowercase_ (				self							)					->    Dict:
										lowerCAmelCase_ :      Optional[Any]    				=				VQModel.from_pretrained('''fusing/vqgan-dummy'''   ,						output_loading_info=__lowercase							)
										self.assertIsNotNone(__lowercase							)
										self.assertEqual(len(loading_info['''missing_keys''']							)   ,						0							)

										model.to(__lowercase							)
										lowerCAmelCase_ :      Tuple    				=				model(**self.dummy_input							)

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




						def 						lowercase_ (				self							)					->    int:
										lowerCAmelCase_ :      List[str]    				=				VQModel.from_pretrained('''fusing/vqgan-dummy'''							)
										model.to(__lowercase							).eval()

										torch.manual_seed(0							)
										if torch.cuda.is_available():
														torch.cuda.manual_seed_all(0							)

										lowerCAmelCase_ :      Tuple    				=				torch.randn(1   ,						model.config.in_channels   ,						model.config.sample_size   ,						model.config.sample_size							)
										lowerCAmelCase_ :      List[Any]    				=				image.to(__lowercase							)
										with torch.no_grad():
														lowerCAmelCase_ :      Any    				=				model(__lowercase							).sample

										lowerCAmelCase_ :      Dict    				=				output[0, -1, -3:, -3:].flatten().cpu()
										# fmt: off
										lowerCAmelCase_ :      Any    				=				torch.tensor([-0.01_53, -0.40_44, -0.18_80, -0.51_61, -0.24_18, -0.40_72, -0.16_12, -0.06_33, -0.01_43]							)
										# fmt: on
										self.assertTrue(torch.allclose(__lowercase   ,						__lowercase   ,						atol=1e-3							)							)



_UpperCAmelCase			:  int     =frozenset(
    [
        """prompt""",
        """height""",
        """width""",
        """guidance_scale""",
        """negative_prompt""",
        """prompt_embeds""",
        """negative_prompt_embeds""",
        """cross_attention_kwargs""",
    ]
)

_UpperCAmelCase			:  List[Any]     =frozenset(["""prompt""", """negative_prompt"""])

_UpperCAmelCase			:  Dict     =frozenset([])

_UpperCAmelCase			:  int     =frozenset(["""image"""])

_UpperCAmelCase			:  Tuple     =frozenset(
    [
        """image""",
        """height""",
        """width""",
        """guidance_scale""",
    ]
)

_UpperCAmelCase			:  int     =frozenset(["""image"""])

_UpperCAmelCase			:  str     =frozenset(
    [
        """prompt""",
        """image""",
        """height""",
        """width""",
        """guidance_scale""",
        """negative_prompt""",
        """prompt_embeds""",
        """negative_prompt_embeds""",
    ]
)

_UpperCAmelCase			:  int     =frozenset(["""prompt""", """image""", """negative_prompt"""])

_UpperCAmelCase			:  Optional[int]     =frozenset(
    [
        # Text guided image variation with an image mask
        """prompt""",
        """image""",
        """mask_image""",
        """height""",
        """width""",
        """guidance_scale""",
        """negative_prompt""",
        """prompt_embeds""",
        """negative_prompt_embeds""",
    ]
)

_UpperCAmelCase			:  Optional[int]     =frozenset(["""prompt""", """image""", """mask_image""", """negative_prompt"""])

_UpperCAmelCase			:  Optional[Any]     =frozenset(
    [
        # image variation with an image mask
        """image""",
        """mask_image""",
        """height""",
        """width""",
        """guidance_scale""",
    ]
)

_UpperCAmelCase			:  Optional[Any]     =frozenset(["""image""", """mask_image"""])

_UpperCAmelCase			:  Union[str, Any]     =frozenset(
    [
        """example_image""",
        """image""",
        """mask_image""",
        """height""",
        """width""",
        """guidance_scale""",
    ]
)

_UpperCAmelCase			:  Tuple     =frozenset(["""example_image""", """image""", """mask_image"""])

_UpperCAmelCase			:  Any     =frozenset(["""class_labels"""])

_UpperCAmelCase			:  List[Any]     =frozenset(["""class_labels"""])

_UpperCAmelCase			:  int     =frozenset(["""batch_size"""])

_UpperCAmelCase			:  str     =frozenset([])

_UpperCAmelCase			:  str     =frozenset(["""batch_size"""])

_UpperCAmelCase			:  Optional[Any]     =frozenset([])

_UpperCAmelCase			:  Tuple     =frozenset(
    [
        """prompt""",
        """audio_length_in_s""",
        """guidance_scale""",
        """negative_prompt""",
        """prompt_embeds""",
        """negative_prompt_embeds""",
        """cross_attention_kwargs""",
    ]
)

_UpperCAmelCase			:  Tuple     =frozenset(["""prompt""", """negative_prompt"""])
_UpperCAmelCase			:  List[str]     =frozenset(["""input_tokens"""])

_UpperCAmelCase			:  Optional[Any]     =frozenset(["""input_tokens"""])