



"""simple docstring"""






from typing import Dict, List, Optional, Union

import numpy as np

from .feature_extraction_utils import BatchFeature, FeatureExtractionMixin
from .utils import PaddingStrategy, TensorType, is_tf_tensor, is_torch_tensor, logging, to_numpy


lowerCAmelCase__      =    logging.get_logger(__name__)






class 	__snake_case	(     _lowercase):




  def __init__(			self					:						Union[str, Any]						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						float						, **__lowerCAmelCase					:						List[str]		):






         """simple docstring"""






         _lowerCamelCase						:    List[Any]         =			feature_size
         _lowerCamelCase						:    str         =			sampling_rate
         _lowerCamelCase						:    Dict         =			padding_value

         _lowerCamelCase						:    int         =			kwargs.pop('''padding_side'''						, '''right'''		)
         _lowerCamelCase						:    Tuple         =			kwargs.pop('''return_attention_mask'''						, __lowerCAmelCase		)

         super().__init__(**__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]						, __lowerCAmelCase					:						Union[
          BatchFeature,
          List[BatchFeature],
          Dict[str, BatchFeature],
          Dict[str, List[BatchFeature]],
          List[Dict[str, BatchFeature]],
      ]						, __lowerCAmelCase					:						Union[bool, str, PaddingStrategy] = True						, __lowerCAmelCase					:						Optional[int] = None						, __lowerCAmelCase					:						bool = False						, __lowerCAmelCase					:						Optional[int] = None						, __lowerCAmelCase					:						Optional[bool] = None						, __lowerCAmelCase					:						Optional[Union[str, TensorType]] = None						, ):






         """simple docstring"""






         if isinstance(__lowerCAmelCase						, (list, tuple)		) and isinstance(processed_features[0]						, (dict, BatchFeature)		):
                _lowerCamelCase						:    Dict         =			{
                    key: [example[key] for example in processed_features] for key in processed_features[0].keys()
                }

         # The model's main input name, usually `input_values`, has be passed for padding
         if self.model_input_names[0] not in processed_features:
                raise ValueError(
                    '''You should supply an instance of `transformers.BatchFeature` or list of `transformers.BatchFeature`'''
                    f''' to this method that includes {self.model_input_names[0]}, but you provided'''
                    f''' {list(processed_features.keys()		)}'''		)

         _lowerCamelCase						:    Optional[int]         =			processed_features[self.model_input_names[0]]
         _lowerCamelCase						:    str         =			(
             return_attention_mask if return_attention_mask is not None else self.return_attention_mask
         )

         if len(__lowerCAmelCase		) == 0:
                if return_attention_mask:
                       _lowerCamelCase						:    str         =			[]
                return processed_features

         # If we have PyTorch/TF tensors or lists as inputs, we cast them as Numpy arrays
         # and rebuild them afterwards if no return_tensors is specified
         # Note that we lose the specific device the tensor may be on for PyTorch

         _lowerCamelCase						:    Tuple         =			required_input[0]
         if isinstance(__lowerCAmelCase						, (list, tuple)		):
                # first_element might be an empty list/tuple in some edge cases so we grab the first non empty element.
                _lowerCamelCase						:    str         =			0
                while len(required_input[index]		) == 0:
                       index += 1
                if index < len(__lowerCAmelCase		):
                       _lowerCamelCase						:    Union[str, Any]         =			required_input[index][0]

         if return_tensors is None:
                if is_tf_tensor(__lowerCAmelCase		):
                       _lowerCamelCase						:    List[str]         =			'''tf'''
                elif is_torch_tensor(__lowerCAmelCase		):
                       _lowerCamelCase						:    Union[str, Any]         =			'''pt'''
                elif isinstance(__lowerCAmelCase						, (int, float, list, tuple, np.ndarray)		):
                       _lowerCamelCase						:    str         =			'''np'''
                else:
                       raise ValueError(
                           f'''type of {first_element} unknown: {type(__lowerCAmelCase		)}. '''
                           '''Should be one of a python, numpy, pytorch or tensorflow object.'''		)

         for key, value in processed_features.items():
                if isinstance(value[0]						, (int, float)		):
                       _lowerCamelCase						:    List[Any]         =			to_numpy(__lowerCAmelCase		)
                else:
                       _lowerCamelCase						:    int         =			[to_numpy(__lowerCAmelCase		) for v in value]

        # Convert padding_strategy in PaddingStrategy
         _lowerCamelCase						:    Tuple         =			self._get_padding_strategies(padding=__lowerCAmelCase						, max_length=__lowerCAmelCase		)

         _lowerCamelCase						:    List[Any]         =			processed_features[self.model_input_names[0]]

         _lowerCamelCase						:    List[Any]         =			len(__lowerCAmelCase		)
         if not all(len(__lowerCAmelCase		) == batch_size for v in processed_features.values()		):
                raise ValueError('''Some items in the output dictionary have a different batch size than others.'''		)

         _lowerCamelCase						:    List[Any]         =			[]
         for i in range(__lowerCAmelCase		):
                _lowerCamelCase						:    Optional[Any]         =			{k: v[i] for k, v in processed_features.items()}
                # truncation
                _lowerCamelCase						:    List[Any]         =			self._truncate(
                    __lowerCAmelCase						, max_length=__lowerCAmelCase						, pad_to_multiple_of=__lowerCAmelCase						, truncation=__lowerCAmelCase						, )
                truncated_inputs.append(__lowerCAmelCase		)

         if padding_strategy == PaddingStrategy.LONGEST:
                # make sure that `max_length` cannot be longer than the longest truncated length
                _lowerCamelCase						:    Dict         =			max(len(input_slice[self.model_input_names[0]]		) for input_slice in truncated_inputs		)
                _lowerCamelCase						:    Optional[Any]         =			PaddingStrategy.MAX_LENGTH

         _lowerCamelCase						:    Optional[Any]         =			{}
         for i in range(__lowerCAmelCase		):
                # padding
                _lowerCamelCase						:    Tuple         =			self._pad(
                    truncated_inputs[i]						, max_length=__lowerCAmelCase						, padding_strategy=__lowerCAmelCase						, pad_to_multiple_of=__lowerCAmelCase						, return_attention_mask=__lowerCAmelCase						, )

                for key, value in outputs.items():
                       if key not in batch_outputs:
                              _lowerCamelCase						:    str         =			[]
                       if value.dtype is np.dtype(np.floataa		):
                              _lowerCamelCase						:    Union[str, Any]         =			value.astype(np.floataa		)
                       batch_outputs[key].append(__lowerCAmelCase		)

         return BatchFeature(__lowerCAmelCase						, tensor_type=__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any						, __lowerCAmelCase					:						Union[Dict[str, np.ndarray], BatchFeature]						, __lowerCAmelCase					:						Optional[int] = None						, __lowerCAmelCase					:						PaddingStrategy = PaddingStrategy.DO_NOT_PAD						, __lowerCAmelCase					:						Optional[int] = None						, __lowerCAmelCase					:						Optional[bool] = None						, ):






         """simple docstring"""






         _lowerCamelCase						:    Union[str, Any]         =			processed_features[self.model_input_names[0]]

         if padding_strategy == PaddingStrategy.LONGEST:
                _lowerCamelCase						:    Optional[int]         =			len(__lowerCAmelCase		)

         if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
                _lowerCamelCase						:    List[Any]         =			((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of

         _lowerCamelCase						:    Any         =			padding_strategy != PaddingStrategy.DO_NOT_PAD and len(__lowerCAmelCase		) < max_length

         if return_attention_mask and "attention_mask" not in processed_features:
                _lowerCamelCase						:    Any         =			np.ones(len(__lowerCAmelCase		)						, dtype=np.intaa		)

         if needs_to_be_padded:
                _lowerCamelCase						:    int         =			max_length - len(__lowerCAmelCase		)
                if self.padding_side == "right":
                       if return_attention_mask:
                              _lowerCamelCase						:    List[Any]         =			np.pad(
                                  processed_features['''attention_mask''']						, (0, difference)		)
                       _lowerCamelCase						:    str         =			((0, difference), (0, 0)) if self.feature_size > 1 else (0, difference)
                       _lowerCamelCase						:    int         =			np.pad(
                           __lowerCAmelCase						, __lowerCAmelCase						, '''constant'''						, constant_values=self.padding_value		)
                elif self.padding_side == "left":
                       if return_attention_mask:
                              _lowerCamelCase						:    Tuple         =			np.pad(
                                  processed_features['''attention_mask''']						, (difference, 0)		)
                       _lowerCamelCase						:    Dict         =			((difference, 0), (0, 0)) if self.feature_size > 1 else (difference, 0)
                       _lowerCamelCase						:    Tuple         =			np.pad(
                           __lowerCAmelCase						, __lowerCAmelCase						, '''constant'''						, constant_values=self.padding_value		)
                else:
                       raise ValueError('''Invalid padding strategy:''' + str(self.padding_side		)		)

         return processed_features




  def 						SCREAMING_SNAKE_CASE    (			self					:						int						, __lowerCAmelCase					:						Union[Dict[str, np.ndarray], BatchFeature]						, __lowerCAmelCase					:						Optional[int] = None						, __lowerCAmelCase					:						Optional[int] = None						, __lowerCAmelCase					:						Optional[bool] = None						, ):






         """simple docstring"""






         if not truncation:
                return processed_features
         elif truncation and max_length is None:
                raise ValueError('''When setting ``truncation=True``, make sure that ``max_length`` is defined.'''		)

         _lowerCamelCase						:    Tuple         =			processed_features[self.model_input_names[0]]

         # find `max_length` that fits `pad_to_multiple_of`
         if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
                _lowerCamelCase						:    str         =			((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of

         _lowerCamelCase						:    List[Any]         =			len(__lowerCAmelCase		) > max_length

         if needs_to_be_truncated:
                _lowerCamelCase						:    Optional[Any]         =			processed_features[self.model_input_names[0]][:max_length]
                if "attention_mask" in processed_features:
                       _lowerCamelCase						:    List[str]         =			processed_features['''attention_mask'''][:max_length]

         return processed_features




  def 						SCREAMING_SNAKE_CASE    (			self					:						int						, __lowerCAmelCase					:						Dict=False						, __lowerCAmelCase					:						Dict=None		):






         """simple docstring"""






         if padding is not False:
                if padding is True:
                       _lowerCamelCase						:    List[Any]         =			PaddingStrategy.LONGEST  # Default to pad to the longest sequence in the batch
                elif not isinstance(__lowerCAmelCase						, __lowerCAmelCase		):
                       _lowerCamelCase						:    int         =			PaddingStrategy(__lowerCAmelCase		)
                elif isinstance(__lowerCAmelCase						, __lowerCAmelCase		):
                       _lowerCamelCase						:    Optional[Any]         =			padding
         else:
                _lowerCamelCase						:    List[Any]         =			PaddingStrategy.DO_NOT_PAD

         # Set max length if needed
         if max_length is None:
                if padding_strategy == PaddingStrategy.MAX_LENGTH:
                       raise ValueError(
                           f'''When setting ``padding={PaddingStrategy.MAX_LENGTH}``, make sure that max_length is defined'''		)

        # Test if we have a padding value
         if padding_strategy != PaddingStrategy.DO_NOT_PAD and (self.padding_value is None):
                raise ValueError(
                    '''Asking to pad but the feature_extractor does not have a padding value. Please select a value to use'''
                    ''' as `padding_value`. For example: `feature_extractor.padding_value = 0.0`.'''		)

         return padding_strategy







"""simple docstring"""






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 	__snake_case	(     _lowercase  ,					unittest.TestCase):
  # TODO: is there an appropriate internal test set?
  snake_case__    :       List[str]        =						"ssube/stable-diffusion-x4-upscaler-onnx"




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]						, __lowerCAmelCase					:						int=0		):






         """simple docstring"""






         _lowerCamelCase						:    Tuple         =			floats_tensor((1, 3, 1_2_8, 1_2_8)						, rng=random.Random(__lowerCAmelCase		)		)
         _lowerCamelCase						:    Union[str, Any]         =			torch.manual_seed(__lowerCAmelCase		)
         _lowerCamelCase						:    Tuple         =			{
             '''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					:						Optional[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    Union[str, Any]         =			OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						, provider='''CPUExecutionProvider'''		)
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

         _lowerCamelCase						:    Union[str, Any]         =			self.get_dummy_inputs()
         _lowerCamelCase						:    Any         =			pipe(**__lowerCAmelCase		).images
         _lowerCamelCase						:    Dict         =			image[0, -3:, -3:, -1].flatten()

         # started as 128, should now be 512
         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         _lowerCamelCase						:    str         =			np.array(
             [0.6_97_47_82, 0.68_90_20_93, 0.70_13_58_85, 0.7_58_36_18, 0.7_80_45_45, 0.7_85_49_12, 0.78_66_74_26, 0.78_74_38_63, 0.78_07_02_23]		)
         assert np.abs(image_slice - expected_slice		).max() < 1E-1




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						, provider='''CPUExecutionProvider'''		)
         _lowerCamelCase						:    List[str]         =			PNDMScheduler.from_config(pipe.scheduler.config						, skip_prk_steps=__lowerCAmelCase		)
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

         _lowerCamelCase						:    int         =			self.get_dummy_inputs()
         _lowerCamelCase						:    Optional[Any]         =			pipe(**__lowerCAmelCase		).images
         _lowerCamelCase						:    str         =			image[0, -3:, -3:, -1]

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         _lowerCamelCase						:    Optional[int]         =			np.array(
             [0.6_89_88_92, 0.59_24_05_56, 0.52_49_95_27, 0.58_86_62_15, 0.52_25_82_35, 0.52_57_27_15, 0.62_41_44_73, 0.6_17_43_87, 0.6_21_49_64]		)
         assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-1




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						, provider='''CPUExecutionProvider'''		)
         _lowerCamelCase						:    Optional[int]         =			DPMSolverMultistepScheduler.from_config(pipe.scheduler.config		)
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

         _lowerCamelCase						:    Tuple         =			self.get_dummy_inputs()
         _lowerCamelCase						:    str         =			pipe(**__lowerCAmelCase		).images
         _lowerCamelCase						:    Tuple         =			image[0, -3:, -3:, -1]

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         _lowerCamelCase						:    str         =			np.array(
             [0.7_65_92_78, 0.76_43_76_64, 0.75_57_91_07, 0.7_69_11_16, 0.77_66_69_86, 0.7_72_76_72, 0.7_75_86_64, 0.7_81_22_26, 0.76_94_25_15]		)

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




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    Any         =			OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						, provider='''CPUExecutionProvider'''		)
         _lowerCamelCase						:    Union[str, Any]         =			EulerDiscreteScheduler.from_config(pipe.scheduler.config		)
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

         _lowerCamelCase						:    Optional[Any]         =			self.get_dummy_inputs()
         _lowerCamelCase						:    Tuple         =			pipe(**__lowerCAmelCase		).images
         _lowerCamelCase						:    Union[str, Any]         =			image[0, -3:, -3:, -1]

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         _lowerCamelCase						:    Union[str, Any]         =			np.array(
             [0.6_97_47_82, 0.68_90_20_93, 0.70_13_58_85, 0.7_58_36_18, 0.7_80_45_45, 0.7_85_49_12, 0.78_66_74_26, 0.78_74_38_63, 0.78_07_02_23]		)
         assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-1




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    Dict         =			OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						, provider='''CPUExecutionProvider'''		)
         _lowerCamelCase						:    int         =			EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config		)
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

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

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         _lowerCamelCase						:    Optional[int]         =			np.array(
             [0.77_42_44_96, 0.77_36_01, 0.7_64_52_88, 0.7_76_95_98, 0.7_77_27_39, 0.7_73_86_88, 0.78_18_72_33, 0.77_87_95_84, 0.76_70_43]		)

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






@nightly
@require_onnxruntime
@require_torch_gpu
class 	__snake_case	(     unittest.TestCase):




  @property
  def 						SCREAMING_SNAKE_CASE    (			self					:						str		):






         """simple docstring"""






         return (
             "CUDAExecutionProvider",
             {
                 "gpu_mem_limit": "15000000000",  # 15GB
                 "arena_extend_strategy": "kSameAsRequested",
             },
         )




  @property
  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			ort.SessionOptions()
         _lowerCamelCase						:    List[str]         =			False
         return options




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/img2img/sketch-mountains-input.jpg'''		)
         _lowerCamelCase						:    Any         =			init_image.resize((1_2_8, 1_2_8)		)
         # using the PNDM scheduler by default
         _lowerCamelCase						:    List[str]         =			OnnxStableDiffusionUpscalePipeline.from_pretrained(
             '''ssube/stable-diffusion-x4-upscaler-onnx'''						, provider=self.gpu_provider						, sess_options=self.gpu_options						, )
         pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

         _lowerCamelCase						:    int         =			'''A fantasy landscape, trending on artstation'''

         _lowerCamelCase						:    List[Any]         =			torch.manual_seed(0		)
         _lowerCamelCase						:    List[str]         =			pipe(
             prompt=__lowerCAmelCase						, image=__lowerCAmelCase						, guidance_scale=7.5						, num_inference_steps=1_0						, generator=__lowerCAmelCase						, output_type='''np'''						, )
         _lowerCamelCase						:    List[Any]         =			output.images
         _lowerCamelCase						:    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)
         _lowerCamelCase						:    str         =			np.array([0.48_83, 0.49_47, 0.49_80, 0.49_75, 0.49_82, 0.49_80, 0.50_00, 0.50_06, 0.49_72]		)
         # 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					:						str		):






         """simple docstring"""






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

         _lowerCamelCase						:    Optional[int]         =			'''A fantasy landscape, trending on artstation'''

         _lowerCamelCase						:    int         =			torch.manual_seed(0		)
         _lowerCamelCase						:    List[str]         =			pipe(
             prompt=__lowerCAmelCase						, image=__lowerCAmelCase						, guidance_scale=7.5						, num_inference_steps=2_0						, generator=__lowerCAmelCase						, output_type='''np'''						, )
         _lowerCamelCase						:    Union[str, Any]         =			output.images
         _lowerCamelCase						:    Optional[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)
         _lowerCamelCase						:    str         =			np.array(
             [0.50_17_37_53, 0.50_22_33_56, 0.50_20_39, 0.50_23_30_36, 0.5_02_37_25, 0.5_02_26_01, 0.5_01_87_58, 0.50_23_40_85, 0.50_24_15_66]		)
         # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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


