





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       UpperCamelCase_		(	snake_case_				,	snake_case_				,	unittest.TestCase       ):
					lowerCAmelCase_				=						UNetaDModel
					lowerCAmelCase_				=						'''sample'''




					@property
					def 						lowerCAmelCase      (		self	)		->				Any:
										_snake_case						=					4
										_snake_case						=					14
										_snake_case						=					16

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

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




					@property
					def 						lowerCAmelCase      (		self	)		->				str:
										return (4, 14, 16)




					@property
					def 						lowerCAmelCase      (		self	)		->				Optional[int]:
										return (4, 14, 16)




					def 						lowerCAmelCase      (		self	)		->				str:
										pass




					def 						lowerCAmelCase      (		self	)		->				Union[str, Any]:
										pass




					@unittest.skipIf(torch_device == 'mps'					,    'mish op not supported in MPS'	)
					def 						lowerCAmelCase      (		self	)		->				str:
										super().test_determinism()




					@unittest.skipIf(torch_device == 'mps'					,    'mish op not supported in MPS'	)
					def 						lowerCAmelCase      (		self	)		->				Tuple:
										super().test_outputs_equivalence()




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




					@unittest.skipIf(torch_device == 'mps'					,    'mish op not supported in MPS'	)
					def 						lowerCAmelCase      (		self	)		->				int:
										super().test_from_save_pretrained_variant()




					@unittest.skipIf(torch_device == 'mps'					,    'mish op not supported in MPS'	)
					def 						lowerCAmelCase      (		self	)		->				Dict:
										super().test_model_from_pretrained()




					@unittest.skipIf(torch_device == 'mps'					,    'mish op not supported in MPS'	)
					def 						lowerCAmelCase      (		self	)		->				str:
										super().test_output()




					def 						lowerCAmelCase      (		self	)		->				Tuple:
										_snake_case						=					{
										    '''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''',
										}
										_snake_case						=					self.dummy_input
										return init_dict, inputs_dict




					@unittest.skipIf(torch_device == 'mps'					,    'mish op not supported in MPS'	)
					def 						lowerCAmelCase      (		self	)		->				Any:
										_snake_case						=					UNetaDModel.from_pretrained(
										    'bglick13/hopper-medium-v2-value-function-hor32'					,    output_loading_info=SCREAMING_SNAKE_CASE_					,    subfolder='unet'	)
										self.assertIsNotNone(SCREAMING_SNAKE_CASE_	)
										self.assertEqual(len(loading_info['missing_keys']	)					,    0	)

										model.to(SCREAMING_SNAKE_CASE_	)
										_snake_case						=					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 						lowerCAmelCase      (		self	)		->				List[Any]:
										_snake_case						=					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	)

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

										with torch.no_grad():
															_snake_case						=					model(SCREAMING_SNAKE_CASE_					,    SCREAMING_SNAKE_CASE_	).sample.permute(0					,    2					,    1	)

										_snake_case						=					output[0, -3:, -3:].flatten()
										# fmt: off
										_snake_case						=					torch.tensor([-2.13_71_72, 1.1_42_60_16, 0.3_68_86_87, -0.76_69_22, 0.7_30_31_46, 0.11_03_88_64, -0.4_76_06_33, 0.13_27_01_72, 0.02_59_13_48]	)
										# fmt: on
										self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE_					,    SCREAMING_SNAKE_CASE_					,    rtol=1E-3	)	)




					def 						lowerCAmelCase      (		self	)		->				Any:
										# Not implemented yet for this UNet
										pass






					@slow
					def 						lowerCAmelCase      (		self	)		->				Optional[int]:
										_snake_case						=					'''harmonai/maestro-150k'''
										_snake_case						=					UNetaDModel.from_pretrained(SCREAMING_SNAKE_CASE_					,    subfolder='unet'	)
										model.to(SCREAMING_SNAKE_CASE_	)

										_snake_case						=					6_5536
										_snake_case						=					torch.sin(torch.arange(SCREAMING_SNAKE_CASE_	)[None, None, :].repeat(1					,    2					,    1	)	).to(SCREAMING_SNAKE_CASE_	)
										_snake_case						=					torch.tensor([1]	).to(SCREAMING_SNAKE_CASE_	)

										with torch.no_grad():
															_snake_case						=					model(SCREAMING_SNAKE_CASE_					,    SCREAMING_SNAKE_CASE_	).sample

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

										assert (output_sum - 2_24.08_96).abs() < 0.5
										assert (output_max - 0.06_07).abs() < 4E-4





class       UpperCamelCase_		(	snake_case_				,	snake_case_				,	unittest.TestCase       ):
					lowerCAmelCase_				=						UNetaDModel
					lowerCAmelCase_				=						'''sample'''




					@property
					def 						lowerCAmelCase      (		self	)		->				Any:
										_snake_case						=					4
										_snake_case						=					14
										_snake_case						=					16

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

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




					@property
					def 						lowerCAmelCase      (		self	)		->				Tuple:
										return (4, 14, 16)




					@property
					def 						lowerCAmelCase      (		self	)		->				List[Any]:
										return (4, 14, 1)




					@unittest.skipIf(torch_device == 'mps'					,    'mish op not supported in MPS'	)
					def 						lowerCAmelCase      (		self	)		->				Tuple:
										super().test_determinism()




					@unittest.skipIf(torch_device == 'mps'					,    'mish op not supported in MPS'	)
					def 						lowerCAmelCase      (		self	)		->				Tuple:
										super().test_outputs_equivalence()




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




					@unittest.skipIf(torch_device == 'mps'					,    'mish op not supported in MPS'	)
					def 						lowerCAmelCase      (		self	)		->				Any:
										super().test_from_save_pretrained_variant()




					@unittest.skipIf(torch_device == 'mps'					,    'mish op not supported in MPS'	)
					def 						lowerCAmelCase      (		self	)		->				Dict:
										super().test_model_from_pretrained()




					@unittest.skipIf(torch_device == 'mps'					,    'mish op not supported in MPS'	)
					def 						lowerCAmelCase      (		self	)		->				Optional[Any]:
										# UNetRL is a value-function is different output shape
										_snake_case						=					self.prepare_init_args_and_inputs_for_common()
										_snake_case						=					self.model_class(**SCREAMING_SNAKE_CASE_	)
										model.to(SCREAMING_SNAKE_CASE_	)
										model.eval()

										with torch.no_grad():
															_snake_case						=					model(**SCREAMING_SNAKE_CASE_	)

															if isinstance(SCREAMING_SNAKE_CASE_					,    SCREAMING_SNAKE_CASE_	):
																				_snake_case						=					output.sample

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




					def 						lowerCAmelCase      (		self	)		->				List[Any]:
										pass




					def 						lowerCAmelCase      (		self	)		->				List[str]:
										pass




					def 						lowerCAmelCase      (		self	)		->				List[Any]:
										_snake_case						=					{
										    '''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''',
										}
										_snake_case						=					self.dummy_input
										return init_dict, inputs_dict




					@unittest.skipIf(torch_device == 'mps'					,    'mish op not supported in MPS'	)
					def 						lowerCAmelCase      (		self	)		->				Optional[Any]:
										_snake_case						=					UNetaDModel.from_pretrained(
										    'bglick13/hopper-medium-v2-value-function-hor32'					,    output_loading_info=SCREAMING_SNAKE_CASE_					,    subfolder='value_function'	)
										self.assertIsNotNone(SCREAMING_SNAKE_CASE_	)
										self.assertEqual(len(vf_loading_info['missing_keys']	)					,    0	)

										value_function.to(SCREAMING_SNAKE_CASE_	)
										_snake_case						=					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 						lowerCAmelCase      (		self	)		->				Union[str, Any]:
										_snake_case						=					UNetaDModel.from_pretrained(
										    'bglick13/hopper-medium-v2-value-function-hor32'					,    output_loading_info=SCREAMING_SNAKE_CASE_					,    subfolder='value_function'	)
										torch.manual_seed(0	)
										if torch.cuda.is_available():
															torch.cuda.manual_seed_all(0	)

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

										with torch.no_grad():
															_snake_case						=					value_function(SCREAMING_SNAKE_CASE_					,    SCREAMING_SNAKE_CASE_	).sample

										# fmt: off
										_snake_case						=					torch.tensor([1_65.25] * seq_len	)
										# fmt: on
										self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE_					,    SCREAMING_SNAKE_CASE_					,    rtol=1E-3	)	)






					def 						lowerCAmelCase      (		self	)		->				Optional[Any]:
										# Not implemented yet for this UNet
										pass









"""simple docstring"""





_lowercase						   =     (1, 1_53, 3_70, 3_71, 16_34, 24_67_80_51, 1_15_13_22_19_01_87_63_99_25_65_09_55_97_97_39_71_52_24_01)
_lowercase						   =     (-1_53, -1, 0, 1.2, 2_00, "A", [], {}, None)

def 		lowerCAmelCase				(			lowerCamelCase		):
						'''simple docstring'''


						if not isinstance(lowerCamelCase   ,      lowerCamelCase		) or n < 1:
												return False

						# Initialization of sum and number of digits.
						snake_case__							:       Optional[Any]				      =       0
						snake_case__							:       Union[str, Any]				      =       0
						snake_case__							:       List[str]				      =       n
						# Calculation of digits of the number
						while temp > 0:
												number_of_digits += 1
												temp //= 10
						# Dividing number into separate digits and find Armstrong number
						snake_case__							:       Optional[int]				      =       n
						while temp > 0:
												snake_case__							:       Optional[int]				      =       temp % 10
												total += rem**number_of_digits
												temp //= 10
						return n == total

def 		lowerCAmelCase				(			lowerCamelCase		):
						'''simple docstring'''


						if not isinstance(lowerCamelCase   ,      lowerCamelCase		) or n < 1:
												return False

						# Init a "histogram" of the digits
						snake_case__							:       Union[str, Any]				      =       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
						snake_case__							:       int				      =       0
						snake_case__							:       int				      =       0
						snake_case__							:       Dict				      =       n
						while temp > 0:
												snake_case__      ,						snake_case__							:       Optional[int]				      =       divmod(lowerCamelCase   ,      10		)
												digit_histogram[rem] += 1
												digit_total += 1

						for cnt, i in zip(lowerCamelCase   ,      range(len(lowerCamelCase		)		)		):
												total += cnt * i**digit_total

						return n == total

def 		lowerCAmelCase				(			lowerCamelCase		):
						'''simple docstring'''


						if not isinstance(lowerCamelCase   ,      lowerCamelCase		) or n < 1:
												return False
						snake_case__							:       Union[str, Any]				      =       len(str(lowerCamelCase		)		)  # the power that all digits will be raised to
						# check if sum of each digit multiplied expo times is equal to number
						return n == sum(int(lowerCamelCase		) ** expo for i in str(lowerCamelCase		)		)

def 		lowerCAmelCase				(			):
						'''simple docstring'''


						snake_case__							:       Dict				      =       int(input('''Enter an integer to see if it is an Armstrong number: '''		).strip()		)
						print(f'''{num} is {'' if armstrong_number(lowerCamelCase		) else 'not '}an Armstrong number.'''		)
						print(f'''{num} is {'' if narcissistic_number(lowerCamelCase		) else 'not '}an Armstrong number.'''		)
						print(f'''{num} is {'' if pluperfect_number(lowerCamelCase		) else 'not '}an Armstrong number.'''		)


if __name__ == "__main__":
							import doctest

							doctest.testmod()
							main()


