


from typing import List, Optional, Union

import numpy as np

from ....audio_utils import mel_filter_bank, optimal_fft_length, spectrogram, window_function
from ....feature_extraction_sequence_utils import SequenceFeatureExtractor
from ....feature_extraction_utils import BatchFeature
from ....file_utils import PaddingStrategy, TensorType
from ....utils import logging


SCREAMING_SNAKE_CASE       :		Tuple         =				logging.get_logger(__name__)






class      UpperCamelCase					(		lowercase__     ):


   '''simple docstring'''







   lowercase					:	List[Any]									=["""input_features""", """attention_mask"""]
   def __init__(     self   ,							UpperCamelCase_=80   ,							UpperCamelCase_=1_6000   ,							UpperCamelCase_=0.0   ,							UpperCamelCase_=10   ,							UpperCamelCase_=25   ,							UpperCamelCase_="hamming_window"   ,							UpperCamelCase_=3_2768.0   ,							UpperCamelCase_=0.97   ,							UpperCamelCase_=1.0   ,							UpperCamelCase_=True   ,							UpperCamelCase_=True   ,							UpperCamelCase_=False   ,							**UpperCamelCase_   ,							):
        super().__init__(feature_size=UpperCamelCase_   ,							sampling_rate=UpperCamelCase_   ,							padding_value=UpperCamelCase_   ,							**UpperCamelCase_		)

        lowercase_       :List[str]     =       feature_size
        lowercase_       :List[Any]     =       sampling_rate
        lowercase_       :Dict     =       padding_value
        lowercase_       :Any     =       hop_length
        lowercase_       :Optional[int]     =       win_length
        lowercase_       :Optional[int]     =       frame_signal_scale
        lowercase_       :List[str]     =       preemphasis_coeff
        lowercase_       :Optional[int]     =       mel_floor
        lowercase_       :Optional[int]     =       normalize_means
        lowercase_       :Optional[Any]     =       normalize_vars
        lowercase_       :List[str]     =       win_function
        lowercase_       :Tuple     =       return_attention_mask

        lowercase_       :Any     =       win_length * sampling_rate // 1000
        lowercase_       :List[Any]     =       hop_length * sampling_rate // 1000

        lowercase_       :List[Any]     =       optimal_fft_length(self.sample_size		)
        lowercase_       :Union[str, Any]     =       (self.n_fft // 2) + 1
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_		):
        if self.win_function == "hamming_window":
             lowercase_       :Optional[Any]     =       window_function(window_length=self.sample_size   ,							name=self.win_function   ,							periodic=UpperCamelCase_		)
        else:
             lowercase_       :List[str]     =       window_function(window_length=self.sample_size   ,							name=self.win_function		)

        lowercase_       :str     =       mel_filter_bank(
            num_frequency_bins=self.n_freqs   ,							num_mel_filters=self.feature_size   ,							min_frequency=0.0   ,							max_frequency=self.sampling_rate / 2.0   ,							sampling_rate=self.sampling_rate   ,							)

        lowercase_       :List[Any]     =       spectrogram(
            one_waveform * self.frame_signal_scale   ,							window=UpperCamelCase_   ,							frame_length=self.sample_size   ,							hop_length=self.sample_stride   ,							fft_length=self.n_fft   ,							center=UpperCamelCase_   ,							preemphasis=self.preemphasis_coeff   ,							mel_filters=UpperCamelCase_   ,							mel_floor=self.mel_floor   ,							log_mel='''log'''   ,							)
        return msfc_features.T
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		):
        # make sure we normalize float32 arrays
        if self.normalize_means:
             lowercase_       :Optional[Any]     =       x[:input_length].mean(axis=0		)
             lowercase_       :Tuple     =       np.subtract(UpperCamelCase_   ,							UpperCamelCase_		)
        if self.normalize_vars:
             lowercase_       :Optional[Any]     =       x[:input_length].std(axis=0		)
             lowercase_       :int     =       np.divide(UpperCamelCase_   ,							UpperCamelCase_		)

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

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

        return x
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_ = None		):
        lowercase_       :Optional[Any]     =       attention_mask.sum(-1		) if attention_mask is not None else [x.shape[0] for x in input_features]
        return [self._normalize_one(UpperCamelCase_   ,							UpperCamelCase_   ,							self.padding_value		) for x, n in zip(UpperCamelCase_   ,							UpperCamelCase_		)]




   def __call__(     self   ,							UpperCamelCase_   ,							UpperCamelCase_ = False   ,							UpperCamelCase_ = None   ,							UpperCamelCase_ = False   ,							UpperCamelCase_ = None   ,							UpperCamelCase_ = None   ,							UpperCamelCase_ = None   ,							UpperCamelCase_ = None   ,							**UpperCamelCase_   ,							):

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

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

        if is_batched:
             lowercase_       :str     =       [np.asarray(UpperCamelCase_   ,							dtype=np.floataa		) for speech in raw_speech]
        elif not is_batched and not isinstance(UpperCamelCase_   ,							np.ndarray		):
             lowercase_       :List[Any]     =       np.asarray(UpperCamelCase_   ,							dtype=np.floataa		)
        elif isinstance(UpperCamelCase_   ,							np.ndarray		) and raw_speech.dtype is np.dtype(np.floataa		):
             lowercase_       :Union[str, Any]     =       raw_speech.astype(np.floataa		)

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

        # extract fbank features
        lowercase_       :Dict     =       [self._extract_mfsc_features(UpperCamelCase_		) for one_waveform in raw_speech]

        # convert into correct format for padding
        lowercase_       :Optional[Any]     =       BatchFeature({'''input_features''': features}		)

        lowercase_       :List[Any]     =       self.pad(
            UpperCamelCase_   ,							padding=UpperCamelCase_   ,							max_length=UpperCamelCase_   ,							truncation=UpperCamelCase_   ,							pad_to_multiple_of=UpperCamelCase_   ,							return_attention_mask=UpperCamelCase_   ,							**UpperCamelCase_   ,							)
        # make sure list is in array format
        lowercase_       :Dict     =       padded_inputs.get('''input_features'''		)
        if isinstance(input_features[0]   ,							UpperCamelCase_		):
             lowercase_       :Tuple     =       [np.asarray(UpperCamelCase_   ,							dtype=np.floataa		) for feature in input_features]

        lowercase_       :Optional[Any]     =       padded_inputs.get('''attention_mask'''		)
        if attention_mask is not None:
             lowercase_       :Union[str, Any]     =       [np.asarray(UpperCamelCase_   ,							dtype=np.intaa		) for array in attention_mask]

        if self.normalize_means or self.normalize_vars:
             lowercase_       :Optional[int]     =       (
                 np.array(UpperCamelCase_   ,							dtype=np.intaa		)
                 if self._get_padding_strategies(UpperCamelCase_   ,							max_length=UpperCamelCase_		) is not PaddingStrategy.DO_NOT_PAD
                 and padding
                 else None
             )
             lowercase_       :List[Any]     =       self.normalize(
                 padded_inputs['''input_features''']   ,							attention_mask=UpperCamelCase_		)

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

        return padded_inputs





import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


SCREAMING_SNAKE_CASE       :		int         =				False






class      UpperCamelCase					(		unittest.TestCase     ):


   '''simple docstring'''







   def 		UpperCamelCase  (     self		):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()
   @property
   def 		UpperCamelCase  (     self		):
        return 12
   @property
   def 		UpperCamelCase  (     self		):
        return 12
   @property
   def 		UpperCamelCase  (     self		):
        return 32
   @property
   def 		UpperCamelCase  (     self		):
        torch.manual_seed(0		)
        lowercase_       :List[Any]     =       VQModel(
            block_out_channels=[32, 64]   ,							in_channels=3   ,							out_channels=3   ,							down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']   ,							up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']   ,							latent_channels=3   ,							num_vq_embeddings=self.num_embed   ,							vq_embed_dim=3   ,							)
        return model
   @property
   def 		UpperCamelCase  (     self		):
        lowercase_       :Optional[Any]     =       CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)
        return tokenizer
   @property
   def 		UpperCamelCase  (     self		):
        torch.manual_seed(0		)
        lowercase_       :int     =       CLIPTextConfig(
            bos_token_id=0   ,							eos_token_id=2   ,							hidden_size=self.text_embedder_hidden_size   ,							intermediate_size=37   ,							layer_norm_eps=1E-05   ,							num_attention_heads=4   ,							num_hidden_layers=5   ,							pad_token_id=1   ,							vocab_size=1000   ,							)
        return CLIPTextModel(UpperCamelCase_		)
   @property
   def 		UpperCamelCase  (     self		):
        torch.manual_seed(0		)

        lowercase_       :int     =       12
        lowercase_       :List[Any]     =       12

        lowercase_       :Dict     =       {
            '''attention_bias''': True,
            '''cross_attention_dim''': 32,
            '''attention_head_dim''': height * width,
            '''num_attention_heads''': 1,
            '''num_vector_embeds''': self.num_embed,
            '''num_embeds_ada_norm''': self.num_embeds_ada_norm,
            '''norm_num_groups''': 32,
            '''sample_size''': width,
            '''activation_fn''': '''geglu-approximate''',
        }

        lowercase_       :int     =       TransformeraDModel(**UpperCamelCase_		)
        return model
   def 		UpperCamelCase  (     self		):
        lowercase_       :List[str]     =       '''cpu'''

        lowercase_       :int     =       self.dummy_vqvae
        lowercase_       :int     =       self.dummy_text_encoder
        lowercase_       :Any     =       self.dummy_tokenizer
        lowercase_       :Optional[int]     =       self.dummy_transformer
        lowercase_       :List[str]     =       VQDiffusionScheduler(self.num_embed		)
        lowercase_       :int     =       LearnedClassifierFreeSamplingEmbeddings(learnable=UpperCamelCase_		)

        lowercase_       :List[Any]     =       VQDiffusionPipeline(
            vqvae=UpperCamelCase_   ,							text_encoder=UpperCamelCase_   ,							tokenizer=UpperCamelCase_   ,							transformer=UpperCamelCase_   ,							scheduler=UpperCamelCase_   ,							learned_classifier_free_sampling_embeddings=UpperCamelCase_   ,							)
        lowercase_       :Union[str, Any]     =       pipe.to(UpperCamelCase_		)
        pipe.set_progress_bar_config(disable=UpperCamelCase_		)

        lowercase_       :Dict     =       '''teddy bear playing in the pool'''

        lowercase_       :Dict     =       torch.Generator(device=UpperCamelCase_		).manual_seed(0		)
        lowercase_       :Optional[int]     =       pipe([prompt]   ,							generator=UpperCamelCase_   ,							num_inference_steps=2   ,							output_type='''np'''		)
        lowercase_       :Any     =       output.images

        lowercase_       :Tuple     =       torch.Generator(device=UpperCamelCase_		).manual_seed(0		)
        lowercase_       :str     =       pipe(
            [prompt]   ,							generator=UpperCamelCase_   ,							output_type='''np'''   ,							return_dict=UpperCamelCase_   ,							num_inference_steps=2		)[0]

        lowercase_       :Optional[Any]     =       image[0, -3:, -3:, -1]
        lowercase_       :Tuple     =       image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (1, 24, 24, 3)

        lowercase_       :str     =       np.array([0.6551, 0.6168, 0.5008, 0.5676, 0.5659, 0.4295, 0.6073, 0.5599, 0.4992]		)

        assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 1E-2




   def 		UpperCamelCase  (     self		):
        lowercase_       :int     =       '''cpu'''

        lowercase_       :Dict     =       self.dummy_vqvae
        lowercase_       :str     =       self.dummy_text_encoder
        lowercase_       :List[Any]     =       self.dummy_tokenizer
        lowercase_       :Any     =       self.dummy_transformer
        lowercase_       :Optional[Any]     =       VQDiffusionScheduler(self.num_embed		)
        lowercase_       :List[str]     =       LearnedClassifierFreeSamplingEmbeddings(
            learnable=UpperCamelCase_   ,							hidden_size=self.text_embedder_hidden_size   ,							length=tokenizer.model_max_length		)

        lowercase_       :Optional[int]     =       VQDiffusionPipeline(
            vqvae=UpperCamelCase_   ,							text_encoder=UpperCamelCase_   ,							tokenizer=UpperCamelCase_   ,							transformer=UpperCamelCase_   ,							scheduler=UpperCamelCase_   ,							learned_classifier_free_sampling_embeddings=UpperCamelCase_   ,							)
        lowercase_       :Dict     =       pipe.to(UpperCamelCase_		)
        pipe.set_progress_bar_config(disable=UpperCamelCase_		)

        lowercase_       :int     =       '''teddy bear playing in the pool'''

        lowercase_       :Tuple     =       torch.Generator(device=UpperCamelCase_		).manual_seed(0		)
        lowercase_       :Optional[int]     =       pipe([prompt]   ,							generator=UpperCamelCase_   ,							num_inference_steps=2   ,							output_type='''np'''		)
        lowercase_       :Optional[Any]     =       output.images

        lowercase_       :Optional[int]     =       torch.Generator(device=UpperCamelCase_		).manual_seed(0		)
        lowercase_       :Dict     =       pipe(
            [prompt]   ,							generator=UpperCamelCase_   ,							output_type='''np'''   ,							return_dict=UpperCamelCase_   ,							num_inference_steps=2		)[0]

        lowercase_       :List[str]     =       image[0, -3:, -3:, -1]
        lowercase_       :Optional[int]     =       image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (1, 24, 24, 3)

        lowercase_       :Dict     =       np.array([0.6693, 0.6075, 0.4959, 0.5701, 0.5583, 0.4333, 0.6171, 0.5684, 0.4988]		)

        assert np.abs(image_slice.flatten() - expected_slice		).max() < 2.0
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 1E-2






@slow
@require_torch_gpu
class      UpperCamelCase					(		unittest.TestCase     ):


   '''simple docstring'''







   def 		UpperCamelCase  (     self		):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()




   def 		UpperCamelCase  (     self		):
        lowercase_       :List[str]     =       load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy'''		)

        lowercase_       :int     =       VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq'''		)
        lowercase_       :Tuple     =       pipeline.to(UpperCamelCase_		)
        pipeline.set_progress_bar_config(disable=UpperCamelCase_		)

        # requires GPU generator for gumbel softmax
        # don't use GPU generator in tests though
        lowercase_       :Dict     =       torch.Generator(device=UpperCamelCase_		).manual_seed(0		)
        lowercase_       :int     =       pipeline(
            '''teddy bear playing in the pool'''   ,							num_images_per_prompt=1   ,							generator=UpperCamelCase_   ,							output_type='''np'''   ,							)

        lowercase_       :List[str]     =       output.images[0]

        assert image.shape == (256, 256, 3)
        assert np.abs(expected_image - image		).max() < 2.0

