

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


UpperCAmelCase_							:					Optional[int]		 =   False






class      lowerCamelCase_      (	unittest.TestCase	):




	def        lowerCAmelCase_       (     self						:     Optional[Any]							):
		super().tearDown()
		gc.collect()
		torch.cuda.empty_cache()




	@property
	def        lowerCAmelCase_       (     self						:     List[Any]							):
		return 12




	@property
	def        lowerCAmelCase_       (     self						:     Tuple							):
		return 12




	@property
	def        lowerCAmelCase_       (     self						:     int							):
		return 32




	@property
	def        lowerCAmelCase_       (     self						:     str							):
		torch.manual_seed(0							)
		__A :			Tuple		    =			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        lowerCAmelCase_       (     self						:     Tuple							):
		__A :			List[str]		    =			CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""							)
		return tokenizer




	@property
	def        lowerCAmelCase_       (     self						:     str							):
		torch.manual_seed(0							)
		__A :			Tuple		    =			CLIPTextConfig(
		    bos_token_id=0					,      eos_token_id=2					,      hidden_size=self.text_embedder_hidden_size					,      intermediate_size=37					,      layer_norm_eps=1e-0_5					,      num_attention_heads=4					,      num_hidden_layers=5					,      pad_token_id=1					,      vocab_size=1000					,      )
		return CLIPTextModel(a__							)




	@property
	def        lowerCAmelCase_       (     self						:     Optional[Any]							):
		torch.manual_seed(0							)

		__A :			int		    =			12
		__A :			Optional[int]		    =			12

		__A :			List[str]		    =			{
		    '''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''',
		}

		__A :			Optional[Any]		    =			TransformeraDModel(**a__							)
		return model




	def        lowerCAmelCase_       (     self						:     Union[str, Any]							):
		__A :			Any		    =			'''cpu'''

		__A :			Dict		    =			self.dummy_vqvae
		__A :			Optional[int]		    =			self.dummy_text_encoder
		__A :			str		    =			self.dummy_tokenizer
		__A :			Dict		    =			self.dummy_transformer
		__A :			List[str]		    =			VQDiffusionScheduler(self.num_embed							)
		__A :			Optional[int]		    =			LearnedClassifierFreeSamplingEmbeddings(learnable=a__							)

		__A :			Any		    =			VQDiffusionPipeline(
		    vqvae=a__					,      text_encoder=a__					,      tokenizer=a__					,      transformer=a__					,      scheduler=a__					,      learned_classifier_free_sampling_embeddings=a__					,      )
		__A :			Tuple		    =			pipe.to(a__							)
		pipe.set_progress_bar_config(disable=a__							)

		__A :			Any		    =			'''teddy bear playing in the pool'''

		__A :			str		    =			torch.Generator(device=a__							).manual_seed(0							)
		__A :			Tuple		    =			pipe([prompt]					,      generator=a__					,      num_inference_steps=2					,      output_type="""np"""							)
		__A :			List[Any]		    =			output.images

		__A :			Optional[int]		    =			torch.Generator(device=a__							).manual_seed(0							)
		__A :			Optional[int]		    =			pipe(
		    [prompt]					,      generator=a__					,      output_type="""np"""					,      return_dict=a__					,      num_inference_steps=2							)[0]

		__A :			List[Any]		    =			image[0, -3:, -3:, -1]
		__A :			List[Any]		    =			image_from_tuple[0, -3:, -3:, -1]

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

		__A :			int		    =			np.array([0.6_5_5_1, 0.6_1_6_8, 0.5_0_0_8, 0.5_6_7_6, 0.5_6_5_9, 0.4_2_9_5, 0.6_0_7_3, 0.5_5_9_9, 0.4_9_9_2]							)

		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        lowerCAmelCase_       (     self						:     int							):
		__A :			Tuple		    =			'''cpu'''

		__A :			Optional[int]		    =			self.dummy_vqvae
		__A :			Union[str, Any]		    =			self.dummy_text_encoder
		__A :			Union[str, Any]		    =			self.dummy_tokenizer
		__A :			Dict		    =			self.dummy_transformer
		__A :			int		    =			VQDiffusionScheduler(self.num_embed							)
		__A :			Optional[int]		    =			LearnedClassifierFreeSamplingEmbeddings(
		    learnable=a__					,      hidden_size=self.text_embedder_hidden_size					,      length=tokenizer.model_max_length							)

		__A :			Any		    =			VQDiffusionPipeline(
		    vqvae=a__					,      text_encoder=a__					,      tokenizer=a__					,      transformer=a__					,      scheduler=a__					,      learned_classifier_free_sampling_embeddings=a__					,      )
		__A :			int		    =			pipe.to(a__							)
		pipe.set_progress_bar_config(disable=a__							)

		__A :			Optional[Any]		    =			'''teddy bear playing in the pool'''

		__A :			Dict		    =			torch.Generator(device=a__							).manual_seed(0							)
		__A :			int		    =			pipe([prompt]					,      generator=a__					,      num_inference_steps=2					,      output_type="""np"""							)
		__A :			Tuple		    =			output.images

		__A :			List[Any]		    =			torch.Generator(device=a__							).manual_seed(0							)
		__A :			str		    =			pipe(
		    [prompt]					,      generator=a__					,      output_type="""np"""					,      return_dict=a__					,      num_inference_steps=2							)[0]

		__A :			Optional[int]		    =			image[0, -3:, -3:, -1]
		__A :			Union[str, Any]		    =			image_from_tuple[0, -3:, -3:, -1]

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

		__A :			Tuple		    =			np.array([0.6_6_9_3, 0.6_0_7_5, 0.4_9_5_9, 0.5_7_0_1, 0.5_5_8_3, 0.4_3_3_3, 0.6_1_7_1, 0.5_6_8_4, 0.4_9_8_8]							)

		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      lowerCamelCase_      (	unittest.TestCase	):




	def        lowerCAmelCase_       (     self						:     Union[str, Any]							):
		super().tearDown()
		gc.collect()
		torch.cuda.empty_cache()




	def        lowerCAmelCase_       (     self						:     List[Any]							):
		__A :			Any		    =			load_numpy(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
		    """/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy"""							)

		__A :			Optional[Any]		    =			VQDiffusionPipeline.from_pretrained("""microsoft/vq-diffusion-ithq"""							)
		__A :			Union[str, Any]		    =			pipeline.to(a__							)
		pipeline.set_progress_bar_config(disable=a__							)

		# requires GPU generator for gumbel softmax
		# don't use GPU generator in tests though
		__A :			Optional[int]		    =			torch.Generator(device=a__							).manual_seed(0							)
		__A :			str		    =			pipeline(
		    """teddy bear playing in the pool"""					,      num_images_per_prompt=1					,      generator=a__					,      output_type="""np"""					,      )

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

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












import copy
from typing import Any, Dict, List, Optional, Union

import numpy as np

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


_lowerCamelCase      							=   logging.get_logger(__name__)







class       __A							(   lowerCamelCase__	):




      """simple docstring"""







      UpperCAmelCase__							=   ["""input_features"""]


      def __init__(		self   ,				a__=80   ,				a__=1_6000   ,				a__=160   ,				a__=30   ,				a__=400   ,				a__=0.0   ,				a__=False   ,				**a__   ,				):



             """simple docstring"""





             super().__init__(
                 feature_size=a__   ,				sampling_rate=a__   ,				padding_value=a__   ,				return_attention_mask=a__   ,				**a__   ,				)
             _lowerCamelCase				:  int						=		n_fft
             _lowerCamelCase				:  List[Any]						=		hop_length
             _lowerCamelCase				:  Tuple						=		chunk_length
             _lowerCamelCase				:  Optional[Any]						=		chunk_length * sampling_rate
             _lowerCamelCase				:  Optional[Any]						=		self.n_samples // hop_length
             _lowerCamelCase				:  List[str]						=		sampling_rate
             _lowerCamelCase				:  Optional[int]						=		mel_filter_bank(
                 num_frequency_bins=1 + n_fft // 2   ,				num_mel_filters=a__   ,				min_frequency=0.0   ,				max_frequency=8000.0   ,				sampling_rate=a__   ,				norm='''slaney'''   ,				mel_scale='''slaney'''   ,				)


      def 					__snake_case      (		self   ,				a__):



             """simple docstring"""





             _lowerCamelCase				:  Any						=		spectrogram(
                 a__   ,				window_function(self.n_fft   ,				'''hann''')   ,				frame_length=self.n_fft   ,				hop_length=self.hop_length   ,				power=2.0   ,				mel_filters=self.mel_filters   ,				log_mel='''log10'''   ,				)
             _lowerCamelCase				:  str						=		log_spec[:, :-1]
             _lowerCamelCase				:  Union[str, Any]						=		np.maximum(a__   ,				log_spec.max() - 8.0)
             _lowerCamelCase				:  Optional[int]						=		(log_spec + 4.0) / 4.0
             return log_spec


      @staticmethod
      # Copied from transformers.models.wav2vec2.feature_extraction_wav2vec2.Wav2Vec2FeatureExtractor.zero_mean_unit_var_norm
      def 					__snake_case      (		a__   ,				a__   ,				a__ = 0.0):



             """simple docstring"""





             if attention_mask is not None:
                    _lowerCamelCase				:  Dict						=		np.array(a__   ,				np.intaa)
                    _lowerCamelCase				:  Any						=		[]

                    for vector, length in zip(a__   ,				attention_mask.sum(-1)):
                           _lowerCamelCase				:  Any						=		(vector - vector[:length].mean()) / np.sqrt(vector[:length].var() + 1e-7)
                           if length < normed_slice.shape[0]:
                                  _lowerCamelCase				:  Optional[int]						=		padding_value

                           normed_input_values.append(a__)
             else:
                    _lowerCamelCase				:  Optional[int]						=		[(x - x.mean()) / np.sqrt(x.var() + 1e-7) for x in input_values]

             return normed_input_values


      def __call__(		self   ,				a__   ,				a__ = True   ,				a__ = None   ,				a__ = None   ,				a__ = None   ,				a__ = "max_length"   ,				a__ = None   ,				a__ = None   ,				a__ = None   ,				**a__   ,				):



             """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.__class__.__name__} was trained using a"""
                               F""" sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input"""
                               F""" was sampled with {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.''')

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

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

             # always return batch
             if not is_batched:
                    _lowerCamelCase				:  List[Any]						=		[np.asarray([raw_speech]).T]

             _lowerCamelCase				:  List[str]						=		BatchFeature({'''input_features''': raw_speech})

             # convert into correct format for padding

             _lowerCamelCase				:  Optional[Any]						=		self.pad(
                 a__   ,				padding=a__   ,				max_length=max_length if max_length else self.n_samples   ,				truncation=a__   ,				pad_to_multiple_of=a__   ,				return_attention_mask=return_attention_mask or do_normalize   ,				)

             # zero-mean and unit-variance normalization
             if do_normalize:
                    _lowerCamelCase				:  Dict						=		self.zero_mean_unit_var_norm(
                        padded_inputs['''input_features''']   ,				attention_mask=padded_inputs['''attention_mask''']   ,				padding_value=self.padding_value   ,				)
                    _lowerCamelCase				:  int						=		np.stack(padded_inputs['''input_features''']   ,				axis=0)

             # make sure list is in array format
             _lowerCamelCase				:  Optional[Any]						=		padded_inputs.get('''input_features''').transpose(2   ,				0   ,				1)

             _lowerCamelCase				:  Tuple						=		[self._np_extract_fbank_features(a__) for waveform in input_features[0]]

             if isinstance(input_features[0]   ,				a__):
                    _lowerCamelCase				:  str						=		[np.asarray(a__   ,				dtype=np.floataa) for feature in input_features]
             else:
                    _lowerCamelCase				:  Dict						=		input_features

             if return_attention_mask:
                    # rescale from sample (48000) to feature (3000)
                    _lowerCamelCase				:  Optional[int]						=		padded_inputs['''attention_mask'''][:, :: self.hop_length]

             if return_tensors is not None:
                    _lowerCamelCase				:  str						=		padded_inputs.convert_to_tensors(a__)

             return padded_inputs



      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  List[Any]						=		copy.deepcopy(self.__dict__)
             _lowerCamelCase				:  int						=		self.__class__.__name__
             if "mel_filters" in output:
                    del output["mel_filters"]
             return output

