




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        _lowerCamelCase  (	_SCREAMING_SNAKE_CASE							,	_SCREAMING_SNAKE_CASE							,	unittest.TestCase      ):

      """simple docstring"""

      lowerCAmelCase__	:       Optional[Any]			=    UNetaDModel
      lowerCAmelCase__	:       Dict			=    "sample"







      @property
      def 					snake_case		(   self    :       Tuple ):
          __UpperCamelCase     						=  4
          __UpperCamelCase     						=  14
          __UpperCamelCase     						=  16

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

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







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







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







      def 					snake_case		(   self    :       Tuple ):
          pass







      def 					snake_case		(   self    :       Union[str, Any] ):
          pass







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







      @unittest.skipIf(torch_device == '''mps'''       ,				'''mish op not supported in MPS''' )
      def 					snake_case		(   self    :       Union[str, Any] ):
          super().test_outputs_equivalence()







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







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







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







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







      def 					snake_case		(   self    :       Tuple ):
          __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 					snake_case		(   self    :       Union[str, Any] ):
          __UpperCamelCase						,					__UpperCamelCase     						=  UNetaDModel.from_pretrained(
              '''bglick13/hopper-medium-v2-value-function-hor32'''       ,				output_loading_info=snake_case       ,				subfolder='''unet''' )
          self.assertIsNotNone(snake_case )
          self.assertEqual(len(loading_info['''missing_keys'''] )       ,				0 )

          model.to(snake_case )
          __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 					snake_case		(   self    :       Tuple ):
          __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(snake_case       ,				snake_case ).sample.permute(0       ,				2       ,				1 )

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







      def 					snake_case		(   self    :       List[Any] ):
          # Not implemented yet for this UNet
          pass







      @slow
      def 					snake_case		(   self    :       Dict ):
          __UpperCamelCase     						=  '''harmonai/maestro-150k'''
          __UpperCamelCase     						=  UNetaDModel.from_pretrained(snake_case       ,				subfolder='''unet''' )
          model.to(snake_case )

          __UpperCamelCase     						=  65536
          __UpperCamelCase     						=  torch.sin(torch.arange(snake_case )[None, None, :].repeat(1       ,				2       ,				1 ) ).to(snake_case )
          __UpperCamelCase     						=  torch.tensor([1] ).to(snake_case )

          with torch.no_grad():
              __UpperCamelCase     						=  model(snake_case       ,				snake_case ).sample

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

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







class        _lowerCamelCase  (	_SCREAMING_SNAKE_CASE							,	_SCREAMING_SNAKE_CASE							,	unittest.TestCase      ):

      """simple docstring"""

      lowerCAmelCase__	:       Optional[int]			=    UNetaDModel
      lowerCAmelCase__	:       str			=    "sample"







      @property
      def 					snake_case		(   self    :       Optional[int] ):
          __UpperCamelCase     						=  4
          __UpperCamelCase     						=  14
          __UpperCamelCase     						=  16

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

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







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







      @property
      def 					snake_case		(   self    :       List[str] ):
          return (4, 14, 1)







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







      @unittest.skipIf(torch_device == '''mps'''       ,				'''mish op not supported in MPS''' )
      def 					snake_case		(   self    :       Union[str, Any] ):
          super().test_outputs_equivalence()







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







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







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







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

          with torch.no_grad():
              __UpperCamelCase     						=  model(**snake_case )

              if isinstance(snake_case       ,				snake_case ):
                  __UpperCamelCase     						=  output.sample

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







      def 					snake_case		(   self    :       List[Any] ):
          pass







      def 					snake_case		(   self    :       Union[str, Any] ):
          pass







      def 					snake_case		(   self    :       Optional[Any] ):
          __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 					snake_case		(   self    :       Any ):
          __UpperCamelCase						,					__UpperCamelCase     						=  UNetaDModel.from_pretrained(
              '''bglick13/hopper-medium-v2-value-function-hor32'''       ,				output_loading_info=snake_case       ,				subfolder='''value_function''' )
          self.assertIsNotNone(snake_case )
          self.assertEqual(len(vf_loading_info['''missing_keys'''] )       ,				0 )

          value_function.to(snake_case )
          __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 					snake_case		(   self    :       List[str] ):
          __UpperCamelCase						,					__UpperCamelCase     						=  UNetaDModel.from_pretrained(
              '''bglick13/hopper-medium-v2-value-function-hor32'''       ,				output_loading_info=snake_case       ,				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(snake_case       ,				snake_case ).sample

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







      def 					snake_case		(   self    :       List[Any] ):
          # Not implemented yet for this UNet
          pass












def     __SCREAMING_SNAKE_CASE      (     lowercase_    ,   lowercase_					)				->		int:

    """simple docstring"""


    __UpperCamelCase     						=  0
    for ele in arr:
        total_sum += product_sum(lowercase_    ,   depth + 1					) if isinstance(lowercase_    ,   lowercase_					) else ele
    return total_sum * depth




def     __SCREAMING_SNAKE_CASE      (     lowercase_					)				->		int:

    """simple docstring"""


    return product_sum(lowercase_    ,   1					)


if __name__ == "__main__":
 import doctest

 doctest.testmod()






