import random
import unittest

import numpy as np

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionImgaImgPipeline,
    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      ):
 __lowerCamelCase			:      List[Any]           =					"""hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"""
 def 		UpperCAmelCase__				(     self						,      snake_case__=0  )		->							int:
       '''simple docstring'''





       UpperCAmelCase			:       Union[str, Any]			     =floats_tensor((1, 3, 128, 128)						,      rng=random.Random(__UpperCamelCase  )  )
       UpperCAmelCase			:       Optional[int]			     =np.random.RandomState(__UpperCamelCase  )
       UpperCAmelCase			:       int			     ={
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''image''': image,
           '''generator''': generator,
           '''num_inference_steps''': 3,
           '''strength''': 0.75,
           '''guidance_scale''': 7.5,
           '''output_type''': '''numpy''',
       }
       return inputs
 def 		UpperCAmelCase__				(     self  )		->							Optional[int]:
       '''simple docstring'''





       UpperCAmelCase			:       Optional[int]			     =OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint						,      provider='''CPUExecutionProvider'''  )
       pipe.set_progress_bar_config(disable=__UpperCamelCase  )

       UpperCAmelCase			:       List[str]			     =self.get_dummy_inputs()
       UpperCAmelCase			:       int			     =pipe(**__UpperCamelCase  ).images
       UpperCAmelCase			:       Tuple			     =image[0, -3:, -3:, -1].flatten()

       assert image.shape == (1, 128, 128, 3)
       UpperCAmelCase			:       Optional[int]			     =np.array([0.6_9643, 0.5_8484, 0.5_0314, 0.5_8760, 0.5_5368, 0.5_9643, 0.5_1529, 0.4_1217, 0.4_9087]  )
       assert np.abs(image_slice - expected_slice  ).max() < 1e-1
 def 		UpperCAmelCase__				(     self  )		->							List[Any]:
       '''simple docstring'''





       UpperCAmelCase			:       Union[str, Any]			     =OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint						,      provider='''CPUExecutionProvider'''  )
       UpperCAmelCase			:       List[Any]			     =PNDMScheduler.from_config(pipe.scheduler.config						,      skip_prk_steps=__UpperCamelCase  )
       pipe.set_progress_bar_config(disable=__UpperCamelCase  )

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

       assert image.shape == (1, 128, 128, 3)
       UpperCAmelCase			:       Union[str, Any]			     =np.array([0.6_1737, 0.5_4642, 0.5_3183, 0.5_4465, 0.5_2742, 0.6_0525, 0.4_9969, 0.4_0655, 0.4_8154]  )

       assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-1
 def 		UpperCAmelCase__				(     self  )		->							List[Any]:
       '''simple docstring'''





       UpperCAmelCase			:       Optional[Any]			     =OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint						,      provider='''CPUExecutionProvider'''  )
       UpperCAmelCase			:       Dict			     =LMSDiscreteScheduler.from_config(pipe.scheduler.config  )
       pipe.set_progress_bar_config(disable=__UpperCamelCase  )

       # warmup pass to apply optimizations
       UpperCAmelCase			:       Dict			     =pipe(**self.get_dummy_inputs()  )

       UpperCAmelCase			:       int			     =self.get_dummy_inputs()
       UpperCAmelCase			:       List[str]			     =pipe(**__UpperCamelCase  ).images
       UpperCAmelCase			:       int			     =image[0, -3:, -3:, -1]

       assert image.shape == (1, 128, 128, 3)
       UpperCAmelCase			:       Union[str, Any]			     =np.array([0.5_2761, 0.5_9977, 0.4_9033, 0.4_9619, 0.5_4282, 0.5_0311, 0.4_7600, 0.4_0918, 0.4_5203]  )

       assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-1
 def 		UpperCAmelCase__				(     self  )		->							Tuple:
       '''simple docstring'''





       UpperCAmelCase			:       Optional[int]			     =OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint						,      provider='''CPUExecutionProvider'''  )
       UpperCAmelCase			:       List[str]			     =EulerDiscreteScheduler.from_config(pipe.scheduler.config  )
       pipe.set_progress_bar_config(disable=__UpperCamelCase  )

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

       assert image.shape == (1, 128, 128, 3)
       UpperCAmelCase			:       Dict			     =np.array([0.5_2911, 0.6_0004, 0.4_9229, 0.4_9805, 0.5_4502, 0.5_0680, 0.4_7777, 0.4_1028, 0.4_5304]  )

       assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-1
 def 		UpperCAmelCase__				(     self  )		->							Dict:
       '''simple docstring'''





       UpperCAmelCase			:       Tuple			     =OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint						,      provider='''CPUExecutionProvider'''  )
       UpperCAmelCase			:       str			     =EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config  )
       pipe.set_progress_bar_config(disable=__UpperCamelCase  )

       UpperCAmelCase			:       List[str]			     =self.get_dummy_inputs()
       UpperCAmelCase			:       List[str]			     =pipe(**__UpperCamelCase  ).images
       UpperCAmelCase			:       List[str]			     =image[0, -3:, -3:, -1]

       assert image.shape == (1, 128, 128, 3)
       UpperCAmelCase			:       Union[str, Any]			     =np.array([0.5_2911, 0.6_0004, 0.4_9229, 0.4_9805, 0.5_4502, 0.5_0680, 0.4_7777, 0.4_1028, 0.4_5304]  )

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






 def 		UpperCAmelCase__				(     self  )		->							Dict:
       '''simple docstring'''





       UpperCAmelCase			:       List[Any]			     =OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint						,      provider='''CPUExecutionProvider'''  )
       UpperCAmelCase			:       Optional[Any]			     =DPMSolverMultistepScheduler.from_config(pipe.scheduler.config  )
       pipe.set_progress_bar_config(disable=__UpperCamelCase  )

       UpperCAmelCase			:       Union[str, Any]			     =self.get_dummy_inputs()
       UpperCAmelCase			:       int			     =pipe(**__UpperCamelCase  ).images
       UpperCAmelCase			:       str			     =image[0, -3:, -3:, -1]

       assert image.shape == (1, 128, 128, 3)
       UpperCAmelCase			:       Any			     =np.array([0.6_5331, 0.5_8277, 0.4_8204, 0.5_6059, 0.5_3665, 0.5_6235, 0.5_0969, 0.4_0009, 0.4_6552]  )

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







@nightly
@require_onnxruntime
@require_torch_gpu
class 	__snake_case      (   unittest.TestCase      ):
 @property
 def 		UpperCAmelCase__				(     self  )		->							Dict:
       '''simple docstring'''





       return (
           "CUDAExecutionProvider",
           {
               "gpu_mem_limit": "15000000000",  # 15GB
               "arena_extend_strategy": "kSameAsRequested",
           },
       )
 @property
 def 		UpperCAmelCase__				(     self  )		->							int:
       '''simple docstring'''





       UpperCAmelCase			:       List[str]			     =ort.SessionOptions()
       UpperCAmelCase			:       Optional[Any]			     =False
       return options
 def 		UpperCAmelCase__				(     self  )		->							int:
       '''simple docstring'''





       UpperCAmelCase			:       Union[str, Any]			     =load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/img2img/sketch-mountains-input.jpg'''  )
       UpperCAmelCase			:       int			     =init_image.resize((768, 512)  )
       # using the PNDM scheduler by default
       UpperCAmelCase			:       Optional[int]			     =OnnxStableDiffusionImgaImgPipeline.from_pretrained(
           '''CompVis/stable-diffusion-v1-4'''						,      revision='''onnx'''						,      safety_checker=__UpperCamelCase						,      feature_extractor=__UpperCamelCase						,      provider=self.gpu_provider						,      sess_options=self.gpu_options						,      )
       pipe.set_progress_bar_config(disable=__UpperCamelCase  )

       UpperCAmelCase			:       str			     ='''A fantasy landscape, trending on artstation'''

       UpperCAmelCase			:       Optional[int]			     =np.random.RandomState(0  )
       UpperCAmelCase			:       Union[str, Any]			     =pipe(
           prompt=__UpperCamelCase						,      image=__UpperCamelCase						,      strength=0.75						,      guidance_scale=7.5						,      num_inference_steps=10						,      generator=__UpperCamelCase						,      output_type='''np'''						,      )
       UpperCAmelCase			:       int			     =output.images
       UpperCAmelCase			:       List[str]			     =images[0, 255:258, 383:386, -1]

       assert images.shape == (1, 512, 768, 3)
       UpperCAmelCase			:       str			     =np.array([0.4909, 0.5059, 0.5372, 0.4623, 0.4876, 0.5049, 0.4820, 0.4956, 0.5019]  )
       # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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






 def 		UpperCAmelCase__				(     self  )		->							int:
       '''simple docstring'''





       UpperCAmelCase			:       Optional[Any]			     =load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/img2img/sketch-mountains-input.jpg'''  )
       UpperCAmelCase			:       Optional[Any]			     =init_image.resize((768, 512)  )
       UpperCAmelCase			:       List[str]			     =LMSDiscreteScheduler.from_pretrained(
           '''runwayml/stable-diffusion-v1-5'''						,      subfolder='''scheduler'''						,      revision='''onnx'''  )
       UpperCAmelCase			:       Dict			     =OnnxStableDiffusionImgaImgPipeline.from_pretrained(
           '''runwayml/stable-diffusion-v1-5'''						,      revision='''onnx'''						,      scheduler=__UpperCamelCase						,      safety_checker=__UpperCamelCase						,      feature_extractor=__UpperCamelCase						,      provider=self.gpu_provider						,      sess_options=self.gpu_options						,      )
       pipe.set_progress_bar_config(disable=__UpperCamelCase  )

       UpperCAmelCase			:       List[str]			     ='''A fantasy landscape, trending on artstation'''

       UpperCAmelCase			:       List[Any]			     =np.random.RandomState(0  )
       UpperCAmelCase			:       Optional[int]			     =pipe(
           prompt=__UpperCamelCase						,      image=__UpperCamelCase						,      strength=0.75						,      guidance_scale=7.5						,      num_inference_steps=20						,      generator=__UpperCamelCase						,      output_type='''np'''						,      )
       UpperCAmelCase			:       Tuple			     =output.images
       UpperCAmelCase			:       str			     =images[0, 255:258, 383:386, -1]

       assert images.shape == (1, 512, 768, 3)
       UpperCAmelCase			:       List[Any]			     =np.array([0.8043, 0.926, 0.9581, 0.8119, 0.8954, 0.913, 0.7209, 0.7463, 0.7431]  )
       # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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







"""simple docstring"""
_snake_case     :		Optional[int]       						= [
    'DownloadConfig',
    'DownloadManager',
    'DownloadMode',
    'StreamingDownloadManager',
]

from .download_config import DownloadConfig
from .download_manager import DownloadManager, DownloadMode
from .streaming_download_manager import StreamingDownloadManager


