



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    SCREAMING_SNAKE_CASE__						(      UpperCAmelCase_   ,    unittest.TestCase				):
   __lowerCAmelCase						:     Union[str, Any]		=		'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'

   def 			SCREAMING_SNAKE_CASE  (     self   ,  _SCREAMING_SNAKE_CASE=0		)						->   Optional[int]:





    '''simple docstring'''




    UpperCAmelCase			:      str			=			floats_tensor((1, 3, 128, 128)   ,  rng=random.Random(_lowercase		)		)
    UpperCAmelCase			:      List[Any]			=			np.random.RandomState(_lowercase		)
    UpperCAmelCase			:      Optional[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 			SCREAMING_SNAKE_CASE  (     self		)						->   Optional[int]:





    '''simple docstring'''




    UpperCAmelCase			:      Union[str, Any]			=			OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint   ,  provider="""CPUExecutionProvider"""		)
    pipe.set_progress_bar_config(disable=_lowercase		)

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

    assert image.shape == (1, 128, 128, 3)
    UpperCAmelCase			:      Union[str, Any]			=			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 			SCREAMING_SNAKE_CASE  (     self		)						->   List[str]:





    '''simple docstring'''




    UpperCAmelCase			:      Dict			=			OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint   ,  provider="""CPUExecutionProvider"""		)
    UpperCAmelCase			:      Dict			=			PNDMScheduler.from_config(pipe.scheduler.config   ,  skip_prk_steps=_lowercase		)
    pipe.set_progress_bar_config(disable=_lowercase		)

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

    assert image.shape == (1, 128, 128, 3)
    UpperCAmelCase			:      int			=			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 			SCREAMING_SNAKE_CASE  (     self		)						->   Union[str, Any]:





    '''simple docstring'''




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

    # warmup pass to apply optimizations
    UpperCAmelCase			:      Union[str, Any]			=			pipe(**self.get_dummy_inputs()		)

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

    assert image.shape == (1, 128, 128, 3)
    UpperCAmelCase			:      List[str]			=			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 			SCREAMING_SNAKE_CASE  (     self		)						->   Union[str, Any]:





    '''simple docstring'''




    UpperCAmelCase			:      int			=			OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint   ,  provider="""CPUExecutionProvider"""		)
    UpperCAmelCase			:      str			=			EulerDiscreteScheduler.from_config(pipe.scheduler.config		)
    pipe.set_progress_bar_config(disable=_lowercase		)

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

    assert image.shape == (1, 128, 128, 3)
    UpperCAmelCase			:      List[str]			=			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 			SCREAMING_SNAKE_CASE  (     self		)						->   Any:





    '''simple docstring'''




    UpperCAmelCase			:      Union[str, Any]			=			OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint   ,  provider="""CPUExecutionProvider"""		)
    UpperCAmelCase			:      str			=			EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config		)
    pipe.set_progress_bar_config(disable=_lowercase		)

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

    assert image.shape == (1, 128, 128, 3)
    UpperCAmelCase			:      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 			SCREAMING_SNAKE_CASE  (     self		)						->   int:





    '''simple docstring'''




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

    UpperCAmelCase			:      Dict			=			self.get_dummy_inputs()
    UpperCAmelCase			:      Tuple			=			pipe(**_lowercase		).images
    UpperCAmelCase			:      Dict			=			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    SCREAMING_SNAKE_CASE__						(      unittest.TestCase				):

   @property
   def 			SCREAMING_SNAKE_CASE  (     self		)						->   int:





    '''simple docstring'''




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

   @property
   def 			SCREAMING_SNAKE_CASE  (     self		)						->   Dict:





    '''simple docstring'''




    UpperCAmelCase			:      Optional[Any]			=			ort.SessionOptions()
    UpperCAmelCase			:      Union[str, Any]			=			False
    return options

   def 			SCREAMING_SNAKE_CASE  (     self		)						->   Optional[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			:      List[str]			=			init_image.resize((768, 512)		)
    # using the PNDM scheduler by default
    UpperCAmelCase			:      Tuple			=			OnnxStableDiffusionImgaImgPipeline.from_pretrained(
        """CompVis/stable-diffusion-v1-4"""   ,  revision="""onnx"""   ,  safety_checker=_lowercase   ,  feature_extractor=_lowercase   ,  provider=self.gpu_provider   ,  sess_options=self.gpu_options   ,  )
    pipe.set_progress_bar_config(disable=_lowercase		)

    UpperCAmelCase			:      int			=			"""A fantasy landscape, trending on artstation"""

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

    assert images.shape == (1, 512, 768, 3)
    UpperCAmelCase			:      Union[str, Any]			=			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 			SCREAMING_SNAKE_CASE  (     self		)						->   Tuple:





    '''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			:      Optional[int]			=			init_image.resize((768, 512)		)
    UpperCAmelCase			:      Optional[int]			=			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=_lowercase   ,  safety_checker=_lowercase   ,  feature_extractor=_lowercase   ,  provider=self.gpu_provider   ,  sess_options=self.gpu_options   ,  )
    pipe.set_progress_bar_config(disable=_lowercase		)

    UpperCAmelCase			:      Optional[Any]			=			"""A fantasy landscape, trending on artstation"""

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

    assert images.shape == (1, 512, 768, 3)
    UpperCAmelCase			:      Optional[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"""
from typing import TYPE_CHECKING

from ...utils import _LazyModule


A:     Dict										=	{"tokenization_bertweet": ["BertweetTokenizer"]}


if TYPE_CHECKING:
      from .tokenization_bertweet import BertweetTokenizer

else:
      import sys

      A:     int										=	_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
