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

    """simple docstring"""



    UpperCAmelCase_    :				Any						  =    UNetaDModel
    UpperCAmelCase_    :				Dict						  =    "sample"


    @property
    def 							A (    self      )   ->   List[str]:
         """simple docstring"""







         SCREAMING_SNAKE_CASE           =     4
         SCREAMING_SNAKE_CASE           =     14
         SCREAMING_SNAKE_CASE           =     16

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

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


    @property
    def 							A (    self      )   ->   Any:
         """simple docstring"""







         return (4, 14, 16)


    @property
    def 							A (    self      )   ->   List[Any]:
         """simple docstring"""







         return (4, 14, 16)


    def 							A (    self      )   ->   Dict:
         """simple docstring"""







         pass


    def 							A (    self      )   ->   List[str]:
         """simple docstring"""







         pass


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







         super().test_determinism()


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







         super().test_outputs_equivalence()


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







         super().test_from_save_pretrained()


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







         super().test_from_save_pretrained_variant()


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







         super().test_model_from_pretrained()


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







         super().test_output()


    def 							A (    self      )   ->   Tuple:
         """simple docstring"""







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


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







         SCREAMING_SNAKE_CASE     ,   SCREAMING_SNAKE_CASE           =     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      )
         SCREAMING_SNAKE_CASE           =     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 							A (    self      )   ->   Dict:
         """simple docstring"""







         SCREAMING_SNAKE_CASE           =     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      )

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

         with torch.no_grad():
              SCREAMING_SNAKE_CASE           =     model(__UpperCamelCase		,     __UpperCamelCase      ).sample.permute(0		,     2		,     1      )

         SCREAMING_SNAKE_CASE           =     output[0, -3:, -3:].flatten()
         # fmt: off
         SCREAMING_SNAKE_CASE           =     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 							A (    self      )   ->   Dict:
         """simple docstring"""







         pass




    @slow
    def 							A (    self      )   ->   str:
         """simple docstring"""







         SCREAMING_SNAKE_CASE           =     'harmonai/maestro-150k'
         SCREAMING_SNAKE_CASE           =     UNetaDModel.from_pretrained(__UpperCamelCase		,     subfolder='unet'      )
         model.to(__UpperCamelCase      )

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

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

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

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






class 					UpperCamelCase     (	lowercase_  ,				lowercase_  ,				unittest.TestCase				):

    """simple docstring"""



    UpperCAmelCase_    :				Any						  =    UNetaDModel
    UpperCAmelCase_    :				str						  =    "sample"


    @property
    def 							A (    self      )   ->   Dict:
         """simple docstring"""







         SCREAMING_SNAKE_CASE           =     4
         SCREAMING_SNAKE_CASE           =     14
         SCREAMING_SNAKE_CASE           =     16

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

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


    @property
    def 							A (    self      )   ->   str:
         """simple docstring"""







         return (4, 14, 16)


    @property
    def 							A (    self      )   ->   str:
         """simple docstring"""







         return (4, 14, 1)


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







         super().test_determinism()


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







         super().test_outputs_equivalence()


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







         super().test_from_save_pretrained()


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







         super().test_from_save_pretrained_variant()


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







         super().test_model_from_pretrained()


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







         SCREAMING_SNAKE_CASE     ,   SCREAMING_SNAKE_CASE           =     self.prepare_init_args_and_inputs_for_common()
         SCREAMING_SNAKE_CASE           =     self.model_class(**__UpperCamelCase      )
         model.to(__UpperCamelCase      )
         model.eval()

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

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

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


    def 							A (    self      )   ->   Dict:
         """simple docstring"""







         pass


    def 							A (    self      )   ->   Tuple:
         """simple docstring"""







         pass


    def 							A (    self      )   ->   Union[str, Any]:
         """simple docstring"""







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


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







         SCREAMING_SNAKE_CASE     ,   SCREAMING_SNAKE_CASE           =     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      )
         SCREAMING_SNAKE_CASE           =     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 							A (    self      )   ->   Any:
         """simple docstring"""







         SCREAMING_SNAKE_CASE     ,   SCREAMING_SNAKE_CASE           =     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      )

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

         with torch.no_grad():
              SCREAMING_SNAKE_CASE           =     value_function(__UpperCamelCase		,     __UpperCamelCase      ).sample

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




    def 							A (    self      )   ->   Any:
         """simple docstring"""







         pass


"""simple docstring"""




import abc
import copy
import dataclasses
from dataclasses import dataclass
from typing import ClassVar, Dict, Type, TypeVar

from ..features import Features


snake_case			       =		TypeVar('T', bound='TaskTemplate')




@dataclass(frozen=__magic_name__				)
class 					UpperCamelCase     (	abc.ABC				):

				"""simple docstring"""



				UpperCAmelCase_    :				str
				UpperCAmelCase_    :				ClassVar[Features]
				UpperCAmelCase_    :				ClassVar[Features]


				def 							A (    self		,     lowercase__      )   ->   T:
									"""simple docstring"""







									return copy.deepcopy(self      )


				@property
				def 							A (    self      )   ->   Features:
									"""simple docstring"""







									return Features(**self.input_schema		,     **self.label_schema      )


				@property
				@abc.abstractmethod
				def 							A (    self      )   ->   Dict[str, str]:
									"""simple docstring"""







									raise NotImplementedError




				@classmethod
				def 							A (    cls		,     lowercase__      )   ->   T:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     {f.name for f in dataclasses.fields(cls      )}
									return cls(**{k: v for k, v in template_dict.items() if k in field_names}      )
