


from typing import List, Optional, Union

import numpy as np
import torch
import torchaudio.compliance.kaldi as ta_kaldi

from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import PaddingStrategy, TensorType, logging


UpperCAmelCase_      : Optional[int]  					=       logging.get_logger(__name__)
class  SCREAMING_SNAKE_CASE__  (				lowercase__			):
      snake_case__				:     str        =						['''input_features''', '''attention_mask''']
      def __init__(       self				:       Any     ,							SCREAMING_SNAKE_CASE__				:       Tuple=8_0     ,							SCREAMING_SNAKE_CASE__				:       Dict=1_6_0_0_0     ,							SCREAMING_SNAKE_CASE__				:       List[Any]=8_0     ,							SCREAMING_SNAKE_CASE__				:       Union[str, Any]=0.0     ,							SCREAMING_SNAKE_CASE__				:       Dict=True     ,							SCREAMING_SNAKE_CASE__				:       List[str]=True     ,							SCREAMING_SNAKE_CASE__				:       int=True     ,							**SCREAMING_SNAKE_CASE__				:       Optional[int]     ,							)		->			List[Any]:
             super().__init__(feature_size=SCREAMING_SNAKE_CASE__     ,							sampling_rate=SCREAMING_SNAKE_CASE__     ,							padding_value=SCREAMING_SNAKE_CASE__     ,							**SCREAMING_SNAKE_CASE__     )
             a_					:		List[Any]							=  num_mel_bins
             a_					:		Any							=  do_ceptral_normalize
             a_					:		Optional[Any]							=  normalize_means
             a_					:		Dict							=  normalize_vars
             a_					:		Union[str, Any]							=  True
      def        SCREAMING_SNAKE_CASE							(       self				:       Any     ,							SCREAMING_SNAKE_CASE__				:       np.ndarray     ,							)		->			np.ndarray:
             a_					:		Optional[int]							=  waveform * (2**1_5)  # Kaldi compliance: 16-bit signed integers
             a_					:		Union[str, Any]							=  torch.from_numpy(SCREAMING_SNAKE_CASE__     ).unsqueeze(0     )
             a_					:		Optional[Any]							=  ta_kaldi.fbank(SCREAMING_SNAKE_CASE__     ,							num_mel_bins=self.num_mel_bins     ,							sample_frequency=self.sampling_rate     )
             return features.numpy()
      @staticmethod
      def        SCREAMING_SNAKE_CASE							(       SCREAMING_SNAKE_CASE__				:       np.ndarray     ,							SCREAMING_SNAKE_CASE__				:       int     ,							SCREAMING_SNAKE_CASE__				:       Optional[bool] = True     ,							SCREAMING_SNAKE_CASE__				:       Optional[bool] = True     ,							SCREAMING_SNAKE_CASE__				:       float = 0.0     ,							)		->			np.ndarray:
             # make sure we normalize float32 arrays
             if normalize_means:
                    a_					:		Union[str, Any]							=  x[:input_length].mean(axis=0     )
                    a_					:		List[Any]							=  np.subtract(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     )
             if normalize_vars:
                    a_					:		Tuple							=  x[:input_length].std(axis=0     )
                    a_					:		str							=  np.divide(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     )

             if input_length < x.shape[0]:
                    a_					:		Tuple							=  padding_value

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

             return x
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     ,							SCREAMING_SNAKE_CASE__				:       List[np.ndarray]     ,							SCREAMING_SNAKE_CASE__				:       Optional[np.ndarray] = None     )		->			List[np.ndarray]:
             a_					:		Union[str, Any]							=  attention_mask.sum(-1     ) if attention_mask is not None else [x.shape[0] for x in input_features]
             return [
                 self.utterance_cmvn(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     ,							self.normalize_means     ,							self.normalize_vars     ,							self.padding_value     )
                 for x, n in zip(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     )
             ]

      def __call__(       self				:       Union[str, Any]     ,							SCREAMING_SNAKE_CASE__				:       Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]     ,							SCREAMING_SNAKE_CASE__				:       Union[bool, str, PaddingStrategy] = False     ,							SCREAMING_SNAKE_CASE__				:       Optional[int] = None     ,							SCREAMING_SNAKE_CASE__				:       bool = False     ,							SCREAMING_SNAKE_CASE__				:       Optional[int] = None     ,							SCREAMING_SNAKE_CASE__				:       Optional[Union[str, TensorType]] = None     ,							SCREAMING_SNAKE_CASE__				:       Optional[int] = None     ,							SCREAMING_SNAKE_CASE__				:       Optional[bool] = None     ,							**SCREAMING_SNAKE_CASE__				:       Any     ,							)		->			BatchFeature:

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

             a_					:		Any							=  isinstance(SCREAMING_SNAKE_CASE__     ,							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}"""     )
             a_					:		Any							=  is_batched_numpy or (
                 isinstance(SCREAMING_SNAKE_CASE__     ,							(list, tuple)     ) and (isinstance(raw_speech[0]     ,							(np.ndarray, tuple, list)     ))
             )

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

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

             # extract fbank features
             a_					:		Any							=  [self._extract_fbank_features(SCREAMING_SNAKE_CASE__     ) for waveform in raw_speech]

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

             a_					:		List[Any]							=  self.pad(
                 SCREAMING_SNAKE_CASE__     ,							padding=SCREAMING_SNAKE_CASE__     ,							max_length=SCREAMING_SNAKE_CASE__     ,							truncation=SCREAMING_SNAKE_CASE__     ,							pad_to_multiple_of=SCREAMING_SNAKE_CASE__     ,							return_attention_mask=SCREAMING_SNAKE_CASE__     ,							**SCREAMING_SNAKE_CASE__     ,							)

             # make sure list is in array format
             a_					:		Any							=  padded_inputs.get('input_features'     )
             if isinstance(input_features[0]     ,							SCREAMING_SNAKE_CASE__     ):
                    a_					:		Union[str, Any]							=  [np.asarray(SCREAMING_SNAKE_CASE__     ,							dtype=np.floataa     ) for feature in input_features]

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

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

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

             return padded_inputs






import random
import unittest

import numpy as np
import torch

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionUpscalePipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
   import onnxruntime as ort
class  SCREAMING_SNAKE_CASE__  (				lowercase__						,    unittest.TestCase			):
      # TODO: is there an appropriate internal test set?
      snake_case__				:     Any        =						'''ssube/stable-diffusion-x4-upscaler-onnx'''
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     ,							SCREAMING_SNAKE_CASE__				:       int=0     )		->			Tuple:
             a_					:		Union[str, Any]							=  floats_tensor((1, 3, 1_2_8, 1_2_8)     ,							rng=random.Random(SCREAMING_SNAKE_CASE__     )     )
             a_					:		Dict							=  torch.manual_seed(SCREAMING_SNAKE_CASE__     )
             a_					:		Union[str, Any]							=  {
                 'prompt': 'A painting of a squirrel eating a burger',
                 'image': image,
                 'generator': generator,
                 'num_inference_steps': 3,
                 'guidance_scale': 7.5,
                 'output_type': 'numpy',
             }
             return inputs
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     )		->			Dict:
             a_					:		Dict							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,							provider='CPUExecutionProvider'     )
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

             a_					:		Optional[int]							=  self.get_dummy_inputs()
             a_					:		int							=  pipe(**SCREAMING_SNAKE_CASE__     ).images
             a_					:		Tuple							=  image[0, -3:, -3:, -1].flatten()

             # started as 128, should now be 512
             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		List[Any]							=  np.array(
                 [0.6974782, 0.68902093, 0.70135885, 0.7583618, 0.7804545, 0.7854912, 0.78667426, 0.78743863, 0.78070223]     )
             assert np.abs(image_slice - expected_slice     ).max() < 1E-1
      def        SCREAMING_SNAKE_CASE							(       self				:       Dict     )		->			Dict:
             a_					:		Tuple							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,							provider='CPUExecutionProvider'     )
             a_					:		int							=  PNDMScheduler.from_config(pipe.scheduler.config     ,							skip_prk_steps=SCREAMING_SNAKE_CASE__     )
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

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

             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		str							=  np.array(
                 [0.6898892, 0.59240556, 0.52499527, 0.58866215, 0.52258235, 0.52572715, 0.62414473, 0.6174387, 0.6214964]     )
             assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1E-1
      def        SCREAMING_SNAKE_CASE							(       self				:       Any     )		->			Dict:
             a_					:		Optional[Any]							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,							provider='CPUExecutionProvider'     )
             a_					:		Optional[Any]							=  DPMSolverMultistepScheduler.from_config(pipe.scheduler.config     )
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

             a_					:		Optional[int]							=  self.get_dummy_inputs()
             a_					:		Dict							=  pipe(**SCREAMING_SNAKE_CASE__     ).images
             a_					:		str							=  image[0, -3:, -3:, -1]

             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		Optional[Any]							=  np.array(
                 [0.7659278, 0.76437664, 0.75579107, 0.7691116, 0.77666986, 0.7727672, 0.7758664, 0.7812226, 0.76942515]     )

             assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1E-1
      def        SCREAMING_SNAKE_CASE							(       self				:       List[Any]     )		->			int:
             a_					:		List[str]							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,							provider='CPUExecutionProvider'     )
             a_					:		int							=  EulerDiscreteScheduler.from_config(pipe.scheduler.config     )
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

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

             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		int							=  np.array(
                 [0.6974782, 0.68902093, 0.70135885, 0.7583618, 0.7804545, 0.7854912, 0.78667426, 0.78743863, 0.78070223]     )
             assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1E-1

      def        SCREAMING_SNAKE_CASE							(       self				:       str     )		->			Optional[Any]:
             a_					:		Tuple							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,							provider='CPUExecutionProvider'     )
             a_					:		Optional[int]							=  EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config     )
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

             a_					:		Union[str, Any]							=  self.get_dummy_inputs()
             a_					:		List[str]							=  pipe(**SCREAMING_SNAKE_CASE__     ).images
             a_					:		List[str]							=  image[0, -3:, -3:, -1]

             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		Union[str, Any]							=  np.array(
                 [0.77424496, 0.773601, 0.7645288, 0.7769598, 0.7772739, 0.7738688, 0.78187233, 0.77879584, 0.767043]     )

             assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1E-1
@nightly
@require_onnxruntime
@require_torch_gpu
class  SCREAMING_SNAKE_CASE__  (				unittest.TestCase			):
      @property
      def        SCREAMING_SNAKE_CASE							(       self				:       Tuple     )		->			Union[str, Any]:
             return (
                 "CUDAExecutionProvider",
                 {
                     "gpu_mem_limit": "15000000000",  # 15GB
                     "arena_extend_strategy": "kSameAsRequested",
                 },
             )
      @property
      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[Any]     )		->			List[str]:
             a_					:		List[str]							=  ort.SessionOptions()
             a_					:		int							=  False
             return options
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     )		->			Tuple:
             a_					:		str							=  load_image(
                 'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
                 '/img2img/sketch-mountains-input.jpg'     )
             a_					:		int							=  init_image.resize((1_2_8, 1_2_8)     )
             # using the PNDM scheduler by default
             a_					:		Dict							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(
                 'ssube/stable-diffusion-x4-upscaler-onnx'     ,							provider=self.gpu_provider     ,							sess_options=self.gpu_options     ,							)
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

             a_					:		Tuple							=  'A fantasy landscape, trending on artstation'

             a_					:		str							=  torch.manual_seed(0     )
             a_					:		List[str]							=  pipe(
                 prompt=SCREAMING_SNAKE_CASE__     ,							image=SCREAMING_SNAKE_CASE__     ,							guidance_scale=7.5     ,							num_inference_steps=1_0     ,							generator=SCREAMING_SNAKE_CASE__     ,							output_type='np'     ,							)
             a_					:		Dict							=  output.images
             a_					:		Any							=  images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

             assert images.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		str							=  np.array([0.4883, 0.4947, 0.4980, 0.4975, 0.4982, 0.4980, 0.5000, 0.5006, 0.4972]     )
             # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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

      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[int]     )		->			List[str]:
             a_					:		Dict							=  load_image(
                 'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
                 '/img2img/sketch-mountains-input.jpg'     )
             a_					:		List[str]							=  init_image.resize((1_2_8, 1_2_8)     )
             a_					:		Dict							=  LMSDiscreteScheduler.from_pretrained(
                 'ssube/stable-diffusion-x4-upscaler-onnx'     ,							subfolder='scheduler'     )
             a_					:		Union[str, Any]							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(
                 'ssube/stable-diffusion-x4-upscaler-onnx'     ,							scheduler=SCREAMING_SNAKE_CASE__     ,							provider=self.gpu_provider     ,							sess_options=self.gpu_options     ,							)
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

             a_					:		Any							=  'A fantasy landscape, trending on artstation'

             a_					:		Tuple							=  torch.manual_seed(0     )
             a_					:		Optional[Any]							=  pipe(
                 prompt=SCREAMING_SNAKE_CASE__     ,							image=SCREAMING_SNAKE_CASE__     ,							guidance_scale=7.5     ,							num_inference_steps=2_0     ,							generator=SCREAMING_SNAKE_CASE__     ,							output_type='np'     ,							)
             a_					:		str							=  output.images
             a_					:		List[Any]							=  images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

             assert images.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		Tuple							=  np.array(
                 [0.50173753, 0.50223356, 0.502039, 0.50233036, 0.5023725, 0.5022601, 0.5018758, 0.50234085, 0.50241566]     )
             # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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


