






import baseaa
import io
import json
import os
from copy import deepcopy

from ..optimizer import AcceleratedOptimizer
from ..scheduler import AcceleratedScheduler
class       _UpperCAmelCase					:

				"""simple docstring"""




				def __init__(     self						:    Union[str, Any],  lowerCamelCase						:    Optional[int]  ):

								'''simple docstring'''
								if isinstance(SCREAMING_SNAKE_CASE__,  SCREAMING_SNAKE_CASE__  ):
												# Don't modify user's data should they want to reuse it (e.g. in tests), because once we
												# modified it, it will not be accepted here again, since `auto` values would have been overridden
												lowercase__    =			deepcopy(SCREAMING_SNAKE_CASE__  )
								elif os.path.exists(SCREAMING_SNAKE_CASE__  ):
												with io.open(SCREAMING_SNAKE_CASE__,  '''r''',  encoding='''utf-8'''  ) as f:
																lowercase__    =			json.load(SCREAMING_SNAKE_CASE__  )
								else:
												try:
																lowercase__    =			baseaa.urlsafe_baadecode(SCREAMING_SNAKE_CASE__  ).decode('''utf-8'''  )
																lowercase__    =			json.loads(SCREAMING_SNAKE_CASE__  )
												except (UnicodeDecodeError, AttributeError, ValueError):
																raise ValueError(
																    F"""Expected a string path to an existing deepspeed config, or a dictionary, or a base64 encoded string. Received: {config_file_or_dict}"""  )

								lowercase__    =			config

								self.set_stage_and_offload()




				def 				lowercase__   (     self						:    List[Any]  ):

								'''simple docstring'''
								# zero stage - this is done as early as possible, before model is created, to allow
								# ``is_deepspeed_zero3_enabled`` query and getting to the early deepspeed config object
								# during ``zero.Init()`` which needs to know the dtype, and some other hparams.
								lowercase__    =			self.get_value('''zero_optimization.stage''',  -1  )

								# offload
								lowercase__    =			False
								if self.is_zeroa() or self.is_zeroa():
												lowercase__    =			set(['''cpu''', '''nvme''']  )
												lowercase__    =			set(
												    [
												        self.get_value('''zero_optimization.offload_optimizer.device'''  ),
												        self.get_value('''zero_optimization.offload_param.device'''  ),
												    ]  )
												if len(offload_devices & offload_devices_valid  ) > 0:
																lowercase__    =			True




				def 				lowercase__   (     self						:    int,  lowerCamelCase						:    Tuple  ):

								'''simple docstring'''
								lowercase__    =			self.config

								# find the config node of interest if it exists
								lowercase__    =			ds_key_long.split('''.'''  )
								lowercase__    =			nodes.pop()
								for node in nodes:
												lowercase__    =			config.get(SCREAMING_SNAKE_CASE__  )
												if config is None:
																return None, ds_key

								return config, ds_key




				def 				lowercase__   (     self						:    Union[str, Any],  lowerCamelCase						:    Tuple,  lowerCamelCase						:    List[Any]=None  ):

								'''simple docstring'''
								lowercase__    =			self.find_config_node(SCREAMING_SNAKE_CASE__  )
								if config is None:
												return default
								return config.get(SCREAMING_SNAKE_CASE__,  SCREAMING_SNAKE_CASE__  )




				def 				lowercase__   (     self						:    Optional[Any],  lowerCamelCase						:    Optional[int],  lowerCamelCase						:    int=False  ):

								'''simple docstring'''
								lowercase__    =			self.config

								# find the config node of interest if it exists
								lowercase__    =			ds_key_long.split('''.'''  )
								for node in nodes:
												lowercase__    =			config
												lowercase__    =			config.get(SCREAMING_SNAKE_CASE__  )
												if config is None:
																if must_exist:
																				raise ValueError(F"""Can't find {ds_key_long} entry in the config: {self.config}"""  )
																else:
																				return

        # if found remove it
								if parent_config is not None:
												parent_config.pop(SCREAMING_SNAKE_CASE__  )




				def 				lowercase__   (     self						:    Union[str, Any],  lowerCamelCase						:    List[Any]  ):

								'''simple docstring'''
								lowercase__    =			self.get_value(SCREAMING_SNAKE_CASE__  )
								return False if value is None else bool(SCREAMING_SNAKE_CASE__  )




				def 				lowercase__   (     self						:    str,  lowerCamelCase						:    Tuple  ):

								'''simple docstring'''
								lowercase__    =			self.get_value(SCREAMING_SNAKE_CASE__  )
								return False if value is None else not bool(SCREAMING_SNAKE_CASE__  )




				def 				lowercase__   (     self						:    Optional[int]  ):

								'''simple docstring'''
								return self._stage == 2




				def 				lowercase__   (     self						:    Optional[Any]  ):

								'''simple docstring'''
								return self._stage == 3




				def 				lowercase__   (     self						:    str  ):

								'''simple docstring'''
								return self._offload



class       _UpperCAmelCase					:

				"""simple docstring"""




				def __init__(     self						:    Optional[Any],  lowerCamelCase						:    Any  ):

								'''simple docstring'''
								lowercase__    =			engine




				def 				lowercase__   (     self						:    Optional[Any],  lowerCamelCase						:    Any,  **lowerCamelCase						:    Union[str, Any]  ):

								'''simple docstring'''
								# runs backpropagation and handles mixed precision
								self.engine.backward(SCREAMING_SNAKE_CASE__,  **SCREAMING_SNAKE_CASE__  )

								# Deepspeed's `engine.step` performs the following operations:
								# - gradient accumulation check
								# - gradient clipping
								# - optimizer step
								# - zero grad
								# - checking overflow
								# - lr_scheduler step (only if engine.lr_scheduler is not None)
								self.engine.step()
								# and this plugin overrides the above calls with no-ops when Accelerate runs under
								# Deepspeed, but allows normal functionality for non-Deepspeed cases thus enabling a simple
								# training loop that works transparently under many training regimes.



class       _UpperCAmelCase					(       lowercase__	):

				"""simple docstring"""




				def __init__(     self						:    Any,  lowerCamelCase						:    str  ):

								'''simple docstring'''
								super().__init__(SCREAMING_SNAKE_CASE__,  device_placement=SCREAMING_SNAKE_CASE__,  scaler=SCREAMING_SNAKE_CASE__  )
								lowercase__    =			hasattr(self.optimizer,  '''overflow'''  )




				def 				lowercase__   (     self						:    Dict,  lowerCamelCase						:    Dict=None  ):

								'''simple docstring'''
								pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed




				def 				lowercase__   (     self						:    Union[str, Any]  ):

								'''simple docstring'''
								pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed




				@property
				def 				lowercase__   (     self						:    Tuple  ):

								'''simple docstring'''
								if self.__has_overflow__:
												return self.optimizer.overflow
								return False



class       _UpperCAmelCase					(       lowercase__	):

				"""simple docstring"""




				def __init__(     self						:    Union[str, Any],  lowerCamelCase						:    Optional[int],  lowerCamelCase						:    Optional[Any]  ):

								'''simple docstring'''
								super().__init__(SCREAMING_SNAKE_CASE__,  SCREAMING_SNAKE_CASE__  )




				def 				lowercase__   (     self						:    Optional[int]  ):

								'''simple docstring'''
								pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed



class       _UpperCAmelCase					:

				"""simple docstring"""




				def __init__(     self						:    Tuple,  lowerCamelCase						:    Tuple,  lowerCamelCase						:    Tuple=0.001,  lowerCamelCase						:    List[Any]=0,  **lowerCamelCase						:    List[Any]  ):

								'''simple docstring'''
								lowercase__    =			params
								lowercase__    =			lr
								lowercase__    =			weight_decay
								lowercase__    =			kwargs



class       _UpperCAmelCase					:

				"""simple docstring"""




				def __init__(     self						:    List[str],  lowerCamelCase						:    Tuple,  lowerCamelCase						:    str=None,  lowerCamelCase						:    Dict=0,  **lowerCamelCase						:    List[Any]  ):

								'''simple docstring'''
								lowercase__    =			optimizer
								lowercase__    =			total_num_steps
								lowercase__    =			warmup_num_steps
								lowercase__    =			kwargs








import random
import unittest

import numpy as np
import torch

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionUpscalePipeline,
    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__  (				lowercase__						,    unittest.TestCase			):
      # TODO: is there an appropriate internal test set?
      snake_case__				:     Any        =						'''ssube/stable-diffusion-x4-upscaler-onnx'''
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     ,							SCREAMING_SNAKE_CASE__				:       int=0     )		->			Tuple:
             a_					:		Union[str, Any]							=  floats_tensor((1, 3, 1_2_8, 1_2_8)     ,							rng=random.Random(SCREAMING_SNAKE_CASE__     )     )
             a_					:		Dict							=  torch.manual_seed(SCREAMING_SNAKE_CASE__     )
             a_					:		Union[str, Any]							=  {
                 'prompt': 'A painting of a squirrel eating a burger',
                 'image': image,
                 'generator': generator,
                 'num_inference_steps': 3,
                 'guidance_scale': 7.5,
                 'output_type': 'numpy',
             }
             return inputs
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     )		->			Dict:
             a_					:		Dict							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,							provider='CPUExecutionProvider'     )
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

             a_					:		Optional[int]							=  self.get_dummy_inputs()
             a_					:		int							=  pipe(**SCREAMING_SNAKE_CASE__     ).images
             a_					:		Tuple							=  image[0, -3:, -3:, -1].flatten()

             # started as 128, should now be 512
             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		List[Any]							=  np.array(
                 [0.6974782, 0.68902093, 0.70135885, 0.7583618, 0.7804545, 0.7854912, 0.78667426, 0.78743863, 0.78070223]     )
             assert np.abs(image_slice - expected_slice     ).max() < 1E-1
      def        SCREAMING_SNAKE_CASE							(       self				:       Dict     )		->			Dict:
             a_					:		Tuple							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,							provider='CPUExecutionProvider'     )
             a_					:		int							=  PNDMScheduler.from_config(pipe.scheduler.config     ,							skip_prk_steps=SCREAMING_SNAKE_CASE__     )
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

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

             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		str							=  np.array(
                 [0.6898892, 0.59240556, 0.52499527, 0.58866215, 0.52258235, 0.52572715, 0.62414473, 0.6174387, 0.6214964]     )
             assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1E-1
      def        SCREAMING_SNAKE_CASE							(       self				:       Any     )		->			Dict:
             a_					:		Optional[Any]							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,							provider='CPUExecutionProvider'     )
             a_					:		Optional[Any]							=  DPMSolverMultistepScheduler.from_config(pipe.scheduler.config     )
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

             a_					:		Optional[int]							=  self.get_dummy_inputs()
             a_					:		Dict							=  pipe(**SCREAMING_SNAKE_CASE__     ).images
             a_					:		str							=  image[0, -3:, -3:, -1]

             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		Optional[Any]							=  np.array(
                 [0.7659278, 0.76437664, 0.75579107, 0.7691116, 0.77666986, 0.7727672, 0.7758664, 0.7812226, 0.76942515]     )

             assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1E-1
      def        SCREAMING_SNAKE_CASE							(       self				:       List[Any]     )		->			int:
             a_					:		List[str]							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,							provider='CPUExecutionProvider'     )
             a_					:		int							=  EulerDiscreteScheduler.from_config(pipe.scheduler.config     )
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

             a_					:		Optional[Any]							=  self.get_dummy_inputs()
             a_					:		Dict							=  pipe(**SCREAMING_SNAKE_CASE__     ).images
             a_					:		List[str]							=  image[0, -3:, -3:, -1]

             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		int							=  np.array(
                 [0.6974782, 0.68902093, 0.70135885, 0.7583618, 0.7804545, 0.7854912, 0.78667426, 0.78743863, 0.78070223]     )
             assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1E-1

      def        SCREAMING_SNAKE_CASE							(       self				:       str     )		->			Optional[Any]:
             a_					:		Tuple							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,							provider='CPUExecutionProvider'     )
             a_					:		Optional[int]							=  EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config     )
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

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

             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		Union[str, Any]							=  np.array(
                 [0.77424496, 0.773601, 0.7645288, 0.7769598, 0.7772739, 0.7738688, 0.78187233, 0.77879584, 0.767043]     )

             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				:       Tuple     )		->			Union[str, Any]:
             return (
                 "CUDAExecutionProvider",
                 {
                     "gpu_mem_limit": "15000000000",  # 15GB
                     "arena_extend_strategy": "kSameAsRequested",
                 },
             )
      @property
      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[Any]     )		->			List[str]:
             a_					:		List[str]							=  ort.SessionOptions()
             a_					:		int							=  False
             return options
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     )		->			Tuple:
             a_					:		str							=  load_image(
                 'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
                 '/img2img/sketch-mountains-input.jpg'     )
             a_					:		int							=  init_image.resize((1_2_8, 1_2_8)     )
             # using the PNDM scheduler by default
             a_					:		Dict							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(
                 'ssube/stable-diffusion-x4-upscaler-onnx'     ,							provider=self.gpu_provider     ,							sess_options=self.gpu_options     ,							)
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

             a_					:		Tuple							=  'A fantasy landscape, trending on artstation'

             a_					:		str							=  torch.manual_seed(0     )
             a_					:		List[str]							=  pipe(
                 prompt=SCREAMING_SNAKE_CASE__     ,							image=SCREAMING_SNAKE_CASE__     ,							guidance_scale=7.5     ,							num_inference_steps=1_0     ,							generator=SCREAMING_SNAKE_CASE__     ,							output_type='np'     ,							)
             a_					:		Dict							=  output.images
             a_					:		Any							=  images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

             assert images.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		str							=  np.array([0.4883, 0.4947, 0.4980, 0.4975, 0.4982, 0.4980, 0.5000, 0.5006, 0.4972]     )
             # 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				:       Optional[int]     )		->			List[str]:
             a_					:		Dict							=  load_image(
                 'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
                 '/img2img/sketch-mountains-input.jpg'     )
             a_					:		List[str]							=  init_image.resize((1_2_8, 1_2_8)     )
             a_					:		Dict							=  LMSDiscreteScheduler.from_pretrained(
                 'ssube/stable-diffusion-x4-upscaler-onnx'     ,							subfolder='scheduler'     )
             a_					:		Union[str, Any]							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(
                 'ssube/stable-diffusion-x4-upscaler-onnx'     ,							scheduler=SCREAMING_SNAKE_CASE__     ,							provider=self.gpu_provider     ,							sess_options=self.gpu_options     ,							)
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

             a_					:		Any							=  'A fantasy landscape, trending on artstation'

             a_					:		Tuple							=  torch.manual_seed(0     )
             a_					:		Optional[Any]							=  pipe(
                 prompt=SCREAMING_SNAKE_CASE__     ,							image=SCREAMING_SNAKE_CASE__     ,							guidance_scale=7.5     ,							num_inference_steps=2_0     ,							generator=SCREAMING_SNAKE_CASE__     ,							output_type='np'     ,							)
             a_					:		str							=  output.images
             a_					:		List[Any]							=  images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

             assert images.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		Tuple							=  np.array(
                 [0.50173753, 0.50223356, 0.502039, 0.50233036, 0.5023725, 0.5022601, 0.5018758, 0.50234085, 0.50241566]     )
             # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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


