




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__			(lowerCAmelCase_					,       lowerCAmelCase_					,       unittest.TestCase  ):

				'''simple docstring'''



				lowerCamelCase_   :			List[Any]			  =   UNetaDModel
				lowerCamelCase_   :			Optional[Any]			  =   """sample"""







				@property
				def 					_lowercase  (					self  )  ->   Dict:
											lowerCamelCase     :   List[str]    				=      4
											lowerCamelCase     :   str    				=      14
											lowerCamelCase     :   int    				=      16

											lowerCamelCase     :   List[str]    				=      floats_tensor((batch_size, num_features, seq_len)  ).to(__a  )
											lowerCamelCase     :   Optional[Any]    				=      torch.tensor([10] * batch_size  ).to(__a  )

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







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







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







				def 					_lowercase  (					self  )  ->   Tuple:
											pass







				def 					_lowercase  (					self  )  ->   Any:
											pass







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







				@unittest.skipIf(torch_device == "mps"  ,						"mish op not supported in MPS"  )
				def 					_lowercase  (					self  )  ->   List[str]:
											super().test_outputs_equivalence()







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







				@unittest.skipIf(torch_device == "mps"  ,						"mish op not supported in MPS"  )
				def 					_lowercase  (					self  )  ->   Optional[Any]:
											super().test_from_save_pretrained_variant()







				@unittest.skipIf(torch_device == "mps"  ,						"mish op not supported in MPS"  )
				def 					_lowercase  (					self  )  ->   Any:
											super().test_model_from_pretrained()







				@unittest.skipIf(torch_device == "mps"  ,						"mish op not supported in MPS"  )
				def 					_lowercase  (					self  )  ->   Union[str, Any]:
											super().test_output()







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







				@unittest.skipIf(torch_device == "mps"  ,						"mish op not supported in MPS"  )
				def 					_lowercase  (					self  )  ->   Union[str, Any]:
											lowerCamelCase     :   List[str]    				=      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  )
											lowerCamelCase     :   Optional[int]    				=      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  )  ->   List[str]:
											lowerCamelCase     :   List[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  )

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

											with torch.no_grad():
																		lowerCamelCase     :   Tuple    				=      model(__a  ,						__a  ).sample.permute(0  ,						2  ,						1  )

											lowerCamelCase     :   int    				=      output[0, -3:, -3:].flatten()
											# fmt: off
											lowerCamelCase     :   List[str]    				=      torch.tensor([-2.137172, 1.1426016, 0.3688687, -0.766922, 0.7303146, 0.11038864, -0.4760633, 0.13270172, 0.02591348]  )
											# fmt: on
											self.assertTrue(torch.allclose(__a  ,						__a  ,						rtol=1e-3  )  )







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







				@slow
				def 					_lowercase  (					self  )  ->   str:
											lowerCamelCase     :   Tuple    				=      'harmonai/maestro-150k'
											lowerCamelCase     :   Any    				=      UNetaDModel.from_pretrained(__a  ,						subfolder="unet"  )
											model.to(__a  )

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

											with torch.no_grad():
																		lowerCamelCase     :   Optional[int]    				=      model(__a  ,						__a  ).sample

											lowerCamelCase     :   Union[str, Any]    				=      output.abs().sum()
											lowerCamelCase     :   Dict    				=      output.abs().max()

											assert (output_sum - 224.0896).abs() < 0.5
											assert (output_max - 0.0607).abs() < 4e-4




class 	UpperCamelCase__			(lowerCAmelCase_					,       lowerCAmelCase_					,       unittest.TestCase  ):

				'''simple docstring'''



				lowerCamelCase_   :			Union[str, Any]			  =   UNetaDModel
				lowerCamelCase_   :			str			  =   """sample"""







				@property
				def 					_lowercase  (					self  )  ->   List[str]:
											lowerCamelCase     :   str    				=      4
											lowerCamelCase     :   str    				=      14
											lowerCamelCase     :   Optional[int]    				=      16

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

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







				@property
				def 					_lowercase  (					self  )  ->   int:
											return (4, 14, 16)







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







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







				@unittest.skipIf(torch_device == "mps"  ,						"mish op not supported in MPS"  )
				def 					_lowercase  (					self  )  ->   List[str]:
											super().test_outputs_equivalence()







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







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







				@unittest.skipIf(torch_device == "mps"  ,						"mish op not supported in MPS"  )
				def 					_lowercase  (					self  )  ->   Tuple:
											super().test_model_from_pretrained()







				@unittest.skipIf(torch_device == "mps"  ,						"mish op not supported in MPS"  )
				def 					_lowercase  (					self  )  ->   Any:
											lowerCamelCase     :   List[Any]    				=      self.prepare_init_args_and_inputs_for_common()
											lowerCamelCase     :   int    				=      self.model_class(**__a  )
											model.to(__a  )
											model.eval()

											with torch.no_grad():
																		lowerCamelCase     :   Optional[Any]    				=      model(**__a  )

																		if isinstance(__a  ,						__a  ):
																									lowerCamelCase     :   Any    				=      output.sample

											self.assertIsNotNone(__a  )
											lowerCamelCase     :   Tuple    				=      torch.Size((inputs_dict["sample"].shape[0], 1)  )
											self.assertEqual(output.shape  ,						__a  ,						"Input and output shapes do not match"  )







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







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







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







				@unittest.skipIf(torch_device == "mps"  ,						"mish op not supported in MPS"  )
				def 					_lowercase  (					self  )  ->   Optional[int]:
											lowerCamelCase     :   List[str]    				=      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  )
											lowerCamelCase     :   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 					_lowercase  (					self  )  ->   Optional[Any]:
											lowerCamelCase     :   Dict    				=      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  )

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

											with torch.no_grad():
																		lowerCamelCase     :   Any    				=      value_function(__a  ,						__a  ).sample

											# fmt: off
											lowerCamelCase     :   str    				=      torch.tensor([165.25] * seq_len  )
											# fmt: on
											self.assertTrue(torch.allclose(__a  ,						__a  ,						rtol=1e-3  )  )







				def 					_lowercase  (					self  )  ->   Dict:
											pass












from __future__ import annotations

def   A   (   _SCREAMING_SNAKE_CASE					)			->				list[int]:
       assert (
           isinstance(_SCREAMING_SNAKE_CASE				,_SCREAMING_SNAKE_CASE					) and mask > 0
       ), f'''mask needs to be positive integer, your input {mask}'''
       lowerCamelCase     :   str    				=      []
       lowerCamelCase     :   Any    				=      mask

       while submask:
              all_submasks.append(_SCREAMING_SNAKE_CASE					)
              lowerCamelCase     :   Union[str, Any]    				=      (submask - 1) & mask

       return all_submasks


if __name__ == "__main__":
       import doctest

       doctest.testmod()






