
"""simple docstring"""



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 		a_    (			_UpperCAmelCase  ,    unittest.TestCase						):
    a   :					Optional[Any]			     =					'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'







    def   _snake_case       (      self      :							Union[str, Any]  ,       __UpperCamelCase      :							Any=0  )		->Dict:






        '''simple docstring'''






        _UpperCAmelCase        = floats_tensor((1, 3, 1_28, 1_28)  ,       rng=random.Random(__UpperCamelCase  )  )
        _UpperCAmelCase        = np.random.RandomState(__UpperCamelCase  )
        _UpperCAmelCase        = {
            """prompt""": """A painting of a squirrel eating a burger""",
            """image""": image,
            """generator""": generator,
            """num_inference_steps""": 3,
            """strength""": 0.7_5,
            """guidance_scale""": 7.5,
            """output_type""": """numpy""",
        }
        return inputs







    def   _snake_case       (      self      :							Tuple  )		->Optional[int]:






        '''simple docstring'''






        _UpperCAmelCase        = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint  ,       provider="""CPUExecutionProvider"""  )
        pipe.set_progress_bar_config(disable=__UpperCamelCase  )

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

        assert image.shape == (1, 1_28, 1_28, 3)
        _UpperCAmelCase        = np.array([0.6_9_6_4_3, 0.5_8_4_8_4, 0.5_0_3_1_4, 0.5_8_7_6_0, 0.5_5_3_6_8, 0.5_9_6_4_3, 0.5_1_5_2_9, 0.4_1_2_1_7, 0.4_9_0_8_7]  )
        assert np.abs(image_slice - expected_slice  ).max() < 1e-1







    def   _snake_case       (      self      :							Tuple  )		->Optional[int]:






        '''simple docstring'''






        _UpperCAmelCase        = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint  ,       provider="""CPUExecutionProvider"""  )
        _UpperCAmelCase        = PNDMScheduler.from_config(pipe.scheduler.config  ,       skip_prk_steps=__UpperCamelCase  )
        pipe.set_progress_bar_config(disable=__UpperCamelCase  )

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

        assert image.shape == (1, 1_28, 1_28, 3)
        _UpperCAmelCase        = np.array([0.6_1_7_3_7, 0.5_4_6_4_2, 0.5_3_1_8_3, 0.5_4_4_6_5, 0.5_2_7_4_2, 0.6_0_5_2_5, 0.4_9_9_6_9, 0.4_0_6_5_5, 0.4_8_1_5_4]  )

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







    def   _snake_case       (      self      :							Dict  )		->int:






        '''simple docstring'''






        _UpperCAmelCase        = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint  ,       provider="""CPUExecutionProvider"""  )
        _UpperCAmelCase        = LMSDiscreteScheduler.from_config(pipe.scheduler.config  )
        pipe.set_progress_bar_config(disable=__UpperCamelCase  )

        # warmup pass to apply optimizations
        _UpperCAmelCase        = pipe(**self.get_dummy_inputs()  )

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

        assert image.shape == (1, 1_28, 1_28, 3)
        _UpperCAmelCase        = np.array([0.5_2_7_6_1, 0.5_9_9_7_7, 0.4_9_0_3_3, 0.4_9_6_1_9, 0.5_4_2_8_2, 0.5_0_3_1_1, 0.4_7_6_0_0, 0.4_0_9_1_8, 0.4_5_2_0_3]  )

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







    def   _snake_case       (      self      :							Dict  )		->Optional[int]:






        '''simple docstring'''






        _UpperCAmelCase        = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint  ,       provider="""CPUExecutionProvider"""  )
        _UpperCAmelCase        = EulerDiscreteScheduler.from_config(pipe.scheduler.config  )
        pipe.set_progress_bar_config(disable=__UpperCamelCase  )

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

        assert image.shape == (1, 1_28, 1_28, 3)
        _UpperCAmelCase        = np.array([0.5_2_9_1_1, 0.6_0_0_0_4, 0.4_9_2_2_9, 0.4_9_8_0_5, 0.5_4_5_0_2, 0.5_0_6_8_0, 0.4_7_7_7_7, 0.4_1_0_2_8, 0.4_5_3_0_4]  )

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







    def   _snake_case       (      self      :							str  )		->Tuple:






        '''simple docstring'''






        _UpperCAmelCase        = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint  ,       provider="""CPUExecutionProvider"""  )
        _UpperCAmelCase        = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config  )
        pipe.set_progress_bar_config(disable=__UpperCamelCase  )

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

        assert image.shape == (1, 1_28, 1_28, 3)
        _UpperCAmelCase        = np.array([0.5_2_9_1_1, 0.6_0_0_0_4, 0.4_9_2_2_9, 0.4_9_8_0_5, 0.5_4_5_0_2, 0.5_0_6_8_0, 0.4_7_7_7_7, 0.4_1_0_2_8, 0.4_5_3_0_4]  )

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







    def   _snake_case       (      self      :							Any  )		->int:






        '''simple docstring'''






        _UpperCAmelCase        = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint  ,       provider="""CPUExecutionProvider"""  )
        _UpperCAmelCase        = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config  )
        pipe.set_progress_bar_config(disable=__UpperCamelCase  )

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

        assert image.shape == (1, 1_28, 1_28, 3)
        _UpperCAmelCase        = np.array([0.6_5_3_3_1, 0.5_8_2_7_7, 0.4_8_2_0_4, 0.5_6_0_5_9, 0.5_3_6_6_5, 0.5_6_2_3_5, 0.5_0_9_6_9, 0.4_0_0_0_9, 0.4_6_5_5_2]  )

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







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







    @property
    def   _snake_case       (      self      :							Any  )		->Optional[int]:






        '''simple docstring'''






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







    @property
    def   _snake_case       (      self      :							Dict  )		->Optional[Any]:






        '''simple docstring'''






        _UpperCAmelCase        = ort.SessionOptions()
        _UpperCAmelCase        = False
        return options







    def   _snake_case       (      self      :							Tuple  )		->Any:






        '''simple docstring'''






        _UpperCAmelCase        = load_image(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
            """/img2img/sketch-mountains-input.jpg"""  )
        _UpperCAmelCase        = init_image.resize((7_68, 5_12)  )
        # using the PNDM scheduler by default
        _UpperCAmelCase        = 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        = """A fantasy landscape, trending on artstation"""

        _UpperCAmelCase        = np.random.RandomState(0  )
        _UpperCAmelCase        = pipe(
            prompt=__UpperCamelCase  ,       image=__UpperCamelCase  ,       strength=0.7_5  ,       guidance_scale=7.5  ,       num_inference_steps=10  ,       generator=__UpperCamelCase  ,       output_type="""np"""  ,       )
        _UpperCAmelCase        = output.images
        _UpperCAmelCase        = images[0, 2_55:2_58, 3_83:3_86, -1]

        assert images.shape == (1, 5_12, 7_68, 3)
        _UpperCAmelCase        = np.array([0.4_9_0_9, 0.5_0_5_9, 0.5_3_7_2, 0.4_6_2_3, 0.4_8_7_6, 0.5_0_4_9, 0.4_8_2_0, 0.4_9_5_6, 0.5_0_1_9]  )
        # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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







    def   _snake_case       (      self      :							Optional[Any]  )		->Optional[int]:






        '''simple docstring'''






        _UpperCAmelCase        = load_image(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
            """/img2img/sketch-mountains-input.jpg"""  )
        _UpperCAmelCase        = init_image.resize((7_68, 5_12)  )
        _UpperCAmelCase        = LMSDiscreteScheduler.from_pretrained(
            """runwayml/stable-diffusion-v1-5"""  ,       subfolder="""scheduler"""  ,       revision="""onnx"""  )
        _UpperCAmelCase        = 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        = """A fantasy landscape, trending on artstation"""

        _UpperCAmelCase        = np.random.RandomState(0  )
        _UpperCAmelCase        = pipe(
            prompt=__UpperCamelCase  ,       image=__UpperCamelCase  ,       strength=0.7_5  ,       guidance_scale=7.5  ,       num_inference_steps=20  ,       generator=__UpperCamelCase  ,       output_type="""np"""  ,       )
        _UpperCAmelCase        = output.images
        _UpperCAmelCase        = images[0, 2_55:2_58, 3_83:3_86, -1]

        assert images.shape == (1, 5_12, 7_68, 3)
        _UpperCAmelCase        = np.array([0.8_0_4_3, 0.9_2_6, 0.9_5_8_1, 0.8_1_1_9, 0.8_9_5_4, 0.9_1_3, 0.7_2_0_9, 0.7_4_6_3, 0.7_4_3_1]  )
        # 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"""



import json
import os
import tempfile

import transformers

import datasets
from utils import generate_example_dataset, get_duration


a       :						int       		=		5_0_0_0_0_0

a		,			a       :						Union[str, Any]       		=		os.path.split(__file__)
a       :						Dict       		=		os.path.join(RESULTS_BASEPATH, '''results''', RESULTS_FILENAME.replace('''.py''', '''.json'''))


@get_duration
def 			_UpperCamelCase						(    _A       ,   **_A   )     ->	Any:





    """simple docstring"""

    _UpperCAmelCase        = dataset.map(**_A   )


@get_duration
def 			_UpperCamelCase						(    _A       ,   **_A   )     ->	Dict:





    """simple docstring"""

    _UpperCAmelCase        = dataset.filter(**_A   )


def 			_UpperCamelCase						(    )     ->	Tuple:





    """simple docstring"""

    _UpperCAmelCase        = {"""num examples""": SPEED_TEST_N_EXAMPLES}
    with tempfile.TemporaryDirectory() as tmp_dir:
        _UpperCAmelCase        = datasets.Features({"""text""": datasets.Value("""string"""   ), """numbers""": datasets.Value("""float32"""   )}   )
        _UpperCAmelCase        = generate_example_dataset(
            os.path.join(_A       ,   """dataset.arrow"""   )       ,   _A       ,   num_examples=_A   )

        _UpperCAmelCase        = transformers.AutoTokenizer.from_pretrained("""bert-base-cased"""       ,   use_fast=_A   )

        def tokenize(_A   ):
            return tokenizer(examples["""text"""]   )

        _UpperCAmelCase        = map(_A   )

        _UpperCAmelCase        = map(_A       ,   batched=_A   )

        _UpperCAmelCase        = map(_A       ,   function=lambda _A   : None       ,   batched=_A   )

        with dataset.formatted_as(type="""numpy"""   ):
            _UpperCAmelCase        = map(_A       ,   function=lambda _A   : None       ,   batched=_A   )

        with dataset.formatted_as(type="""pandas"""   ):
            _UpperCAmelCase        = map(_A       ,   function=lambda _A   : None       ,   batched=_A   )

        with dataset.formatted_as(type="""torch"""       ,   columns="""numbers"""   ):
            _UpperCAmelCase        = map(_A       ,   function=lambda _A   : None       ,   batched=_A   )

        with dataset.formatted_as(type="""tensorflow"""       ,   columns="""numbers"""   ):
            _UpperCAmelCase        = map(_A       ,   function=lambda _A   : None       ,   batched=_A   )

        _UpperCAmelCase        = map(_A       ,   function=_A       ,   batched=_A   )

        _UpperCAmelCase        = filter(_A   )

        # Activate later when tokenizer support batched inputs
        # with dataset.formatted_as(type='numpy'):
        #     times[func.__name__ + " fast-tokenizer batched numpy"] = func(dataset, function=tokenize, batched=True)

    with open(_A       ,   """wb"""   ) as f:
        f.write(json.dumps(_A   ).encode("""utf-8"""   )   )


if __name__ == "__main__":  # useful to run the profiler
 benchmark_map_filter()