"""simple docstring"""







import gc
import random
import unittest

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

import diffusers
from diffusers import (
    AutoencoderKL,
    EulerDiscreteScheduler,
    StableDiffusionLatentUpscalePipeline,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


def      __lowerCAmelCase(		__UpperCAmelCase			):






	"""simple docstring"""





	_lowercase :    Dict							=							[tensor.shape for tensor in tensor_list]
	return all(shape == shapes[0] for shape in shapes[1:]			)



class  _lowerCamelCase			(__lowerCamelCase						,							__lowerCamelCase						,							__lowerCamelCase						,							unittest.TestCase     ):
				_snake_case     =  StableDiffusionLatentUpscalePipeline
				_snake_case     =  TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
				    "height",
				    "width",
				    "cross_attention_kwargs",
				    "negative_prompt_embeds",
				    "prompt_embeds",
				}
				_snake_case     =  PipelineTesterMixin.required_optional_params - {"num_images_per_prompt"}
				_snake_case     =  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
				_snake_case     =  frozenset(
				    []     )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
				_snake_case     =  frozenset([]     )

				_snake_case     =  True







				@property
				def 							__UpperCAmelCase					( self    :       Optional[Any]						):

					"""simple docstring"""






					_lowercase :    str							=							1
					_lowercase :    Any							=							4
					_lowercase :    Optional[int]							=							(1_6, 1_6)

					_lowercase :    Dict							=							floats_tensor((batch_size, num_channels) + sizes			,					rng=random.Random(0						)						).to(lowerCamelCase_						)
					return image







				def 							__UpperCAmelCase					( self    :       Tuple						):

					"""simple docstring"""






					torch.manual_seed(0						)
					_lowercase :    Tuple							=							UNetaDConditionModel(
					    act_fn='gelu'			,					attention_head_dim=8			,					norm_num_groups=lowerCamelCase_			,					block_out_channels=[3_2, 3_2, 6_4, 6_4]			,					time_cond_proj_dim=1_6_0			,					conv_in_kernel=1			,					conv_out_kernel=1			,					cross_attention_dim=3_2			,					down_block_types=(
					        'KDownBlock2D',
					        'KCrossAttnDownBlock2D',
					        'KCrossAttnDownBlock2D',
					        'KCrossAttnDownBlock2D',
					    )			,					in_channels=8			,					mid_block_type=lowerCamelCase_			,					only_cross_attention=lowerCamelCase_			,					out_channels=5			,					resnet_time_scale_shift='scale_shift'			,					time_embedding_type='fourier'			,					timestep_post_act='gelu'			,					up_block_types=('KCrossAttnUpBlock2D', 'KCrossAttnUpBlock2D', 'KCrossAttnUpBlock2D', 'KUpBlock2D')			,					)
					_lowercase :    Optional[Any]							=							AutoencoderKL(
					    block_out_channels=[3_2, 3_2, 6_4, 6_4]			,					in_channels=3			,					out_channels=3			,					down_block_types=[
					        'DownEncoderBlock2D',
					        'DownEncoderBlock2D',
					        'DownEncoderBlock2D',
					        'DownEncoderBlock2D',
					    ]			,					up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D']			,					latent_channels=4			,					)
					_lowercase :    List[Any]							=							EulerDiscreteScheduler(prediction_type='sample'						)
					_lowercase :    str							=							CLIPTextConfig(
					    bos_token_id=0			,					eos_token_id=2			,					hidden_size=3_2			,					intermediate_size=3_7			,					layer_norm_eps=1E-05			,					num_attention_heads=4			,					num_hidden_layers=5			,					pad_token_id=1			,					vocab_size=1_0_0_0			,					hidden_act='quick_gelu'			,					projection_dim=5_1_2			,					)
					_lowercase :    int							=							CLIPTextModel(lowerCamelCase_						)
					_lowercase :    Tuple							=							CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'						)

					_lowercase :    int							=							{
					    'unet': model.eval(),
					    'vae': vae.eval(),
					    'scheduler': scheduler,
					    'text_encoder': text_encoder,
					    'tokenizer': tokenizer,
					}

					return components







				def 							__UpperCAmelCase					( self    :       Any			,					lowerCamelCase_    :       str			,					lowerCamelCase_    :       List[str]=0						):

					"""simple docstring"""






					if str(lowerCamelCase_						).startswith('mps'						):
						_lowercase :    Optional[int]							=							torch.manual_seed(lowerCamelCase_						)
					else:
						_lowercase :    Dict							=							torch.Generator(device=lowerCamelCase_						).manual_seed(lowerCamelCase_						)
					_lowercase :    Optional[int]							=							{
					    'prompt': 'A painting of a squirrel eating a burger',
					    'image': self.dummy_image.cpu(),
					    'generator': generator,
					    'num_inference_steps': 2,
					    'output_type': 'numpy',
					}
					return inputs







				def 							__UpperCAmelCase					( self    :       Union[str, Any]						):

					"""simple docstring"""






					_lowercase :    int							=							'cpu'

					_lowercase :    Dict							=							self.get_dummy_components()
					_lowercase :    Optional[int]							=							self.pipeline_class(**lowerCamelCase_						)
					pipe.to(lowerCamelCase_						)
					pipe.set_progress_bar_config(disable=lowerCamelCase_						)

					_lowercase :    List[Any]							=							self.get_dummy_inputs(lowerCamelCase_						)
					_lowercase :    Optional[Any]							=							pipe(**lowerCamelCase_						).images
					_lowercase :    List[str]							=							image[0, -3:, -3:, -1]

					self.assertEqual(image.shape			,					(1, 2_5_6, 2_5_6, 3)						)
					_lowercase :    List[Any]							=							np.array(
					    [0.4722_2412, 0.4192_1633, 0.4471_7434, 0.4687_4192, 0.4258_8258, 0.4615_0726, 0.467_7534, 0.4558_3832, 0.4857_9055]						)
					_lowercase :    str							=							np.abs(image_slice.flatten() - expected_slice						).max()
					self.assertLessEqual(lowerCamelCase_			,					1E-3						)







				def 							__UpperCAmelCase					( self    :       Optional[int]						):

					"""simple docstring"""






					super().test_attention_slicing_forward_pass(expected_max_diff=7E-3						)







				def 							__UpperCAmelCase					( self    :       Optional[Any]						):

					"""simple docstring"""






					super().test_cpu_offload_forward_pass(expected_max_diff=3E-3						)







				def 							__UpperCAmelCase					( self    :       Any						):

					"""simple docstring"""






					super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3						)







				def 							__UpperCAmelCase					( self    :       int						):

					"""simple docstring"""






					super().test_inference_batch_single_identical(expected_max_diff=7E-3						)







				def 							__UpperCAmelCase					( self    :       Any						):

					"""simple docstring"""






					super().test_pt_np_pil_outputs_equivalent(expected_max_diff=3E-3						)







				def 							__UpperCAmelCase					( self    :       int						):

					"""simple docstring"""






					super().test_save_load_local(expected_max_difference=3E-3						)







				def 							__UpperCAmelCase					( self    :       str						):

					"""simple docstring"""






					super().test_save_load_optional_components(expected_max_difference=3E-3						)







				def 							__UpperCAmelCase					( self    :       str						):

					"""simple docstring"""






					_lowercase :    Dict							=							[
					    'DDIMScheduler',
					    'DDPMScheduler',
					    'PNDMScheduler',
					    'HeunDiscreteScheduler',
					    'EulerAncestralDiscreteScheduler',
					    'KDPM2DiscreteScheduler',
					    'KDPM2AncestralDiscreteScheduler',
					    'DPMSolverSDEScheduler',
					]
					_lowercase :    Optional[Any]							=							self.get_dummy_components()
					_lowercase :    Dict							=							self.pipeline_class(**lowerCamelCase_						)

					# make sure that PNDM does not need warm-up
					pipe.scheduler.register_to_config(skip_prk_steps=lowerCamelCase_						)

					pipe.to(lowerCamelCase_						)
					pipe.set_progress_bar_config(disable=lowerCamelCase_						)
					_lowercase :    Optional[Any]							=							self.get_dummy_inputs(lowerCamelCase_						)
					_lowercase :    Optional[int]							=							2

					_lowercase :    List[str]							=							[]
					for scheduler_enum in KarrasDiffusionSchedulers:
						if scheduler_enum.name in skip_schedulers:
							# no sigma schedulers are not supported
							# no schedulers
							continue

						_lowercase :    Dict							=							getattr(lowerCamelCase_			,					scheduler_enum.name						)
						_lowercase :    int							=							scheduler_cls.from_config(pipe.scheduler.config						)
						_lowercase :    Optional[Any]							=							pipe(**lowerCamelCase_						)[0]
						outputs.append(lowerCamelCase_						)

					assert check_same_shape(lowerCamelCase_						)



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







				def 							__UpperCAmelCase					( self    :       Optional[Any]						):

					"""simple docstring"""






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







				def 							__UpperCAmelCase					( self    :       Any						):

					"""simple docstring"""






					_lowercase :    Optional[int]							=							torch.manual_seed(3_3						)

					_lowercase :    Tuple							=							StableDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4'			,					torch_dtype=torch.floataa						)
					pipe.to('cuda'						)

					_lowercase :    Any							=							StableDiffusionLatentUpscalePipeline.from_pretrained(
					    'stabilityai/sd-x2-latent-upscaler'			,					torch_dtype=torch.floataa						)
					upscaler.to('cuda'						)

					_lowercase :    Dict							=							'a photo of an astronaut high resolution, unreal engine, ultra realistic'

					_lowercase :    Tuple							=							pipe(lowerCamelCase_			,					generator=lowerCamelCase_			,					output_type='latent'						).images

					_lowercase :    str							=							upscaler(
					    prompt=lowerCamelCase_			,					image=lowerCamelCase_			,					num_inference_steps=2_0			,					guidance_scale=0			,					generator=lowerCamelCase_			,					output_type='np'			,					).images[0]

					_lowercase :    Union[str, Any]							=							load_numpy(
					    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/astronaut_1024.npy'						)
					assert np.abs((expected_image - image).mean()						) < 5E-2







				def 							__UpperCAmelCase					( self    :       int						):

					"""simple docstring"""






					_lowercase :    List[Any]							=							torch.manual_seed(3_3						)

					_lowercase :    str							=							StableDiffusionLatentUpscalePipeline.from_pretrained(
					    'stabilityai/sd-x2-latent-upscaler'			,					torch_dtype=torch.floataa						)
					upscaler.to('cuda'						)

					_lowercase :    Optional[int]							=							'the temple of fire by Ross Tran and Gerardo Dottori, oil on canvas'

					_lowercase :    int							=							load_image(
					    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_512.png'						)

					_lowercase :    List[str]							=							upscaler(
					    prompt=lowerCamelCase_			,					image=lowerCamelCase_			,					num_inference_steps=2_0			,					guidance_scale=0			,					generator=lowerCamelCase_			,					output_type='np'			,					).images[0]

					_lowercase :    Optional[Any]							=							load_numpy(
					    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_1024.npy'						)
					assert np.abs((expected_image - image).max()						) < 5E-2




"""simple docstring"""







from typing import List, Optional, Union

import numpy as np

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


SCREAMING_SNAKE_CASE									=						logging.get_logger(__name__)



class  _lowerCamelCase			(__lowerCamelCase     ):
				_snake_case     =  ["input_features", "attention_mask"]







				def __init__( self    :       int			,					lowerCamelCase_    :       List[str]=8_0			,					lowerCamelCase_    :       Tuple=1_6_0_0_0			,					lowerCamelCase_    :       Union[str, Any]=0.0			,					lowerCamelCase_    :       List[Any]=1_0			,					lowerCamelCase_    :       List[str]=2_5			,					lowerCamelCase_    :       List[Any]="hamming_window"			,					lowerCamelCase_    :       Tuple=3_2768.0			,					lowerCamelCase_    :       int=0.97			,					lowerCamelCase_    :       Optional[int]=1.0			,					lowerCamelCase_    :       Optional[Any]=True			,					lowerCamelCase_    :       Union[str, Any]=True			,					lowerCamelCase_    :       Optional[int]=False			,					**lowerCamelCase_    :       Optional[Any]			,					):

					"""simple docstring"""






					super().__init__(feature_size=lowerCamelCase_			,					sampling_rate=lowerCamelCase_			,					padding_value=lowerCamelCase_			,					**lowerCamelCase_						)

					_lowercase :    Dict							=							feature_size
					_lowercase :    Dict							=							sampling_rate
					_lowercase :    Tuple							=							padding_value
					_lowercase :    int							=							hop_length
					_lowercase :    Any							=							win_length
					_lowercase :    Union[str, Any]							=							frame_signal_scale
					_lowercase :    Tuple							=							preemphasis_coeff
					_lowercase :    Tuple							=							mel_floor
					_lowercase :    Tuple							=							normalize_means
					_lowercase :    List[Any]							=							normalize_vars
					_lowercase :    List[str]							=							win_function
					_lowercase :    int							=							return_attention_mask

					_lowercase :    Optional[Any]							=							win_length * sampling_rate // 1_0_0_0
					_lowercase :    Tuple							=							hop_length * sampling_rate // 1_0_0_0

					_lowercase :    str							=							optimal_fft_length(self.sample_size						)
					_lowercase :    Dict							=							(self.n_fft // 2) + 1







				def 							__UpperCAmelCase					( self    :       List[str]			,					lowerCamelCase_    :       np.array						):

					"""simple docstring"""






					if self.win_function == "hamming_window":
						_lowercase :    List[Any]							=							window_function(window_length=self.sample_size			,					name=self.win_function			,					periodic=lowerCamelCase_						)
					else:
						_lowercase :    Union[str, Any]							=							window_function(window_length=self.sample_size			,					name=self.win_function						)

					_lowercase :    Tuple							=							mel_filter_bank(
					    num_frequency_bins=self.n_freqs			,					num_mel_filters=self.feature_size			,					min_frequency=0.0			,					max_frequency=self.sampling_rate / 2.0			,					sampling_rate=self.sampling_rate			,					)

					_lowercase :    Tuple							=							spectrogram(
					    one_waveform * self.frame_signal_scale			,					window=lowerCamelCase_			,					frame_length=self.sample_size			,					hop_length=self.sample_stride			,					fft_length=self.n_fft			,					center=lowerCamelCase_			,					preemphasis=self.preemphasis_coeff			,					mel_filters=lowerCamelCase_			,					mel_floor=self.mel_floor			,					log_mel='log'			,					)
					return msfc_features.T







				def 							__UpperCAmelCase					( self    :       Optional[Any]			,					lowerCamelCase_    :       Dict			,					lowerCamelCase_    :       Dict			,					lowerCamelCase_    :       Tuple						):

					"""simple docstring"""






					if self.normalize_means:
						_lowercase :    Optional[int]							=							x[:input_length].mean(axis=0						)
						_lowercase :    int							=							np.subtract(lowerCamelCase_			,					lowerCamelCase_						)
					if self.normalize_vars:
						_lowercase :    int							=							x[:input_length].std(axis=0						)
						_lowercase :    Optional[Any]							=							np.divide(lowerCamelCase_			,					lowerCamelCase_						)

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

					# make sure array is in float32
					_lowercase :    Tuple							=							x.astype(np.floataa						)

					return x







				def 							__UpperCAmelCase					( self    :       List[Any]			,					lowerCamelCase_    :       List[np.ndarray]			,					lowerCamelCase_    :       Optional[np.ndarray] = None						):

					"""simple docstring"""






					_lowercase :    Dict							=							attention_mask.sum(-1						) if attention_mask is not None else [x.shape[0] for x in input_features]
					return [self._normalize_one(lowerCamelCase_			,					lowerCamelCase_			,					self.padding_value						) for x, n in zip(lowerCamelCase_			,					lowerCamelCase_						)]







				def __call__( self    :       Dict			,					lowerCamelCase_    :       Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]			,					lowerCamelCase_    :       Union[bool, str, PaddingStrategy] = False			,					lowerCamelCase_    :       Optional[int] = None			,					lowerCamelCase_    :       bool = False			,					lowerCamelCase_    :       Optional[int] = None			,					lowerCamelCase_    :       Optional[bool] = None			,					lowerCamelCase_    :       Optional[Union[str, TensorType]] = None			,					lowerCamelCase_    :       Optional[int] = None			,					**lowerCamelCase_    :       Optional[int]			,					):

					"""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.'						)

					_lowercase :    Optional[Any]							=							isinstance(lowerCamelCase_			,					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}'''						)
					_lowercase :    Optional[int]							=							is_batched_numpy or (
					    isinstance(lowerCamelCase_			,					(list, tuple)						) and (isinstance(raw_speech[0]			,					(np.ndarray, tuple, list)						))
					)

					if is_batched:
						_lowercase :    str							=							[np.asarray(lowerCamelCase_			,					dtype=np.floataa						) for speech in raw_speech]
					elif not is_batched and not isinstance(lowerCamelCase_			,					np.ndarray						):
						_lowercase :    Tuple							=							np.asarray(lowerCamelCase_			,					dtype=np.floataa						)
					elif isinstance(lowerCamelCase_			,					np.ndarray						) and raw_speech.dtype is np.dtype(np.floataa						):
						_lowercase :    Optional[int]							=							raw_speech.astype(np.floataa						)

					# always return batch
					if not is_batched:
						_lowercase :    str							=							[raw_speech]

					# extract fbank features
					_lowercase :    Optional[Any]							=							[self._extract_mfsc_features(lowerCamelCase_						) for one_waveform in raw_speech]

					# convert into correct format for padding
					_lowercase :    Optional[int]							=							BatchFeature({'input_features': features}						)

					_lowercase :    Tuple							=							self.pad(
					    lowerCamelCase_			,					padding=lowerCamelCase_			,					max_length=lowerCamelCase_			,					truncation=lowerCamelCase_			,					pad_to_multiple_of=lowerCamelCase_			,					return_attention_mask=lowerCamelCase_			,					**lowerCamelCase_			,					)
					# make sure list is in array format
					_lowercase :    Dict							=							padded_inputs.get('input_features'						)
					if isinstance(input_features[0]			,					lowerCamelCase_						):
						_lowercase :    List[str]							=							[np.asarray(lowerCamelCase_			,					dtype=np.floataa						) for feature in input_features]

					_lowercase :    List[Any]							=							padded_inputs.get('attention_mask'						)
					if attention_mask is not None:
						_lowercase :    Union[str, Any]							=							[np.asarray(lowerCamelCase_			,					dtype=np.intaa						) for array in attention_mask]

					if self.normalize_means or self.normalize_vars:
						_lowercase :    int							=							(
						    np.array(lowerCamelCase_			,					dtype=np.intaa						)
						    if self._get_padding_strategies(lowerCamelCase_			,					max_length=lowerCamelCase_						) is not PaddingStrategy.DO_NOT_PAD
						    and padding
						    else None
						)
						_lowercase :    List[Any]							=							self.normalize(
						    padded_inputs['input_features']			,					attention_mask=lowerCamelCase_						)

					if return_tensors is not None:
						_lowercase :    Union[str, Any]							=							padded_inputs.convert_to_tensors(lowerCamelCase_						)

					return padded_inputs



