




from math import acos, sin
from typing import List, Tuple, Union

import numpy as np
import torch
from PIL import Image

from ...models import AutoencoderKL, UNetaDConditionModel
from ...schedulers import DDIMScheduler, DDPMScheduler
from ...utils import randn_tensor
from ..pipeline_utils import AudioPipelineOutput, BaseOutput, DiffusionPipeline, ImagePipelineOutput
from .mel import Mel







class     __lowerCAmelCase (       __magic_name__ ):

			"""simple docstring"""







			snake_case_   		=       ['''vqvae''']


			def __init__(      self			,						lowerCamelCase__			,						lowerCamelCase__			,						lowerCamelCase__			,						lowerCamelCase__			,						)    ->      Optional[Any]:



							'''simple docstring'''
							super().__init__()
							self.register_modules(unet=lowerCamelCase__			,						scheduler=lowerCamelCase__			,						mel=lowerCamelCase__			,						vqvae=lowerCamelCase__							)


			def       lowercase_    (      self							)    ->      int:



							'''simple docstring'''
							return 50 if isinstance(self.scheduler			,						lowerCamelCase__							) else 1_000


			@torch.no_grad()
			def __call__(      self			,						lowerCamelCase__ = 1			,						lowerCamelCase__ = None			,						lowerCamelCase__ = None			,						lowerCamelCase__ = 0			,						lowerCamelCase__ = 0			,						lowerCamelCase__ = None			,						lowerCamelCase__ = None			,						lowerCamelCase__ = 0			,						lowerCamelCase__ = 0			,						lowerCamelCase__ = None			,						lowerCamelCase__ = 0			,						lowerCamelCase__ = None			,						lowerCamelCase__ = None			,						lowerCamelCase__=True			,						)    ->      Union[
			    Union[AudioPipelineOutput, ImagePipelineOutput],
			    Tuple[List[Image.Image], Tuple[int, List[np.ndarray]]],
			]:



							'''simple docstring'''
							__lowerCamelCase		       =    steps or self.get_default_steps()
							self.scheduler.set_timesteps(lowerCamelCase__							)
							__lowerCamelCase		       =    step_generator or generator
							# For backwards compatibility
							if type(self.unet.config.sample_size							) == int:
											__lowerCamelCase		       =    (self.unet.config.sample_size, self.unet.config.sample_size)
							if noise is None:
											__lowerCamelCase		       =    randn_tensor(
											    (
											        batch_size,
											        self.unet.config.in_channels,
											        self.unet.config.sample_size[0],
											        self.unet.config.sample_size[1],
											    )			,						generator=lowerCamelCase__			,						device=self.device			,						)
							__lowerCamelCase		       =    noise
							__lowerCamelCase		       =    None

							if audio_file is not None or raw_audio is not None:
											self.mel.load_audio(lowerCamelCase__			,						lowerCamelCase__							)
											__lowerCamelCase		       =    self.mel.audio_slice_to_image(lowerCamelCase__							)
											__lowerCamelCase		       =    np.frombuffer(input_image.tobytes()			,						dtype='uint8'							).reshape(
											    (input_image.height, input_image.width)							)
											__lowerCamelCase		       =    (input_image / 255) * 2 - 1
											__lowerCamelCase		       =    torch.tensor(input_image[np.newaxis, :, :]			,						dtype=torch.float							).to(self.device							)

											if self.vqvae is not None:
															__lowerCamelCase		       =    self.vqvae.encode(torch.unsqueeze(lowerCamelCase__			,						0							)							).latent_dist.sample(
															    generator=lowerCamelCase__							)[0]
															__lowerCamelCase		       =    self.vqvae.config.scaling_factor * input_images

											if start_step > 0:
															__lowerCamelCase		       =    self.scheduler.add_noise(lowerCamelCase__			,						lowerCamelCase__			,						self.scheduler.timesteps[start_step - 1]							)

											__lowerCamelCase		       =    (
											    self.unet.config.sample_size[1] * self.mel.get_sample_rate() / self.mel.x_res / self.mel.hop_length
											)
											__lowerCamelCase		       =    int(mask_start_secs * pixels_per_second							)
											__lowerCamelCase		       =    int(mask_end_secs * pixels_per_second							)
											__lowerCamelCase		       =    self.scheduler.add_noise(lowerCamelCase__			,						lowerCamelCase__			,						torch.tensor(self.scheduler.timesteps[start_step:]							)							)

							for step, t in enumerate(self.progress_bar(self.scheduler.timesteps[start_step:]							)							):
											if isinstance(self.unet			,						lowerCamelCase__							):
															__lowerCamelCase		       =    self.unet(lowerCamelCase__			,						lowerCamelCase__			,						lowerCamelCase__							)['sample']
											else:
															__lowerCamelCase		       =    self.unet(lowerCamelCase__			,						lowerCamelCase__							)['sample']

											if isinstance(self.scheduler			,						lowerCamelCase__							):
															__lowerCamelCase		       =    self.scheduler.step(
															    model_output=lowerCamelCase__			,						timestep=lowerCamelCase__			,						sample=lowerCamelCase__			,						eta=lowerCamelCase__			,						generator=lowerCamelCase__			,						)['prev_sample']
											else:
															__lowerCamelCase		       =    self.scheduler.step(
															    model_output=lowerCamelCase__			,						timestep=lowerCamelCase__			,						sample=lowerCamelCase__			,						generator=lowerCamelCase__			,						)['prev_sample']

											if mask is not None:
															if mask_start > 0:
																			__lowerCamelCase		       =    mask[:, step, :, :mask_start]
															if mask_end > 0:
																			__lowerCamelCase		       =    mask[:, step, :, -mask_end:]

							if self.vqvae is not None:
											# 0.18215 was scaling factor used in training to ensure unit variance
											__lowerCamelCase		       =    1 / self.vqvae.config.scaling_factor * images
											__lowerCamelCase		       =    self.vqvae.decode(lowerCamelCase__							)['sample']

							__lowerCamelCase		       =    (images / 2 + 0.5).clamp(0			,						1							)
							__lowerCamelCase		       =    images.cpu().permute(0			,						2			,						3			,						1							).numpy()
							__lowerCamelCase		       =    (images * 255).round().astype('uint8'							)
							__lowerCamelCase		       =    list(
							    (Image.fromarray(_[:, :, 0]							) for _ in images)
							    if images.shape[3] == 1
							    else (Image.fromarray(lowerCamelCase__			,						mode='RGB'							).convert('L'							) for _ in images)							)

							__lowerCamelCase		       =    [self.mel.image_to_audio(lowerCamelCase__							) for _ in images]
							if not return_dict:
											return images, (self.mel.get_sample_rate(), audios)

							return BaseOutput(**AudioPipelineOutput(np.array(lowerCamelCase__							)[:, np.newaxis, :]							)			,						**ImagePipelineOutput(lowerCamelCase__							)							)


			@torch.no_grad()
			def       lowercase_    (      self			,						lowerCamelCase__			,						lowerCamelCase__ = 50							)    ->      np.ndarray:



							'''simple docstring'''
							assert isinstance(self.scheduler			,						lowerCamelCase__							)
							self.scheduler.set_timesteps(lowerCamelCase__							)
							__lowerCamelCase		       =    np.array(
							    [np.frombuffer(image.tobytes()			,						dtype='uint8'							).reshape((1, image.height, image.width)							) for image in images]							)
							__lowerCamelCase		       =    (sample / 255) * 2 - 1
							__lowerCamelCase		       =    torch.Tensor(lowerCamelCase__							).to(self.device							)

							for t in self.progress_bar(torch.flip(self.scheduler.timesteps			,						(0,)							)							):
											__lowerCamelCase		       =    t - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
											__lowerCamelCase		       =    self.scheduler.alphas_cumprod[t]
											__lowerCamelCase		       =    (
											    self.scheduler.alphas_cumprod[prev_timestep]
											    if prev_timestep >= 0
											    else self.scheduler.final_alpha_cumprod
											)
											__lowerCamelCase		       =    1 - alpha_prod_t
											__lowerCamelCase		       =    self.unet(lowerCamelCase__			,						lowerCamelCase__							)['sample']
											__lowerCamelCase		       =    (1 - alpha_prod_t_prev) ** 0.5 * model_output
											__lowerCamelCase		       =    (sample - pred_sample_direction) * alpha_prod_t_prev ** (-0.5)
											__lowerCamelCase		       =    sample * alpha_prod_t ** 0.5 + beta_prod_t ** 0.5 * model_output

							return sample






			@staticmethod
			def       lowercase_    (      lowerCamelCase__			,						lowerCamelCase__			,						lowerCamelCase__							)    ->      torch.Tensor:



							'''simple docstring'''
							__lowerCamelCase		       =    acos(torch.dot(torch.flatten(lowerCamelCase__							)			,						torch.flatten(lowerCamelCase__							)							) / torch.norm(lowerCamelCase__							) / torch.norm(lowerCamelCase__							)							)
							return sin((1 - alpha) * theta							) * xa / sin(lowerCamelCase__							) + sin(alpha * theta							) * xa / sin(lowerCamelCase__							)








import itertools
import os
import random
import tempfile
import unittest

import numpy as np
from datasets import load_dataset

from transformers import is_speech_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
					from transformers import WhisperFeatureExtractor

if is_torch_available():
					import torch

__A        =    random.Random()







def 			lowerCamelCase_		(       UpperCamelCase__				:      Optional[int] ,     UpperCamelCase__				:      List[str]=1.0 ,     UpperCamelCase__				:      Union[str, Any]=None ,     UpperCamelCase__				:      Optional[int]=None      )     ->       Optional[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     __lowerCAmelCase (       unittest.TestCase ):

			"""simple docstring"""







			def __init__(      self			,						lowerCamelCase__			,						lowerCamelCase__=7			,						lowerCamelCase__=400			,						lowerCamelCase__=2_000			,						lowerCamelCase__=10			,						lowerCamelCase__=160			,						lowerCamelCase__=8			,						lowerCamelCase__=0.0			,						lowerCamelCase__=4_000			,						lowerCamelCase__=False			,						lowerCamelCase__=True			,						)    ->      List[str]:



							'''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		       =    padding_value
							__lowerCamelCase		       =    sampling_rate
							__lowerCamelCase		       =    return_attention_mask
							__lowerCamelCase		       =    do_normalize
							__lowerCamelCase		       =    feature_size
							__lowerCamelCase		       =    chunk_length
							__lowerCamelCase		       =    hop_length


			def       lowercase_    (      self							)    ->      Any:



							'''simple docstring'''
							return {
							    "feature_size": self.feature_size,
							    "hop_length": self.hop_length,
							    "chunk_length": self.chunk_length,
							    "padding_value": self.padding_value,
							    "sampling_rate": self.sampling_rate,
							    "return_attention_mask": self.return_attention_mask,
							    "do_normalize": self.do_normalize,
							}






			def       lowercase_    (      self			,						lowerCamelCase__=False			,						lowerCamelCase__=False							)    ->      Optional[int]:



							'''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     __lowerCAmelCase (       __magic_name__			,				unittest.TestCase ):

			"""simple docstring"""







			snake_case_   		=       WhisperFeatureExtractor if is_speech_available() else None


			def       lowercase_    (      self							)    ->      Any:



							'''simple docstring'''
							__lowerCamelCase		       =    WhisperFeatureExtractionTester(self							)


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



							'''simple docstring'''
							__lowerCamelCase		       =    self.feature_extraction_class(**self.feat_extract_dict							)

							with tempfile.TemporaryDirectory() as tmpdirname:
											__lowerCamelCase		       =    feat_extract_first.save_pretrained(lowerCamelCase__							)[0]
											check_json_file_has_correct_format(lowerCamelCase__							)
											__lowerCamelCase		       =    self.feature_extraction_class.from_pretrained(lowerCamelCase__							)

							__lowerCamelCase		       =    feat_extract_first.to_dict()
							__lowerCamelCase		       =    feat_extract_second.to_dict()
							__lowerCamelCase		       =    feat_extract_first.mel_filters
							__lowerCamelCase		       =    feat_extract_second.mel_filters
							self.assertTrue(np.allclose(lowerCamelCase__			,						lowerCamelCase__							)							)
							self.assertEqual(lowerCamelCase__			,						lowerCamelCase__							)


			def       lowercase_    (      self							)    ->      Dict:



							'''simple docstring'''
							__lowerCamelCase		       =    self.feature_extraction_class(**self.feat_extract_dict							)

							with tempfile.TemporaryDirectory() as tmpdirname:
											__lowerCamelCase		       =    os.path.join(lowerCamelCase__			,						'feat_extract.json'							)
											feat_extract_first.to_json_file(lowerCamelCase__							)
											__lowerCamelCase		       =    self.feature_extraction_class.from_json_file(lowerCamelCase__							)

							__lowerCamelCase		       =    feat_extract_first.to_dict()
							__lowerCamelCase		       =    feat_extract_second.to_dict()
							__lowerCamelCase		       =    feat_extract_first.mel_filters
							__lowerCamelCase		       =    feat_extract_second.mel_filters
							self.assertTrue(np.allclose(lowerCamelCase__			,						lowerCamelCase__							)							)
							self.assertEqual(lowerCamelCase__			,						lowerCamelCase__							)


			def       lowercase_    (      self							)    ->      Dict:



							'''simple docstring'''
							# Tests that all call wrap to encode_plus and batch_encode_plus
							__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(800			,						1_400			,						200							)]
							__lowerCamelCase		       =    [np.asarray(lowerCamelCase__							) for speech_input in speech_inputs]

							# Test feature size
							__lowerCamelCase		       =    feature_extractor(lowerCamelCase__			,						padding='max_length'			,						return_tensors='np'							).input_features
							self.assertTrue(input_features.ndim == 3							)
							self.assertTrue(input_features.shape[-1] == feature_extractor.nb_max_frames							)
							self.assertTrue(input_features.shape[-2] == 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 (800, 800, 800)]
							__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							)							)

							# Test truncation required
							__lowerCamelCase		       =    [floats_list((1, x)							)[0] for x in range(200			,						(feature_extractor.n_samples + 500)			,						200							)]
							__lowerCamelCase		       =    [np.asarray(lowerCamelCase__							) for speech_input in speech_inputs]

							__lowerCamelCase		       =    [x[: feature_extractor.n_samples] for x in speech_inputs]
							__lowerCamelCase		       =    [np.asarray(lowerCamelCase__							) for speech_input in speech_inputs_truncated]

							__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       lowercase_    (      self							)    ->      List[str]:



							'''simple docstring'''
							import torch

							__lowerCamelCase		       =    self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
							__lowerCamelCase		       =    np.random.rand(100			,						32							).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       lowercase_    (      self			,						lowerCamelCase__							)    ->      List[str]:



							'''simple docstring'''
							__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       lowercase_    (      self							)    ->      Tuple:



							'''simple docstring'''
							# fmt: off
							__lowerCamelCase		       =    torch.tensor(
							    [
							        0.11_93, -0.09_46, -0.10_98, -0.01_96, 0.02_25, -0.06_90, -0.17_36, 0.09_51,
							        0.09_71, -0.08_17, -0.07_02, 0.01_62, 0.02_60, 0.00_17, -0.01_92, -0.16_78,
							        0.07_09, -0.18_67, -0.06_55, -0.02_74, -0.02_34, -0.18_84, -0.05_16, -0.05_54,
							        -0.02_74, -0.14_25, -0.14_23, 0.08_37, 0.03_77, -0.08_54
							    ]							)
							# fmt: on

							__lowerCamelCase		       =    self._load_datasamples(1							)
							__lowerCamelCase		       =    WhisperFeatureExtractor()
							__lowerCamelCase		       =    feature_extractor(lowerCamelCase__			,						return_tensors='pt'							).input_features
							self.assertEqual(input_features.shape			,						(1, 80, 3_000)							)
							self.assertTrue(torch.allclose(input_features[0, 0, :30]			,						lowerCamelCase__			,						atol=1e-4							)							)






			def       lowercase_    (      self							)    ->      List[Any]:



							'''simple docstring'''
							__lowerCamelCase		       =    self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
							__lowerCamelCase		       =    self._load_datasamples(1							)[0]
							__lowerCamelCase		       =    ((audio - audio.min()) / (audio.max() - audio.min())) * 65_535  # Rescale to [0, 65535] to show issue
							__lowerCamelCase		       =    feat_extract.zero_mean_unit_var_norm([audio]			,						attention_mask=lowerCamelCase__							)[0]

							self.assertTrue(np.all(np.mean(lowerCamelCase__							) < 1e-3							)							)
							self.assertTrue(np.all(np.abs(np.var(lowerCamelCase__							) - 1							) < 1e-3							)							)


