





'''simple docstring'''






import gc
import unittest

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

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference




class  _A (					__lowercase						,					unittest.TestCase		):
		lowercase__:   Dict				     =   ShapEPipeline
		lowercase__:   Tuple				     =   ['''prompt''']
		lowercase__:   str				     =   ['''prompt''']
		lowercase__:   str				     =   [
		    '''num_images_per_prompt''',
		    '''num_inference_steps''',
		    '''generator''',
		    '''latents''',
		    '''guidance_scale''',
		    '''frame_size''',
		    '''output_type''',
		    '''return_dict''',
		]
		lowercase__:   Dict				     =   False

		@property
		def 	lowercase__     (   self					:   str			)					->    int:


					"""simple docstring"""

					return 32

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


					"""simple docstring"""

					return 32

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


					"""simple docstring"""

					return self.time_input_dim * 4

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


					"""simple docstring"""

					return 8

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


					"""simple docstring"""

					__snake_case					:      Tuple							      =     CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""			)
					return tokenizer

		@property
		def 	lowercase__     (   self					:   Any			)					->    Any:


					"""simple docstring"""

					torch.manual_seed(0			)
					__snake_case					:      int							      =     CLIPTextConfig(
					    bos_token_id=0  ,      eos_token_id=2  ,      hidden_size=self.text_embedder_hidden_size  ,      projection_dim=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=10_00  ,      )
					return CLIPTextModelWithProjection(__magic_name__			)

		@property
		def 	lowercase__     (   self					:   Optional[int]			)					->    Optional[int]:


					"""simple docstring"""

					torch.manual_seed(0			)

					__snake_case					:      Tuple							      =     {
					    """num_attention_heads""": 2,
					    """attention_head_dim""": 16,
					    """embedding_dim""": self.time_input_dim,
					    """num_embeddings""": 32,
					    """embedding_proj_dim""": self.text_embedder_hidden_size,
					    """time_embed_dim""": self.time_embed_dim,
					    """num_layers""": 1,
					    """clip_embed_dim""": self.time_input_dim * 2,
					    """additional_embeddings""": 0,
					    """time_embed_act_fn""": """gelu""",
					    """norm_in_type""": """layer""",
					    """encoder_hid_proj_type""": None,
					    """added_emb_type""": None,
					}

					__snake_case					:      Union[str, Any]							      =     PriorTransformer(**__magic_name__			)
					return model

		@property
		def 	lowercase__     (   self					:   Optional[int]			)					->    str:


					"""simple docstring"""

					torch.manual_seed(0			)

					__snake_case					:      Dict							      =     {
					    """param_shapes""": (
					        (self.renderer_dim, 93),
					        (self.renderer_dim, 8),
					        (self.renderer_dim, 8),
					        (self.renderer_dim, 8),
					    ),
					    """d_latent""": self.time_input_dim,
					    """d_hidden""": self.renderer_dim,
					    """n_output""": 12,
					    """background""": (
					        0.1,
					        0.1,
					        0.1,
					    ),
					}
					__snake_case					:      Dict							      =     ShapERenderer(**__magic_name__			)
					return model

		def 	lowercase__     (   self					:   Optional[int]			)					->    Any:


					"""simple docstring"""

					__snake_case					:      List[str]							      =     self.dummy_prior
					__snake_case					:      Any							      =     self.dummy_text_encoder
					__snake_case					:      Tuple							      =     self.dummy_tokenizer
					__snake_case					:      int							      =     self.dummy_renderer

					__snake_case					:      Tuple							      =     HeunDiscreteScheduler(
					    beta_schedule="""exp"""  ,      num_train_timesteps=10_24  ,      prediction_type="""sample"""  ,      use_karras_sigmas=__magic_name__  ,      clip_sample=__magic_name__  ,      clip_sample_range=1.0  ,      )
					__snake_case					:      int							      =     {
					    """prior""": prior,
					    """text_encoder""": text_encoder,
					    """tokenizer""": tokenizer,
					    """renderer""": renderer,
					    """scheduler""": scheduler,
					}

					return components

		def 	lowercase__     (   self					:   Optional[Any]  ,      __magic_name__					:   List[Any]  ,      __magic_name__					:   Any=0			)					->    Union[str, Any]:


					"""simple docstring"""

					if str(__magic_name__			).startswith("""mps"""			):
								__snake_case					:      List[Any]							      =     torch.manual_seed(__magic_name__			)
					else:
								__snake_case					:      Optional[int]							      =     torch.Generator(device=__magic_name__			).manual_seed(__magic_name__			)
					__snake_case					:      str							      =     {
					    """prompt""": """horse""",
					    """generator""": generator,
					    """num_inference_steps""": 1,
					    """frame_size""": 32,
					    """output_type""": """np""",
					}
					return inputs

		def 	lowercase__     (   self					:   List[Any]			)					->    Optional[Any]:


					"""simple docstring"""

					__snake_case					:      Tuple							      =     """cpu"""

					__snake_case					:      Dict							      =     self.get_dummy_components()

					__snake_case					:      Dict							      =     self.pipeline_class(**__magic_name__			)
					__snake_case					:      List[Any]							      =     pipe.to(__magic_name__			)

					pipe.set_progress_bar_config(disable=__magic_name__			)

					__snake_case					:      Optional[Any]							      =     pipe(**self.get_dummy_inputs(__magic_name__			)			)
					__snake_case					:      str							      =     output.images[0]
					__snake_case					:      int							      =     image[0, -3:, -3:, -1]

					assert image.shape == (20, 32, 32, 3)

					__snake_case					:      int							      =     np.array(
					    [
					        0.00039216,
					        0.00039216,
					        0.00039216,
					        0.00039216,
					        0.00039216,
					        0.00039216,
					        0.00039216,
					        0.00039216,
					        0.00039216,
					    ]			)

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

		def 	lowercase__     (   self					:   Any			)					->    str:


					"""simple docstring"""

					self._test_inference_batch_consistent(batch_sizes=[1, 2]			)

		def 	lowercase__     (   self					:   int			)					->    List[str]:


					"""simple docstring"""

					__snake_case					:      Dict							      =     torch_device == """cpu"""
					__snake_case					:      int							      =     True

					self._test_inference_batch_single_identical(
					    batch_size=2  ,      test_max_difference=__magic_name__  ,      relax_max_difference=__magic_name__  ,      )

		def 	lowercase__     (   self					:   Tuple			)					->    Tuple:


					"""simple docstring"""

					__snake_case					:      int							      =     self.get_dummy_components()
					__snake_case					:      List[str]							      =     self.pipeline_class(**__magic_name__			)
					__snake_case					:      List[Any]							      =     pipe.to(__magic_name__			)
					pipe.set_progress_bar_config(disable=__magic_name__			)

					__snake_case					:      str							      =     1
					__snake_case					:      Optional[int]							      =     2

					__snake_case					:      int							      =     self.get_dummy_inputs(__magic_name__			)

					for key in inputs.keys():
								if key in self.batch_params:
											__snake_case					:      List[str]							      =     batch_size * [inputs[key]]

					__snake_case					:      Optional[Any]							      =     pipe(**__magic_name__  ,      num_images_per_prompt=__magic_name__			)[0]

					assert images.shape[0] == batch_size * num_images_per_prompt





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

		def 	lowercase__     (   self					:   List[Any]			)					->    Optional[int]:


					"""simple docstring"""

					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()

		def 	lowercase__     (   self					:   List[str]			)					->    List[str]:


					"""simple docstring"""

					__snake_case					:      List[Any]							      =     load_numpy(
					    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
					    """/shap_e/test_shap_e_np_out.npy"""			)
					__snake_case					:      List[Any]							      =     ShapEPipeline.from_pretrained("""openai/shap-e"""			)
					__snake_case					:      Optional[int]							      =     pipe.to(__magic_name__			)
					pipe.set_progress_bar_config(disable=__magic_name__			)

					__snake_case					:      Optional[int]							      =     torch.Generator(device=__magic_name__			).manual_seed(0			)

					__snake_case					:      Dict							      =     pipe(
					    """a shark"""  ,      generator=__magic_name__  ,      guidance_scale=15.0  ,      num_inference_steps=64  ,      frame_size=64  ,      output_type="""np"""  ,      ).images[0]

					assert images.shape == (20, 64, 64, 3)

					assert_mean_pixel_difference(__magic_name__  ,      __magic_name__			)









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


__lowerCamelCase			:	List[str]					    =	logging.get_logger(__name__)



class    A__  (			__snake_case     ):
      _UpperCAmelCase      :List[str]		       =  ['input_features', 'attention_mask']





      def __init__(					self			,				A_=80			,				A_=1_6000			,				A_=80			,				A_=0.0			,				A_=True			,				A_=True			,				A_=True			,				**A_			,				):
       '''simple docstring'''




       super().__init__(feature_size=A_			,				sampling_rate=A_			,				padding_value=A_			,				**A_  )
       UpperCamelCase      :   List[Any]			      =  num_mel_bins
       UpperCamelCase      :   Any			      =  do_ceptral_normalize
       UpperCamelCase      :   int			      =  normalize_means
       UpperCamelCase      :   Tuple			      =  normalize_vars
       UpperCamelCase      :   List[Any]			      =  True





      def 			__UpperCamelCase(					self			,				A_			,				):
       '''simple docstring'''




       UpperCamelCase      :   Optional[int]			      =  waveform * (2**15)  # Kaldi compliance: 16-bit signed integers
       UpperCamelCase      :   Optional[Any]			      =  torch.from_numpy(A_  ).unsqueeze(0  )
       UpperCamelCase      :   Dict			      =  ta_kaldi.fbank(A_			,				num_mel_bins=self.num_mel_bins			,				sample_frequency=self.sampling_rate  )
       return features.numpy()





      @staticmethod
      def 			__UpperCamelCase(					A_			,				A_			,				A_ = True			,				A_ = True			,				A_ = 0.0			,				):
       '''simple docstring'''




       if normalize_means:
        UpperCamelCase      :   str			      =  x[:input_length].mean(axis=0  )
        UpperCamelCase      :   Any			      =  np.subtract(A_			,				A_  )
       if normalize_vars:
        UpperCamelCase      :   List[str]			      =  x[:input_length].std(axis=0  )
        UpperCamelCase      :   Optional[int]			      =  np.divide(A_			,				A_  )

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

       # make sure array is in float32
       UpperCamelCase      :   Union[str, Any]			      =  x.astype(np.floataa  )

       return x





      def 			__UpperCamelCase(					self			,				A_			,				A_ = None  ):
       '''simple docstring'''




       UpperCamelCase      :   Optional[Any]			      =  attention_mask.sum(-1  ) if attention_mask is not None else [x.shape[0] for x in input_features]
       return [
           self.utterance_cmvn(A_			,				A_			,				self.normalize_means			,				self.normalize_vars			,				self.padding_value  )
           for x, n in zip(A_			,				A_  )
       ]





      def __call__(					self			,				A_			,				A_ = False			,				A_ = None			,				A_ = False			,				A_ = None			,				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} 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      :   str			      =  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}"""  )
       UpperCamelCase      :   List[Any]			      =  is_batched_numpy or (
           isinstance(A_			,				(list, tuple)  ) and (isinstance(raw_speech[0]			,				(np.ndarray, tuple, list)  ))
       )

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

       # always return batch
       if not is_batched:
        UpperCamelCase      :   Optional[int]			      =  [raw_speech]

       # extract fbank features
       UpperCamelCase      :   Optional[int]			      =  [self._extract_fbank_features(A_  ) for waveform in raw_speech]

       # convert into correct format for padding
       UpperCamelCase      :   List[Any]			      =  BatchFeature({"input_features": features}  )

       UpperCamelCase      :   Any			      =  self.pad(
           A_			,				padding=A_			,				max_length=A_			,				truncation=A_			,				pad_to_multiple_of=A_			,				return_attention_mask=A_			,				**A_			,				)

       # make sure list is in array format
       UpperCamelCase      :   Optional[int]			      =  padded_inputs.get("input_features"  )
       if isinstance(input_features[0]			,				A_  ):
        UpperCamelCase      :   str			      =  [np.asarray(A_			,				dtype=np.floataa  ) for feature in input_features]

       UpperCamelCase      :   Optional[int]			      =  padded_inputs.get("attention_mask"  )
       if attention_mask is not None:
        UpperCamelCase      :   Optional[int]			      =  [np.asarray(A_			,				dtype=np.intaa  ) for array in attention_mask]

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

       if return_tensors is not None:
        UpperCamelCase      :   List[str]			      =  padded_inputs.convert_to_tensors(A_  )

       return padded_inputs
