

import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, LDMSuperResolutionPipeline, UNetaDModel, VQModel
from diffusers.utils import PIL_INTERPOLATION, floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch


enable_full_determinism()





class    __UpperCAmelCase	(			unittest.TestCase		):



   @property
   def 			__magic_name__       (				self				:	Union[str, Any]		):
     UpperCAmelCase			:				List[str]										=					1
     UpperCAmelCase			:				Union[str, Any]										=					3
     UpperCAmelCase			:				Optional[Any]										=					(3_2, 3_2)

     UpperCAmelCase			:				List[Any]										=					floats_tensor((batch_size, num_channels) + sizes,						rng=random.Random(0		)		).to(__A		)
     return image



   @property
   def 			__magic_name__       (				self				:	Tuple		):
     torch.manual_seed(0		)
     UpperCAmelCase			:				Tuple										=					UNetaDModel(
         block_out_channels=(3_2, 6_4),						layers_per_block=2,						sample_size=3_2,						in_channels=6,						out_channels=3,						down_block_types=('''DownBlock2D''', '''AttnDownBlock2D'''),						up_block_types=('''AttnUpBlock2D''', '''UpBlock2D'''),						)
     return model



   @property
   def 			__magic_name__       (				self				:	Optional[int]		):
     torch.manual_seed(0		)
     UpperCAmelCase			:				Any										=					VQModel(
         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,						)
     return model



   def 			__magic_name__       (				self				:	List[Any]		):
     UpperCAmelCase			:				Optional[int]										=					'''cpu'''
     UpperCAmelCase			:				int										=					self.dummy_uncond_unet
     UpperCAmelCase			:				Tuple										=					DDIMScheduler()
     UpperCAmelCase			:				List[str]										=					self.dummy_vq_model

     UpperCAmelCase			:				Optional[int]										=					LDMSuperResolutionPipeline(unet=__A,						vqvae=__A,						scheduler=__A		)
     ldm.to(__A		)
     ldm.set_progress_bar_config(disable=__A		)

     UpperCAmelCase			:				Optional[int]										=					self.dummy_image.to(__A		)

     UpperCAmelCase			:				str										=					torch.Generator(device=__A		).manual_seed(0		)
     UpperCAmelCase			:				Optional[Any]										=					ldm(image=__A,						generator=__A,						num_inference_steps=2,						output_type='''numpy'''		).images

     UpperCAmelCase			:				Tuple										=					image[0, -3:, -3:, -1]

     assert image.shape == (1, 6_4, 6_4, 3)
     UpperCAmelCase			:				List[str]										=					np.array([0.8_6_7_8, 0.8_2_4_5, 0.6_3_8_1, 0.6_8_3_0, 0.4_3_8_5, 0.5_5_9_9, 0.4_6_4_1, 0.6_2_0_1, 0.5_1_5_0]		)

     assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2





   @unittest.skipIf(torch_device != '''cuda''',						'''This test requires a GPU'''		)
   def 			__magic_name__       (				self				:	str		):
     UpperCAmelCase			:				List[str]										=					self.dummy_uncond_unet
     UpperCAmelCase			:				str										=					DDIMScheduler()
     UpperCAmelCase			:				str										=					self.dummy_vq_model

     # put models in fp16
     UpperCAmelCase			:				List[Any]										=					unet.half()
     UpperCAmelCase			:				Union[str, Any]										=					vqvae.half()

     UpperCAmelCase			:				List[Any]										=					LDMSuperResolutionPipeline(unet=__A,						vqvae=__A,						scheduler=__A		)
     ldm.to(__A		)
     ldm.set_progress_bar_config(disable=__A		)

     UpperCAmelCase			:				int										=					self.dummy_image.to(__A		)

     UpperCAmelCase			:				Tuple										=					ldm(__A,						num_inference_steps=2,						output_type='''numpy'''		).images

     assert image.shape == (1, 6_4, 6_4, 3)






@slow
@require_torch
class    __UpperCAmelCase	(			unittest.TestCase		):



   def 			__magic_name__       (				self				:	Union[str, Any]		):
     UpperCAmelCase			:				str										=					load_image(
         '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
         '''/vq_diffusion/teddy_bear_pool.png'''		)
     UpperCAmelCase			:				List[str]										=					init_image.resize((6_4, 6_4),						resample=PIL_INTERPOLATION['''lanczos''']		)

     UpperCAmelCase			:				str										=					LDMSuperResolutionPipeline.from_pretrained('''duongna/ldm-super-resolution''',						device_map='''auto'''		)
     ldm.set_progress_bar_config(disable=__A		)

     UpperCAmelCase			:				List[Any]										=					torch.manual_seed(0		)
     UpperCAmelCase			:				Any										=					ldm(image=__A,						generator=__A,						num_inference_steps=2_0,						output_type='''numpy'''		).images

     UpperCAmelCase			:				str										=					image[0, -3:, -3:, -1]

     assert image.shape == (1, 2_5_6, 2_5_6, 3)
     UpperCAmelCase			:				Optional[int]										=					np.array([0.7_6_4_4, 0.7_6_7_9, 0.7_6_4_2, 0.7_6_3_3, 0.7_6_6_6, 0.7_5_6_0, 0.7_4_2_5, 0.7_2_5_7, 0.6_9_0_7]		)

     assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2













"""simple docstring"""







import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, slow, torch_device

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin







class 							_SCREAMING_SNAKE_CASE    (	__snake_case							,			__snake_case							,			unittest.TestCase				):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE     =UNetaDModel
							_SCREAMING_SNAKE_CASE     ='sample'




							@property
							def 			lowercase		( self:  Any     ):

												'''simple docstring'''

												a__ 							=	4
												a__ 							=	14
												a__ 							=	16

												a__ 							=	floats_tensor((batch_size, num_features, seq_len)     ).to(__A     )
												a__ 							=	torch.tensor([10] * batch_size     ).to(__A     )

												return {"sample": noise, "timestep": time_step}




							@property
							def 			lowercase		( self:  List[Any]     ):

												'''simple docstring'''

												return (4, 14, 16)




							@property
							def 			lowercase		( self:  List[Any]     ):

												'''simple docstring'''

												return (4, 14, 16)




							def 			lowercase		( self:  Optional[Any]     ):

												'''simple docstring'''

												pass




							def 			lowercase		( self:  Any     ):

												'''simple docstring'''

												pass




							@unittest.skipIf(torch_device == '''mps'''					,		'''mish op not supported in MPS'''     )
							def 			lowercase		( self:  Optional[int]     ):

												'''simple docstring'''

												super().test_determinism()




							@unittest.skipIf(torch_device == '''mps'''					,		'''mish op not supported in MPS'''     )
							def 			lowercase		( self:  Optional[Any]     ):

												'''simple docstring'''

												super().test_outputs_equivalence()




							@unittest.skipIf(torch_device == '''mps'''					,		'''mish op not supported in MPS'''     )
							def 			lowercase		( self:  Tuple     ):

												'''simple docstring'''

												super().test_from_save_pretrained()




							@unittest.skipIf(torch_device == '''mps'''					,		'''mish op not supported in MPS'''     )
							def 			lowercase		( self:  int     ):

												'''simple docstring'''

												super().test_from_save_pretrained_variant()




							@unittest.skipIf(torch_device == '''mps'''					,		'''mish op not supported in MPS'''     )
							def 			lowercase		( self:  int     ):

												'''simple docstring'''

												super().test_model_from_pretrained()




							@unittest.skipIf(torch_device == '''mps'''					,		'''mish op not supported in MPS'''     )
							def 			lowercase		( self:  Tuple     ):

												'''simple docstring'''

												super().test_output()




							def 			lowercase		( self:  Optional[Any]     ):

												'''simple docstring'''

												a__ 							=	{
												    '''block_out_channels''': (32, 64, 128, 256),
												    '''in_channels''': 14,
												    '''out_channels''': 14,
												    '''time_embedding_type''': '''positional''',
												    '''use_timestep_embedding''': True,
												    '''flip_sin_to_cos''': False,
												    '''freq_shift''': 1.0,
												    '''out_block_type''': '''OutConv1DBlock''',
												    '''mid_block_type''': '''MidResTemporalBlock1D''',
												    '''down_block_types''': ('''DownResnetBlock1D''', '''DownResnetBlock1D''', '''DownResnetBlock1D''', '''DownResnetBlock1D'''),
												    '''up_block_types''': ('''UpResnetBlock1D''', '''UpResnetBlock1D''', '''UpResnetBlock1D'''),
												    '''act_fn''': '''mish''',
												}
												a__ 							=	self.dummy_input
												return init_dict, inputs_dict




							@unittest.skipIf(torch_device == '''mps'''					,		'''mish op not supported in MPS'''     )
							def 			lowercase		( self:  str     ):

												'''simple docstring'''

												a__							,a__ 							=	UNetaDModel.from_pretrained(
												    '''bglick13/hopper-medium-v2-value-function-hor32'''					,		output_loading_info=__A					,		subfolder='''unet'''     )
												self.assertIsNotNone(__A     )
												self.assertEqual(len(loading_info['''missing_keys''']     )					,		0     )

												model.to(__A     )
												a__ 							=	model(**self.dummy_input     )

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




							@unittest.skipIf(torch_device == '''mps'''					,		'''mish op not supported in MPS'''     )
							def 			lowercase		( self:  str     ):

												'''simple docstring'''

												a__ 							=	UNetaDModel.from_pretrained('''bglick13/hopper-medium-v2-value-function-hor32'''					,		subfolder='''unet'''     )
												torch.manual_seed(0     )
												if torch.cuda.is_available():
																	torch.cuda.manual_seed_all(0     )

												a__ 							=	model.config.in_channels
												a__ 							=	16
												a__ 							=	torch.randn((1, seq_len, num_features)     ).permute(
												    0					,		2					,		1     )  # match original, we can update values and remove
												a__ 							=	torch.full((num_features,)					,		0     )

												with torch.no_grad():
																	a__ 							=	model(__A					,		__A     ).sample.permute(0					,		2					,		1     )

												a__ 							=	output[0, -3:, -3:].flatten()
												# fmt: off
												a__ 							=	torch.tensor([-2.1_3_7_1_7_2, 1.1_4_2_6_0_1_6, 0.3_6_8_8_6_8_7, -0.7_6_6_9_2_2, 0.7_3_0_3_1_4_6, 0.1_1_0_3_8_8_6_4, -0.4_7_6_0_6_3_3, 0.1_3_2_7_0_1_7_2, 0.0_2_5_9_1_3_4_8]     )
												# fmt: on
												self.assertTrue(torch.allclose(__A					,		__A					,		rtol=1e-3     )     )




							def 			lowercase		( self:  int     ):

												'''simple docstring'''

												pass





							@slow
							def 			lowercase		( self:  Any     ):

												'''simple docstring'''

												a__ 							=	'''harmonai/maestro-150k'''
												a__ 							=	UNetaDModel.from_pretrained(__A					,		subfolder='''unet'''     )
												model.to(__A     )

												a__ 							=	65536
												a__ 							=	torch.sin(torch.arange(__A     )[None, None, :].repeat(1					,		2					,		1     )     ).to(__A     )
												a__ 							=	torch.tensor([1]     ).to(__A     )

												with torch.no_grad():
																	a__ 							=	model(__A					,		__A     ).sample

												a__ 							=	output.abs().sum()
												a__ 							=	output.abs().max()

												assert (output_sum - 2_2_4.0_8_9_6).abs() < 0.5
												assert (output_max - 0.0_6_0_7).abs() < 4e-4







class 							_SCREAMING_SNAKE_CASE    (	__snake_case							,			__snake_case							,			unittest.TestCase				):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE     =UNetaDModel
							_SCREAMING_SNAKE_CASE     ='sample'




							@property
							def 			lowercase		( self:  List[str]     ):

												'''simple docstring'''

												a__ 							=	4
												a__ 							=	14
												a__ 							=	16

												a__ 							=	floats_tensor((batch_size, num_features, seq_len)     ).to(__A     )
												a__ 							=	torch.tensor([10] * batch_size     ).to(__A     )

												return {"sample": noise, "timestep": time_step}




							@property
							def 			lowercase		( self:  Dict     ):

												'''simple docstring'''

												return (4, 14, 16)




							@property
							def 			lowercase		( self:  int     ):

												'''simple docstring'''

												return (4, 14, 1)




							@unittest.skipIf(torch_device == '''mps'''					,		'''mish op not supported in MPS'''     )
							def 			lowercase		( self:  int     ):

												'''simple docstring'''

												super().test_determinism()




							@unittest.skipIf(torch_device == '''mps'''					,		'''mish op not supported in MPS'''     )
							def 			lowercase		( self:  str     ):

												'''simple docstring'''

												super().test_outputs_equivalence()




							@unittest.skipIf(torch_device == '''mps'''					,		'''mish op not supported in MPS'''     )
							def 			lowercase		( self:  Optional[Any]     ):

												'''simple docstring'''

												super().test_from_save_pretrained()




							@unittest.skipIf(torch_device == '''mps'''					,		'''mish op not supported in MPS'''     )
							def 			lowercase		( self:  List[str]     ):

												'''simple docstring'''

												super().test_from_save_pretrained_variant()




							@unittest.skipIf(torch_device == '''mps'''					,		'''mish op not supported in MPS'''     )
							def 			lowercase		( self:  Optional[int]     ):

												'''simple docstring'''

												super().test_model_from_pretrained()




							@unittest.skipIf(torch_device == '''mps'''					,		'''mish op not supported in MPS'''     )
							def 			lowercase		( self:  int     ):

												'''simple docstring'''

												a__							,a__ 							=	self.prepare_init_args_and_inputs_for_common()
												a__ 							=	self.model_class(**__A     )
												model.to(__A     )
												model.eval()

												with torch.no_grad():
																	a__ 							=	model(**__A     )

																	if isinstance(__A					,		__A     ):
																						a__ 							=	output.sample

												self.assertIsNotNone(__A     )
												a__ 							=	torch.Size((inputs_dict['''sample'''].shape[0], 1)     )
												self.assertEqual(output.shape					,		__A					,		'''Input and output shapes do not match'''     )




							def 			lowercase		( self:  Dict     ):

												'''simple docstring'''

												pass




							def 			lowercase		( self:  Dict     ):

												'''simple docstring'''

												pass




							def 			lowercase		( self:  Any     ):

												'''simple docstring'''

												a__ 							=	{
												    '''in_channels''': 14,
												    '''out_channels''': 14,
												    '''down_block_types''': ['''DownResnetBlock1D''', '''DownResnetBlock1D''', '''DownResnetBlock1D''', '''DownResnetBlock1D'''],
												    '''up_block_types''': [],
												    '''out_block_type''': '''ValueFunction''',
												    '''mid_block_type''': '''ValueFunctionMidBlock1D''',
												    '''block_out_channels''': [32, 64, 128, 256],
												    '''layers_per_block''': 1,
												    '''downsample_each_block''': True,
												    '''use_timestep_embedding''': True,
												    '''freq_shift''': 1.0,
												    '''flip_sin_to_cos''': False,
												    '''time_embedding_type''': '''positional''',
												    '''act_fn''': '''mish''',
												}
												a__ 							=	self.dummy_input
												return init_dict, inputs_dict




							@unittest.skipIf(torch_device == '''mps'''					,		'''mish op not supported in MPS'''     )
							def 			lowercase		( self:  List[Any]     ):

												'''simple docstring'''

												a__							,a__ 							=	UNetaDModel.from_pretrained(
												    '''bglick13/hopper-medium-v2-value-function-hor32'''					,		output_loading_info=__A					,		subfolder='''value_function'''     )
												self.assertIsNotNone(__A     )
												self.assertEqual(len(vf_loading_info['''missing_keys''']     )					,		0     )

												value_function.to(__A     )
												a__ 							=	value_function(**self.dummy_input     )

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




							@unittest.skipIf(torch_device == '''mps'''					,		'''mish op not supported in MPS'''     )
							def 			lowercase		( self:  Union[str, Any]     ):

												'''simple docstring'''

												a__							,a__ 							=	UNetaDModel.from_pretrained(
												    '''bglick13/hopper-medium-v2-value-function-hor32'''					,		output_loading_info=__A					,		subfolder='''value_function'''     )
												torch.manual_seed(0     )
												if torch.cuda.is_available():
																	torch.cuda.manual_seed_all(0     )

												a__ 							=	value_function.config.in_channels
												a__ 							=	14
												a__ 							=	torch.randn((1, seq_len, num_features)     ).permute(
												    0					,		2					,		1     )  # match original, we can update values and remove
												a__ 							=	torch.full((num_features,)					,		0     )

												with torch.no_grad():
																	a__ 							=	value_function(__A					,		__A     ).sample

												# fmt: off
												a__ 							=	torch.tensor([1_6_5.2_5] * seq_len     )
												# fmt: on
												self.assertTrue(torch.allclose(__A					,		__A					,		rtol=1e-3     )     )





							def 			lowercase		( self:  List[str]     ):

												'''simple docstring'''

												pass


