




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


   '''simple docstring'''







   UpperCamelCase__	:       Tuple     =      UNetaDModel
   UpperCamelCase__	:       Dict     =      'sample'




   @property
   def 	UpperCAmelCase       (					self     :	Optional[int]			)					->Any:
     '''simple docstring'''







     UpperCAmelCase__    						=   4
     UpperCAmelCase__    						=   14
     UpperCAmelCase__    						=   16

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

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




   @property
   def 	UpperCAmelCase       (					self     :	Any			)					->List[Any]:
     '''simple docstring'''







     return (4, 14, 16)




   @property
   def 	UpperCAmelCase       (					self     :	str			)					->List[Any]:
     '''simple docstring'''







     return (4, 14, 16)




   def 	UpperCAmelCase       (					self     :	Tuple			)					->Optional[int]:
     '''simple docstring'''







     pass




   def 	UpperCAmelCase       (					self     :	Union[str, Any]			)					->Optional[Any]:
     '''simple docstring'''







     pass




   @unittest.skipIf(torch_device == """mps"""						,    """mish op not supported in MPS"""			)
   def 	UpperCAmelCase       (					self     :	Optional[int]			)					->Any:
     '''simple docstring'''







     super().test_determinism()




   @unittest.skipIf(torch_device == """mps"""						,    """mish op not supported in MPS"""			)
   def 	UpperCAmelCase       (					self     :	Dict			)					->Dict:
     '''simple docstring'''







     super().test_outputs_equivalence()




   @unittest.skipIf(torch_device == """mps"""						,    """mish op not supported in MPS"""			)
   def 	UpperCAmelCase       (					self     :	List[Any]			)					->Optional[Any]:
     '''simple docstring'''







     super().test_from_save_pretrained()




   @unittest.skipIf(torch_device == """mps"""						,    """mish op not supported in MPS"""			)
   def 	UpperCAmelCase       (					self     :	Optional[Any]			)					->Tuple:
     '''simple docstring'''







     super().test_from_save_pretrained_variant()




   @unittest.skipIf(torch_device == """mps"""						,    """mish op not supported in MPS"""			)
   def 	UpperCAmelCase       (					self     :	Dict			)					->Optional[Any]:
     '''simple docstring'''







     super().test_model_from_pretrained()




   @unittest.skipIf(torch_device == """mps"""						,    """mish op not supported in MPS"""			)
   def 	UpperCAmelCase       (					self     :	Dict			)					->Optional[int]:
     '''simple docstring'''







     super().test_output()




   def 	UpperCAmelCase       (					self     :	List[Any]			)					->List[Any]:
     '''simple docstring'''







     UpperCAmelCase__    						=   {
         '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__    						=   self.dummy_input
     return init_dict, inputs_dict




   @unittest.skipIf(torch_device == """mps"""						,    """mish op not supported in MPS"""			)
   def 	UpperCAmelCase       (					self     :	int			)					->Dict:
     '''simple docstring'''







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

     model.to(lowerCamelCase_			)
     UpperCAmelCase__    						=   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 	UpperCAmelCase       (					self     :	int			)					->Any:
     '''simple docstring'''







     UpperCAmelCase__    						=   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__    						=   model.config.in_channels
     UpperCAmelCase__    						=   16
     UpperCAmelCase__    						=   torch.randn((1, seq_len, num_features)			).permute(
         0						,    2						,    1			)  # match original, we can update values and remove
     UpperCAmelCase__    						=   torch.full((num_features,)						,    0			)

     with torch.no_grad():
       UpperCAmelCase__    						=   model(lowerCamelCase_						,    lowerCamelCase_			).sample.permute(0						,    2						,    1			)

     UpperCAmelCase__    						=   output[0, -3:, -3:].flatten()
     # fmt: off
     UpperCAmelCase__    						=   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(lowerCamelCase_						,    lowerCamelCase_						,    rtol=1E-3			)			)




   def 	UpperCAmelCase       (					self     :	Dict			)					->Union[str, Any]:
     '''simple docstring'''







     pass







   @slow
   def 	UpperCAmelCase       (					self     :	int			)					->List[Any]:
     '''simple docstring'''







     UpperCAmelCase__    						=   'harmonai/maestro-150k'
     UpperCAmelCase__    						=   UNetaDModel.from_pretrained(lowerCamelCase_						,    subfolder="""unet"""			)
     model.to(lowerCamelCase_			)

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

     with torch.no_grad():
       UpperCAmelCase__    						=   model(lowerCamelCase_						,    lowerCamelCase_			).sample

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

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





class     snake_case  (	snake_case__ ,  snake_case__ ,  unittest.TestCase							):


   '''simple docstring'''







   UpperCamelCase__	:       Tuple     =      UNetaDModel
   UpperCamelCase__	:       Dict     =      'sample'




   @property
   def 	UpperCAmelCase       (					self     :	Optional[Any]			)					->int:
     '''simple docstring'''







     UpperCAmelCase__    						=   4
     UpperCAmelCase__    						=   14
     UpperCAmelCase__    						=   16

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

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




   @property
   def 	UpperCAmelCase       (					self     :	int			)					->Any:
     '''simple docstring'''







     return (4, 14, 16)




   @property
   def 	UpperCAmelCase       (					self     :	List[Any]			)					->List[Any]:
     '''simple docstring'''







     return (4, 14, 1)




   @unittest.skipIf(torch_device == """mps"""						,    """mish op not supported in MPS"""			)
   def 	UpperCAmelCase       (					self     :	Any			)					->str:
     '''simple docstring'''







     super().test_determinism()




   @unittest.skipIf(torch_device == """mps"""						,    """mish op not supported in MPS"""			)
   def 	UpperCAmelCase       (					self     :	Any			)					->List[Any]:
     '''simple docstring'''







     super().test_outputs_equivalence()




   @unittest.skipIf(torch_device == """mps"""						,    """mish op not supported in MPS"""			)
   def 	UpperCAmelCase       (					self     :	List[str]			)					->Dict:
     '''simple docstring'''







     super().test_from_save_pretrained()




   @unittest.skipIf(torch_device == """mps"""						,    """mish op not supported in MPS"""			)
   def 	UpperCAmelCase       (					self     :	Optional[int]			)					->Dict:
     '''simple docstring'''







     super().test_from_save_pretrained_variant()




   @unittest.skipIf(torch_device == """mps"""						,    """mish op not supported in MPS"""			)
   def 	UpperCAmelCase       (					self     :	Dict			)					->int:
     '''simple docstring'''







     super().test_model_from_pretrained()




   @unittest.skipIf(torch_device == """mps"""						,    """mish op not supported in MPS"""			)
   def 	UpperCAmelCase       (					self     :	Optional[int]			)					->Any:
     '''simple docstring'''







     UpperCAmelCase__    						=   self.prepare_init_args_and_inputs_for_common()
     UpperCAmelCase__    						=   self.model_class(**lowerCamelCase_			)
     model.to(lowerCamelCase_			)
     model.eval()

     with torch.no_grad():
       UpperCAmelCase__    						=   model(**lowerCamelCase_			)

       if isinstance(lowerCamelCase_						,    lowerCamelCase_			):
         UpperCAmelCase__    						=   output.sample

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




   def 	UpperCAmelCase       (					self     :	Tuple			)					->Any:
     '''simple docstring'''







     pass




   def 	UpperCAmelCase       (					self     :	List[Any]			)					->Any:
     '''simple docstring'''







     pass




   def 	UpperCAmelCase       (					self     :	Dict			)					->Optional[int]:
     '''simple docstring'''







     UpperCAmelCase__    						=   {
         '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__    						=   self.dummy_input
     return init_dict, inputs_dict




   @unittest.skipIf(torch_device == """mps"""						,    """mish op not supported in MPS"""			)
   def 	UpperCAmelCase       (					self     :	List[Any]			)					->str:
     '''simple docstring'''







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

     value_function.to(lowerCamelCase_			)
     UpperCAmelCase__    						=   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 	UpperCAmelCase       (					self     :	Optional[Any]			)					->List[Any]:
     '''simple docstring'''







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

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

     with torch.no_grad():
       UpperCAmelCase__    						=   value_function(lowerCamelCase_						,    lowerCamelCase_			).sample

     # fmt: off
     UpperCAmelCase__    						=   torch.tensor([165.25] * seq_len			)
     # fmt: on
     self.assertTrue(torch.allclose(lowerCamelCase_						,    lowerCamelCase_						,    rtol=1E-3			)			)







   def 	UpperCAmelCase       (					self     :	Dict			)					->Optional[Any]:
     '''simple docstring'''







     pass










def 						lowercase		(     _lowerCAmelCase      ):
  UpperCAmelCase__    						=   5381
  for x in s:
    UpperCAmelCase__    						=   ((hash_value << 5) + hash_value) + ord(_lowerCAmelCase      )
  return hash_value & 0xffffffff




