

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

			"""simple docstring"""

			SCREAMING_SNAKE_CASE_      =    UNetaDModel
			SCREAMING_SNAKE_CASE_      =    'sample'






			@property
			def  _lowerCAmelCase      (				self  :       Optional[int]    ):
								UpperCamelCase__		:     Union[str, Any]								=4
								UpperCamelCase__		:     Optional[Any]								=14
								UpperCamelCase__		:     List[str]								=16

								UpperCamelCase__		:     str								=floats_tensor((batch_size, num_features, seq_len)    ).to(lowercase_    )
								UpperCamelCase__		:     Optional[int]								=torch.tensor([10] * batch_size    ).to(lowercase_    )

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






			@property
			def  _lowerCAmelCase      (				self  :       int    ):
								return (4, 14, 16)






			@property
			def  _lowerCAmelCase      (				self  :       List[str]    ):
								return (4, 14, 16)






			def  _lowerCAmelCase      (				self  :       Dict    ):
								pass






			def  _lowerCAmelCase      (				self  :       Dict    ):
								pass






			@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  :       Union[str, Any]    ):
								super().test_outputs_equivalence()






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






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






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






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






			def  _lowerCAmelCase      (				self  :       Any    ):
								UpperCamelCase__		:     str								={
								    '''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''',
								}
								UpperCamelCase__		:     int								=self.dummy_input
								return init_dict, inputs_dict






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

								model.to(lowercase_    )
								UpperCamelCase__		:     Optional[Any]								=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  :       Optional[Any]    ):
								UpperCamelCase__		:     Any								=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    )

								UpperCamelCase__		:     List[str]								=model.config.in_channels
								UpperCamelCase__		:     List[Any]								=16
								UpperCamelCase__		:     List[str]								=torch.randn((1, seq_len, num_features)    ).permute(
								    0    ,	2    ,	1    )  # match original, we can update values and remove
								UpperCamelCase__		:     Optional[int]								=torch.full((num_features,)    ,	0    )

								with torch.no_grad():
													UpperCamelCase__		:     Dict								=model(lowercase_    ,	lowercase_    ).sample.permute(0    ,	2    ,	1    )

								UpperCamelCase__		:     int								=output[0, -3:, -3:].flatten()
								# fmt: off
								UpperCamelCase__		:     List[Any]								=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(lowercase_    ,	lowercase_    ,	rtol=1e-3    )    )






			def  _lowerCAmelCase      (				self  :       Tuple    ):
								# Not implemented yet for this UNet
								pass






			@slow
			def  _lowerCAmelCase      (				self  :       Any    ):
								UpperCamelCase__		:     List[Any]								='''harmonai/maestro-150k'''
								UpperCamelCase__		:     Dict								=UNetaDModel.from_pretrained(lowercase_    ,	subfolder='''unet'''    )
								model.to(lowercase_    )

								UpperCamelCase__		:     Optional[Any]								=6_5536
								UpperCamelCase__		:     int								=torch.sin(torch.arange(lowercase_    )[None, None, :].repeat(1    ,	2    ,	1    )    ).to(lowercase_    )
								UpperCamelCase__		:     str								=torch.tensor([1]    ).to(lowercase_    )

								with torch.no_grad():
													UpperCamelCase__		:     Tuple								=model(lowercase_    ,	lowercase_    ).sample

								UpperCamelCase__		:     str								=output.abs().sum()
								UpperCamelCase__		:     List[Any]								=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      __a (   snake_case__,     snake_case__,     unittest.TestCase     ):

			"""simple docstring"""

			SCREAMING_SNAKE_CASE_      =    UNetaDModel
			SCREAMING_SNAKE_CASE_      =    'sample'






			@property
			def  _lowerCAmelCase      (				self  :       Any    ):
								UpperCamelCase__		:     List[Any]								=4
								UpperCamelCase__		:     Any								=14
								UpperCamelCase__		:     Union[str, Any]								=16

								UpperCamelCase__		:     Union[str, Any]								=floats_tensor((batch_size, num_features, seq_len)    ).to(lowercase_    )
								UpperCamelCase__		:     Optional[Any]								=torch.tensor([10] * batch_size    ).to(lowercase_    )

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






			@property
			def  _lowerCAmelCase      (				self  :       str    ):
								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  :       int    ):
								super().test_determinism()






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






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






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






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






			@unittest.skipIf(torch_device == '''mps'''    ,	'''mish op not supported in MPS'''    )
			def  _lowerCAmelCase      (				self  :       Union[str, Any]    ):
								# UNetRL is a value-function is different output shape
								UpperCamelCase__	,      UpperCamelCase__		:     int								=self.prepare_init_args_and_inputs_for_common()
								UpperCamelCase__		:     Any								=self.model_class(**lowercase_    )
								model.to(lowercase_    )
								model.eval()

								with torch.no_grad():
													UpperCamelCase__		:     int								=model(**lowercase_    )

													if isinstance(lowercase_    ,	lowercase_    ):
																		UpperCamelCase__		:     List[Any]								=output.sample

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






			def  _lowerCAmelCase      (				self  :       Dict    ):
								pass






			def  _lowerCAmelCase      (				self  :       List[str]    ):
								pass






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






			@unittest.skipIf(torch_device == '''mps'''    ,	'''mish op not supported in MPS'''    )
			def  _lowerCAmelCase      (				self  :       Tuple    ):
								UpperCamelCase__	,      UpperCamelCase__		:     int								=UNetaDModel.from_pretrained(
								    '''bglick13/hopper-medium-v2-value-function-hor32'''    ,	output_loading_info=lowercase_    ,	subfolder='''value_function'''    )
								self.assertIsNotNone(lowercase_    )
								self.assertEqual(len(vf_loading_info['''missing_keys''']    )    ,	0    )

								value_function.to(lowercase_    )
								UpperCamelCase__		:     Optional[Any]								=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]    ):
								UpperCamelCase__	,      UpperCamelCase__		:     Dict								=UNetaDModel.from_pretrained(
								    '''bglick13/hopper-medium-v2-value-function-hor32'''    ,	output_loading_info=lowercase_    ,	subfolder='''value_function'''    )
								torch.manual_seed(0    )
								if torch.cuda.is_available():
													torch.cuda.manual_seed_all(0    )

								UpperCamelCase__		:     int								=value_function.config.in_channels
								UpperCamelCase__		:     Any								=14
								UpperCamelCase__		:     Tuple								=torch.randn((1, seq_len, num_features)    ).permute(
								    0    ,	2    ,	1    )  # match original, we can update values and remove
								UpperCamelCase__		:     Union[str, Any]								=torch.full((num_features,)    ,	0    )

								with torch.no_grad():
													UpperCamelCase__		:     int								=value_function(lowercase_    ,	lowercase_    ).sample

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






			def  _lowerCAmelCase      (				self  :       Union[str, Any]    ):
								# Not implemented yet for this UNet
								pass





"""simple docstring"""





def 							_lowerCAmelCase   (    UpperCAmelCase  :				int   ,				UpperCAmelCase  :				int   ,				UpperCAmelCase  :				int					):


					'''simple docstring'''






					if power < 0:
										return -1
					base %= mod
					UpperCamelCase__		:     Any								=1

					while power > 0:
										if power & 1:
															UpperCamelCase__		:     Optional[Any]								=(result * base) % mod
										UpperCamelCase__		:     Union[str, Any]								=power >> 1
										UpperCamelCase__		:     Union[str, Any]								=(base * base) % mod

					return result



def 							_lowerCAmelCase   (    ):


					'''simple docstring'''






					print(modular_exponential(3   ,				200   ,				13					)					)


if __name__ == "__main__":
				import doctest

				doctest.testmod()

				main()


