





'''simple docstring'''


import inspect
from typing import Optional, Union

import numpy as np
import PIL
import torch
from torch.nn import functional as F
from torchvision import transforms
from transformers import CLIPFeatureExtractor, CLIPModel, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    DPMSolverMultistepScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput
from diffusers.utils import (
    PIL_INTERPOLATION,
    randn_tensor,
)


def 		__a(SCREAMING_SNAKE_CASE_     :    int						,				SCREAMING_SNAKE_CASE_     :    Tuple						,				SCREAMING_SNAKE_CASE_     :    int			):



    '''simple docstring'''





    if isinstance(__snake_case						,				torch.Tensor			):
        return image
    elif isinstance(__snake_case						,				PIL.Image.Image			):
        _lowerCAmelCase			=					[image]

    if isinstance(image[0]						,				PIL.Image.Image			):
        _lowerCAmelCase			=					[np.array(i.resize((w, h)						,				resample=PIL_INTERPOLATION["lanczos"]			)			)[None, :] for i in image]
        _lowerCAmelCase			=					np.concatenate(__snake_case						,				axis=0			)
        _lowerCAmelCase			=					np.array(__snake_case			).astype(np.floataa			) / 255.0
        _lowerCAmelCase			=					image.transpose(0						,				3						,				1						,				2			)
        _lowerCAmelCase			=					2.0 * image - 1.0
        _lowerCAmelCase			=					torch.from_numpy(__snake_case			)
    elif isinstance(image[0]						,				torch.Tensor			):
        _lowerCAmelCase			=					torch.cat(__snake_case						,				dim=0			)
    return image


def 		__a(SCREAMING_SNAKE_CASE_     :    Tuple						,				SCREAMING_SNAKE_CASE_     :    Optional[Any]						,				SCREAMING_SNAKE_CASE_     :    Any						,				SCREAMING_SNAKE_CASE_     :    str=0.9995			):



    '''simple docstring'''





    if not isinstance(__snake_case						,				np.ndarray			):
        _lowerCAmelCase			=					True
        _lowerCAmelCase			=					va.device
        _lowerCAmelCase			=					va.cpu().numpy()
        _lowerCAmelCase			=					va.cpu().numpy()

    _lowerCAmelCase			=					np.sum(va * va / (np.linalg.norm(__snake_case			) * np.linalg.norm(__snake_case			))			)
    if np.abs(__snake_case			) > DOT_THRESHOLD:
        _lowerCAmelCase			=					(1 - t) * va + t * va
    else:
        _lowerCAmelCase			=					np.arccos(__snake_case			)
        _lowerCAmelCase			=					np.sin(__snake_case			)
        _lowerCAmelCase			=					theta_a * t
        _lowerCAmelCase			=					np.sin(__snake_case			)
        _lowerCAmelCase			=					np.sin(theta_a - theta_t			) / sin_theta_a
        _lowerCAmelCase			=					sin_theta_t / sin_theta_a
        _lowerCAmelCase			=					sa * va + sa * va

    if inputs_are_torch:
        _lowerCAmelCase			=					torch.from_numpy(__snake_case			).to(__snake_case			)

    return va


def 		__a(SCREAMING_SNAKE_CASE_     :    Optional[Any]						,				SCREAMING_SNAKE_CASE_     :    str			):



    '''simple docstring'''





    _lowerCAmelCase			=					F.normalize(__snake_case						,				dim=-1			)
    _lowerCAmelCase			=					F.normalize(__snake_case						,				dim=-1			)
    return (x - y).norm(dim=-1			).div(2			).arcsin().pow(2			).mul(2			)


def 		__a(SCREAMING_SNAKE_CASE_     :    Union[str, Any]						,				SCREAMING_SNAKE_CASE_     :    Tuple			):



    '''simple docstring'''





    for param in model.parameters():
        _lowerCAmelCase			=					value





class       lowerCAmelCase_						(						__magic_name__						):




 def __init__(		self	,   _lowerCAmelCase	,   _lowerCAmelCase	,   _lowerCAmelCase	,   _lowerCAmelCase	,   _lowerCAmelCase	,   _lowerCAmelCase	,   _lowerCAmelCase	,   _lowerCAmelCase=None	,   _lowerCAmelCase=None	,   _lowerCAmelCase=None	,   )			->							Tuple:
     super().__init__()
     self.register_modules(
         vae=__lowerCAmelCase	,   text_encoder=__lowerCAmelCase	,   clip_model=__lowerCAmelCase	,   tokenizer=__lowerCAmelCase	,   unet=__lowerCAmelCase	,   scheduler=__lowerCAmelCase	,   feature_extractor=__lowerCAmelCase	,   coca_model=__lowerCAmelCase	,   coca_tokenizer=__lowerCAmelCase	,   coca_transform=__lowerCAmelCase	,   )
     _lowerCAmelCase			=					(
         feature_extractor.size
         if isinstance(feature_extractor.size	,   __lowerCAmelCase						)
         else feature_extractor.size["shortest_edge"]
     )
     _lowerCAmelCase			=					transforms.Normalize(mean=feature_extractor.image_mean	,   std=feature_extractor.image_std						)
     set_requires_grad(self.text_encoder	,   __lowerCAmelCase						)
     set_requires_grad(self.clip_model	,   __lowerCAmelCase						)




 def 							_snake_case      (		self	,   _lowerCAmelCase = "auto"						)			->							Optional[int]:







     if slice_size == "auto":
         # half the attention head size is usually a good trade-off between
         # speed and memory
         _lowerCAmelCase			=					self.unet.config.attention_head_dim // 2
     self.unet.set_attention_slice(__lowerCAmelCase						)




 def 							_snake_case      (		self						)			->							Optional[int]:
     self.enable_attention_slicing(__lowerCAmelCase						)




 def 							_snake_case      (		self						)			->							Dict:
     set_requires_grad(self.vae	,   __lowerCAmelCase						)




 def 							_snake_case      (		self						)			->							Any:
     set_requires_grad(self.vae	,   __lowerCAmelCase						)




 def 							_snake_case      (		self						)			->							List[Any]:
     set_requires_grad(self.unet	,   __lowerCAmelCase						)




 def 							_snake_case      (		self						)			->							Optional[Any]:
     set_requires_grad(self.unet	,   __lowerCAmelCase						)




 def 							_snake_case      (		self	,   _lowerCAmelCase	,   _lowerCAmelCase	,   _lowerCAmelCase						)			->							Union[str, Any]:
     _lowerCAmelCase			=					min(int(num_inference_steps * strength						)	,   __lowerCAmelCase						)

     _lowerCAmelCase			=					max(num_inference_steps - init_timestep	,   0						)
     _lowerCAmelCase			=					self.scheduler.timesteps[t_start:]

     return timesteps, num_inference_steps - t_start




 def 							_snake_case      (		self	,   _lowerCAmelCase	,   _lowerCAmelCase	,   _lowerCAmelCase	,   _lowerCAmelCase	,   _lowerCAmelCase	,   _lowerCAmelCase=None						)			->							List[str]:







     if not isinstance(__lowerCAmelCase	,   torch.Tensor						):
         raise ValueError(f'''`image` has to be of type `torch.Tensor` but is {type(__lowerCAmelCase						)}'''						)

     _lowerCAmelCase			=					image.to(device=__lowerCAmelCase	,   dtype=__lowerCAmelCase						)

     if isinstance(__lowerCAmelCase	,   __lowerCAmelCase						):
         _lowerCAmelCase			=					[
             self.vae.encode(image[i : i + 1]						).latent_dist.sample(generator[i]						) for i in range(__lowerCAmelCase						)
         ]
         _lowerCAmelCase			=					torch.cat(__lowerCAmelCase	,   dim=0						)
     else:
         _lowerCAmelCase			=					self.vae.encode(__lowerCAmelCase						).latent_dist.sample(__lowerCAmelCase						)

     # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
     _lowerCAmelCase			=					0.18215 * init_latents
     _lowerCAmelCase			=					init_latents.repeat_interleave(__lowerCAmelCase	,   dim=0						)

     _lowerCAmelCase			=					randn_tensor(init_latents.shape	,   generator=__lowerCAmelCase	,   device=__lowerCAmelCase	,   dtype=__lowerCAmelCase						)

     # get latents
     _lowerCAmelCase			=					self.scheduler.add_noise(__lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase						)
     _lowerCAmelCase			=					init_latents

     return latents




 def 							_snake_case      (		self	,   _lowerCAmelCase						)			->							Any:
     _lowerCAmelCase			=					self.coca_transform(__lowerCAmelCase						).unsqueeze(0						)
     with torch.no_grad(), torch.cuda.amp.autocast():
         _lowerCAmelCase			=					self.coca_model.generate(transformed_image.to(device=self.device	,   dtype=self.coca_model.dtype						)						)
     _lowerCAmelCase			=					self.coca_tokenizer.decode(generated[0].cpu().numpy()						)
     return generated.split("<end_of_text>"						)[0].replace("<start_of_text>"	,   ""						).rstrip(" .,"						)




 def 							_snake_case      (		self	,   _lowerCAmelCase	,   _lowerCAmelCase						)			->							Union[str, Any]:
     _lowerCAmelCase			=					self.feature_extractor.preprocess(__lowerCAmelCase						)
     _lowerCAmelCase			=					torch.from_numpy(clip_image_input["pixel_values"][0]						).unsqueeze(0						).to(self.device						).half()
     _lowerCAmelCase			=					self.clip_model.get_image_features(__lowerCAmelCase						)
     _lowerCAmelCase			=					image_embeddings_clip / image_embeddings_clip.norm(p=2	,   dim=-1	,   keepdim=__lowerCAmelCase						)
     _lowerCAmelCase			=					image_embeddings_clip.repeat_interleave(__lowerCAmelCase	,   dim=0						)
     return image_embeddings_clip




 @torch.enable_grad()
 def 							_snake_case      (		self	,   _lowerCAmelCase	,   _lowerCAmelCase	,   _lowerCAmelCase	,   _lowerCAmelCase	,   _lowerCAmelCase	,   _lowerCAmelCase	,   _lowerCAmelCase	,   )			->							Tuple:
     _lowerCAmelCase			=					latents.detach().requires_grad_()

     _lowerCAmelCase			=					self.scheduler.scale_model_input(__lowerCAmelCase	,   __lowerCAmelCase						)

     # predict the noise residual
     _lowerCAmelCase			=					self.unet(__lowerCAmelCase	,   __lowerCAmelCase	,   encoder_hidden_states=__lowerCAmelCase						).sample

     if isinstance(self.scheduler	,   (PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)						):
         _lowerCAmelCase			=					self.scheduler.alphas_cumprod[timestep]
         _lowerCAmelCase			=					1 - alpha_prod_t
         # compute predicted original sample from predicted noise also called
         # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
         _lowerCAmelCase			=					(latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

         _lowerCAmelCase			=					torch.sqrt(__lowerCAmelCase						)
         _lowerCAmelCase			=					pred_original_sample * (fac) + latents * (1 - fac)
     elif isinstance(self.scheduler	,   __lowerCAmelCase						):
         _lowerCAmelCase			=					self.scheduler.sigmas[index]
         _lowerCAmelCase			=					latents - sigma * noise_pred
     else:
         raise ValueError(f'''scheduler type {type(self.scheduler						)} not supported'''						)

     # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
     _lowerCAmelCase			=					1 / 0.18215 * sample
     _lowerCAmelCase			=					self.vae.decode(__lowerCAmelCase						).sample
     _lowerCAmelCase			=					(image / 2 + 0.5).clamp(0	,   1						)

     _lowerCAmelCase			=					transforms.Resize(self.feature_extractor_size						)(__lowerCAmelCase						)
     _lowerCAmelCase			=					self.normalize(__lowerCAmelCase						).to(latents.dtype						)

     _lowerCAmelCase			=					self.clip_model.get_image_features(__lowerCAmelCase						)
     _lowerCAmelCase			=					image_embeddings_clip / image_embeddings_clip.norm(p=2	,   dim=-1	,   keepdim=__lowerCAmelCase						)

     _lowerCAmelCase			=					spherical_dist_loss(__lowerCAmelCase	,   __lowerCAmelCase						).mean() * clip_guidance_scale

     _lowerCAmelCase			=					-torch.autograd.grad(__lowerCAmelCase	,   __lowerCAmelCase						)[0]

     if isinstance(self.scheduler	,   __lowerCAmelCase						):
         _lowerCAmelCase			=					latents.detach() + grads * (sigma**2)
         _lowerCAmelCase			=					noise_pred_original
     else:
         _lowerCAmelCase			=					noise_pred_original - torch.sqrt(__lowerCAmelCase						) * grads
     return noise_pred, latents





 @torch.no_grad()
 def __call__(		self	,   _lowerCAmelCase	,   _lowerCAmelCase	,   _lowerCAmelCase = None	,   _lowerCAmelCase = None	,   _lowerCAmelCase = 512	,   _lowerCAmelCase = 512	,   _lowerCAmelCase = 0.6	,   _lowerCAmelCase = 50	,   _lowerCAmelCase = 7.5	,   _lowerCAmelCase = 1	,   _lowerCAmelCase = 0.0	,   _lowerCAmelCase = 100	,   _lowerCAmelCase = None	,   _lowerCAmelCase = "pil"	,   _lowerCAmelCase = True	,   _lowerCAmelCase = 0.8	,   _lowerCAmelCase = 0.1	,   _lowerCAmelCase = 0.1	,   )			->							Tuple:







     if isinstance(__lowerCAmelCase	,   __lowerCAmelCase						) and len(__lowerCAmelCase						) != batch_size:
         raise ValueError(f'''You have passed {batch_size} batch_size, but only {len(__lowerCAmelCase						)} generators.'''						)

     if height % 8 != 0 or width % 8 != 0:
         raise ValueError(f'''`height` and `width` have to be divisible by 8 but are {height} and {width}.'''						)

     if isinstance(__lowerCAmelCase	,   torch.Generator						) and batch_size > 1:
         _lowerCAmelCase			=					[generator] + [None] * (batch_size - 1)

     _lowerCAmelCase			=					[
         ("model", self.coca_model is None),
         ("tokenizer", self.coca_tokenizer is None),
         ("transform", self.coca_transform is None),
     ]
     _lowerCAmelCase			=					[x[0] for x in coca_is_none if x[1]]
     _lowerCAmelCase			=					", ".join(__lowerCAmelCase						)
     # generate prompts with coca model if prompt is None
     if content_prompt is None:
         if len(__lowerCAmelCase						):
             raise ValueError(
                 f'''Content prompt is None and CoCa [{coca_is_none_str}] is None.'''
                 f'''Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline.'''						)
         _lowerCAmelCase			=					self.get_image_description(__lowerCAmelCase						)
     if style_prompt is None:
         if len(__lowerCAmelCase						):
             raise ValueError(
                 f'''Style prompt is None and CoCa [{coca_is_none_str}] is None.'''
                 f''' Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline.'''						)
         _lowerCAmelCase			=					self.get_image_description(__lowerCAmelCase						)

     # get prompt text embeddings for content and style
     _lowerCAmelCase			=					self.tokenizer(
         __lowerCAmelCase	,   padding="max_length"	,   max_length=self.tokenizer.model_max_length	,   truncation=__lowerCAmelCase	,   return_tensors="pt"	,   )
     _lowerCAmelCase			=					self.text_encoder(content_text_input.input_ids.to(self.device						)						)[0]

     _lowerCAmelCase			=					self.tokenizer(
         __lowerCAmelCase	,   padding="max_length"	,   max_length=self.tokenizer.model_max_length	,   truncation=__lowerCAmelCase	,   return_tensors="pt"	,   )
     _lowerCAmelCase			=					self.text_encoder(style_text_input.input_ids.to(self.device						)						)[0]

     _lowerCAmelCase			=					slerp(__lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase						)

     # duplicate text embeddings for each generation per prompt
     _lowerCAmelCase			=					text_embeddings.repeat_interleave(__lowerCAmelCase	,   dim=0						)

     # set timesteps
     _lowerCAmelCase			=					"offset" in set(inspect.signature(self.scheduler.set_timesteps						).parameters.keys()						)
     _lowerCAmelCase			=					{}
     if accepts_offset:
         _lowerCAmelCase			=					1

     self.scheduler.set_timesteps(__lowerCAmelCase	,   **__lowerCAmelCase						)
     # Some schedulers like PNDM have timesteps as arrays
     # It's more optimized to move all timesteps to correct device beforehand
     self.scheduler.timesteps.to(self.device						)

     _lowerCAmelCase	,	_lowerCAmelCase			=					self.get_timesteps(__lowerCAmelCase	,   __lowerCAmelCase	,   self.device						)
     _lowerCAmelCase			=					timesteps[:1].repeat(__lowerCAmelCase						)

     # Preprocess image
     _lowerCAmelCase			=					preprocess(__lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase						)
     _lowerCAmelCase			=					self.prepare_latents(
         __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   text_embeddings.dtype	,   self.device	,   __lowerCAmelCase						)

     _lowerCAmelCase			=					preprocess(__lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase						)
     _lowerCAmelCase			=					self.prepare_latents(
         __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   text_embeddings.dtype	,   self.device	,   __lowerCAmelCase						)

     _lowerCAmelCase			=					slerp(__lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase						)

     if clip_guidance_scale > 0:
         _lowerCAmelCase			=					self.get_clip_image_embeddings(__lowerCAmelCase	,   __lowerCAmelCase						)
         _lowerCAmelCase			=					self.get_clip_image_embeddings(__lowerCAmelCase	,   __lowerCAmelCase						)
         _lowerCAmelCase			=					slerp(
             __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase						)

     # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
     # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
     # corresponds to doing no classifier free guidance.
     _lowerCAmelCase			=					guidance_scale > 1.0
     # get unconditional embeddings for classifier free guidance
     if do_classifier_free_guidance:
         _lowerCAmelCase			=					content_text_input.input_ids.shape[-1]
         _lowerCAmelCase			=					self.tokenizer([""]	,   padding="max_length"	,   max_length=__lowerCAmelCase	,   return_tensors="pt"						)
         _lowerCAmelCase			=					self.text_encoder(uncond_input.input_ids.to(self.device						)						)[0]
         # duplicate unconditional embeddings for each generation per prompt
         _lowerCAmelCase			=					uncond_embeddings.repeat_interleave(__lowerCAmelCase	,   dim=0						)

         # 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
         _lowerCAmelCase			=					torch.cat([uncond_embeddings, text_embeddings]						)

     # get the initial random noise unless the user supplied it

     # Unlike in other pipelines, latents need to be generated in the target device
     # for 1-to-1 results reproducibility with the CompVis implementation.
     # However this currently doesn't work in `mps`.
     _lowerCAmelCase			=					(batch_size, self.unet.config.in_channels, height // 8, width // 8)
     _lowerCAmelCase			=					text_embeddings.dtype
     if latents is None:
         if self.device.type == "mps":
             # randn does not work reproducibly on mps
             _lowerCAmelCase			=					torch.randn(__lowerCAmelCase	,   generator=__lowerCAmelCase	,   device="cpu"	,   dtype=__lowerCAmelCase						).to(
                 self.device						)
         else:
             _lowerCAmelCase			=					torch.randn(__lowerCAmelCase	,   generator=__lowerCAmelCase	,   device=self.device	,   dtype=__lowerCAmelCase						)
     else:
         if latents.shape != latents_shape:
             raise ValueError(f'''Unexpected latents shape, got {latents.shape}, expected {latents_shape}'''						)
         _lowerCAmelCase			=					latents.to(self.device						)

     # scale the initial noise by the standard deviation required by the scheduler
     _lowerCAmelCase			=					latents * self.scheduler.init_noise_sigma

     # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
     # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
     # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
     # and should be between [0, 1]
     _lowerCAmelCase			=					"eta" in set(inspect.signature(self.scheduler.step						).parameters.keys()						)
     _lowerCAmelCase			=					{}
     if accepts_eta:
         _lowerCAmelCase			=					eta

     # check if the scheduler accepts generator
     _lowerCAmelCase			=					"generator" in set(inspect.signature(self.scheduler.step						).parameters.keys()						)
     if accepts_generator:
         _lowerCAmelCase			=					generator

     with self.progress_bar(total=__lowerCAmelCase						):
         for i, t in enumerate(__lowerCAmelCase						):
             # expand the latents if we are doing classifier free guidance
             _lowerCAmelCase			=					torch.cat([latents] * 2						) if do_classifier_free_guidance else latents
             _lowerCAmelCase			=					self.scheduler.scale_model_input(__lowerCAmelCase	,   __lowerCAmelCase						)

             # predict the noise residual
             _lowerCAmelCase			=					self.unet(__lowerCAmelCase	,   __lowerCAmelCase	,   encoder_hidden_states=__lowerCAmelCase						).sample

             # perform classifier free guidance
             if do_classifier_free_guidance:
                 _lowerCAmelCase	,	_lowerCAmelCase			=					noise_pred.chunk(2						)
                 _lowerCAmelCase			=					noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

             # perform clip guidance
             if clip_guidance_scale > 0:
                 _lowerCAmelCase			=					(
                     text_embeddings.chunk(2						)[1] if do_classifier_free_guidance else text_embeddings
                 )
                 _lowerCAmelCase	,	_lowerCAmelCase			=					self.cond_fn(
                     __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   )

             # compute the previous noisy sample x_t -> x_t-1
             _lowerCAmelCase			=					self.scheduler.step(__lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   **__lowerCAmelCase						).prev_sample

        # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
     _lowerCAmelCase			=					1 / 0.18215 * latents
     _lowerCAmelCase			=					self.vae.decode(__lowerCAmelCase						).sample

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

     if output_type == "pil":
         _lowerCAmelCase			=					self.numpy_to_pil(__lowerCAmelCase						)

     if not return_dict:
         return (image, None)

     return StableDiffusionPipelineOutput(images=__lowerCAmelCase	,   nsfw_content_detected=__lowerCAmelCase						)









import itertools
import random
import unittest

import numpy as np

from transformers import is_speech_available
from transformers.testing_utils import require_torch, require_torchaudio

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
			from transformers import SpeechaTextFeatureExtractor

_a 						=      random.Random()





def lowerCAmelCase__(__snake_case       ,__snake_case=1.0       ,__snake_case=None       ,__snake_case=None							)				->      List[Any]:







	'''simple docstring'''







	if rng is None:
		lowerCamelCase__								=		global_rng

	lowerCamelCase__								=		[]
	for batch_idx in range(shape[0]							):
		values.append([]							)
		for _ in range(shape[1]							):
			values[-1].append(rng.random() * scale							)

	return values






@require_torch
@require_torchaudio
class 					__A    (    unittest.TestCase ):

			'''simple docstring'''







			def __init__(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase=7       ,  __lowerCAmelCase=4_0_0       ,  __lowerCAmelCase=2_0_0_0       ,  __lowerCAmelCase=2_4       ,  __lowerCAmelCase=2_4       ,  __lowerCAmelCase=0.0       ,  __lowerCAmelCase=1_6_0_0_0       ,  __lowerCAmelCase=True       ,  __lowerCAmelCase=True       ,  ):



				'''simple docstring'''







				lowerCamelCase__								=		parent
				lowerCamelCase__								=		batch_size
				lowerCamelCase__								=		min_seq_length
				lowerCamelCase__								=		max_seq_length
				lowerCamelCase__								=		(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
				lowerCamelCase__								=		feature_size
				lowerCamelCase__								=		num_mel_bins
				lowerCamelCase__								=		padding_value
				lowerCamelCase__								=		sampling_rate
				lowerCamelCase__								=		return_attention_mask
				lowerCamelCase__								=		do_normalize





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				return {
				    "feature_size": self.feature_size,
				    "num_mel_bins": self.num_mel_bins,
				    "padding_value": self.padding_value,
				    "sampling_rate": self.sampling_rate,
				    "return_attention_mask": self.return_attention_mask,
				    "do_normalize": self.do_normalize,
				}





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase=False       ,  __lowerCAmelCase=False							):



				'''simple docstring'''







				def _flatten(__lowerCAmelCase							):
					return list(itertools.chain(*__lowerCAmelCase							)							)

				if equal_length:
					lowerCamelCase__								=		[floats_list((self.max_seq_length, self.feature_size)							) for _ in range(self.batch_size							)]
				else:
					# make sure that inputs increase in size
					lowerCamelCase__								=		[
					    floats_list((x, self.feature_size)							)
					    for x in range(self.min_seq_length       ,  self.max_seq_length       ,  self.seq_length_diff							)
					]
				if numpify:
					lowerCamelCase__								=		[np.asarray(__lowerCAmelCase							) for x in speech_inputs]
				return speech_inputs






@require_torch
@require_torchaudio
class 					__A    (    lowerCAmelCase			,				unittest.TestCase ):

			'''simple docstring'''







			lowerCAmelCase_							     =						SpeechaTextFeatureExtractor if is_speech_available() else None





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		SpeechaTextFeatureExtractionTester(self							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase							):



				'''simple docstring'''







				self.assertTrue(np.all(np.mean(__lowerCAmelCase       ,  axis=0							) < 1E-3							)							)
				self.assertTrue(np.all(np.abs(np.var(__lowerCAmelCase       ,  axis=0							) - 1							) < 1E-3							)							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
				# create three inputs of length 800, 1000, and 1200
				lowerCamelCase__								=		[floats_list((1, x)							)[0] for x in range(8_0_0       ,  1_4_0_0       ,  2_0_0							)]
				lowerCamelCase__								=		[np.asarray(__lowerCAmelCase							) for speech_input in speech_inputs]

				# Test feature size
				lowerCamelCase__								=		feature_extractor(__lowerCAmelCase       ,  padding=__lowerCAmelCase       ,  return_tensors='''np'''							).input_features
				self.assertTrue(input_features.ndim == 3							)
				self.assertTrue(input_features.shape[-1] == feature_extractor.feature_size							)

				# Test not batched input
				lowerCamelCase__								=		feature_extractor(speech_inputs[0]       ,  return_tensors='''np'''							).input_features
				lowerCamelCase__								=		feature_extractor(np_speech_inputs[0]       ,  return_tensors='''np'''							).input_features
				self.assertTrue(np.allclose(__lowerCAmelCase       ,  __lowerCAmelCase       ,  atol=1E-3							)							)

				# Test batched
				lowerCamelCase__								=		feature_extractor(__lowerCAmelCase       ,  return_tensors='''np'''							).input_features
				lowerCamelCase__								=		feature_extractor(__lowerCAmelCase       ,  return_tensors='''np'''							).input_features
				for enc_seq_a, enc_seq_a in zip(__lowerCAmelCase       ,  __lowerCAmelCase							):
					self.assertTrue(np.allclose(__lowerCAmelCase       ,  __lowerCAmelCase       ,  atol=1E-3							)							)

				# Test 2-D numpy arrays are batched.
				lowerCamelCase__								=		[floats_list((1, x)							)[0] for x in (8_0_0, 8_0_0, 8_0_0)]
				lowerCamelCase__								=		np.asarray(__lowerCAmelCase							)
				lowerCamelCase__								=		feature_extractor(__lowerCAmelCase       ,  return_tensors='''np'''							).input_features
				lowerCamelCase__								=		feature_extractor(__lowerCAmelCase       ,  return_tensors='''np'''							).input_features
				for enc_seq_a, enc_seq_a in zip(__lowerCAmelCase       ,  __lowerCAmelCase							):
					self.assertTrue(np.allclose(__lowerCAmelCase       ,  __lowerCAmelCase       ,  atol=1E-3							)							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
				lowerCamelCase__								=		[floats_list((1, x)							)[0] for x in range(8_0_0       ,  1_4_0_0       ,  2_0_0							)]

				lowerCamelCase__								=		['''longest''', '''max_length''', '''do_not_pad''']
				lowerCamelCase__								=		[None, 1_6, None]
				for max_length, padding in zip(__lowerCAmelCase       ,  __lowerCAmelCase							):
					lowerCamelCase__								=		feature_extractor(
					    __lowerCAmelCase       ,  padding=__lowerCAmelCase       ,  max_length=__lowerCAmelCase       ,  return_attention_mask=__lowerCAmelCase							)
					lowerCamelCase__								=		inputs.input_features
					lowerCamelCase__								=		inputs.attention_mask
					lowerCamelCase__								=		[np.sum(__lowerCAmelCase							) for x in attention_mask]

					self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]							)
					self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]							)
					self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
				lowerCamelCase__								=		[floats_list((1, x)							)[0] for x in range(8_0_0       ,  1_4_0_0       ,  2_0_0							)]

				lowerCamelCase__								=		['''longest''', '''max_length''', '''do_not_pad''']
				lowerCamelCase__								=		[None, 1_6, None]
				for max_length, padding in zip(__lowerCAmelCase       ,  __lowerCAmelCase							):
					lowerCamelCase__								=		feature_extractor(
					    __lowerCAmelCase       ,  max_length=__lowerCAmelCase       ,  padding=__lowerCAmelCase       ,  return_tensors='''np'''       ,  return_attention_mask=__lowerCAmelCase							)
					lowerCamelCase__								=		inputs.input_features
					lowerCamelCase__								=		inputs.attention_mask
					lowerCamelCase__								=		[np.sum(__lowerCAmelCase							) for x in attention_mask]

					self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]							)
					self.assertTrue(input_features[0][fbank_feat_lengths[0] :].sum() < 1E-6							)
					self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]							)
					self.assertTrue(input_features[0][fbank_feat_lengths[1] :].sum() < 1E-6							)
					self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
				lowerCamelCase__								=		[floats_list((1, x)							)[0] for x in range(8_0_0       ,  1_4_0_0       ,  2_0_0							)]
				lowerCamelCase__								=		feature_extractor(
				    __lowerCAmelCase       ,  padding='''max_length'''       ,  max_length=4       ,  truncation=__lowerCAmelCase       ,  return_tensors='''np'''       ,  return_attention_mask=__lowerCAmelCase       ,  )
				lowerCamelCase__								=		inputs.input_features
				lowerCamelCase__								=		inputs.attention_mask
				lowerCamelCase__								=		np.sum(attention_mask == 1       ,  axis=1							)

				self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]							)
				self._check_zero_mean_unit_variance(input_features[1]							)
				self._check_zero_mean_unit_variance(input_features[2]							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
				lowerCamelCase__								=		[floats_list((1, x)							)[0] for x in range(8_0_0       ,  1_4_0_0       ,  2_0_0							)]
				lowerCamelCase__								=		feature_extractor(
				    __lowerCAmelCase       ,  padding='''longest'''       ,  max_length=4       ,  truncation=__lowerCAmelCase       ,  return_tensors='''np'''       ,  return_attention_mask=__lowerCAmelCase       ,  )
				lowerCamelCase__								=		inputs.input_features
				lowerCamelCase__								=		inputs.attention_mask
				lowerCamelCase__								=		np.sum(attention_mask == 1       ,  axis=1							)

				self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]							)
				self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]							)
				self._check_zero_mean_unit_variance(input_features[2]							)

				# make sure that if max_length < longest -> then pad to max_length
				self.assertEqual(input_features.shape       ,  (3, 4, 2_4)							)

				lowerCamelCase__								=		[floats_list((1, x)							)[0] for x in range(8_0_0       ,  1_4_0_0       ,  2_0_0							)]
				lowerCamelCase__								=		feature_extractor(
				    __lowerCAmelCase       ,  padding='''longest'''       ,  max_length=1_6       ,  truncation=__lowerCAmelCase       ,  return_tensors='''np'''       ,  return_attention_mask=__lowerCAmelCase       ,  )
				lowerCamelCase__								=		inputs.input_features
				lowerCamelCase__								=		inputs.attention_mask
				lowerCamelCase__								=		np.sum(attention_mask == 1       ,  axis=1							)

				self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]							)
				self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]							)
				self._check_zero_mean_unit_variance(input_features[2]							)

				# make sure that if max_length < longest -> then pad to max_length
				self.assertEqual(input_features.shape       ,  (3, 6, 2_4)							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				import torch

				lowerCamelCase__								=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
				lowerCamelCase__								=		np.random.rand(1_0_0       ,  3_2							).astype(np.floataa							)
				lowerCamelCase__								=		np_speech_inputs.tolist()

				for inputs in [py_speech_inputs, np_speech_inputs]:
					lowerCamelCase__								=		feature_extractor.pad([{'''input_features''': inputs}]       ,  return_tensors='''np'''							)
					self.assertTrue(np_processed.input_features.dtype == np.floataa							)
					lowerCamelCase__								=		feature_extractor.pad([{'''input_features''': inputs}]       ,  return_tensors='''pt'''							)
					self.assertTrue(pt_processed.input_features.dtype == torch.floataa							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase							):



				'''simple docstring'''







				from datasets import load_dataset

				lowerCamelCase__								=		load_dataset('''hf-internal-testing/librispeech_asr_dummy'''       ,  '''clean'''       ,  split='''validation'''							)
				# automatic decoding with librispeech
				lowerCamelCase__								=		ds.sort('''id'''							).select(range(__lowerCAmelCase							)							)[:num_samples]['''audio''']

				return [x["array"] for x in speech_samples]





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		np.array([
				    -1.5745, -1.7713, -1.7020, -1.6069, -1.2250, -1.1105, -0.9072, -0.8241,
				    -1.2310, -0.8098, -0.3320, -0.4101, -0.7985, -0.4996, -0.8213, -0.9128,
				    -1.0420, -1.1286, -1.0440, -0.7999, -0.8405, -1.2275, -1.5443, -1.4625,
				]							)
				# fmt: on

				lowerCamelCase__								=		self._load_datasamples(1							)
				lowerCamelCase__								=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
				lowerCamelCase__								=		feature_extractor(__lowerCAmelCase       ,  return_tensors='''pt'''							).input_features
				self.assertEquals(input_features.shape       ,  (1, 5_8_4, 2_4)							)
				self.assertTrue(np.allclose(input_features[0, 0, :3_0]       ,  __lowerCAmelCase       ,  atol=1E-4							)							)



