

"""simple docstring"""





from copy import deepcopy
from typing import Optional, Union

import numpy as np

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding
from ...utils import TensorType, is_tf_available, is_torch_available


if is_torch_available():
   import torch

if is_tf_available():
   import tensorflow as tf




class 					_lowerCAmelCase  (	lowercase   ):







 """simple docstring"""







 __UpperCAmelCase	:					Union[str, Any]   		=	["image_processor"]
 __UpperCAmelCase	:					Optional[Any]   		=	"SamImageProcessor"


 def __init__(			self							:			Any,					UpperCAmelCase__							:			Dict	):
    super().__init__(UpperCAmelCase__	)
    __lowercase								=						self.image_processor
    __lowercase								=						-1_0
    __lowercase								=						self.image_processor.size["longest_edge"]


 def __call__(			self							:			Dict,					UpperCAmelCase__							:			int=None,					UpperCAmelCase__							:			List[str]=None,					UpperCAmelCase__							:			List[str]=None,					UpperCAmelCase__							:			Union[str, Any]=None,					UpperCAmelCase__							:			Optional[Union[str, TensorType]] = None,					**UpperCAmelCase__							:			Tuple,					):
    __lowercase								=						self.image_processor(
        UpperCAmelCase__,					return_tensors=UpperCAmelCase__,					**UpperCAmelCase__,					)

    # pop arguments that are not used in the foward but used nevertheless
    __lowercase								=						encoding_image_processor["original_sizes"]

    if hasattr(UpperCAmelCase__,					"numpy"	):  # Checks if Torch or TF tensor
       __lowercase								=						original_sizes.numpy()

    __lowercase		,__lowercase		,__lowercase								=						self._check_and_preprocess_points(
        input_points=UpperCAmelCase__,					input_labels=UpperCAmelCase__,					input_boxes=UpperCAmelCase__,					)

    __lowercase								=						self._normalize_and_convert(
        UpperCAmelCase__,					UpperCAmelCase__,					input_points=UpperCAmelCase__,					input_labels=UpperCAmelCase__,					input_boxes=UpperCAmelCase__,					return_tensors=UpperCAmelCase__,					)

    return encoding_image_processor


 def 					_lowercase (			self							:			Union[str, Any],					UpperCAmelCase__							:			Tuple,					UpperCAmelCase__							:			List[Any],					UpperCAmelCase__							:			List[Any]=None,					UpperCAmelCase__							:			Optional[int]=None,					UpperCAmelCase__							:			Optional[Any]=None,					UpperCAmelCase__							:			int="pt",					):
    if input_points is not None:
       if len(UpperCAmelCase__	) != len(UpperCAmelCase__	):
          __lowercase								=						[
              self._normalize_coordinates(self.target_size,					UpperCAmelCase__,					original_sizes[0]	) for point in input_points
          ]
       else:
          __lowercase								=						[
              self._normalize_coordinates(self.target_size,					UpperCAmelCase__,					UpperCAmelCase__	)
              for point, original_size in zip(UpperCAmelCase__,					UpperCAmelCase__	)
          ]
       # check that all arrays have the same shape
       if not all(point.shape == input_points[0].shape for point in input_points	):
          if input_labels is not None:
             __lowercase		,__lowercase								=						self._pad_points_and_labels(UpperCAmelCase__,					UpperCAmelCase__	)

       __lowercase								=						np.array(UpperCAmelCase__	)

    if input_labels is not None:
       __lowercase								=						np.array(UpperCAmelCase__	)

    if input_boxes is not None:
       if len(UpperCAmelCase__	) != len(UpperCAmelCase__	):
          __lowercase								=						[
              self._normalize_coordinates(self.target_size,					UpperCAmelCase__,					original_sizes[0],					is_bounding_box=UpperCAmelCase__	)
              for box in input_boxes
          ]
       else:
          __lowercase								=						[
              self._normalize_coordinates(self.target_size,					UpperCAmelCase__,					UpperCAmelCase__,					is_bounding_box=UpperCAmelCase__	)
              for box, original_size in zip(UpperCAmelCase__,					UpperCAmelCase__	)
          ]
       __lowercase								=						np.array(UpperCAmelCase__	)

    if input_boxes is not None:
       if return_tensors == "pt":
          __lowercase								=						torch.from_numpy(UpperCAmelCase__	)
          # boxes batch size of 1 by default
          __lowercase								=						input_boxes.unsqueeze(1	) if len(input_boxes.shape	) != 3 else input_boxes
       elif return_tensors == "tf":
          __lowercase								=						tf.convert_to_tensor(UpperCAmelCase__	)
          # boxes batch size of 1 by default
          __lowercase								=						tf.expand_dims(UpperCAmelCase__,					1	) if len(input_boxes.shape	) != 3 else input_boxes
       encoding_image_processor.update({"input_boxes": input_boxes}	)
    if input_points is not None:
       if return_tensors == "pt":
          __lowercase								=						torch.from_numpy(UpperCAmelCase__	)
          # point batch size of 1 by default
          __lowercase								=						input_points.unsqueeze(1	) if len(input_points.shape	) != 4 else input_points
       elif return_tensors == "tf":
          __lowercase								=						tf.convert_to_tensor(UpperCAmelCase__	)
          # point batch size of 1 by default
          __lowercase								=						tf.expand_dims(UpperCAmelCase__,					1	) if len(input_points.shape	) != 4 else input_points
       encoding_image_processor.update({"input_points": input_points}	)
    if input_labels is not None:
       if return_tensors == "pt":
          __lowercase								=						torch.from_numpy(UpperCAmelCase__	)
          # point batch size of 1 by default
          __lowercase								=						input_labels.unsqueeze(1	) if len(input_labels.shape	) != 3 else input_labels
       elif return_tensors == "tf":
          __lowercase								=						tf.convert_to_tensor(UpperCAmelCase__	)
          # point batch size of 1 by default
          __lowercase								=						tf.expand_dims(UpperCAmelCase__,					1	) if len(input_labels.shape	) != 3 else input_labels
       encoding_image_processor.update({"input_labels": input_labels}	)

    return encoding_image_processor


 def 					_lowercase (			self							:			Tuple,					UpperCAmelCase__							:			int,					UpperCAmelCase__							:			Optional[Any]	):
    __lowercase								=						max([point.shape[0] for point in input_points]	)
    __lowercase								=						[]
    for i, point in enumerate(UpperCAmelCase__	):
       if point.shape[0] != expected_nb_points:
          __lowercase								=						np.concatenate(
              [point, np.zeros((expected_nb_points - point.shape[0], 2)	) + self.point_pad_value],					axis=0	)
          __lowercase								=						np.append(input_labels[i],					[self.point_pad_value]	)
       processed_input_points.append(UpperCAmelCase__	)
    __lowercase								=						processed_input_points
    return input_points, input_labels


 def 					_lowercase (			self							:			List[Any],					UpperCAmelCase__							:			int,					UpperCAmelCase__							:			np.ndarray,					UpperCAmelCase__							:			Any,					UpperCAmelCase__							:			Tuple=False	):
    __lowercase		,__lowercase								=						original_size
    __lowercase		,__lowercase								=						self.image_processor._get_preprocess_shape(UpperCAmelCase__,					longest_edge=UpperCAmelCase__	)
    __lowercase								=						deepcopy(UpperCAmelCase__	).astype(UpperCAmelCase__	)

    if is_bounding_box:
       __lowercase								=						coords.reshape(-1,					2,					2	)

    __lowercase								=						coords[..., 0] * (new_w / old_w)
    __lowercase								=						coords[..., 1] * (new_h / old_h)

    if is_bounding_box:
       __lowercase								=						coords.reshape(-1,					4	)

    return coords


 def 					_lowercase (			self							:			Tuple,					UpperCAmelCase__							:			List[Any]=None,					UpperCAmelCase__							:			Optional[int]=None,					UpperCAmelCase__							:			List[Any]=None,					):
    if input_points is not None:
       if hasattr(UpperCAmelCase__,					"numpy"	):  # Checks for TF or Torch tensor
          __lowercase								=						input_points.numpy().tolist()

       if not isinstance(UpperCAmelCase__,					UpperCAmelCase__	) or not isinstance(input_points[0],					UpperCAmelCase__	):
          raise ValueError("Input points must be a list of list of floating points."	)
       __lowercase								=						[np.array(UpperCAmelCase__	) for input_point in input_points]
    else:
       __lowercase								=						None

    if input_labels is not None:
       if hasattr(UpperCAmelCase__,					"numpy"	):
          __lowercase								=						input_labels.numpy().tolist()

       if not isinstance(UpperCAmelCase__,					UpperCAmelCase__	) or not isinstance(input_labels[0],					UpperCAmelCase__	):
          raise ValueError("Input labels must be a list of list integers."	)
       __lowercase								=						[np.array(UpperCAmelCase__	) for label in input_labels]
    else:
       __lowercase								=						None

    if input_boxes is not None:
       if hasattr(UpperCAmelCase__,					"numpy"	):
          __lowercase								=						input_boxes.numpy().tolist()

       if (
           not isinstance(UpperCAmelCase__,					UpperCAmelCase__	)
           or not isinstance(input_boxes[0],					UpperCAmelCase__	)
           or not isinstance(input_boxes[0][0],					UpperCAmelCase__	)
       ):
          raise ValueError("Input boxes must be a list of list of list of floating points."	)
       __lowercase								=						[np.array(UpperCAmelCase__	).astype(np.floataa	) for box in input_boxes]
    else:
       __lowercase								=						None

    return input_points, input_labels, input_boxes


 @property
 def 					_lowercase (			self							:			Optional[Any]	):
    __lowercase								=						self.image_processor.model_input_names
    return list(dict.fromkeys(UpperCAmelCase__	)	)



 def 					_lowercase (			self							:			Optional[int],					*UpperCAmelCase__							:			List[str],					**UpperCAmelCase__							:			Union[str, Any]	):
    return self.image_processor.post_process_masks(*UpperCAmelCase__,					**UpperCAmelCase__	)







"""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 					_lowerCAmelCase  (	lowercase     ,unittest.TestCase   ):







 """simple docstring"""







 __UpperCAmelCase	:					str   		=	"ssube/stable-diffusion-x4-upscaler-onnx"


 def 					_lowercase (			self							:			Union[str, Any],					UpperCAmelCase__							:			List[str]=0	):
    __lowercase								=						floats_tensor((1, 3, 1_2_8, 1_2_8),					rng=random.Random(UpperCAmelCase__	)	)
    __lowercase								=						torch.manual_seed(UpperCAmelCase__	)
    __lowercase								=						{
        "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 					_lowercase (			self							:			Any	):
    __lowercase								=						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint,					provider="CPUExecutionProvider"	)
    pipe.set_progress_bar_config(disable=UpperCAmelCase__	)

    __lowercase								=						self.get_dummy_inputs()
    __lowercase								=						pipe(**UpperCAmelCase__	).images
    __lowercase								=						image[0, -3:, -3:, -1].flatten()

    # started as 128, should now be 512
    assert image.shape == (1, 5_1_2, 5_1_2, 3)
    __lowercase								=						np.array(
        [0.6_974_782, 0.68_902_093, 0.70_135_885, 0.7_583_618, 0.7_804_545, 0.7_854_912, 0.78_667_426, 0.78_743_863, 0.78_070_223]	)
    assert np.abs(image_slice - expected_slice	).max() < 1E-1


 def 					_lowercase (			self							:			Optional[Any]	):
    __lowercase								=						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint,					provider="CPUExecutionProvider"	)
    __lowercase								=						PNDMScheduler.from_config(pipe.scheduler.config,					skip_prk_steps=UpperCAmelCase__	)
    pipe.set_progress_bar_config(disable=UpperCAmelCase__	)

    __lowercase								=						self.get_dummy_inputs()
    __lowercase								=						pipe(**UpperCAmelCase__	).images
    __lowercase								=						image[0, -3:, -3:, -1]

    assert image.shape == (1, 5_1_2, 5_1_2, 3)
    __lowercase								=						np.array(
        [0.6_898_892, 0.59_240_556, 0.52_499_527, 0.58_866_215, 0.52_258_235, 0.52_572_715, 0.62_414_473, 0.6_174_387, 0.6_214_964]	)
    assert np.abs(image_slice.flatten() - expected_slice	).max() < 1E-1


 def 					_lowercase (			self							:			int	):
    __lowercase								=						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint,					provider="CPUExecutionProvider"	)
    __lowercase								=						DPMSolverMultistepScheduler.from_config(pipe.scheduler.config	)
    pipe.set_progress_bar_config(disable=UpperCAmelCase__	)

    __lowercase								=						self.get_dummy_inputs()
    __lowercase								=						pipe(**UpperCAmelCase__	).images
    __lowercase								=						image[0, -3:, -3:, -1]

    assert image.shape == (1, 5_1_2, 5_1_2, 3)
    __lowercase								=						np.array(
        [0.7_659_278, 0.76_437_664, 0.75_579_107, 0.7_691_116, 0.77_666_986, 0.7_727_672, 0.7_758_664, 0.7_812_226, 0.76_942_515]	)

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


 def 					_lowercase (			self							:			str	):
    __lowercase								=						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint,					provider="CPUExecutionProvider"	)
    __lowercase								=						EulerDiscreteScheduler.from_config(pipe.scheduler.config	)
    pipe.set_progress_bar_config(disable=UpperCAmelCase__	)

    __lowercase								=						self.get_dummy_inputs()
    __lowercase								=						pipe(**UpperCAmelCase__	).images
    __lowercase								=						image[0, -3:, -3:, -1]

    assert image.shape == (1, 5_1_2, 5_1_2, 3)
    __lowercase								=						np.array(
        [0.6_974_782, 0.68_902_093, 0.70_135_885, 0.7_583_618, 0.7_804_545, 0.7_854_912, 0.78_667_426, 0.78_743_863, 0.78_070_223]	)
    assert np.abs(image_slice.flatten() - expected_slice	).max() < 1E-1



 def 					_lowercase (			self							:			Any	):
    __lowercase								=						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint,					provider="CPUExecutionProvider"	)
    __lowercase								=						EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config	)
    pipe.set_progress_bar_config(disable=UpperCAmelCase__	)

    __lowercase								=						self.get_dummy_inputs()
    __lowercase								=						pipe(**UpperCAmelCase__	).images
    __lowercase								=						image[0, -3:, -3:, -1]

    assert image.shape == (1, 5_1_2, 5_1_2, 3)
    __lowercase								=						np.array(
        [0.77_424_496, 0.773_601, 0.7_645_288, 0.7_769_598, 0.7_772_739, 0.7_738_688, 0.78_187_233, 0.77_879_584, 0.767_043]	)

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







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







 """simple docstring"""







 @property
 def 					_lowercase (			self							:			Tuple	):
    return (
        "CUDAExecutionProvider",
        {
            "gpu_mem_limit": "15000000000",  # 15GB
            "arena_extend_strategy": "kSameAsRequested",
        },
    )


 @property
 def 					_lowercase (			self							:			Dict	):
    __lowercase								=						ort.SessionOptions()
    __lowercase								=						False
    return options


 def 					_lowercase (			self							:			Dict	):
    __lowercase								=						load_image(
        "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
        "/img2img/sketch-mountains-input.jpg"	)
    __lowercase								=						init_image.resize((1_2_8, 1_2_8)	)
    # using the PNDM scheduler by default
    __lowercase								=						OnnxStableDiffusionUpscalePipeline.from_pretrained(
        "ssube/stable-diffusion-x4-upscaler-onnx",					provider=self.gpu_provider,					sess_options=self.gpu_options,					)
    pipe.set_progress_bar_config(disable=UpperCAmelCase__	)

    __lowercase								=						"A fantasy landscape, trending on artstation"

    __lowercase								=						torch.manual_seed(0	)
    __lowercase								=						pipe(
        prompt=UpperCAmelCase__,					image=UpperCAmelCase__,					guidance_scale=7.5,					num_inference_steps=1_0,					generator=UpperCAmelCase__,					output_type="np",					)
    __lowercase								=						output.images
    __lowercase								=						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)
    __lowercase								=						np.array([0.4_883, 0.4_947, 0.4_980, 0.4_975, 0.4_982, 0.4_980, 0.5_000, 0.5_006, 0.4_972]	)
    # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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



 def 					_lowercase (			self							:			str	):
    __lowercase								=						load_image(
        "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
        "/img2img/sketch-mountains-input.jpg"	)
    __lowercase								=						init_image.resize((1_2_8, 1_2_8)	)
    __lowercase								=						LMSDiscreteScheduler.from_pretrained(
        "ssube/stable-diffusion-x4-upscaler-onnx",					subfolder="scheduler"	)
    __lowercase								=						OnnxStableDiffusionUpscalePipeline.from_pretrained(
        "ssube/stable-diffusion-x4-upscaler-onnx",					scheduler=UpperCAmelCase__,					provider=self.gpu_provider,					sess_options=self.gpu_options,					)
    pipe.set_progress_bar_config(disable=UpperCAmelCase__	)

    __lowercase								=						"A fantasy landscape, trending on artstation"

    __lowercase								=						torch.manual_seed(0	)
    __lowercase								=						pipe(
        prompt=UpperCAmelCase__,					image=UpperCAmelCase__,					guidance_scale=7.5,					num_inference_steps=2_0,					generator=UpperCAmelCase__,					output_type="np",					)
    __lowercase								=						output.images
    __lowercase								=						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)
    __lowercase								=						np.array(
        [0.50_173_753, 0.50_223_356, 0.502_039, 0.50_233_036, 0.5_023_725, 0.5_022_601, 0.5_018_758, 0.50_234_085, 0.50_241_566]	)
    # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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




