

'''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_					(							_A			,  _A			,  unittest.TestCase				):






     '''simple docstring'''



     UpperCAmelCase_		:      Optional[Any]    						=    UNetaDModel
     UpperCAmelCase_		:      Any    						=    """sample"""


     @property
     def   UpperCAmelCase_			(						self					:	str			)		->      Tuple:
      UpperCAmelCase			:      List[Any]		   =					4
      UpperCAmelCase			:      Union[str, Any]		   =					14
      UpperCAmelCase			:      Union[str, Any]		   =					16

      UpperCAmelCase			:      Optional[int]		   =					floats_tensor((batch_size, num_features, seq_len)			).to(__A			)
      UpperCAmelCase			:      Optional[int]		   =					torch.tensor([10] * batch_size			).to(__A			)

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


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


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


     def   UpperCAmelCase_			(						self					:	str			)		->      Dict:
      pass


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


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


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


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


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


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


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


     def   UpperCAmelCase_			(						self					:	Optional[int]			)		->      Any:
      UpperCAmelCase			:      int		   =					{
          '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			:      Tuple		   =					self.dummy_input
      return init_dict, inputs_dict


     @unittest.skipIf(torch_device == 'mps'	,     'mish op not supported in MPS'			)
     def   UpperCAmelCase_			(						self					:	Optional[Any]			)		->      str:
      UpperCAmelCase   ,					UpperCAmelCase			:      Tuple		   =					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			)
      UpperCAmelCase			:      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   UpperCAmelCase_			(						self					:	Optional[int]			)		->      List[str]:
      UpperCAmelCase			:      List[str]		   =					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			:      Optional[int]		   =					model.config.in_channels
      UpperCAmelCase			:      Dict		   =					16
      UpperCAmelCase			:      str		   =					torch.randn((1, seq_len, num_features)			).permute(
          0	,     2	,     1			)  # match original, we can update values and remove
      UpperCAmelCase			:      str		   =					torch.full((num_features,)	,     0			)

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

      UpperCAmelCase			:      str		   =					output[0, -3:, -3:].flatten()
      # fmt: off
      UpperCAmelCase			:      List[str]		   =					torch.tensor([-2.13_7172, 1.142_6016, 0.368_8687, -0.76_6922, 0.730_3146, 0.1103_8864, -0.476_0633, 0.1327_0172, 0.0259_1348]			)
      # fmt: on
      self.assertTrue(torch.allclose(__A	,     __A	,     rtol=1E-3			)			)


     def   UpperCAmelCase_			(						self					:	Any			)		->      Any:
      pass




     @slow
     def   UpperCAmelCase_			(						self					:	Optional[int]			)		->      List[Any]:
      UpperCAmelCase			:      Optional[int]		   =					'harmonai/maestro-150k'
      UpperCAmelCase			:      List[str]		   =					UNetaDModel.from_pretrained(__A	,     subfolder='unet'			)
      model.to(__A			)

      UpperCAmelCase			:      Optional[int]		   =					65_536
      UpperCAmelCase			:      Any		   =					torch.sin(torch.arange(__A			)[None, None, :].repeat(1	,     2	,     1			)			).to(__A			)
      UpperCAmelCase			:      Any		   =					torch.tensor([1]			).to(__A			)

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

      UpperCAmelCase			:      int		   =					output.abs().sum()
      UpperCAmelCase			:      Optional[int]		   =					output.abs().max()

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




class 			A_					(							_A			,  _A			,  unittest.TestCase				):






     '''simple docstring'''



     UpperCAmelCase_		:      Union[str, Any]    						=    UNetaDModel
     UpperCAmelCase_		:      Dict    						=    """sample"""


     @property
     def   UpperCAmelCase_			(						self					:	str			)		->      Any:
      UpperCAmelCase			:      Any		   =					4
      UpperCAmelCase			:      int		   =					14
      UpperCAmelCase			:      List[Any]		   =					16

      UpperCAmelCase			:      str		   =					floats_tensor((batch_size, num_features, seq_len)			).to(__A			)
      UpperCAmelCase			:      Any		   =					torch.tensor([10] * batch_size			).to(__A			)

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


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


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


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


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


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


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


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


     @unittest.skipIf(torch_device == 'mps'	,     'mish op not supported in MPS'			)
     def   UpperCAmelCase_			(						self					:	Union[str, Any]			)		->      List[str]:
      UpperCAmelCase   ,					UpperCAmelCase			:      Optional[Any]		   =					self.prepare_init_args_and_inputs_for_common()
      UpperCAmelCase			:      Optional[Any]		   =					self.model_class(**__A			)
      model.to(__A			)
      model.eval()

      with torch.no_grad():
       UpperCAmelCase			:      Tuple		   =					model(**__A			)

       if isinstance(__A	,     __A			):
        UpperCAmelCase			:      Optional[int]		   =					output.sample

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


     def   UpperCAmelCase_			(						self					:	Tuple			)		->      Tuple:
      pass


     def   UpperCAmelCase_			(						self					:	Optional[int]			)		->      Dict:
      pass


     def   UpperCAmelCase_			(						self					:	str			)		->      int:
      UpperCAmelCase			:      Union[str, 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   UpperCAmelCase_			(						self					:	str			)		->      Any:
      UpperCAmelCase   ,					UpperCAmelCase			:      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			)
      UpperCAmelCase			:      Dict		   =					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					:	Any			)		->      Optional[Any]:
      UpperCAmelCase   ,					UpperCAmelCase			:      Tuple		   =					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			)

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

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

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




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









'''simple docstring'''







from ..utils import DummyObject, requires_backends

class      lowerCamelCase__     (     metaclass=_A				):






	'''simple docstring'''

	A__		=	['''torch''', '''librosa''']




	def __init__(  self		:  Optional[Any]   ,		*__A		:  Tuple   ,		**__A		:  Tuple					)     ->			Union[str, Any]:




								'''simple docstring'''





								requires_backends(self   ,		["""torch""", """librosa"""]					)




	@classmethod
	def      lowercase__							(  cls		:  Dict   ,		*__A		:  List[Any]   ,		**__A		:  List[str]					)     ->			Union[str, Any]:




								'''simple docstring'''





								requires_backends(cls   ,		["""torch""", """librosa"""]					)




	@classmethod
	def      lowercase__							(  cls		:  Optional[int]   ,		*__A		:  List[str]   ,		**__A		:  Tuple					)     ->			Any:




								'''simple docstring'''





								requires_backends(cls   ,		["""torch""", """librosa"""]					)





class      lowerCamelCase__     (     metaclass=_A				):






	'''simple docstring'''

	A__		=	['''torch''', '''librosa''']




	def __init__(  self		:  int   ,		*__A		:  Dict   ,		**__A		:  List[Any]					)     ->			List[Any]:




								'''simple docstring'''





								requires_backends(self   ,		["""torch""", """librosa"""]					)




	@classmethod
	def      lowercase__							(  cls		:  Optional[int]   ,		*__A		:  Optional[int]   ,		**__A		:  List[Any]					)     ->			Union[str, Any]:




								'''simple docstring'''





								requires_backends(cls   ,		["""torch""", """librosa"""]					)




	@classmethod
	def      lowercase__							(  cls		:  List[str]   ,		*__A		:  List[str]   ,		**__A		:  int					)     ->			List[str]:




								'''simple docstring'''





								requires_backends(cls   ,		["""torch""", """librosa"""]					)





