



'''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  lowerCamelCase_   (    SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       unittest.TestCase					):
  """simple docstring"""




  a_							 =UNetaDModel
  a_							 ="""sample"""




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

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

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




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




  @property
  def 						_lowercase					(			self				:     int       ) -> str:
    return (4, 14, 16)




  def 						_lowercase					(			self				:     List[Any]       ) -> Dict:
    pass




  def 						_lowercase					(			self				:     List[str]       ) -> Dict:
    pass




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




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




  @unittest.skipIf(torch_device == 'mps'      ,     'mish op not supported in MPS'       )
  def 						_lowercase					(			self				:     Dict       ) -> Any:
    super().test_from_save_pretrained()




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




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




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




  def 						_lowercase					(			self				:     Tuple       ) -> Dict:
    __lowerCamelCase	:	Tuple   	=			{
        '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',
    }
    __lowerCamelCase	:	Tuple   	=			self.dummy_input
    return init_dict, inputs_dict




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

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

    with torch.no_grad():
      __lowerCamelCase	:	Optional[Any]   	=			model(_a      ,     _a       ).sample.permute(0      ,     2      ,     1       )

    __lowerCamelCase	:	Union[str, Any]   	=			output[0, -3:, -3:].flatten()
    # fmt: off
    __lowerCamelCase	:	List[str]   	=			torch.tensor([-2.137172, 1.1426016, 0.3688687, -0.766922, 0.7303146, 0.11038864, -0.4760633, 0.13270172, 0.02591348]       )
    # fmt: on
    self.assertTrue(torch.allclose(_a      ,     _a      ,     rtol=1e-3       )       )




  def 						_lowercase					(			self				:     Optional[int]       ) -> Optional[int]:
    # Not implemented yet for this UNet
    pass




  @slow
  def 						_lowercase					(			self				:     int       ) -> Optional[Any]:
    __lowerCamelCase	:	int   	=			'harmonai/maestro-150k'
    __lowerCamelCase	:	Union[str, Any]   	=			UNetaDModel.from_pretrained(_a      ,     subfolder='unet'       )
    model.to(_a       )

    __lowerCamelCase	:	int   	=			6_5536
    __lowerCamelCase	:	Dict   	=			torch.sin(torch.arange(_a       )[None, None, :].repeat(1      ,     2      ,     1       )       ).to(_a       )
    __lowerCamelCase	:	List[str]   	=			torch.tensor([1]       ).to(_a       )

    with torch.no_grad():
      __lowerCamelCase	:	List[str]   	=			model(_a      ,     _a       ).sample

    __lowerCamelCase	:	Optional[int]   	=			output.abs().sum()
    __lowerCamelCase	:	List[Any]   	=			output.abs().max()

    assert (output_sum - 224.0896).abs() < 0.5
    assert (output_max - 0.0607).abs() < 4e-4






class  lowerCamelCase_   (    SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       unittest.TestCase					):
  """simple docstring"""




  a_							 =UNetaDModel
  a_							 ="""sample"""




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

    __lowerCamelCase	:	Tuple   	=			floats_tensor((batch_size, num_features, seq_len)       ).to(_a       )
    __lowerCamelCase	:	Optional[Any]   	=			torch.tensor([10] * batch_size       ).to(_a       )

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




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




  @property
  def 						_lowercase					(			self				:     Dict       ) -> List[str]:
    return (4, 14, 1)




  @unittest.skipIf(torch_device == 'mps'      ,     'mish op not supported in MPS'       )
  def 						_lowercase					(			self				:     List[str]       ) -> Union[str, Any]:
    super().test_determinism()




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




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




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




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




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

    with torch.no_grad():
      __lowerCamelCase	:	Tuple   	=			model(**_a       )

      if isinstance(_a      ,     _a       ):
        __lowerCamelCase	:	Optional[int]   	=			output.sample

    self.assertIsNotNone(_a       )
    __lowerCamelCase	:	Union[str, Any]   	=			torch.Size((inputs_dict['sample'].shape[0], 1)       )
    self.assertEqual(output.shape      ,     _a      ,     'Input and output shapes do not match'       )




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




  def 						_lowercase					(			self				:     str       ) -> Optional[int]:
    pass




  def 						_lowercase					(			self				:     Any       ) -> List[Any]:
    __lowerCamelCase	:	Tuple   	=			{
        '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',
    }
    __lowerCamelCase	:	int   	=			self.dummy_input
    return init_dict, inputs_dict




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

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

    with torch.no_grad():
      __lowerCamelCase	:	Any   	=			value_function(_a      ,     _a       ).sample

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




  def 						_lowercase					(			self				:     Optional[Any]       ) -> int:
    # Not implemented yet for this UNet
    pass











'''simple docstring'''





import unittest

from diffusers import __version__
from diffusers.utils import deprecate






class  lowerCamelCase_   (    unittest.TestCase					):
  """simple docstring"""




  a_							 =""".""".join([str(int(__version__.split("""."""					)[0]					) + 1					)] + __version__.split("""."""					)[1:]					)
  a_							 ="""0.0.1"""




  def 						_lowercase					(			self				:     Optional[Any]       ) -> int:
    __lowerCamelCase	:	Any   	=			{'deprecated_arg': 4}

    with self.assertWarns(_a       ) as warning:
      __lowerCamelCase	:	Optional[int]   	=			deprecate('deprecated_arg'      ,     self.higher_version      ,     'message'      ,     take_from=_a       )

    assert output == 4
    assert (
        str(warning.warning       )
        == f'The `deprecated_arg` argument is deprecated and will be removed in version {self.higher_version}.'
        " message"
    )




  def 						_lowercase					(			self				:     Union[str, Any]       ) -> Optional[int]:
    __lowerCamelCase	:	int   	=			{'deprecated_arg': 4}

    with self.assertWarns(_a       ) as warning:
      __lowerCamelCase	:	Optional[Any]   	=			deprecate(('deprecated_arg', self.higher_version, 'message')      ,     take_from=_a       )

    assert output == 4
    assert (
        str(warning.warning       )
        == f'The `deprecated_arg` argument is deprecated and will be removed in version {self.higher_version}.'
        " message"
    )




  def 						_lowercase					(			self				:     int       ) -> Optional[Any]:
    __lowerCamelCase	:	Optional[int]   	=			{'deprecated_arg_1': 4, 'deprecated_arg_2': 8}
    with self.assertWarns(_a       ) as warning:
      __lowerCamelCase	,__lowerCamelCase	:	Any   	=			deprecate(
          ('deprecated_arg_1', self.higher_version, 'Hey')      ,     ('deprecated_arg_2', self.higher_version, 'Hey')      ,     take_from=_a      ,     )
    assert output_a == 4
    assert output_a == 8
    assert (
        str(warning.warnings[0].message       )
        == "The `deprecated_arg_1` argument is deprecated and will be removed in version"
        f' {self.higher_version}. Hey'
    )
    assert (
        str(warning.warnings[1].message       )
        == "The `deprecated_arg_2` argument is deprecated and will be removed in version"
        f' {self.higher_version}. Hey'
    )




  def 						_lowercase					(			self				:     str       ) -> Optional[Any]:
    __lowerCamelCase	:	Optional[int]   	=			{'deprecated_arg': 4}

    with self.assertRaises(_a       ) as error:
      deprecate(('wrong_arg', self.higher_version, 'message')      ,     take_from=_a       )

    assert "test_deprecate_function_incorrect_arg in" in str(error.exception       )
    assert "line" in str(error.exception       )
    assert "got an unexpected keyword argument `deprecated_arg`" in str(error.exception       )




  def 						_lowercase					(			self				:     Optional[Any]       ) -> Union[str, Any]:
    with self.assertWarns(_a       ) as warning:
      deprecate(('deprecated_arg', self.higher_version, 'message')       )

    assert (
        str(warning.warning       )
        == f'`deprecated_arg` is deprecated and will be removed in version {self.higher_version}. message'
    )




  def 						_lowercase					(			self				:     Optional[Any]       ) -> Union[str, Any]:
    with self.assertWarns(_a       ) as warning:
      deprecate(
          ('deprecated_arg_1', self.higher_version, 'Hey')      ,     ('deprecated_arg_2', self.higher_version, 'Hey')      ,     )
    assert (
        str(warning.warnings[0].message       )
        == f'`deprecated_arg_1` is deprecated and will be removed in version {self.higher_version}. Hey'
    )
    assert (
        str(warning.warnings[1].message       )
        == f'`deprecated_arg_2` is deprecated and will be removed in version {self.higher_version}. Hey'
    )




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






    class  lowerCamelCase_   :
      """simple docstring"""




      a_							 =5

    with self.assertWarns(_a       ) as warning:
      __lowerCamelCase	:	int   	=			deprecate(('arg', self.higher_version, 'message')      ,     take_from=Args()       )

    assert arg == 5
    assert (
        str(warning.warning       )
        == f'The `arg` attribute is deprecated and will be removed in version {self.higher_version}. message'
    )




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






    class  lowerCamelCase_   :
      """simple docstring"""




      a_							 =5
      a_							 =7

    with self.assertWarns(_a       ) as warning:
      __lowerCamelCase	,__lowerCamelCase	:	List[str]   	=			deprecate(
          ('arg', self.higher_version, 'message')      ,     ('foo', self.higher_version, 'message')      ,     ('does not exist', self.higher_version, 'message')      ,     take_from=Args()      ,     )

    assert arg_a == 5
    assert arg_a == 7
    assert (
        str(warning.warning       )
        == f'The `arg` attribute is deprecated and will be removed in version {self.higher_version}. message'
    )
    assert (
        str(warning.warnings[0].message       )
        == f'The `arg` attribute is deprecated and will be removed in version {self.higher_version}. message'
    )
    assert (
        str(warning.warnings[1].message       )
        == f'The `foo` attribute is deprecated and will be removed in version {self.higher_version}. message'
    )




  def 						_lowercase					(			self				:     Union[str, Any]       ) -> List[str]:
    __lowerCamelCase	:	Union[str, Any]   	=			{'deprecated_arg': 4}

    with self.assertRaises(_a       ) as error:
      deprecate(('wrong_arg', self.lower_version, 'message')      ,     take_from=_a       )

    assert (
        str(error.exception       )
        == "The deprecation tuple ('wrong_arg', '0.0.1', 'message') should be removed since diffusers' version"
        f' {__version__} is >= {self.lower_version}'
    )




  def 						_lowercase					(			self				:     Optional[int]       ) -> str:
    with self.assertWarns(_a       ) as warning:
      deprecate(('deprecated_arg', self.higher_version, 'This message is better!!!')      ,     standard_warn=_a       )

    assert str(warning.warning       ) == "This message is better!!!"




  def 						_lowercase					(			self				:     List[Any]       ) -> Union[str, Any]:
    with self.assertWarns(_a       ) as warning:
      deprecate(('deprecated_arg', self.higher_version, 'This message is better!!!')      ,     standard_warn=_a       )
    assert str(warning.warning       ) == "This message is better!!!"
    assert "diffusers/tests/others/test_utils.py" in warning.filename






