


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







    """simple docstring"""




    UpperCAmelCase			:				List[Any]      =			UNetaDModel
    UpperCAmelCase			:				int      =			"""sample"""






    @property
    def       __snake_case				(   self							:       Dict):
      a		:						Dict 				=     4
      a		:						Optional[Any] 				=     14
      a		:						List[str] 				=     16

      a		:						List[str] 				=     floats_tensor((batch_size, num_features, seq_len)).to(a__)
      a		:						str 				=     torch.tensor([10] * batch_size).to(a__)

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






    @property
    def       __snake_case				(   self							:       List[Any]):
      return (4, 14, 16)






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






    def       __snake_case				(   self							:       Any):
      pass






    def       __snake_case				(   self							:       str):
      pass






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






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






    @unittest.skipIf(torch_device == "mps"							,					"mish op not supported in MPS")
    def       __snake_case				(   self							:       List[Any]):
      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							:       Dict):
      super().test_model_from_pretrained()






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






    def       __snake_case				(   self							:       List[Any]):
      a		:						List[str] 				=     {
          "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		:						int 				=     self.dummy_input
      return init_dict, inputs_dict






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

      with torch.no_grad():
        a		:						Union[str, Any] 				=     model(a__							,					a__).sample.permute(0							,					2							,					1)

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






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






    @slow
    def       __snake_case				(   self							:       List[str]):
      a		:						Dict 				=     "harmonai/maestro-150k"
      a		:						Optional[Any] 				=     UNetaDModel.from_pretrained(a__							,					subfolder="unet")
      model.to(a__)

      a		:						List[str] 				=     65536
      a		:						Optional[int] 				=     torch.sin(torch.arange(a__)[None, None, :].repeat(1							,					2							,					1)).to(a__)
      a		:						List[Any] 				=     torch.tensor([1]).to(a__)

      with torch.no_grad():
        a		:						List[Any] 				=     model(a__							,					a__).sample

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

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






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







    """simple docstring"""




    UpperCAmelCase			:				Optional[int]      =			UNetaDModel
    UpperCAmelCase			:				int      =			"""sample"""






    @property
    def       __snake_case				(   self							:       Tuple):
      a		:						Any 				=     4
      a		:						List[Any] 				=     14
      a		:						Optional[Any] 				=     16

      a		:						Optional[Any] 				=     floats_tensor((batch_size, num_features, seq_len)).to(a__)
      a		:						Optional[Any] 				=     torch.tensor([10] * batch_size).to(a__)

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






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






    @property
    def       __snake_case				(   self							:       Union[str, Any]):
      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							:       Dict):
      super().test_outputs_equivalence()






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






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






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






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

      with torch.no_grad():
        a		:						Union[str, Any] 				=     model(**a__)

        if isinstance(a__							,					a__):
          a		:						str 				=     output.sample

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






    def       __snake_case				(   self							:       Optional[Any]):
      pass






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






    def       __snake_case				(   self							:       Union[str, Any]):
      a		:						List[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",
      }
      a		:						Tuple 				=     self.dummy_input
      return init_dict, inputs_dict






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

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

      with torch.no_grad():
        a		:						List[str] 				=     value_function(a__							,					a__).sample

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






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








class 		__magic_name__				:


  """simple docstring"""





  class 		__magic_name__				:


    """simple docstring"""





    _UpperCamelCase						 = "_prev", "_data", "_next"


    def __init__(					self	,		a__	,		a__	,		a__		):
      _lowerCamelCase      =						link_p
      _lowerCamelCase      =						element
      _lowerCamelCase      =						link_n


    def     _UpperCAmelCase  (					self		):
      return (
          f''' Prev -> {self._prev is not None}, Next -> {self._next is not None}'''
      )







  def __init__(					self		):
    _lowerCamelCase      =						self._Node(a__	,		a__	,		a__		)
    _lowerCamelCase      =						self._Node(a__	,		a__	,		a__		)
    _lowerCamelCase      =						self._trailer
    _lowerCamelCase      =						self._header
    _lowerCamelCase      =						0


  def __len__(					self		):
    return self._size


  def     _UpperCAmelCase  (					self		):
    return self.__len__() == 0


  def     _UpperCAmelCase  (					self	,		a__	,		a__	,		a__		):
    # Create new_node by setting it's prev.link -> header
    # setting it's next.link -> trailer
    _lowerCamelCase      =						self._Node(a__	,		a__	,		a__		)
    _lowerCamelCase      =						new_node
    _lowerCamelCase      =						new_node
    self._size += 1
    return self


  def     _UpperCAmelCase  (					self	,		a__		):
    _lowerCamelCase      =						node._prev
    _lowerCamelCase      =						node._next

    _lowerCamelCase      =						successor
    _lowerCamelCase      =						predecessor
    self._size -= 1
    _lowerCamelCase      =						node._data
    _lowerCamelCase      =						_lowerCamelCase      =						_lowerCamelCase      =						None
    del node
    return temp







class 		__magic_name__				( _DoublyLinkedBase					):


  """simple docstring"""





  def     _UpperCAmelCase  (					self		):
    if self.is_empty():
      raise Exception('''List is empty'''		)
    return self._header._next._data


  def     _UpperCAmelCase  (					self		):
    if self.is_empty():
      raise Exception('''List is empty'''		)
    return self._trailer._prev._data


  def     _UpperCAmelCase  (					self	,		a__		):
    return self._insert(self._header	,		a__	,		self._header._next		)


  def     _UpperCAmelCase  (					self	,		a__		):
    return self._insert(self._trailer._prev	,		a__	,		self._trailer		)


  def     _UpperCAmelCase  (					self		):
    if self.is_empty():
      raise IndexError('''remove_first from empty list'''		)
    return self._delete(self._header._next		)


  def     _UpperCAmelCase  (					self		):
    if self.is_empty():
      raise IndexError('''remove_first from empty list'''		)
    return self._delete(self._trailer._prev		)




