






from typing import Callable, List, Optional, Tuple, Union

import torch
from transformers import CLIPTextModel, CLIPTokenizer

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin, TransformeraDModel, VQModel
from ...schedulers import VQDiffusionScheduler
from ...utils import logging
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput


lowerCAmelCase_           = logging.get_logger(__name__)  # pylint: disable=invalid-name
class      UpperCamelCase (				snake_case__			,       snake_case__ ):







		"""simple docstring"""



		@register_to_config
		def __init__(  self		:       Optional[int]						,_SCREAMING_SNAKE_CASE		:       bool						,_SCREAMING_SNAKE_CASE		:       Optional[int] = None						,_SCREAMING_SNAKE_CASE		:       Optional[int] = None		)    ->   str:


						'''simple docstring'''
						super().__init__()

						A	 =      learnable

						if self.learnable:
										assert hidden_size is not None, "learnable=True requires `hidden_size` to be set"
										assert length is not None, "learnable=True requires `length` to be set"

										A	 =      torch.zeros(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		)
						else:
										A	 =      None

						A	 =      torch.nn.Parameter(_SCREAMING_SNAKE_CASE		)


class      UpperCamelCase (				snake_case__ ):







		"""simple docstring"""



		snake_case    =   42
		snake_case    =   42
		snake_case    =   42
		snake_case    =   42
		snake_case    =   42
		snake_case    =   42
		def __init__(  self		:       Union[str, Any]						,_SCREAMING_SNAKE_CASE		:       VQModel						,_SCREAMING_SNAKE_CASE		:       CLIPTextModel						,_SCREAMING_SNAKE_CASE		:       CLIPTokenizer						,_SCREAMING_SNAKE_CASE		:       TransformeraDModel						,_SCREAMING_SNAKE_CASE		:       VQDiffusionScheduler						,_SCREAMING_SNAKE_CASE		:       LearnedClassifierFreeSamplingEmbeddings						,)    ->   Dict:


						'''simple docstring'''
						super().__init__()

						self.register_modules(
						    vqvae=_SCREAMING_SNAKE_CASE						,transformer=_SCREAMING_SNAKE_CASE						,text_encoder=_SCREAMING_SNAKE_CASE						,tokenizer=_SCREAMING_SNAKE_CASE						,scheduler=_SCREAMING_SNAKE_CASE						,learned_classifier_free_sampling_embeddings=_SCREAMING_SNAKE_CASE						,)
		def       A(  self		:       str						,_SCREAMING_SNAKE_CASE		:       Tuple						,_SCREAMING_SNAKE_CASE		:       Optional[int]						,_SCREAMING_SNAKE_CASE		:       Union[str, Any]		)    ->   List[str]:


						'''simple docstring'''
						A	 =      len(_SCREAMING_SNAKE_CASE		) if isinstance(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		) else 1

						# get prompt text embeddings
						A	 =      self.tokenizer(
						    _SCREAMING_SNAKE_CASE						,padding='max_length'						,max_length=self.tokenizer.model_max_length						,return_tensors='pt'						,)
						A	 =      text_inputs.input_ids

						if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
										A	 =      self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]		)
										logger.warning(
										    'The following part of your input was truncated because CLIP can only handle sequences up to'
										    f' {self.tokenizer.model_max_length} tokens: {removed_text}'		)
										A	 =      text_input_ids[:, : self.tokenizer.model_max_length]
						A	 =      self.text_encoder(text_input_ids.to(self.device		)		)[0]

						# NOTE: This additional step of normalizing the text embeddings is from VQ-Diffusion.
						# While CLIP does normalize the pooled output of the text transformer when combining
						# the image and text embeddings, CLIP does not directly normalize the last hidden state.
						#
						# CLIP normalizing the pooled output.
						# https://github.com/huggingface/transformers/blob/d92e22d1f28324f513f3080e5c47c071a3916721/src/transformers/models/clip/modeling_clip.py#L1052-L1053
						A	 =      prompt_embeds / prompt_embeds.norm(dim=-1						,keepdim=_SCREAMING_SNAKE_CASE		)

						# duplicate text embeddings for each generation per prompt
						A	 =      prompt_embeds.repeat_interleave(_SCREAMING_SNAKE_CASE						,dim=0		)

						if do_classifier_free_guidance:
										if self.learned_classifier_free_sampling_embeddings.learnable:
														A	 =      self.learned_classifier_free_sampling_embeddings.embeddings
														A	 =      negative_prompt_embeds.unsqueeze(0		).repeat(_SCREAMING_SNAKE_CASE						,1						,1		)
										else:
														A	 =      [''] * batch_size

														A	 =      text_input_ids.shape[-1]
														A	 =      self.tokenizer(
														    _SCREAMING_SNAKE_CASE						,padding='max_length'						,max_length=_SCREAMING_SNAKE_CASE						,truncation=_SCREAMING_SNAKE_CASE						,return_tensors='pt'						,)
														A	 =      self.text_encoder(uncond_input.input_ids.to(self.device		)		)[0]
														# See comment for normalizing text embeddings
														A	 =      negative_prompt_embeds / negative_prompt_embeds.norm(dim=-1						,keepdim=_SCREAMING_SNAKE_CASE		)

										# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
										A	 =      negative_prompt_embeds.shape[1]
										A	 =      negative_prompt_embeds.repeat(1						,_SCREAMING_SNAKE_CASE						,1		)
										A	 =      negative_prompt_embeds.view(batch_size * num_images_per_prompt						,_SCREAMING_SNAKE_CASE						,-1		)

										# For classifier free guidance, we need to do two forward passes.
										# Here we concatenate the unconditional and text embeddings into a single batch
										# to avoid doing two forward passes
										A	 =      torch.cat([negative_prompt_embeds, prompt_embeds]		)

						return prompt_embeds
		@torch.no_grad()
		def __call__(  self		:       Dict						,_SCREAMING_SNAKE_CASE		:       Union[str, List[str]]						,_SCREAMING_SNAKE_CASE		:       int = 1_0_0						,_SCREAMING_SNAKE_CASE		:       float = 5.0						,_SCREAMING_SNAKE_CASE		:       float = 1.0						,_SCREAMING_SNAKE_CASE		:       int = 1						,_SCREAMING_SNAKE_CASE		:       Optional[Union[torch.Generator, List[torch.Generator]]] = None						,_SCREAMING_SNAKE_CASE		:       Optional[torch.FloatTensor] = None						,_SCREAMING_SNAKE_CASE		:       Optional[str] = "pil"						,_SCREAMING_SNAKE_CASE		:       bool = True						,_SCREAMING_SNAKE_CASE		:       Optional[Callable[[int, int, torch.FloatTensor], None]] = None						,_SCREAMING_SNAKE_CASE		:       int = 1						,)    ->   Union[ImagePipelineOutput, Tuple]:


						'''simple docstring'''
						if isinstance(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		):
										A	 =      1
						elif isinstance(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		):
										A	 =      len(_SCREAMING_SNAKE_CASE		)
						else:
										raise ValueError(f'`prompt` has to be of type `str` or `list` but is {type(_SCREAMING_SNAKE_CASE		)}'		)

						A	 =      batch_size * num_images_per_prompt

						A	 =      guidance_scale > 1.0

						A	 =      self._encode_prompt(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		)

						if (callback_steps is None) or (
						    callback_steps is not None and (not isinstance(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		) or callback_steps <= 0)
						):
										raise ValueError(
										    f'`callback_steps` has to be a positive integer but is {callback_steps} of type'
										    f' {type(_SCREAMING_SNAKE_CASE		)}.'		)

						# get the initial completely masked latents unless the user supplied it

						A	 =      (batch_size, self.transformer.num_latent_pixels)
						if latents is None:
										A	 =      self.transformer.num_vector_embeds - 1
										A	 =      torch.full(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		).to(self.device		)
						else:
										if latents.shape != latents_shape:
														raise ValueError(f'Unexpected latents shape, got {latents.shape}, expected {latents_shape}'		)
										if (latents < 0).any() or (latents >= self.transformer.num_vector_embeds).any():
														raise ValueError(
														    'Unexpected latents value(s). All latents be valid embedding indices i.e. in the range 0,'
														    f' {self.transformer.num_vector_embeds - 1} (inclusive).'		)
										A	 =      latents.to(self.device		)

						# set timesteps
						self.scheduler.set_timesteps(_SCREAMING_SNAKE_CASE						,device=self.device		)

						A	 =      self.scheduler.timesteps.to(self.device		)

						A	 =      latents

						for i, t in enumerate(self.progress_bar(_SCREAMING_SNAKE_CASE		)		):
										# expand the sample if we are doing classifier free guidance
										A	 =      torch.cat([sample] * 2		) if do_classifier_free_guidance else sample

										# predict the un-noised image
										# model_output == `log_p_x_0`
										A	 =      self.transformer(_SCREAMING_SNAKE_CASE						,encoder_hidden_states=_SCREAMING_SNAKE_CASE						,timestep=_SCREAMING_SNAKE_CASE		).sample

										if do_classifier_free_guidance:
														A ,		A	 =      model_output.chunk(2		)
														A	 =      model_output_uncond + guidance_scale * (model_output_text - model_output_uncond)
														model_output -= torch.logsumexp(_SCREAMING_SNAKE_CASE						,dim=1						,keepdim=_SCREAMING_SNAKE_CASE		)

										A	 =      self.truncate(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		)

										# remove `log(0)`'s (`-inf`s)
										A	 =      model_output.clamp(-7_0		)

										# compute the previous noisy sample x_t -> x_t-1
										A	 =      self.scheduler.step(_SCREAMING_SNAKE_CASE						,timestep=_SCREAMING_SNAKE_CASE						,sample=_SCREAMING_SNAKE_CASE						,generator=_SCREAMING_SNAKE_CASE		).prev_sample

										# call the callback, if provided
										if callback is not None and i % callback_steps == 0:
														callback(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		)

						A	 =      self.vqvae.config.vq_embed_dim
						A	 =      (batch_size, self.transformer.height, self.transformer.width, embedding_channels)
						A	 =      self.vqvae.quantize.get_codebook_entry(_SCREAMING_SNAKE_CASE						,shape=_SCREAMING_SNAKE_CASE		)
						A	 =      self.vqvae.decode(_SCREAMING_SNAKE_CASE						,force_not_quantize=_SCREAMING_SNAKE_CASE		).sample

						A	 =      (image / 2 + 0.5).clamp(0						,1		)
						A	 =      image.cpu().permute(0						,2						,3						,1		).numpy()

						if output_type == "pil":
										A	 =      self.numpy_to_pil(_SCREAMING_SNAKE_CASE		)

						if not return_dict:
										return (image,)

						return ImagePipelineOutput(images=_SCREAMING_SNAKE_CASE		)







		def       A(  self		:       Any						,_SCREAMING_SNAKE_CASE		:       torch.FloatTensor						,_SCREAMING_SNAKE_CASE		:       float		)    ->   torch.FloatTensor:


						'''simple docstring'''
						A ,		A	 =      torch.sort(_SCREAMING_SNAKE_CASE						,1						,descending=_SCREAMING_SNAKE_CASE		)
						A	 =      torch.exp(_SCREAMING_SNAKE_CASE		)
						A	 =      sorted_p_x_0.cumsum(dim=1		) < truncation_rate

						# Ensure that at least the largest probability is not zeroed out
						A	 =      torch.full_like(keep_mask[:, 0:1, :]						,_SCREAMING_SNAKE_CASE		)
						A	 =      torch.cat((all_true, keep_mask)						,dim=1		)
						A	 =      keep_mask[:, :-1, :]

						A	 =      keep_mask.gather(1						,indices.argsort(1		)		)

						A	 =      log_p_x_0.clone()

						A	 =      -torch.inf  # -inf = log(0)

						return rv













import unittest

import numpy as np

from transformers import MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING, TF_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING
from transformers.pipelines import AudioClassificationPipeline, pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_torchaudio,
    slow,
)

from .test_pipelines_common import ANY
@is_pipeline_test
class      UpperCamelCase (				unittest.TestCase ):







		"""simple docstring"""



		snake_case    =   MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING
		snake_case    =   TF_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING
		def       A(  self		:       int						,_SCREAMING_SNAKE_CASE		:       Any						,_SCREAMING_SNAKE_CASE		:       List[Any]						,_SCREAMING_SNAKE_CASE		:       Tuple		)    ->   Dict:


						'''simple docstring'''
						A	 =      AudioClassificationPipeline(model=_SCREAMING_SNAKE_CASE						,feature_extractor=_SCREAMING_SNAKE_CASE		)

						# test with a raw waveform
						A	 =      np.zeros((3_4_0_0_0,)		)
						A	 =      np.zeros((1_4_0_0_0,)		)
						return audio_classifier, [audioa, audio]
		def       A(  self		:       int						,_SCREAMING_SNAKE_CASE		:       Tuple						,_SCREAMING_SNAKE_CASE		:       List[Any]		)    ->   Union[str, Any]:


						'''simple docstring'''
						A ,		A	 =      examples
						A	 =      audio_classifier(_SCREAMING_SNAKE_CASE		)
						# by default a model is initialized with num_labels=2
						self.assertEqual(
						    _SCREAMING_SNAKE_CASE						,[
						        {'score': ANY(_SCREAMING_SNAKE_CASE		), 'label': ANY(_SCREAMING_SNAKE_CASE		)},
						        {'score': ANY(_SCREAMING_SNAKE_CASE		), 'label': ANY(_SCREAMING_SNAKE_CASE		)},
						    ]						,)
						A	 =      audio_classifier(_SCREAMING_SNAKE_CASE						,top_k=1		)
						self.assertEqual(
						    _SCREAMING_SNAKE_CASE						,[
						        {'score': ANY(_SCREAMING_SNAKE_CASE		), 'label': ANY(_SCREAMING_SNAKE_CASE		)},
						    ]						,)

						self.run_torchaudio(_SCREAMING_SNAKE_CASE		)
		@require_torchaudio
		def       A(  self		:       Optional[int]						,_SCREAMING_SNAKE_CASE		:       int		)    ->   Optional[Any]:


						'''simple docstring'''
						import datasets

						# test with a local file
						A	 =      datasets.load_dataset('hf-internal-testing/librispeech_asr_dummy'						,'clean'						,split='validation'		)
						A	 =      dataset[0]['audio']['array']
						A	 =      audio_classifier(_SCREAMING_SNAKE_CASE		)
						self.assertEqual(
						    _SCREAMING_SNAKE_CASE						,[
						        {'score': ANY(_SCREAMING_SNAKE_CASE		), 'label': ANY(_SCREAMING_SNAKE_CASE		)},
						        {'score': ANY(_SCREAMING_SNAKE_CASE		), 'label': ANY(_SCREAMING_SNAKE_CASE		)},
						    ]						,)
		@require_torch
		def       A(  self		:       int		)    ->   str:


						'''simple docstring'''
						A	 =      'anton-l/wav2vec2-random-tiny-classifier'

						A	 =      pipeline('audio-classification'						,model=_SCREAMING_SNAKE_CASE		)

						A	 =      np.ones((8_0_0_0,)		)
						A	 =      audio_classifier(_SCREAMING_SNAKE_CASE						,top_k=4		)

						A	 =      [
						    {'score': 0.0842, 'label': 'no'},
						    {'score': 0.0838, 'label': 'up'},
						    {'score': 0.0837, 'label': 'go'},
						    {'score': 0.0834, 'label': 'right'},
						]
						A	 =      [
						    {'score': 0.0845, 'label': 'stop'},
						    {'score': 0.0844, 'label': 'on'},
						    {'score': 0.0841, 'label': 'right'},
						    {'score': 0.0834, 'label': 'left'},
						]
						self.assertIn(nested_simplify(_SCREAMING_SNAKE_CASE						,decimals=4		)						,[EXPECTED_OUTPUT, EXPECTED_OUTPUT_PT_2]		)

						A	 =      {'array': np.ones((8_0_0_0,)		), 'sampling_rate': audio_classifier.feature_extractor.sampling_rate}
						A	 =      audio_classifier(_SCREAMING_SNAKE_CASE						,top_k=4		)
						self.assertIn(nested_simplify(_SCREAMING_SNAKE_CASE						,decimals=4		)						,[EXPECTED_OUTPUT, EXPECTED_OUTPUT_PT_2]		)
		@require_torch
		@slow
		def       A(  self		:       Optional[int]		)    ->   Union[str, Any]:


						'''simple docstring'''
						import datasets

						A	 =      'superb/wav2vec2-base-superb-ks'

						A	 =      pipeline('audio-classification'						,model=_SCREAMING_SNAKE_CASE		)
						A	 =      datasets.load_dataset('anton-l/superb_dummy'						,'ks'						,split='test'		)

						A	 =      np.array(dataset[3]['speech']						,dtype=np.floataa		)
						A	 =      audio_classifier(_SCREAMING_SNAKE_CASE						,top_k=4		)
						self.assertEqual(
						    nested_simplify(_SCREAMING_SNAKE_CASE						,decimals=3		)						,[
						        {'score': 0.981, 'label': 'go'},
						        {'score': 0.007, 'label': 'up'},
						        {'score': 0.006, 'label': '_unknown_'},
						        {'score': 0.001, 'label': 'down'},
						    ]						,)







		@require_tf
		@unittest.skip('Audio classification is not implemented for TF'		)
		def       A(  self		:       Any		)    ->   int:


						'''simple docstring'''
						pass





