


'''simple docstring'''






import gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()






class    __magic_name__				(						lowerCAmelCase   ,unittest.TestCase							):
       UpperCAmelCase		 =PriorTransformer
       UpperCAmelCase		 ="hidden_states"






       @property
       def      lowerCAmelCase						(     self)		->    Dict:




          '''simple docstring'''
          _UpperCAmelCase  :  Optional[int]				       =4
          _UpperCAmelCase  :  Optional[int]				       =8
          _UpperCAmelCase  :  List[str]				       =7

          _UpperCAmelCase  :  Optional[int]				       =floats_tensor((batch_size, embedding_dim)).to(snake_case)

          _UpperCAmelCase  :  int				       =floats_tensor((batch_size, embedding_dim)).to(snake_case)
          _UpperCAmelCase  :  Dict				       =floats_tensor((batch_size, num_embeddings, embedding_dim)).to(snake_case)

          return {
              "hidden_states": hidden_states,
              "timestep": 2,
              "proj_embedding": proj_embedding,
              "encoder_hidden_states": encoder_hidden_states,
          }






       def      lowerCAmelCase						(     self    ,				snake_case=0)		->    int:




          '''simple docstring'''
          torch.manual_seed(snake_case)
          _UpperCAmelCase  :  Optional[Any]				       =4
          _UpperCAmelCase  :  Union[str, Any]				       =8
          _UpperCAmelCase  :  Dict				       =7

          _UpperCAmelCase  :  Optional[Any]				       =torch.randn((batch_size, embedding_dim)).to(snake_case)

          _UpperCAmelCase  :  List[Any]				       =torch.randn((batch_size, embedding_dim)).to(snake_case)
          _UpperCAmelCase  :  Union[str, Any]				       =torch.randn((batch_size, num_embeddings, embedding_dim)).to(snake_case)

          return {
              "hidden_states": hidden_states,
              "timestep": 2,
              "proj_embedding": proj_embedding,
              "encoder_hidden_states": encoder_hidden_states,
          }






       @property
       def      lowerCAmelCase						(     self)		->    Dict:




          '''simple docstring'''
          return (4, 8)






       @property
       def      lowerCAmelCase						(     self)		->    Any:




          '''simple docstring'''
          return (4, 8)






       def      lowerCAmelCase						(     self)		->    Tuple:




          '''simple docstring'''
          _UpperCAmelCase  :  Tuple				       ={
              'num_attention_heads': 2,
              'attention_head_dim': 4,
              'num_layers': 2,
              'embedding_dim': 8,
              'num_embeddings': 7,
              'additional_embeddings': 4,
          }
          _UpperCAmelCase  :  Dict				       =self.dummy_input
          return init_dict, inputs_dict






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




          '''simple docstring'''
          _UpperCAmelCase       ,			_UpperCAmelCase  :  Optional[int]				       =PriorTransformer.from_pretrained(
              'hf-internal-testing/prior-dummy'    ,				output_loading_info=snake_case)
          self.assertIsNotNone(snake_case)
          self.assertEqual(len(loading_info['missing_keys'])    ,				0)

          model.to(snake_case)
          _UpperCAmelCase  :  Dict				       =model(**self.dummy_input)[0]

          assert hidden_states is not None, "Make sure output is not None"






       def      lowerCAmelCase						(     self)		->    str:




          '''simple docstring'''
          _UpperCAmelCase       ,			_UpperCAmelCase  :  Dict				       =self.prepare_init_args_and_inputs_for_common()

          _UpperCAmelCase  :  Dict				       =self.model_class(**snake_case)
          _UpperCAmelCase  :  Optional[int]				       =inspect.signature(model.forward)
          # signature.parameters is an OrderedDict => so arg_names order is deterministic
          _UpperCAmelCase  :  Optional[int]				       =[*signature.parameters.keys()]

          _UpperCAmelCase  :  int				       =['hidden_states', 'timestep']
          self.assertListEqual(arg_names[:2]    ,				snake_case)






       def      lowerCAmelCase						(     self)		->    List[Any]:




          '''simple docstring'''
          _UpperCAmelCase  :  List[str]				       =PriorTransformer.from_pretrained('hf-internal-testing/prior-dummy')
          _UpperCAmelCase  :  int				       =model.to(snake_case)

          if hasattr(snake_case    ,				'set_default_attn_processor'):
             model.set_default_attn_processor()

          _UpperCAmelCase  :  Tuple				       =self.get_dummy_seed_input()

          with torch.no_grad():
             _UpperCAmelCase  :  Union[str, Any]				       =model(**snake_case)[0]

          _UpperCAmelCase  :  str				       =output[0, :5].flatten().cpu()
          print(snake_case)

          # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
          # the expected output slices are not the same for CPU and GPU.
          _UpperCAmelCase  :  List[Any]				       =torch.tensor([-1.34_36, -0.28_70, 0.75_38, 0.43_68, -0.02_39])
          self.assertTrue(torch_all_close(snake_case    ,				snake_case    ,				rtol=1E-2))






@slow
class    __magic_name__				(						unittest.TestCase							):






       def      lowerCAmelCase						(     self    ,				snake_case=1    ,				snake_case=7_6_8    ,				snake_case=7_7    ,				snake_case=0)		->    Union[str, Any]:




          '''simple docstring'''
          torch.manual_seed(snake_case)
          _UpperCAmelCase  :  Optional[Any]				       =batch_size
          _UpperCAmelCase  :  int				       =embedding_dim
          _UpperCAmelCase  :  Union[str, Any]				       =num_embeddings

          _UpperCAmelCase  :  List[str]				       =torch.randn((batch_size, embedding_dim)).to(snake_case)

          _UpperCAmelCase  :  Tuple				       =torch.randn((batch_size, embedding_dim)).to(snake_case)
          _UpperCAmelCase  :  Tuple				       =torch.randn((batch_size, num_embeddings, embedding_dim)).to(snake_case)

          return {
              "hidden_states": hidden_states,
              "timestep": 2,
              "proj_embedding": proj_embedding,
              "encoder_hidden_states": encoder_hidden_states,
          }






       def      lowerCAmelCase						(     self)		->    int:




          '''simple docstring'''
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()






       @parameterized.expand(
           [
               # fmt: off
               [1_3, [-0.58_61,  0.12_83, -0.09_31,  0.08_82,  0.44_76,  0.13_29, -0.04_98,  0.06_40]],
               [3_7, [-0.49_13,  0.01_10, -0.04_83,  0.05_41,  0.49_54, -0.01_70,  0.03_54,  0.16_51]],
               # fmt: on
           ])
       def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case)		->    Dict:




          '''simple docstring'''
          _UpperCAmelCase  :  int				       =PriorTransformer.from_pretrained('kandinsky-community/kandinsky-2-1-prior'    ,				subfolder='prior')
          model.to(snake_case)
          _UpperCAmelCase  :  str				       =self.get_dummy_seed_input(seed=snake_case)

          with torch.no_grad():
             _UpperCAmelCase  :  int				       =model(**snake_case)[0]

          assert list(sample.shape) == [1, 7_6_8]

          _UpperCAmelCase  :  Any				       =sample[0, :8].flatten().cpu()
          print(snake_case)
          _UpperCAmelCase  :  Tuple				       =torch.tensor(snake_case)

          assert torch_all_close(snake_case    ,				snake_case    ,				atol=1E-3)





'''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 PaddingStrategy, TensorType, logging


lowercase				=logging.get_logger(__name__)






class    __magic_name__				(						lowerCAmelCase							):
       UpperCAmelCase		 =["input_features", "attention_mask"]






       def __init__(     self    ,				snake_case=8_0    ,				snake_case=1_6_0_0_0    ,				snake_case=8_0    ,				snake_case=0.0    ,				snake_case=True    ,				snake_case=True    ,				snake_case=True    ,				**snake_case    ,				)		->    str:




          '''simple docstring'''
          super().__init__(feature_size=snake_case    ,				sampling_rate=snake_case    ,				padding_value=snake_case    ,				**snake_case)
          _UpperCAmelCase  :  Optional[Any]				       =num_mel_bins
          _UpperCAmelCase  :  Optional[int]				       =do_ceptral_normalize
          _UpperCAmelCase  :  Optional[Any]				       =normalize_means
          _UpperCAmelCase  :  Optional[Any]				       =normalize_vars
          _UpperCAmelCase  :  Tuple				       =True






       def      lowerCAmelCase						(     self    ,				snake_case    ,				)		->    np.ndarray:




          '''simple docstring'''
          _UpperCAmelCase  :  List[Any]				       =waveform * (2**1_5)  # Kaldi compliance: 16-bit signed integers
          _UpperCAmelCase  :  Dict				       =torch.from_numpy(snake_case).unsqueeze(0)
          _UpperCAmelCase  :  Union[str, Any]				       =ta_kaldi.fbank(snake_case    ,				num_mel_bins=self.num_mel_bins    ,				sample_frequency=self.sampling_rate)
          return features.numpy()






       @staticmethod
       def      lowerCAmelCase						(     snake_case    ,				snake_case    ,				snake_case = True    ,				snake_case = True    ,				snake_case = 0.0    ,				)		->    np.ndarray:




          '''simple docstring'''
          # make sure we normalize float32 arrays
          if normalize_means:
             _UpperCAmelCase  :  int				       =x[:input_length].mean(axis=0)
             _UpperCAmelCase  :  List[Any]				       =np.subtract(snake_case    ,				snake_case)
          if normalize_vars:
             _UpperCAmelCase  :  List[Any]				       =x[:input_length].std(axis=0)
             _UpperCAmelCase  :  Optional[int]				       =np.divide(snake_case    ,				snake_case)

          if input_length < x.shape[0]:
             _UpperCAmelCase  :  Dict				       =padding_value

          # make sure array is in float32
          _UpperCAmelCase  :  str				       =x.astype(np.floataa)

          return x






       def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case = None)		->    List[np.ndarray]:




          '''simple docstring'''
          _UpperCAmelCase  :  str				       =attention_mask.sum(-1) if attention_mask is not None else [x.shape[0] for x in input_features]
          return [
              self.utterance_cmvn(snake_case    ,				snake_case    ,				self.normalize_means    ,				self.normalize_vars    ,				self.padding_value)
              for x, n in zip(snake_case    ,				snake_case)
          ]






       def __call__(     self    ,				snake_case    ,				snake_case = False    ,				snake_case = None    ,				snake_case = False    ,				snake_case = None    ,				snake_case = None    ,				snake_case = None    ,				snake_case = None    ,				**snake_case    ,				)		->    BatchFeature:




          '''simple docstring'''

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

          _UpperCAmelCase  :  Optional[int]				       =isinstance(snake_case    ,				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}")
          _UpperCAmelCase  :  List[Any]				       =is_batched_numpy or (
              isinstance(snake_case    ,				(list, tuple)) and (isinstance(raw_speech[0]    ,				(np.ndarray, tuple, list)))
          )

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

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

          # extract fbank features
          _UpperCAmelCase  :  Optional[Any]				       =[self._extract_fbank_features(snake_case) for waveform in raw_speech]

          # convert into correct format for padding
          _UpperCAmelCase  :  List[str]				       =BatchFeature({'input_features': features})

          _UpperCAmelCase  :  Any				       =self.pad(
              snake_case    ,				padding=snake_case    ,				max_length=snake_case    ,				truncation=snake_case    ,				pad_to_multiple_of=snake_case    ,				return_attention_mask=snake_case    ,				**snake_case    ,				)

          # make sure list is in array format
          _UpperCAmelCase  :  Dict				       =padded_inputs.get('input_features')
          if isinstance(input_features[0]    ,				snake_case):
             _UpperCAmelCase  :  Any				       =[np.asarray(snake_case    ,				dtype=np.floataa) for feature in input_features]

          _UpperCAmelCase  :  int				       =padded_inputs.get('attention_mask')
          if attention_mask is not None:
             _UpperCAmelCase  :  Tuple				       =[np.asarray(snake_case    ,				dtype=np.intaa) for array in attention_mask]

          # Utterance-level cepstral mean and variance normalization
          if self.do_ceptral_normalize:
             _UpperCAmelCase  :  Optional[Any]				       =(
                 np.array(snake_case    ,				dtype=np.intaa)
                 if self._get_padding_strategies(snake_case    ,				max_length=snake_case) is not PaddingStrategy.DO_NOT_PAD
                 else None
             )
             _UpperCAmelCase  :  Optional[int]				       =self.normalize(
                 padded_inputs['input_features']    ,				attention_mask=snake_case)

          if return_tensors is not None:
             _UpperCAmelCase  :  List[Any]				       =padded_inputs.convert_to_tensors(snake_case)

          return padded_inputs

