

"""simple docstring"""





import unittest
from typing import Tuple

import torch

from diffusers.utils import floats_tensor, randn_tensor, torch_all_close, torch_device
from diffusers.utils.testing_utils import require_torch







@require_torch
class 				SCREAMING_SNAKE_CASE					:






					"""simple docstring"""



					@property
					def        __lowerCAmelCase (       self			:				Tuple		):
						return self.get_dummy_input()







					@property
					def        __lowerCAmelCase (       self			:				Dict		):
						if self.block_type == "down":
							return (4, 3_2, 1_6, 1_6)
						elif self.block_type == "mid":
							return (4, 3_2, 3_2, 3_2)
						elif self.block_type == "up":
							return (4, 3_2, 6_4, 6_4)

						raise ValueError(F'\'{self.block_type}\' is not a supported block_type. Set it to \'up\', \'mid\', or \'down\'.'		)







					def        __lowerCAmelCase (       self			:				int   ,lowercase_			:				str=True   ,lowercase_			:				List[str]=False   ,lowercase_			:				str=False   ,lowercase_			:				List[str]=False   ,):
						lowerCAmelCase__	:						Any  				=	4
						lowerCAmelCase__	:						Union[str, Any]  				=	3_2
						lowerCAmelCase__	:						int  				=	(3_2, 3_2)

						lowerCAmelCase__	:						str  				=	torch.manual_seed(0		)
						lowerCAmelCase__	:						List[Any]  				=	torch.device(lowercase_		)
						lowerCAmelCase__	:						Optional[int]  				=	(batch_size, num_channels) + sizes
						lowerCAmelCase__	:						Union[str, Any]  				=	randn_tensor(lowercase_   ,generator=lowercase_   ,device=lowercase_		)
						lowerCAmelCase__	:						int  				=	{'''hidden_states''': hidden_states}

						if include_temb:
							lowerCAmelCase__	:						List[str]  				=	1_2_8
							lowerCAmelCase__	:						Tuple  				=	randn_tensor((batch_size, temb_channels)   ,generator=lowercase_   ,device=lowercase_		)

						if include_res_hidden_states_tuple:
							lowerCAmelCase__	:						Optional[int]  				=	torch.manual_seed(1		)
							lowerCAmelCase__	:						Optional[Any]  				=	(randn_tensor(lowercase_   ,generator=lowercase_   ,device=lowercase_		),)

						if include_encoder_hidden_states:
							lowerCAmelCase__	:						int  				=	floats_tensor((batch_size, 3_2, 3_2)		).to(lowercase_		)

						if include_skip_sample:
							lowerCAmelCase__	:						List[str]  				=	randn_tensor(((batch_size, 3) + sizes)   ,generator=lowercase_   ,device=lowercase_		)

						return dummy_input







					def        __lowerCAmelCase (       self			:				Optional[int]		):
						lowerCAmelCase__	:						List[Any]  				=	{
						    '''in_channels''': 3_2,
						    '''out_channels''': 3_2,
						    '''temb_channels''': 1_2_8,
						}
						if self.block_type == "up":
							lowerCAmelCase__	:						str  				=	3_2

						if self.block_type == "mid":
							init_dict.pop('''out_channels'''		)

						lowerCAmelCase__	:						Union[str, Any]  				=	self.dummy_input
						return init_dict, inputs_dict







					def        __lowerCAmelCase (       self			:				List[str]   ,lowercase_			:				Optional[int]		):
						lowerCAmelCase__				,lowerCAmelCase__	:						Optional[int]  				=	self.prepare_init_args_and_inputs_for_common()
						lowerCAmelCase__	:						Dict  				=	self.block_class(**lowercase_		)
						unet_block.to(lowercase_		)
						unet_block.eval()

						with torch.no_grad():
							lowerCAmelCase__	:						Any  				=	unet_block(**lowercase_		)

						if isinstance(lowercase_   ,lowercase_		):
							lowerCAmelCase__	:						Optional[Any]  				=	output[0]

						self.assertEqual(output.shape   ,self.output_shape		)

						lowerCAmelCase__	:						List[str]  				=	output[0, -1, -3:, -3:]
						lowerCAmelCase__	:						List[str]  				=	torch.tensor(lowercase_		).to(lowercase_		)
						assert torch_all_close(output_slice.flatten()   ,lowercase_   ,atol=5E-3		)







					@unittest.skipIf(torch_device == '''mps'''   ,'''Training is not supported in mps'''		)
					def        __lowerCAmelCase (       self			:				Optional[int]		):
						lowerCAmelCase__				,lowerCAmelCase__	:						Optional[Any]  				=	self.prepare_init_args_and_inputs_for_common()
						lowerCAmelCase__	:						Optional[int]  				=	self.block_class(**lowercase_		)
						model.to(lowercase_		)
						model.train()
						lowerCAmelCase__	:						Dict  				=	model(**lowercase_		)

						if isinstance(lowercase_   ,lowercase_		):
							lowerCAmelCase__	:						Union[str, Any]  				=	output[0]

						lowerCAmelCase__	:						int  				=	torch.device(lowercase_		)
						lowerCAmelCase__	:						List[str]  				=	randn_tensor(output.shape   ,device=lowercase_		)
						lowerCAmelCase__	:						int  				=	torch.nn.functional.mse_loss(lowercase_   ,lowercase_		)
						loss.backward()








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





    """simple docstring"""






    lowerCamelCase  :					Union[str, Any]							     =			VQModel
    lowerCamelCase  :					Union[str, Any]							     =			'sample'







    @property
    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=(32, 32)	)				->     Any:
         __lowerCamelCase   :	Tuple								=		4
         __lowerCamelCase   :	Union[str, Any]								=		3

         __lowerCamelCase   :	List[Any]								=		floats_tensor((batch_size, num_channels) + sizes	).to(SCREAMING_SNAKE_CASE_	)

         return {"sample": image}







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







    @property
    def       lowercase_				(							self	)				->     List[Any]:
         return (3, 32, 32)







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







    def       lowercase_				(							self	)				->     str:
         pass







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







    def       lowercase_				(							self	)				->     Any:
         __lowerCamelCase				,       __lowerCamelCase   :	Union[str, Any]								=		VQModel.from_pretrained('fusing/vqgan-dummy'			,     output_loading_info=SCREAMING_SNAKE_CASE_	)
         self.assertIsNotNone(SCREAMING_SNAKE_CASE_	)
         self.assertEqual(len(loading_info['missing_keys']	)			,     0	)

         model.to(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	List[Any]								=		model(**self.dummy_input	)

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







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

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

         __lowerCamelCase   :	List[str]								=		torch.randn(1			,     model.config.in_channels			,     model.config.sample_size			,     model.config.sample_size	)
         __lowerCamelCase   :	Optional[int]								=		image.to(SCREAMING_SNAKE_CASE_	)
         with torch.no_grad():
              __lowerCamelCase   :	int								=		model(SCREAMING_SNAKE_CASE_	).sample

         __lowerCamelCase   :	List[str]								=		output[0, -1, -3:, -3:].flatten().cpu()
         # fmt: off
         __lowerCamelCase   :	Union[str, Any]								=		torch.tensor([-0.0_1_5_3, -0.4_0_4_4, -0.1_8_8_0, -0.5_1_6_1, -0.2_4_1_8, -0.4_0_7_2, -0.1_6_1_2, -0.0_6_3_3, -0.0_1_4_3]	)
         # fmt: on
         self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     atol=1E-3	)	)






