
'''simple docstring'''







import unittest

from diffusers.models.unet_ad_blocks import *  # noqa F403
from diffusers.utils import torch_device

from .test_unet_blocks_common import UNetBlockTesterMixin


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					Optional[int]       =					DownBlockaD  # noqa F405
		a							:					int       =					"""down"""




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [-0.02_32, -0.98_69, 0.80_54, -0.06_37, -0.16_88, -1.42_64, 0.44_70, -1.33_94, 0.09_04]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					Union[str, Any]       =					ResnetDownsampleBlockaD  # noqa F405
		a							:					Optional[int]       =					"""down"""




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [0.07_10, 0.24_10, -0.73_20, -1.07_57, -1.13_43, 0.35_40, -0.01_33, -0.25_76, 0.09_48]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					str       =					AttnDownBlockaD  # noqa F405
		a							:					Tuple       =					"""down"""




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [0.06_36, 0.89_64, -0.62_34, -1.01_31, 0.08_44, 0.49_35, 0.34_37, 0.09_11, -0.29_57]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					Any       =					CrossAttnDownBlockaD  # noqa F405
		a							:					Tuple       =					"""down"""




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE,							_SCREAMING_SNAKE_CASE				= super().prepare_init_args_and_inputs_for_common()
								_SCREAMING_SNAKE_CASE				= 32
								return init_dict, inputs_dict




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [0.22_38, -0.73_96, -0.22_55, -0.38_29, 0.19_25, 1.16_65, 0.06_03, -0.72_95, 0.19_83]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					Any       =					SimpleCrossAttnDownBlockaD  # noqa F405
		a							:					Union[str, Any]       =					"""down"""




		@property
		def     lowercase       (				self      ):
								return super().get_dummy_input(include_encoder_hidden_states=UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE,							_SCREAMING_SNAKE_CASE				= super().prepare_init_args_and_inputs_for_common()
								_SCREAMING_SNAKE_CASE				= 32
								return init_dict, inputs_dict




		@unittest.skipIf(torch_device == "mps"					,						"MPS result is not consistent"      )
		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [0.79_21, -0.09_92, -0.19_62, -0.76_95, -0.42_42, 0.78_04, 0.47_37, 0.27_65, 0.33_38]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					Dict       =					SkipDownBlockaD  # noqa F405
		a							:					Union[str, Any]       =					"""down"""




		@property
		def     lowercase       (				self      ):
								return super().get_dummy_input(include_skip_sample=UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [-0.08_45, -0.20_87, -0.24_65, 0.09_71, 0.19_00, -0.04_84, 0.26_64, 0.41_79, 0.50_69]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					Union[str, Any]       =					AttnSkipDownBlockaD  # noqa F405
		a							:					str       =					"""down"""




		@property
		def     lowercase       (				self      ):
								return super().get_dummy_input(include_skip_sample=UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [0.55_39, 0.16_09, 0.49_24, 0.05_37, -0.19_95, 0.40_50, 0.09_79, -0.27_21, -0.06_42]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					Tuple       =					DownEncoderBlockaD  # noqa F405
		a							:					Optional[int]       =					"""down"""




		@property
		def     lowercase       (				self      ):
								return super().get_dummy_input(include_temb=UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= {
								    "in_channels": 32,
								    "out_channels": 32,
								}
								_SCREAMING_SNAKE_CASE				= self.dummy_input
								return init_dict, inputs_dict




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [1.11_02, 0.53_02, 0.48_72, -0.00_23, -0.80_42, 0.04_83, -0.34_89, -0.56_32, 0.76_26]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					Tuple       =					AttnDownEncoderBlockaD  # noqa F405
		a							:					List[str]       =					"""down"""




		@property
		def     lowercase       (				self      ):
								return super().get_dummy_input(include_temb=UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= {
								    "in_channels": 32,
								    "out_channels": 32,
								}
								_SCREAMING_SNAKE_CASE				= self.dummy_input
								return init_dict, inputs_dict




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [0.89_66, -0.14_86, 0.85_68, 0.81_41, -0.90_46, -0.13_42, -0.09_72, -0.74_17, 0.15_38]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					Any       =					UNetMidBlockaD  # noqa F405
		a							:					Tuple       =					"""mid"""




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= {
								    "in_channels": 32,
								    "temb_channels": 128,
								}
								_SCREAMING_SNAKE_CASE				= self.dummy_input
								return init_dict, inputs_dict




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [-0.10_62, 1.72_48, 0.34_94, 1.45_69, -0.09_10, -1.24_21, -0.99_84, 0.67_36, 1.00_28]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					Union[str, Any]       =					UNetMidBlockaDCrossAttn  # noqa F405
		a							:					Tuple       =					"""mid"""




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE,							_SCREAMING_SNAKE_CASE				= super().prepare_init_args_and_inputs_for_common()
								_SCREAMING_SNAKE_CASE				= 32
								return init_dict, inputs_dict




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [0.01_87, 2.42_20, 0.44_84, 1.12_03, -0.61_21, -1.51_22, -0.82_70, 0.78_51, 1.83_35]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					List[Any]       =					UNetMidBlockaDSimpleCrossAttn  # noqa F405
		a							:					Dict       =					"""mid"""




		@property
		def     lowercase       (				self      ):
								return super().get_dummy_input(include_encoder_hidden_states=UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE,							_SCREAMING_SNAKE_CASE				= super().prepare_init_args_and_inputs_for_common()
								_SCREAMING_SNAKE_CASE				= 32
								return init_dict, inputs_dict




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [0.71_43, 1.99_74, 0.54_48, 1.39_77, 0.12_82, -1.12_37, -1.42_38, 0.55_30, 0.88_80]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					Union[str, Any]       =					UpBlockaD  # noqa F405
		a							:					List[str]       =					"""up"""




		@property
		def     lowercase       (				self      ):
								return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [-0.20_41, -0.41_65, -0.30_22, 0.00_41, -0.66_28, -0.70_53, 0.19_28, -0.03_25, 0.05_23]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					Tuple       =					ResnetUpsampleBlockaD  # noqa F405
		a							:					Optional[int]       =					"""up"""




		@property
		def     lowercase       (				self      ):
								return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [0.22_87, 0.35_49, -0.13_46, 0.47_97, -0.17_15, -0.96_49, 0.73_05, -0.58_64, -0.62_44]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					int       =					CrossAttnUpBlockaD  # noqa F405
		a							:					Any       =					"""up"""




		@property
		def     lowercase       (				self      ):
								return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE,							_SCREAMING_SNAKE_CASE				= super().prepare_init_args_and_inputs_for_common()
								_SCREAMING_SNAKE_CASE				= 32
								return init_dict, inputs_dict




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [-0.14_03, -0.35_15, -0.04_20, -0.14_25, 0.31_67, 0.50_94, -0.21_81, 0.59_31, 0.55_82]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					str       =					SimpleCrossAttnUpBlockaD  # noqa F405
		a							:					Optional[Any]       =					"""up"""




		@property
		def     lowercase       (				self      ):
								return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase					,						include_encoder_hidden_states=UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE,							_SCREAMING_SNAKE_CASE				= super().prepare_init_args_and_inputs_for_common()
								_SCREAMING_SNAKE_CASE				= 32
								return init_dict, inputs_dict




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [0.26_45, 0.14_80, 0.09_09, 0.80_44, -0.97_58, -0.90_83, 0.09_94, -1.14_53, -0.74_02]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					Optional[Any]       =					AttnUpBlockaD  # noqa F405
		a							:					int       =					"""up"""




		@property
		def     lowercase       (				self      ):
								return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase      )




		@unittest.skipIf(torch_device == "mps"					,						"MPS result is not consistent"      )
		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [0.09_79, 0.13_26, 0.00_21, 0.06_59, 0.22_49, 0.00_59, 0.11_32, 0.59_52, 0.10_33]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					str       =					SkipUpBlockaD  # noqa F405
		a							:					Union[str, Any]       =					"""up"""




		@property
		def     lowercase       (				self      ):
								return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [-0.08_93, -0.12_34, -0.15_06, -0.03_32, 0.01_23, -0.02_11, 0.05_66, 0.01_43, 0.03_62]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					Dict       =					AttnSkipUpBlockaD  # noqa F405
		a							:					int       =					"""up"""




		@property
		def     lowercase       (				self      ):
								return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [0.03_61, 0.06_17, 0.27_87, -0.03_50, 0.03_42, 0.34_21, -0.08_43, 0.09_13, 0.30_15]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					str       =					UpDecoderBlockaD  # noqa F405
		a							:					Tuple       =					"""up"""




		@property
		def     lowercase       (				self      ):
								return super().get_dummy_input(include_temb=UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= {"in_channels": 32, "out_channels": 32}

								_SCREAMING_SNAKE_CASE				= self.dummy_input
								return init_dict, inputs_dict




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [0.44_04, 0.19_98, -0.98_86, -0.33_20, -0.31_28, -0.70_34, -0.69_55, -0.23_38, -0.31_37]
								super().test_output(UpperCamelCase      )


class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					str       =					AttnUpDecoderBlockaD  # noqa F405
		a							:					int       =					"""up"""




		@property
		def     lowercase       (				self      ):
								return super().get_dummy_input(include_temb=UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= {"in_channels": 32, "out_channels": 32}

								_SCREAMING_SNAKE_CASE				= self.dummy_input
								return init_dict, inputs_dict




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= [0.67_38, 0.44_91, 0.10_55, 1.07_10, 0.73_16, 0.33_39, 0.33_52, 0.10_23, 0.35_68]
								super().test_output(UpperCamelCase      )

'''simple docstring'''







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

from ..optimizer import AcceleratedOptimizer
from ..scheduler import AcceleratedScheduler


class 			lowerCAmelCase :




		def __init__(				self					,						UpperCamelCase      ):
								if isinstance(UpperCamelCase					,						UpperCamelCase      ):
														# 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
														_SCREAMING_SNAKE_CASE				= deepcopy(UpperCamelCase      )
								elif os.path.exists(UpperCamelCase      ):
														with io.open(UpperCamelCase					,						"r"					,						encoding="utf-8"      ) as f:
																				_SCREAMING_SNAKE_CASE				= json.load(UpperCamelCase      )
								else:
														try:
																				_SCREAMING_SNAKE_CASE				= baseaa.urlsafe_baadecode(UpperCamelCase      ).decode("utf-8"      )
																				_SCREAMING_SNAKE_CASE				= json.loads(UpperCamelCase      )
														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}'      )

								_SCREAMING_SNAKE_CASE				= config

								self.set_stage_and_offload()




		def     lowercase       (				self      ):
								# 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.
								_SCREAMING_SNAKE_CASE				= self.get_value("zero_optimization.stage"					,						-1      )

								# offload
								_SCREAMING_SNAKE_CASE				= False
								if self.is_zeroa() or self.is_zeroa():
														_SCREAMING_SNAKE_CASE				= set(["cpu", "nvme"]      )
														_SCREAMING_SNAKE_CASE				= 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:
																				_SCREAMING_SNAKE_CASE				= True




		def     lowercase       (				self					,						UpperCamelCase      ):
								_SCREAMING_SNAKE_CASE				= self.config

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

								return config, ds_key




		def     lowercase       (				self					,						UpperCamelCase					,						UpperCamelCase=None      ):
								_SCREAMING_SNAKE_CASE,							_SCREAMING_SNAKE_CASE				= self.find_config_node(UpperCamelCase      )
								if config is None:
														return default
								return config.get(UpperCamelCase					,						UpperCamelCase      )




		def     lowercase       (				self					,						UpperCamelCase					,						UpperCamelCase=False      ):
								_SCREAMING_SNAKE_CASE				= self.config

								# find the config node of interest if it exists
								_SCREAMING_SNAKE_CASE				= ds_key_long.split("."      )
								for node in nodes:
														_SCREAMING_SNAKE_CASE				= config
														_SCREAMING_SNAKE_CASE				= config.get(UpperCamelCase      )
														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(UpperCamelCase      )




		def     lowercase       (				self					,						UpperCamelCase      ):
								_SCREAMING_SNAKE_CASE				= self.get_value(UpperCamelCase      )
								return False if value is None else bool(UpperCamelCase      )




		def     lowercase       (				self					,						UpperCamelCase      ):
								_SCREAMING_SNAKE_CASE				= self.get_value(UpperCamelCase      )
								return False if value is None else not bool(UpperCamelCase      )




		def     lowercase       (				self      ):
								return self._stage == 2




		def     lowercase       (				self      ):
								return self._stage == 3




		def     lowercase       (				self      ):
								return self._offload


class 			lowerCAmelCase :




		def __init__(				self					,						UpperCamelCase      ):
								_SCREAMING_SNAKE_CASE				= engine




		def     lowercase       (				self					,						UpperCamelCase					,						**UpperCamelCase      ):
								# runs backpropagation and handles mixed precision
								self.engine.backward(UpperCamelCase					,						**UpperCamelCase      )

								# 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 			lowerCAmelCase (				__UpperCAmelCase			):




		def __init__(				self					,						UpperCamelCase      ):
								super().__init__(UpperCamelCase					,						device_placement=UpperCamelCase					,						scaler=UpperCamelCase      )
								_SCREAMING_SNAKE_CASE				= hasattr(self.optimizer					,						"overflow"      )




		def     lowercase       (				self					,						UpperCamelCase=None      ):
								pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed




		def     lowercase       (				self      ):
								pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed




		@property
		def     lowercase       (				self      ):
								if self.__has_overflow__:
														return self.optimizer.overflow
								return False


class 			lowerCAmelCase (				__UpperCAmelCase			):




		def __init__(				self					,						UpperCamelCase					,						UpperCamelCase      ):
								super().__init__(UpperCamelCase					,						UpperCamelCase      )




		def     lowercase       (				self      ):
								pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed


class 			lowerCAmelCase :




		def __init__(				self					,						UpperCamelCase					,						UpperCamelCase=0.0_01					,						UpperCamelCase=0					,						**UpperCamelCase      ):
								_SCREAMING_SNAKE_CASE				= params
								_SCREAMING_SNAKE_CASE				= lr
								_SCREAMING_SNAKE_CASE				= weight_decay
								_SCREAMING_SNAKE_CASE				= kwargs


class 			lowerCAmelCase :




		def __init__(				self					,						UpperCamelCase					,						UpperCamelCase=None					,						UpperCamelCase=0					,						**UpperCamelCase      ):
								_SCREAMING_SNAKE_CASE				= optimizer
								_SCREAMING_SNAKE_CASE				= total_num_steps
								_SCREAMING_SNAKE_CASE				= warmup_num_steps
								_SCREAMING_SNAKE_CASE				= kwargs