



"""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    __snake_case(					__A			,			__A			,			unittest.TestCase       ):
    _A       = UNetaDModel
    _A       = '''sample'''




    @property
    def   A   (   self			):





      '''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(lowercase__			)
      _SCREAMING_SNAKE_CASE        =			torch.tensor([10] * batch_size			).to(lowercase__			)

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




    @property
    def   A   (   self			):





      '''simple docstring'''



      return (4, 14, 16)




    @property
    def   A   (   self			):





      '''simple docstring'''



      return (4, 14, 16)




    def   A   (   self			):





      '''simple docstring'''



      pass




    def   A   (   self			):





      '''simple docstring'''



      pass




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





      '''simple docstring'''



      super().test_determinism()




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





      '''simple docstring'''



      super().test_outputs_equivalence()




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





      '''simple docstring'''



      super().test_from_save_pretrained()




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





      '''simple docstring'''



      super().test_from_save_pretrained_variant()




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





      '''simple docstring'''



      super().test_model_from_pretrained()




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





      '''simple docstring'''



      super().test_output()




    def   A   (   self			):





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





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			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__			)
      _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			):





      '''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(lowercase__    ,							lowercase__			).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(lowercase__    ,							lowercase__    ,							rtol=1e-3			)			)




    def   A   (   self			):





      '''simple docstring'''



      pass




    @slow
    def   A   (   self			):





      '''simple docstring'''



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

      _SCREAMING_SNAKE_CASE        =			65_536
      _SCREAMING_SNAKE_CASE        =			torch.sin(torch.arange(lowercase__			)[None, None, :].repeat(1    ,							2    ,							1			)			).to(lowercase__			)
      _SCREAMING_SNAKE_CASE        =			torch.tensor([1]			).to(lowercase__			)

      with torch.no_grad():
        _SCREAMING_SNAKE_CASE        =			model(lowercase__    ,							lowercase__			).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    __snake_case(					__A			,			__A			,			unittest.TestCase       ):
    _A       = UNetaDModel
    _A       = '''sample'''




    @property
    def   A   (   self			):





      '''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(lowercase__			)
      _SCREAMING_SNAKE_CASE        =			torch.tensor([10] * batch_size			).to(lowercase__			)

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




    @property
    def   A   (   self			):





      '''simple docstring'''



      return (4, 14, 16)




    @property
    def   A   (   self			):





      '''simple docstring'''



      return (4, 14, 1)




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





      '''simple docstring'''



      super().test_determinism()




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





      '''simple docstring'''



      super().test_outputs_equivalence()




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





      '''simple docstring'''



      super().test_from_save_pretrained()




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





      '''simple docstring'''



      super().test_from_save_pretrained_variant()




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





      '''simple docstring'''



      super().test_model_from_pretrained()




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





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			self.prepare_init_args_and_inputs_for_common()
      _SCREAMING_SNAKE_CASE        =			self.model_class(**lowercase__			)
      model.to(lowercase__			)
      model.eval()

      with torch.no_grad():
        _SCREAMING_SNAKE_CASE        =			model(**lowercase__			)

        if isinstance(lowercase__    ,							lowercase__			):
          _SCREAMING_SNAKE_CASE        =			output.sample

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




    def   A   (   self			):





      '''simple docstring'''



      pass




    def   A   (   self			):





      '''simple docstring'''



      pass




    def   A   (   self			):





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





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			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__			)
      _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			):





      '''simple docstring'''



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

      _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(lowercase__    ,							lowercase__			).sample

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




    def   A   (   self			):





      '''simple docstring'''



      pass








'''simple docstring'''







from typing import List, Optional, Union

import numpy as np
import torch
import torchaudio.compliance.kaldi as ta_kaldi

from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import TensorType, logging


__A   :							int     	=						logging.get_logger(__name__)






class        a__     (      lowerCamelCase    ):
     '''simple docstring'''







     _snake_case							   =	['input_values', 'attention_mask']






     def __init__(self     :						Optional[int]					,					lowercase__     :						Tuple=1					,					lowercase__     :						Optional[int]=16000					,					lowercase__     :						Any=128					,					lowercase__     :						Optional[Any]=1024					,					lowercase__     :						int=0.0					,					lowercase__     :						List[str]=True					,					lowercase__     :						Dict=-4.2_67_73_93					,					lowercase__     :						Tuple=4.5_68_99_74					,					lowercase__     :						Optional[int]=False					,					**lowercase__     :						Optional[int]					,					)	->							Tuple:
           super().__init__(feature_size=lowercase__					,					sampling_rate=lowercase__					,					padding_value=lowercase__					,					**lowercase__       )
           __A		 =   num_mel_bins
           __A		 =   max_length
           __A		 =   do_normalize
           __A		 =   mean
           __A		 =   std
           __A		 =   return_attention_mask






     def      _UpperCamelCase (self     :						List[Any]					,					lowercase__     :						np.ndarray					,					lowercase__     :						int					,					)	->							np.ndarray:
           __A		 =   torch.from_numpy(lowercase__       ).unsqueeze(0       )
           __A		 =   ta_kaldi.fbank(
               lowercase__					,					htk_compat=lowercase__					,					sample_frequency=self.sampling_rate					,					use_energy=lowercase__					,					window_type="hanning"					,					num_mel_bins=self.num_mel_bins					,					dither=0.0					,					frame_shift=10					,					)

           __A		 =   fbank.shape[0]
           __A		 =   max_length - n_frames

           # pad or truncate, depending on difference
           if difference > 0:
                 __A		 =   torch.nn.ZeroPadad((0, 0, 0, difference)       )
                 __A		 =   pad_module(lowercase__       )
           elif difference < 0:
                 __A		 =   fbank[0:max_length, :]

           __A		 =   fbank.numpy()

           return fbank






     def      _UpperCamelCase (self     :						int					,					lowercase__     :						np.ndarray       )	->							np.ndarray:
           return (input_values - (self.mean)) / (self.std * 2)






     def __call__(self     :						str					,					lowercase__     :						Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]					,					lowercase__     :						Optional[int] = None					,					lowercase__     :						Optional[Union[str, TensorType]] = None					,					**lowercase__     :						List[Any]					,					)	->							BatchFeature:

           if sampling_rate is not None:
                 if sampling_rate != self.sampling_rate:
                       raise ValueError(
                           f"The model corresponding to this feature extractor: {self} was trained using a sampling rate of"
                           f" {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled with"
                           f" {self.sampling_rate} and not {sampling_rate}."       )
           else:
                 logger.warning(
                     "It is strongly recommended to pass the `sampling_rate` argument to this function. "
                     "Failing to do so can result in silent errors that might be hard to debug."       )

           __A		 =   isinstance(lowercase__					,					np.ndarray       ) and len(raw_speech.shape       ) > 1
           if is_batched_numpy and len(raw_speech.shape       ) > 2:
                 raise ValueError(f"Only mono-channel audio is supported for input to {self}"       )
           __A		 =   is_batched_numpy or (
               isinstance(lowercase__					,					(list, tuple)       ) and (isinstance(raw_speech[0]					,					(np.ndarray, tuple, list)       ))
           )

           if is_batched:
                 __A		 =   [np.asarray(lowercase__					,					dtype=np.floataa       ) for speech in raw_speech]
           elif not is_batched and not isinstance(lowercase__					,					np.ndarray       ):
                 __A		 =   np.asarray(lowercase__					,					dtype=np.floataa       )
           elif isinstance(lowercase__					,					np.ndarray       ) and raw_speech.dtype is np.dtype(np.floataa       ):
                 __A		 =   raw_speech.astype(np.floataa       )

           # always return batch
           if not is_batched:
                 __A		 =   [raw_speech]

           # extract fbank features and pad/truncate to max_length
           __A		 =   [self._extract_fbank_features(lowercase__					,					max_length=self.max_length       ) for waveform in raw_speech]

           # convert into BatchFeature
           __A		 =   BatchFeature({"input_values": features}       )

           # make sure list is in array format
           __A		 =   padded_inputs.get("input_values"       )
           if isinstance(input_values[0]					,					lowercase__       ):
                 __A		 =   [np.asarray(lowercase__					,					dtype=np.floataa       ) for feature in input_values]

           # normalization
           if self.do_normalize:
                 __A		 =   [self.normalize(lowercase__       ) for feature in input_values]

           if return_tensors is not None:
                 __A		 =   padded_inputs.convert_to_tensors(lowercase__       )

           return padded_inputs


