




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			(			__SCREAMING_SNAKE_CASE ,  unittest.TestCase		):
   UpperCAmelCase__  :    Any          =  '''hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'''







   def lowerCamelCase						(      self		:						List[Any]					,						_snake_case		:						Optional[int]=0):

        """simple docstring"""


        UpperCAmelCase_          =							floats_tensor((1, 3, 128, 128)					,						rng=random.Random(_SCREAMING_SNAKE_CASE))
        UpperCAmelCase_          =							np.random.RandomState(_SCREAMING_SNAKE_CASE)
        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 lowerCamelCase						(      self		:						Dict):

        """simple docstring"""


        UpperCAmelCase_          =							OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint					,						provider='''CPUExecutionProvider''')
        pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)

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

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

        """simple docstring"""


        UpperCAmelCase_          =							OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint					,						provider='''CPUExecutionProvider''')
        UpperCAmelCase_          =							PNDMScheduler.from_config(pipe.scheduler.config					,						skip_prk_steps=_SCREAMING_SNAKE_CASE)
        pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)

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

        assert image.shape == (1, 128, 128, 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 lowerCamelCase						(      self		:						Any):

        """simple docstring"""


        UpperCAmelCase_          =							OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint					,						provider='''CPUExecutionProvider''')
        UpperCAmelCase_          =							LMSDiscreteScheduler.from_config(pipe.scheduler.config)
        pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)

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

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

        assert image.shape == (1, 128, 128, 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 lowerCamelCase						(      self		:						Dict):

        """simple docstring"""


        UpperCAmelCase_          =							OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint					,						provider='''CPUExecutionProvider''')
        UpperCAmelCase_          =							EulerDiscreteScheduler.from_config(pipe.scheduler.config)
        pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)

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

        assert image.shape == (1, 128, 128, 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 lowerCamelCase						(      self		:						str):

        """simple docstring"""


        UpperCAmelCase_          =							OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint					,						provider='''CPUExecutionProvider''')
        UpperCAmelCase_          =							EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
        pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)

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

        assert image.shape == (1, 128, 128, 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 lowerCamelCase						(      self		:						Optional[Any]):

        """simple docstring"""


        UpperCAmelCase_          =							OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint					,						provider='''CPUExecutionProvider''')
        UpperCAmelCase_          =							DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
        pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)

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

        assert image.shape == (1, 128, 128, 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 __snake_case			(			unittest.TestCase		):







   @property
   def lowerCamelCase						(      self		:						Tuple):

        """simple docstring"""


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







   @property
   def lowerCamelCase						(      self		:						Union[str, Any]):

        """simple docstring"""


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







   def lowerCamelCase						(      self		:						List[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((768, 512))
        # using the PNDM scheduler by default
        UpperCAmelCase_          =							OnnxStableDiffusionImgaImgPipeline.from_pretrained(
            '''CompVis/stable-diffusion-v1-4'''					,						revision='''onnx'''					,						safety_checker=_SCREAMING_SNAKE_CASE					,						feature_extractor=_SCREAMING_SNAKE_CASE					,						provider=self.gpu_provider					,						sess_options=self.gpu_options					,						)
        pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)

        UpperCAmelCase_          =							"A fantasy landscape, trending on artstation"

        UpperCAmelCase_          =							np.random.RandomState(0)
        UpperCAmelCase_          =							pipe(
            prompt=_SCREAMING_SNAKE_CASE					,						image=_SCREAMING_SNAKE_CASE					,						strength=0.7_5					,						guidance_scale=7.5					,						num_inference_steps=10					,						generator=_SCREAMING_SNAKE_CASE					,						output_type='''np'''					,						)
        UpperCAmelCase_          =							output.images
        UpperCAmelCase_          =							images[0, 255:258, 383:386, -1]

        assert images.shape == (1, 512, 768, 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 lowerCamelCase						(      self		:						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((768, 512))
        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=_SCREAMING_SNAKE_CASE					,						safety_checker=_SCREAMING_SNAKE_CASE					,						feature_extractor=_SCREAMING_SNAKE_CASE					,						provider=self.gpu_provider					,						sess_options=self.gpu_options					,						)
        pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)

        UpperCAmelCase_          =							"A fantasy landscape, trending on artstation"

        UpperCAmelCase_          =							np.random.RandomState(0)
        UpperCAmelCase_          =							pipe(
            prompt=_SCREAMING_SNAKE_CASE					,						image=_SCREAMING_SNAKE_CASE					,						strength=0.7_5					,						guidance_scale=7.5					,						num_inference_steps=20					,						generator=_SCREAMING_SNAKE_CASE					,						output_type='''np'''					,						)
        UpperCAmelCase_          =							output.images
        UpperCAmelCase_          =							images[0, 255:258, 383:386, -1]

        assert images.shape == (1, 512, 768, 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 subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
	from sagemaker import Session, TrainingJobAnalytics
	from sagemaker.huggingface import HuggingFace



@pytest.mark.skipif(
    literal_eval(os.getenv('TEST_SAGEMAKER'  ,		'False'							)							) is not True  ,		reason='Skipping test because should only be run when releasing minor transformers version'  ,		)
@pytest.mark.usefixtures('sm_env'							)
@parameterized_class(
    [
        {
            'framework': 'pytorch',
            'script': 'run_glue.py',
            'model_name_or_path': 'distilbert-base-cased',
            'instance_type': 'ml.g4dn.xlarge',
            'results': {'train_runtime': 6_5_0, 'eval_accuracy': 0.6, 'eval_loss': 0.9},
        },
        {
            'framework': 'tensorflow',
            'script': 'run_tf.py',
            'model_name_or_path': 'distilbert-base-cased',
            'instance_type': 'ml.g4dn.xlarge',
            'results': {'train_runtime': 6_0_0, 'eval_accuracy': 0.3, 'eval_loss': 0.9},
        },
    ]							)
class 			A__   (  unittest.TestCase							):




	'''simple docstring'''





	def 							_SCREAMING_SNAKE_CASE   (     self:		int)				->       Tuple:

				"""simple docstring"""

				if self.framework == "pytorch":
							subprocess.run(
							    F"""cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py""".split()       ,			encoding="utf-8"       ,			check=_SCREAMING_SNAKE_CASE       ,			)
				assert hasattr(self       ,			"env")

	def 							_SCREAMING_SNAKE_CASE   (     self:		Optional[Any]       ,			_SCREAMING_SNAKE_CASE:		Any=1)				->       Dict:

				"""simple docstring"""

				return HuggingFace(
				    entry_point=self.script       ,			source_dir=self.env.test_path       ,			role=self.env.role       ,			image_uri=self.env.image_uri       ,			base_job_name=F"""{self.env.base_job_name}-single"""       ,			instance_count=_SCREAMING_SNAKE_CASE       ,			instance_type=self.instance_type       ,			debugger_hook_config=_SCREAMING_SNAKE_CASE       ,			hyperparameters={**self.env.hyperparameters, "model_name_or_path": self.model_name_or_path}       ,			metric_definitions=self.env.metric_definitions       ,			py_version="py36"       ,			)

	def 							_SCREAMING_SNAKE_CASE   (     self:		str       ,			_SCREAMING_SNAKE_CASE:		List[Any])				->       Optional[Any]:

				"""simple docstring"""

				TrainingJobAnalytics(_SCREAMING_SNAKE_CASE).export_csv(F"""{self.env.test_path}/{job_name}_metrics.csv""")







	def 							_SCREAMING_SNAKE_CASE   (     self:		str)				->       Optional[Any]:

				"""simple docstring"""

				__lowerCAmelCase      :  Tuple												=       self.create_estimator()

				# run training
				estimator.fit()

				# result dataframe
				__lowerCAmelCase      :  Tuple												=       TrainingJobAnalytics(estimator.latest_training_job.name).dataframe()

				# extract kpis
				__lowerCAmelCase      :  Optional[Any]												=       list(result_metrics_df[result_metrics_df.metric_name == "eval_accuracy"]["value"])
				__lowerCAmelCase      :  Union[str, Any]												=       list(result_metrics_df[result_metrics_df.metric_name == "eval_loss"]["value"])
				# get train time from SageMaker job, this includes starting, preprocessing, stopping
				__lowerCAmelCase      :  Tuple												=       (
				    Session().describe_training_job(estimator.latest_training_job.name).get("TrainingTimeInSeconds"       ,			99_9999)
				)

				# assert kpis
				assert train_runtime <= self.results["train_runtime"]
				assert all(t >= self.results["eval_accuracy"] for t in eval_accuracy)
				assert all(t <= self.results["eval_loss"] for t in eval_loss)

				# dump tests result into json file to share in PR
				with open(F"""{estimator.latest_training_job.name}.json"""       ,			"w") as outfile:
							json.dump({"train_time": train_runtime, "eval_accuracy": eval_accuracy, "eval_loss": eval_loss}       ,			_SCREAMING_SNAKE_CASE)