


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





							'''simple docstring'''

							__lowercase:						Union[str, Any]							 =      """hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"""



							def 						lowerCAmelCase    (			self							: Dict						,    UpperCAmelCase_							: List[str]=0    )   ->List[Any]:



														"""simple docstring"""



														snake_case_          =						floats_tensor((1, 3, 128, 128)						,    rng=random.Random(UpperCAmelCase_    )    )
														snake_case_          =						np.random.RandomState(UpperCAmelCase_    )
														snake_case_          =						{
														    """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 						lowerCAmelCase    (			self							: List[Any]    )   ->str:



														"""simple docstring"""



														snake_case_          =						OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint						,    provider="""CPUExecutionProvider"""    )
														pipe.set_progress_bar_config(disable=UpperCAmelCase_    )

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

														assert image.shape == (1, 128, 128, 3)
														snake_case_          =						np.array([0.69_643, 0.58_484, 0.50_314, 0.58_760, 0.55_368, 0.59_643, 0.51_529, 0.41_217, 0.49_087]    )
														assert np.abs(image_slice - expected_slice    ).max() < 1E-1



							def 						lowerCAmelCase    (			self							: List[str]    )   ->int:



														"""simple docstring"""



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

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

														assert image.shape == (1, 128, 128, 3)
														snake_case_          =						np.array([0.61_737, 0.54_642, 0.53_183, 0.54_465, 0.52_742, 0.60_525, 0.49_969, 0.40_655, 0.48_154]    )

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



							def 						lowerCAmelCase    (			self							: List[str]    )   ->int:



														"""simple docstring"""



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

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

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

														assert image.shape == (1, 128, 128, 3)
														snake_case_          =						np.array([0.52_761, 0.59_977, 0.49_033, 0.49_619, 0.54_282, 0.50_311, 0.47_600, 0.40_918, 0.45_203]    )

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



							def 						lowerCAmelCase    (			self							: Tuple    )   ->Tuple:



														"""simple docstring"""



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

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

														assert image.shape == (1, 128, 128, 3)
														snake_case_          =						np.array([0.52_911, 0.60_004, 0.49_229, 0.49_805, 0.54_502, 0.50_680, 0.47_777, 0.41_028, 0.45_304]    )

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



							def 						lowerCAmelCase    (			self							: Optional[Any]    )   ->Tuple:



														"""simple docstring"""



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

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

														assert image.shape == (1, 128, 128, 3)
														snake_case_          =						np.array([0.52_911, 0.60_004, 0.49_229, 0.49_805, 0.54_502, 0.50_680, 0.47_777, 0.41_028, 0.45_304]    )

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



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



														"""simple docstring"""



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

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

														assert image.shape == (1, 128, 128, 3)
														snake_case_          =						np.array([0.65_331, 0.58_277, 0.48_204, 0.56_059, 0.53_665, 0.56_235, 0.50_969, 0.40_009, 0.46_552]    )

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





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





							'''simple docstring'''

							@property
							def 						lowerCAmelCase    (			self							: Dict    )   ->Tuple:



														"""simple docstring"""



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



							@property
							def 						lowerCAmelCase    (			self							: Optional[int]    )   ->Tuple:



														"""simple docstring"""



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



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



														"""simple docstring"""



														snake_case_          =						load_image(
														    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
														    """/img2img/sketch-mountains-input.jpg"""    )
														snake_case_          =						init_image.resize((768, 512)    )
														# using the PNDM scheduler by default
														snake_case_          =						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_    )

														snake_case_          =						"""A fantasy landscape, trending on artstation"""

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

														assert images.shape == (1, 512, 768, 3)
														snake_case_          =						np.array([0.4_909, 0.5_059, 0.5_372, 0.4_623, 0.4_876, 0.5_049, 0.4_820, 0.4_956, 0.5_019]    )
														# 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							: Dict    )   ->List[Any]:



														"""simple docstring"""



														snake_case_          =						load_image(
														    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
														    """/img2img/sketch-mountains-input.jpg"""    )
														snake_case_          =						init_image.resize((768, 512)    )
														snake_case_          =						LMSDiscreteScheduler.from_pretrained(
														    """runwayml/stable-diffusion-v1-5"""						,    subfolder="""scheduler"""						,    revision="""onnx"""    )
														snake_case_          =						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_    )

														snake_case_          =						"""A fantasy landscape, trending on artstation"""

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

														assert images.shape == (1, 512, 768, 3)
														snake_case_          =						np.array([0.8_043, 0.926, 0.9_581, 0.8_119, 0.8_954, 0.913, 0.7_209, 0.7_463, 0.7_431]    )
														# 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 ast
import os
import re
import shutil
import tempfile
import unittest
from unittest import mock

import torch

from accelerate.test_utils.examples import compare_against_test
from accelerate.test_utils.testing import TempDirTestCase, require_trackers, run_command, slow
from accelerate.utils import write_basic_config


# DataLoaders built from `test_samples/MRPC` for quick testing
# Should mock `{script_name}.get_dataloaders` via:
# @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)

__SCREAMING_SNAKE_CASE							:      Tuple						=      [
    'cross_validation.py',
    'gradient_accumulation.py',
    'local_sgd.py',
    'multi_process_metrics.py',
    'memory.py',
    'automatic_gradient_accumulation.py',
    'fsdp_with_peak_mem_tracking.py',
    'deepspeed_with_config_support.py',
    'megatron_lm_gpt_pretraining.py',
]





class        __A       (unittest.TestCase):





							'''simple docstring'''

							def 						lowerCAmelCase    (			self							: str						,    UpperCAmelCase_							: str						,    UpperCAmelCase_							: bool						,    UpperCAmelCase_							: str = None						,    UpperCAmelCase_							: list = None    )   ->List[Any]:



														"""simple docstring"""



														snake_case_          =						None
														snake_case_          =						os.path.abspath(os.path.join("""examples"""						,    """by_feature"""    )    )
														snake_case_          =						os.path.abspath("""examples"""    )
														for item in os.listdir(UpperCAmelCase_    ):
																					if item not in EXCLUDE_EXAMPLES:
																												snake_case_          =						os.path.join(UpperCAmelCase_						,    UpperCAmelCase_    )
																												if os.path.isfile(UpperCAmelCase_    ) and ".py" in item_path:
																																			with self.subTest(
																																			    tested_script=UpperCAmelCase_						,    feature_script=UpperCAmelCase_						,    tested_section="""main()""" if parser_only else """training_function()"""						,    ):
																																										snake_case_          =						compare_against_test(
																																										    os.path.join(UpperCAmelCase_						,    UpperCAmelCase_    )						,    UpperCAmelCase_						,    UpperCAmelCase_						,    UpperCAmelCase_    )
																																										snake_case_          =						"""\n""".join(UpperCAmelCase_    )
																																										if special_strings is not None:
																																																	for string in special_strings:
																																																								snake_case_          =						diff.replace(UpperCAmelCase_						,    """"""    )
																																										self.assertEqual(UpperCAmelCase_						,    """"""    )



							def 						lowerCAmelCase    (			self							: Tuple    )   ->List[Any]:



														"""simple docstring"""



														self.one_complete_example("""complete_nlp_example.py"""						,    UpperCAmelCase_    )
														self.one_complete_example("""complete_nlp_example.py"""						,    UpperCAmelCase_    )



							def 						lowerCAmelCase    (			self							: Tuple    )   ->Dict:



														"""simple docstring"""



														snake_case_          =						os.path.abspath(os.path.join("""examples"""						,    """cv_example.py"""    )    )
														snake_case_          =						[
														    """ """ * 16 + """{\n\n""",
														    """ """ * 20 + """\"accuracy\": eval_metric[\"accuracy\"],\n\n""",
														    """ """ * 20 + """\"f1\": eval_metric[\"f1\"],\n\n""",
														    """ """ * 20 + """\"train_loss\": total_loss.item() / len(train_dataloader),\n\n""",
														    """ """ * 20 + """\"epoch\": epoch,\n\n""",
														    """ """ * 16 + """},\n\n""",
														    """ """ * 16 + """step=epoch,\n""",
														    """ """ * 12,
														    """ """ * 8 + """for step, batch in enumerate(active_dataloader):\n""",
														]
														self.one_complete_example("""complete_cv_example.py"""						,    UpperCAmelCase_						,    UpperCAmelCase_						,    UpperCAmelCase_    )
														self.one_complete_example("""complete_cv_example.py"""						,    UpperCAmelCase_						,    UpperCAmelCase_						,    UpperCAmelCase_    )





@mock.patch.dict(os.environ  ,						{"""TESTING_MOCKED_DATALOADERS""": """1"""})
class        __A       (snake_case__):





							'''simple docstring'''

							__lowercase:						str							 =      False



							@classmethod
							def 						lowerCAmelCase    (			cls							: Any    )   ->List[str]:



														"""simple docstring"""



														super().setUpClass()
														snake_case_          =						tempfile.mkdtemp()
														snake_case_          =						os.path.join(cls._tmpdir						,    """default_config.yml"""    )

														write_basic_config(save_location=cls.configPath    )
														snake_case_          =						["""accelerate""", """launch""", """--config_file""", cls.configPath]



							@classmethod
							def 						lowerCAmelCase    (			cls							: List[str]    )   ->int:



														"""simple docstring"""



														super().tearDownClass()
														shutil.rmtree(cls._tmpdir    )



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



														"""simple docstring"""



														snake_case_          =						F"""
        examples/by_feature/checkpointing.py
        --checkpointing_steps epoch
        --output_dir {self.tmpdir}
        """.split()
														run_command(self._launch_args + testargs    )
														self.assertTrue(os.path.exists(os.path.join(self.tmpdir						,    """epoch_0"""    )    )    )



							def 						lowerCAmelCase    (			self							: Any    )   ->Optional[Any]:



														"""simple docstring"""



														snake_case_          =						F"""
        examples/by_feature/checkpointing.py
        --checkpointing_steps 1
        --output_dir {self.tmpdir}
        """.split()
														snake_case_          =						run_command(self._launch_args + testargs    )
														self.assertTrue(os.path.exists(os.path.join(self.tmpdir						,    """step_2"""    )    )    )



							def 						lowerCAmelCase    (			self							: Union[str, Any]    )   ->str:



														"""simple docstring"""



														snake_case_          =						F"""
        examples/by_feature/checkpointing.py
        --resume_from_checkpoint {os.path.join(self.tmpdir						,    "epoch_0"    )}
        """.split()
														snake_case_          =						run_command(self._launch_args + testargs						,    return_stdout=UpperCAmelCase_    )
														self.assertNotIn("""epoch 0:"""						,    UpperCAmelCase_    )
														self.assertIn("""epoch 1:"""						,    UpperCAmelCase_    )



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



														"""simple docstring"""



														snake_case_          =						F"""
        examples/by_feature/checkpointing.py
        --resume_from_checkpoint {os.path.join(self.tmpdir						,    "step_2"    )}
        """.split()
														snake_case_          =						run_command(self._launch_args + testargs						,    return_stdout=UpperCAmelCase_    )
														if torch.cuda.is_available():
																					snake_case_          =						torch.cuda.device_count()
														else:
																					snake_case_          =						1
														if num_processes > 1:
																					self.assertNotIn("""epoch 0:"""						,    UpperCAmelCase_    )
																					self.assertIn("""epoch 1:"""						,    UpperCAmelCase_    )
														else:
																					self.assertIn("""epoch 0:"""						,    UpperCAmelCase_    )
																					self.assertIn("""epoch 1:"""						,    UpperCAmelCase_    )



							@slow
							def 						lowerCAmelCase    (			self							: Dict    )   ->Dict:



														"""simple docstring"""



														snake_case_          =						"""
        examples/by_feature/cross_validation.py
        --num_folds 2
        """.split()
														with mock.patch.dict(os.environ						,    {"""TESTING_MOCKED_DATALOADERS""": """0"""}    ):
																					snake_case_          =						run_command(self._launch_args + testargs						,    return_stdout=UpperCAmelCase_    )
																					snake_case_          =						re.findall("""({.+})"""						,    UpperCAmelCase_    )
																					snake_case_          =						[r for r in results if """accuracy""" in r][-1]
																					snake_case_          =						ast.literal_eval(UpperCAmelCase_    )
																					self.assertGreaterEqual(results["""accuracy"""]						,    0.75    )



							def 						lowerCAmelCase    (			self							: Optional[int]    )   ->int:



														"""simple docstring"""



														snake_case_          =						["""examples/by_feature/multi_process_metrics.py"""]
														run_command(self._launch_args + testargs    )



							@require_trackers
							@mock.patch.dict(os.environ						,    {"""WANDB_MODE""": """offline"""}    )
							def 						lowerCAmelCase    (			self							: Dict    )   ->Optional[Any]:



														"""simple docstring"""



														with tempfile.TemporaryDirectory() as tmpdir:
																					snake_case_          =						F"""
            examples/by_feature/tracking.py
            --with_tracking
            --project_dir {tmpdir}
            """.split()
																					run_command(self._launch_args + testargs    )
																					self.assertTrue(os.path.exists(os.path.join(UpperCAmelCase_						,    """tracking"""    )    )    )



							def 						lowerCAmelCase    (			self							: List[str]    )   ->Union[str, Any]:



														"""simple docstring"""



														snake_case_          =						["""examples/by_feature/gradient_accumulation.py"""]
														run_command(self._launch_args + testargs    )



							def 						lowerCAmelCase    (			self							: str    )   ->Dict:



														"""simple docstring"""



														snake_case_          =						["""examples/by_feature/local_sgd.py"""]
														run_command(self._launch_args + testargs    )



