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



    """simple docstring"""

    _UpperCAmelCase			:int							=				UNetaDModel
    _UpperCAmelCase			:int							=				"sample"






    @property
    def _snake_case    (     self							):
       lowercase__:     Optional[int]      =   4
       lowercase__:     Optional[int]      =   14
       lowercase__:     Dict      =   16

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

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






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






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






    def _snake_case    (     self							):
       pass






    def _snake_case    (     self							):
       pass






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






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






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






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






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






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






    def _snake_case    (     self							):
       lowercase__:     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''',
       }
       lowercase__:     Tuple      =   self.dummy_input
       return init_dict, inputs_dict






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

       model.to(_UpperCAmelCase							)
       lowercase__:     List[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 _snake_case    (     self							):
       lowercase__:     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							)

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

       with torch.no_grad():
          lowercase__:     int      =   model(_UpperCAmelCase					,    _UpperCAmelCase							).sample.permute(0					,    2					,    1							)

       lowercase__:     List[Any]      =   output[0, -3:, -3:].flatten()
       # fmt: off
       lowercase__:     Dict      =   torch.tensor([-2.137_172, 1.1_426_016, 0.3_688_687, -0.766_922, 0.7_303_146, 0.11_038_864, -0.4_760_633, 0.13_270_172, 0.02_591_348]							)
       # fmt: on
       self.assertTrue(torch.allclose(_UpperCAmelCase					,    _UpperCAmelCase					,    rtol=1e-3							)							)






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






    @slow
    def _snake_case    (     self							):
       lowercase__:     Optional[int]      =   '''harmonai/maestro-150k'''
       lowercase__:     Tuple      =   UNetaDModel.from_pretrained(_UpperCAmelCase					,    subfolder='''unet'''							)
       model.to(_UpperCAmelCase							)

       lowercase__:     str      =   65536
       lowercase__:     Union[str, Any]      =   torch.sin(torch.arange(_UpperCAmelCase							)[None, None, :].repeat(1					,    2					,    1							)							).to(_UpperCAmelCase							)
       lowercase__:     int      =   torch.tensor([1]							).to(_UpperCAmelCase							)

       with torch.no_grad():
          lowercase__:     Optional[int]      =   model(_UpperCAmelCase					,    _UpperCAmelCase							).sample

       lowercase__:     str      =   output.abs().sum()
       lowercase__:     Tuple      =   output.abs().max()

       assert (output_sum - 224.0_896).abs() < 0.5
       assert (output_max - 0.0_607).abs() < 4e-4



class 					UpperCAmelCase		(_UpperCAmelCase					,_UpperCAmelCase					,unittest.TestCase						):



    """simple docstring"""

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






    @property
    def _snake_case    (     self							):
       lowercase__:     str      =   4
       lowercase__:     Optional[int]      =   14
       lowercase__:     Any      =   16

       lowercase__:     Dict      =   floats_tensor((batch_size, num_features, seq_len)							).to(_UpperCAmelCase							)
       lowercase__:     Dict      =   torch.tensor([10] * batch_size							).to(_UpperCAmelCase							)

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






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






    @property
    def _snake_case    (     self							):
       return (4, 14, 1)






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






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






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






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






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






    @unittest.skipIf(torch_device == '''mps'''					,    '''mish op not supported in MPS'''							)
    def _snake_case    (     self							):
       # UNetRL is a value-function is different output shape
       lowercase__,				lowercase__:     Dict      =   self.prepare_init_args_and_inputs_for_common()
       lowercase__:     List[str]      =   self.model_class(**_UpperCAmelCase							)
       model.to(_UpperCAmelCase							)
       model.eval()

       with torch.no_grad():
          lowercase__:     Any      =   model(**_UpperCAmelCase							)

          if isinstance(_UpperCAmelCase					,    _UpperCAmelCase							):
             lowercase__:     str      =   output.sample

       self.assertIsNotNone(_UpperCAmelCase							)
       lowercase__:     Optional[int]      =   torch.Size((inputs_dict['''sample'''].shape[0], 1)							)
       self.assertEqual(output.shape					,    _UpperCAmelCase					,    '''Input and output shapes do not match'''							)






    def _snake_case    (     self							):
       pass






    def _snake_case    (     self							):
       pass






    def _snake_case    (     self							):
       lowercase__:     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''',
       }
       lowercase__:     Any      =   self.dummy_input
       return init_dict, inputs_dict






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

       value_function.to(_UpperCAmelCase							)
       lowercase__:     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 _snake_case    (     self							):
       lowercase__,				lowercase__:     str      =   UNetaDModel.from_pretrained(
           '''bglick13/hopper-medium-v2-value-function-hor32'''					,    output_loading_info=_UpperCAmelCase					,    subfolder='''value_function'''							)
       torch.manual_seed(0							)
       if torch.cuda.is_available():
          torch.cuda.manual_seed_all(0							)

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

       with torch.no_grad():
          lowercase__:     Any      =   value_function(_UpperCAmelCase					,    _UpperCAmelCase							).sample

       # fmt: off
       lowercase__:     List[Any]      =   torch.tensor([165.25] * seq_len							)
       # fmt: on
       self.assertTrue(torch.allclose(_UpperCAmelCase					,    _UpperCAmelCase					,    rtol=1e-3							)							)






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

"""simple docstring"""






def SCREAMING_SNAKE_CASE__    (     __UpperCAmelCase = 1_0_0_0				)						->						int:
   lowercase__:     Union[str, Any]      =   0
   lowercase__:     Any      =   0
   while 1:
      num += 3
      if num >= n:
         break
      total += num
      num += 2
      if num >= n:
         break
      total += num
      num += 1
      if num >= n:
         break
      total += num
      num += 3
      if num >= n:
         break
      total += num
      num += 1
      if num >= n:
         break
      total += num
      num += 2
      if num >= n:
         break
      total += num
      num += 3
      if num >= n:
         break
      total += num
   return total


if __name__ == "__main__":
       print(f'''{solution() = }''')
