





'''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    (a__    ,   a__    ,   unittest.TestCase       ):
  _lowercase	:						str							       =					UNetaDModel
  _lowercase	:						Optional[int]							       =					"""sample"""







  @property
  def 					UpperCAmelCase_					(		self		)		->    Union[str, Any]:






        """simple docstring"""





        _snake_case		:						str         =  4
        _snake_case		:						Dict         =  14
        _snake_case		:						List[Any]         =  16

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

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







  @property
  def 					UpperCAmelCase_					(		self		)		->    List[Any]:






        """simple docstring"""





        return (4, 14, 16)







  @property
  def 					UpperCAmelCase_					(		self		)		->    List[str]:






        """simple docstring"""





        return (4, 14, 16)







  def 					UpperCAmelCase_					(		self		)		->    Dict:






        """simple docstring"""





        pass







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






        """simple docstring"""





        pass







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






        """simple docstring"""





        super().test_determinism()







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






        """simple docstring"""





        super().test_outputs_equivalence()







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






        """simple docstring"""





        super().test_from_save_pretrained()







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






        """simple docstring"""





        super().test_from_save_pretrained_variant()







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






        """simple docstring"""





        super().test_model_from_pretrained()







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






        """simple docstring"""





        super().test_output()







  def 					UpperCAmelCase_					(		self		)		->    Dict:






        """simple docstring"""





        _snake_case		:						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''',
        }
        _snake_case		:						Union[str, Any]         =  self.dummy_input
        return init_dict, inputs_dict







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






        """simple docstring"""





        _snake_case  ,				_snake_case		:						Optional[Any]         =  UNetaDModel.from_pretrained(
            '''bglick13/hopper-medium-v2-value-function-hor32'''  ,  output_loading_info=lowercase__  ,  subfolder='''unet'''		)
        self.assertIsNotNone(lowercase__		)
        self.assertEqual(len(loading_info['''missing_keys''']		)  ,  0		)

        model.to(lowercase__		)
        _snake_case		:						Optional[int]         =  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 					UpperCAmelCase_					(		self		)		->    Dict:






        """simple docstring"""





        _snake_case		:						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		)

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

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

        _snake_case		:						Optional[int]         =  output[0, -3:, -3:].flatten()
        # fmt: off
        _snake_case		:						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(lowercase__  ,  lowercase__  ,  rtol=1E-3		)		)







  def 					UpperCAmelCase_					(		self		)		->    Optional[Any]:






        """simple docstring"""





        pass







  @slow
  def 					UpperCAmelCase_					(		self		)		->    Tuple:






        """simple docstring"""





        _snake_case		:						Tuple         =  '''harmonai/maestro-150k'''
        _snake_case		:						List[str]         =  UNetaDModel.from_pretrained(lowercase__  ,  subfolder='''unet'''		)
        model.to(lowercase__		)

        _snake_case		:						List[Any]         =  65_536
        _snake_case		:						str         =  torch.sin(torch.arange(lowercase__		)[None, None, :].repeat(1  ,  2  ,  1		)		).to(lowercase__		)
        _snake_case		:						Optional[Any]         =  torch.tensor([1]		).to(lowercase__		)

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

        _snake_case		:						int         =  output.abs().sum()
        _snake_case		:						List[str]         =  output.abs().max()

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





class lowerCamelCase    (a__    ,   a__    ,   unittest.TestCase       ):
  _lowercase	:						List[Any]							       =					UNetaDModel
  _lowercase	:						Optional[Any]							       =					"""sample"""







  @property
  def 					UpperCAmelCase_					(		self		)		->    List[Any]:






        """simple docstring"""





        _snake_case		:						List[Any]         =  4
        _snake_case		:						str         =  14
        _snake_case		:						Optional[int]         =  16

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

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







  @property
  def 					UpperCAmelCase_					(		self		)		->    List[Any]:






        """simple docstring"""





        return (4, 14, 16)







  @property
  def 					UpperCAmelCase_					(		self		)		->    Any:






        """simple docstring"""





        return (4, 14, 1)







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






        """simple docstring"""





        super().test_determinism()







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






        """simple docstring"""





        super().test_outputs_equivalence()







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






        """simple docstring"""





        super().test_from_save_pretrained()







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






        """simple docstring"""





        super().test_from_save_pretrained_variant()







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






        """simple docstring"""





        super().test_model_from_pretrained()







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






        """simple docstring"""





        _snake_case  ,				_snake_case		:						List[Any]         =  self.prepare_init_args_and_inputs_for_common()
        _snake_case		:						Optional[int]         =  self.model_class(**lowercase__		)
        model.to(lowercase__		)
        model.eval()

        with torch.no_grad():
              _snake_case		:						Dict         =  model(**lowercase__		)

              if isinstance(lowercase__  ,  lowercase__		):
                    _snake_case		:						Optional[Any]         =  output.sample

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







  def 					UpperCAmelCase_					(		self		)		->    Optional[int]:






        """simple docstring"""





        pass







  def 					UpperCAmelCase_					(		self		)		->    str:






        """simple docstring"""





        pass







  def 					UpperCAmelCase_					(		self		)		->    str:






        """simple docstring"""





        _snake_case		:						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''',
        }
        _snake_case		:						Union[str, Any]         =  self.dummy_input
        return init_dict, inputs_dict







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






        """simple docstring"""





        _snake_case  ,				_snake_case		:						Any         =  UNetaDModel.from_pretrained(
            '''bglick13/hopper-medium-v2-value-function-hor32'''  ,  output_loading_info=lowercase__  ,  subfolder='''value_function'''		)
        self.assertIsNotNone(lowercase__		)
        self.assertEqual(len(vf_loading_info['''missing_keys''']		)  ,  0		)

        value_function.to(lowercase__		)
        _snake_case		:						Union[str, 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 					UpperCAmelCase_					(		self		)		->    str:






        """simple docstring"""





        _snake_case  ,				_snake_case		:						List[Any]         =  UNetaDModel.from_pretrained(
            '''bglick13/hopper-medium-v2-value-function-hor32'''  ,  output_loading_info=lowercase__  ,  subfolder='''value_function'''		)
        torch.manual_seed(0		)
        if torch.cuda.is_available():
              torch.cuda.manual_seed_all(0		)

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

        with torch.no_grad():
              _snake_case		:						int         =  value_function(lowercase__  ,  lowercase__		).sample

        # fmt: off
        _snake_case		:						Union[str, Any]         =  torch.tensor([165.25] * seq_len		)
        # fmt: on
        self.assertTrue(torch.allclose(lowercase__  ,  lowercase__  ,  rtol=1E-3		)		)







  def 					UpperCAmelCase_					(		self		)		->    Dict:






        """simple docstring"""





        pass









'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCAmelCase       :		Dict		=				{
    'configuration_marian': ['MARIAN_PRETRAINED_CONFIG_ARCHIVE_MAP', 'MarianConfig', 'MarianOnnxConfig'],
}

try:
   if not is_sentencepiece_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase       :		Tuple		=				['MarianTokenizer']

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase       :		List[str]		=				[
       'MARIAN_PRETRAINED_MODEL_ARCHIVE_LIST',
       'MarianForCausalLM',
       'MarianModel',
       'MarianMTModel',
       'MarianPreTrainedModel',
   ]

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase       :		Any		=				['TFMarianModel', 'TFMarianMTModel', 'TFMarianPreTrainedModel']

try:
   if not is_flax_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase       :		Optional[Any]		=				['FlaxMarianModel', 'FlaxMarianMTModel', 'FlaxMarianPreTrainedModel']

if TYPE_CHECKING:
   from .configuration_marian import MARIAN_PRETRAINED_CONFIG_ARCHIVE_MAP, MarianConfig, MarianOnnxConfig

   try:
      if not is_sentencepiece_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_marian import MarianTokenizer

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_marian import (
          MARIAN_PRETRAINED_MODEL_ARCHIVE_LIST,
          MarianForCausalLM,
          MarianModel,
          MarianMTModel,
          MarianPreTrainedModel,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_marian import TFMarianModel, TFMarianMTModel, TFMarianPreTrainedModel

   try:
      if not is_flax_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_flax_marian import FlaxMarianModel, FlaxMarianMTModel, FlaxMarianPreTrainedModel

else:
   import sys

   UpperCAmelCase       :		int		=				_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)


