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_					(				_a  ,  _a  ,  unittest.TestCase     ):
    UpperCAmelCase_      :							int    			=				UNetaDModel
    UpperCAmelCase_      :							Optional[Any]    			=				'sample'




    @property
    def     __lowerCamelCase			(       self	: Tuple )						->			int:
      a__  		=  4
      a__  		=  14
      a__  		=  16

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

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




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




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




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




    def     __lowerCamelCase			(       self	: Optional[Any] )						->			Optional[int]:
      pass




    @unittest.skipIf(torch_device == "mps"   ,   "mish op not supported in MPS" )
    def     __lowerCamelCase			(       self	: Tuple )						->			str:
      super().test_determinism()




    @unittest.skipIf(torch_device == "mps"   ,   "mish op not supported in MPS" )
    def     __lowerCamelCase			(       self	: Optional[int] )						->			List[str]:
      super().test_outputs_equivalence()




    @unittest.skipIf(torch_device == "mps"   ,   "mish op not supported in MPS" )
    def     __lowerCamelCase			(       self	: Union[str, Any] )						->			Optional[Any]:
      super().test_from_save_pretrained()




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




    @unittest.skipIf(torch_device == "mps"   ,   "mish op not supported in MPS" )
    def     __lowerCamelCase			(       self	: int )						->			Tuple:
      super().test_model_from_pretrained()




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




    def     __lowerCamelCase			(       self	: Any )						->			str:
      a__  		=  {
          "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",
      }
      a__  		=  self.dummy_input
      return init_dict, inputs_dict




    @unittest.skipIf(torch_device == "mps"   ,   "mish op not supported in MPS" )
    def     __lowerCamelCase			(       self	: int )						->			Union[str, Any]:
      a__		,a__  		=  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 )
      a__  		=  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     __lowerCamelCase			(       self	: int )						->			int:
      a__  		=  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 )

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

      with torch.no_grad():
        a__  		=  model(lowerCAmelCase   ,   lowerCAmelCase ).sample.permute(0   ,   2   ,   1 )

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




    def     __lowerCamelCase			(       self	: Union[str, Any] )						->			Union[str, Any]:
      # Not implemented yet for this UNet
      pass




    @slow
    def     __lowerCamelCase			(       self	: Tuple )						->			Tuple:
      a__  		=  "harmonai/maestro-150k"
      a__  		=  UNetaDModel.from_pretrained(lowerCAmelCase   ,   subfolder="unet" )
      model.to(lowerCAmelCase )

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

      with torch.no_grad():
        a__  		=  model(lowerCAmelCase   ,   lowerCAmelCase ).sample

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

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







class        UpperCAmelCase_					(				_a  ,  _a  ,  unittest.TestCase     ):
    UpperCAmelCase_      :							Optional[Any]    			=				UNetaDModel
    UpperCAmelCase_      :							int    			=				'sample'




    @property
    def     __lowerCamelCase			(       self	: List[str] )						->			Optional[int]:
      a__  		=  4
      a__  		=  14
      a__  		=  16

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

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




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




    @property
    def     __lowerCamelCase			(       self	: Tuple )						->			Union[str, Any]:
      return (4, 14, 1)




    @unittest.skipIf(torch_device == "mps"   ,   "mish op not supported in MPS" )
    def     __lowerCamelCase			(       self	: Any )						->			Tuple:
      super().test_determinism()




    @unittest.skipIf(torch_device == "mps"   ,   "mish op not supported in MPS" )
    def     __lowerCamelCase			(       self	: Dict )						->			List[Any]:
      super().test_outputs_equivalence()




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




    @unittest.skipIf(torch_device == "mps"   ,   "mish op not supported in MPS" )
    def     __lowerCamelCase			(       self	: Tuple )						->			List[Any]:
      super().test_from_save_pretrained_variant()




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




    @unittest.skipIf(torch_device == "mps"   ,   "mish op not supported in MPS" )
    def     __lowerCamelCase			(       self	: Tuple )						->			Union[str, Any]:
      # UNetRL is a value-function is different output shape
      a__		,a__  		=  self.prepare_init_args_and_inputs_for_common()
      a__  		=  self.model_class(**lowerCAmelCase )
      model.to(lowerCAmelCase )
      model.eval()

      with torch.no_grad():
        a__  		=  model(**lowerCAmelCase )

        if isinstance(lowerCAmelCase   ,   lowerCAmelCase ):
          a__  		=  output.sample

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




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




    def     __lowerCamelCase			(       self	: List[Any] )						->			Tuple:
      pass




    def     __lowerCamelCase			(       self	: int )						->			Any:
      a__  		=  {
          "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",
      }
      a__  		=  self.dummy_input
      return init_dict, inputs_dict




    @unittest.skipIf(torch_device == "mps"   ,   "mish op not supported in MPS" )
    def     __lowerCamelCase			(       self	: Union[str, Any] )						->			Tuple:
      a__		,a__  		=  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 )
      a__  		=  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     __lowerCamelCase			(       self	: Optional[int] )						->			str:
      a__		,a__  		=  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 )

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

      with torch.no_grad():
        a__  		=  value_function(lowerCAmelCase   ,   lowerCAmelCase ).sample

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




    def     __lowerCamelCase			(       self	: List[str] )						->			Optional[Any]:
      # Not implemented yet for this UNet
      pass


def 					__lowercase						(    lowerCamelCase__       )					->int:

  '''simple docstring'''




  if not isinstance(lowerCamelCase__			,					lowerCamelCase__       ):
    raise ValueError("double_factorial() only accepts integral values"       )
  if n < 0:
    raise ValueError("double_factorial() not defined for negative values"       )
  return 1 if n <= 1 else n * double_factorial(n - 2       )


if __name__ == "__main__":
     import doctest

     doctest.testmod()

