



import importlib.util
import weakref
from functools import partialmethod

from .dependency_versions_check import dep_version_check
from .utils import is_accelerate_available, is_torch_available, logging


if is_torch_available():
							import torch

__SCREAMING_SNAKE_CASE     =	logging.get_logger(__name__)




def 						UpperCAmelCase (    ):
			return importlib.util.find_spec("deepspeed"     ) is not None


if is_accelerate_available() and is_deepspeed_available():
							from accelerate.utils.deepspeed import HfDeepSpeedConfig as DeepSpeedConfig
else:
							# Inherits from a dummy `object` if accelerate is not available, so that python succeeds to import this file.
							# Deepspeed glue code will never inherit this dummy object as it checks if accelerate is available.
							from builtins import object as DeepSpeedConfig

class 						lowerCamelCase_     (			lowerCamelCase__    ):






						'''simple docstring'''







						def __init__(						self   :      Dict						,						__lowerCamelCase   :      str			)   ->					Tuple:
									# set global weakref object
									set_hf_deepspeed_config(self			)
									dep_version_check("accelerate"			)
									dep_version_check("deepspeed"			)
									super().__init__(_UpperCAmelCase			)







class 						lowerCamelCase_     (			lowerCamelCase__    ):






						'''simple docstring'''







						def __init__(						self   :      Any						,						__lowerCamelCase   :      int			)   ->					str:
									super().__init__(_UpperCAmelCase			)
									A				:       str														=    None
									A				:       Union[str, Any]														=    []







						def       SCREAMING_SNAKE_CASE__		(						self   :      List[str]			)   ->					str:
									if self._dtype is None:
												raise ValueError("trainer_config_process() wasn\'t called yet to tell dtype"			)
									return self._dtype







						def       SCREAMING_SNAKE_CASE__		(						self   :      Any						,						__lowerCamelCase   :      List[str]			)   ->					Optional[int]:
									A				:       Dict														=    self.get_value(_UpperCAmelCase			)
									if val is None:
												return False
									else:
												return val == "auto"







						def       SCREAMING_SNAKE_CASE__		(						self   :      str						,						__lowerCamelCase   :      Optional[int]						,						__lowerCamelCase   :      Tuple						,						__lowerCamelCase   :      Dict=None						,						__lowerCamelCase   :      Optional[int]=True			)   ->					List[Any]:
									A							,			A				:       List[Any]														=    self.find_config_node(_UpperCAmelCase			)
									if config is None:
												return

									if config.get(_UpperCAmelCase			) == "auto":
												A				:       Optional[Any]														=    hf_val
												return

									if not must_match:
												return

									A				:       Dict														=    config.get(_UpperCAmelCase			)
									if ds_val is not None and ds_val != hf_val:
												self.mismatches.append(F"""- ds {ds_key_long}={ds_val} vs hf {hf_key}={hf_val}"""			)

						a__ =	partialmethod(lowerCamelCase__    ,must_match=lowerCamelCase__    )







						def       SCREAMING_SNAKE_CASE__		(						self   :      Optional[Any]						,						__lowerCamelCase   :      Optional[Any]			)   ->					List[Any]:
									A				:       str														=    args.world_size * args.per_device_train_batch_size * args.gradient_accumulation_steps
									self.fill_match(
									    "train_micro_batch_size_per_gpu"						,						args.per_device_train_batch_size						,						"per_device_train_batch_size"			)
									self.fill_match("gradient_accumulation_steps"						,						args.gradient_accumulation_steps						,						"gradient_accumulation_steps"			)
									self.fill_match("train_batch_size"						,						_UpperCAmelCase						,						"train_batch_size (calculated)"			)
									self.fill_match("gradient_clipping"						,						args.max_grad_norm						,						"max_grad_norm"			)

									self.fill_match("optimizer.params.lr"						,						args.learning_rate						,						"learning_rate"			)
									self.fill_match("optimizer.params.betas"						,						[args.adam_betaa, args.adam_betaa]						,						"adam_beta1+adam_beta2"			)
									self.fill_match("optimizer.params.eps"						,						args.adam_epsilon						,						"adam_epsilon"			)
									self.fill_match("optimizer.params.weight_decay"						,						args.weight_decay						,						"weight_decay"			)

									self.fill_only("scheduler.params.warmup_min_lr"						,						0			)  # not a trainer arg
									self.fill_match("scheduler.params.warmup_max_lr"						,						args.learning_rate						,						"learning_rate"			)
									# total_num_steps - will get set in trainer_config_finalize

									# fp16
									if args.fpaa or args.fpaa_full_eval:
												A				:       Any														=    "apex" if args.fpaa_backend == "apex" else "amp"
									else:
												A				:       Tuple														=    None

									if args.save_on_each_node:
												# deepspeed uses shared storage by default. Let's override this setting if save_on_each_node == True
												A				:       Union[str, Any]														=    self.config.get("checkpoint"						,						{}			)
												A				:       Any														=    args.save_on_each_node

									# amp: similar to the pytorch native amp - it has a bunch of optional params but we won't set
									# any here unless the user did the work
									self.fill_match(
									    "fp16.enabled"						,						((args.fpaa or args.fpaa_full_eval) and fpaa_backend == "amp")						,						"fp16|fp16_full_eval+fp16_backend(amp)"						,						)

									# apex: delegates amp work to apex (which needs to be available), but it cannot be used with any
									# ZeRO features
									self.fill_match("amp.enabled"						,						fpaa_backend == "apex"						,						"fp16+fp16_backend(apex)"			)
									self.fill_match("amp.opt_level"						,						args.fpaa_opt_level						,						"fp16_opt_level"			)

									self.fill_match("bf16.enabled"						,						(args.bfaa or args.bfaa_full_eval)						,						"bf16|bf16_full_eval"			)

									# deepspeed's default mode is fp16 unless there is a config that says differently
									if self.is_true("bf16.enabled"			):
												A				:       str														=    torch.bfloataa
									elif self.is_false("fp16.enabled"			):
												A				:       Dict														=    torch.floataa
									else:
												A				:       List[str]														=    torch.floataa







						def       SCREAMING_SNAKE_CASE__		(						self   :      Dict						,						__lowerCamelCase   :      List[Any]						,						__lowerCamelCase   :      Any						,						__lowerCamelCase   :      Tuple			)   ->					Optional[int]:
									A				:       int														=    [
									    "zero_optimization.reduce_bucket_size",
									    "zero_optimization.stage3_prefetch_bucket_size",
									    "zero_optimization.stage3_param_persistence_threshold",
									]
									A				:       Tuple														=    [x for x in hidden_size_based_keys if self.is_auto(_UpperCAmelCase			)]

									if len(_UpperCAmelCase			) > 0:
												if hasattr(model.config						,						"hidden_size"			):
															A				:       int														=    model.config.hidden_size
												elif hasattr(model.config						,						"hidden_sizes"			):
															# if there are many hidden sizes pick the largest one
															A				:       Optional[Any]														=    max(model.config.hidden_sizes			)
												else:
															raise ValueError(
															    "The model\'s config file has neither `hidden_size` nor `hidden_sizes` entry, "
															    "therefore it\'s not possible to automatically fill out the following `auto` entries "
															    F"""in the DeepSpeed config file: {hidden_size_auto_keys}. You can fix that by replacing """
															    "`auto` values for these keys with an integer value of your choice."			)

												self.fill_only("zero_optimization.reduce_bucket_size"						,						hidden_size * hidden_size			)
												if self.is_zeroa():
															# automatically assign the optimal config values based on model config
															self.fill_only("zero_optimization.stage3_prefetch_bucket_size"						,						0.9 * hidden_size * hidden_size			)
															self.fill_only("zero_optimization.stage3_param_persistence_threshold"						,						10 * hidden_size			)

        # scheduler
									self.fill_match("scheduler.params.total_num_steps"						,						_UpperCAmelCase						,						"num_training_steps (calculated)"			)
									self.fill_match("scheduler.params.warmup_num_steps"						,						args.get_warmup_steps(_UpperCAmelCase			)						,						"warmup_steps"			)

									if len(self.mismatches			) > 0:
												A				:       Optional[Any]														=    "\n".join(self.mismatches			)
												raise ValueError(
												    "Please correct the following DeepSpeed config values that mismatch TrainingArguments"
												    F""" values:\n{mismatches}\nThe easiest method is to set these DeepSpeed config values to \'auto\'."""			)







__SCREAMING_SNAKE_CASE     =	None




def 						UpperCAmelCase (    _lowerCamelCase     ):
			global _hf_deepspeed_config_weak_ref
			# will go away automatically when HfDeepSpeedConfig is destroyed (when TrainingArguments is destroyed)
			A				:       Dict														=    weakref.ref(_lowerCamelCase     )




def 						UpperCAmelCase (    ):
			global _hf_deepspeed_config_weak_ref
			A				:       Union[str, Any]														=    None




def 						UpperCAmelCase (    ):

			if _hf_deepspeed_config_weak_ref is not None and _hf_deepspeed_config_weak_ref() is not None:
						return _hf_deepspeed_config_weak_ref().is_zeroa()
			else:
						return False




def 						UpperCAmelCase (    ):

			if _hf_deepspeed_config_weak_ref is not None and _hf_deepspeed_config_weak_ref() is not None:
						return _hf_deepspeed_config_weak_ref().config
			else:
						return None




def 						UpperCAmelCase (    _lowerCamelCase			,     _lowerCamelCase			,     _lowerCamelCase			,     _lowerCamelCase			,     _lowerCamelCase     ):
			from accelerate.utils import DummyOptim, DummyScheduler

			A				:       str														=    hf_deepspeed_config.config

			# Optimizer + Scheduler
			# Currently supported combos:
			# 1. DS scheduler + DS optimizer: Yes
			# 2. HF scheduler + HF optimizer: Yes
			# 3. DS scheduler + HF optimizer: Yes
			# 4. HF scheduler + DS optimizer: No
			#
			# Unless Offload is enabled in which case it's:
			# 1. DS scheduler + DS optimizer: Yes
			# 2. HF scheduler + HF optimizer: Mostly*
			# 3. DS scheduler + HF optimizer: Mostly*
			# 4. HF scheduler + DS optimizer: No
			#
			# Mostly*: All non-native DeepSpeed optimizers that have both CPU and GPU implementation should work (except LAMB)

			A				:       Union[str, Any]														=    None
			if "optimizer" in config:
						if args.adafactor:
									raise ValueError(
									    "--adafactor was passed, but also found `optimizer` configured in the DeepSpeed config. "
									    "Only one optimizer can be configured."     )
						A				:       List[Any]														=    DummyOptim(params=_lowerCamelCase     )
			else:
						if hf_deepspeed_config.is_offload():
									logger.info(
									    "Detected ZeRO Offload and non-DeepSpeed optimizers: This combination should work as long as the"
									    " custom optimizer has both CPU and GPU implementation (except LAMB)"     )

						# ds supports Adam, OneBitAdam, and Lamb optimizers and can import other optimizers from torch.
						# But trainer uses AdamW by default.
						A				:       Tuple														=    trainer.create_optimizer()
						# To use other optimizers requires voiding warranty with: `zero_allow_untested_optimizer`
						A				:       Union[str, Any]														=    True

			A				:       Optional[int]														=    None
			if "scheduler" in config:
						A				:       Optional[int]														=    DummyScheduler(_lowerCamelCase     )
			else:
						if isinstance(_lowerCamelCase			,     _lowerCamelCase     ):
									raise ValueError(
									    "Found `optimizer` configured in the DeepSpeed config, but no `scheduler`. "
									    "Please configure a scheduler in the DeepSpeed config."     )
						A				:       Dict														=    trainer.create_scheduler(num_training_steps=_lowerCamelCase			,     optimizer=_lowerCamelCase     )

			return optimizer, lr_scheduler




def 						UpperCAmelCase (    _lowerCamelCase			,     _lowerCamelCase			,     _lowerCamelCase=False     ):
			from deepspeed.utils import logger as ds_logger

			A				:       int														=    trainer.model
			A				:       List[str]														=    trainer.args

			A				:       Any														=    trainer.accelerator.state.deepspeed_plugin.hf_ds_config

			# resume config update - some bits like `model` and `num_training_steps` only become available during train
			hf_deepspeed_config.trainer_config_finalize(_lowerCamelCase			,     _lowerCamelCase			,     _lowerCamelCase     )

			# set the Deepspeed log level consistent with the Trainer
			ds_logger.setLevel(args.get_process_log_level()     )

			if inference:
						# only Z3 makes sense for the inference
						if not hf_deepspeed_config.is_zeroa():
									raise ValueError("ZeRO inference only makes sense with ZeRO Stage 3 - please adjust your config"     )

						# in case the training config is re-used for inference
						hf_deepspeed_config.del_config_sub_tree("optimizer"     )
						hf_deepspeed_config.del_config_sub_tree("lr_scheduler"     )
						A							,			A				:       List[Any]														=    None, None
						A				:       int														=    None
			else:
						A				:       Tuple														=    None  # important for when deepspeed_init is used as re-init
						A				:       Optional[int]														=    list(filter(lambda _lowerCamelCase     : p.requires_grad			,     model.parameters()     )     )
						A							,			A				:       str														=    deepspeed_optim_sched(
						    _lowerCamelCase			,     _lowerCamelCase			,     _lowerCamelCase			,     _lowerCamelCase			,     _lowerCamelCase     )

			# keep for quick debug:
			# from pprint import pprint; pprint(config)

			return optimizer, lr_scheduler






def 						UpperCAmelCase (    _lowerCamelCase			,     _lowerCamelCase     ):
			import glob

			A				:       Union[str, Any]														=    sorted(glob.glob(f"""{checkpoint_path}/global_step*"""     )     )

			if len(_lowerCamelCase     ) > 0:
						logger.info(f"""Attempting to resume from {checkpoint_path}"""     )
						# this magically updates self.optimizer and self.lr_scheduler
						A							,			A				:       List[Any]														=    deepspeed_engine.load_checkpoint(
						    _lowerCamelCase			,     load_optimizer_states=_lowerCamelCase			,     load_lr_scheduler_states=_lowerCamelCase     )
						if load_path is None:
									raise ValueError(f"""[deepspeed] failed to resume from checkpoint {checkpoint_path}"""     )
			else:
						raise ValueError(f"""Can\'t find a valid checkpoint at {checkpoint_path}"""     )



import random
import unittest

import torch

from diffusers import IFImgaImgPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin







@skip_mps
class lowerCAmelCase_		(					lowerCamelCase__   ,		lowerCamelCase__   ,		unittest.TestCase			):

 '''simple docstring'''

 __snake_case				=    IFImgaImgPipeline
 __snake_case				=    TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"width", "height"}
 __snake_case				=    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 __snake_case				=    PipelineTesterMixin.required_optional_params - {"latents"}



 def 		UpperCamelCase__	(      self   ):
       return self._get_dummy_components()



 def 		UpperCamelCase__	(      self						, _UpperCAmelCase						, _UpperCAmelCase=0   ):
       if str(_UpperCAmelCase   ).startswith('''mps'''   ):
             snake_case_			= torch.manual_seed(_UpperCAmelCase   )
       else:
             snake_case_			= torch.Generator(device=_UpperCAmelCase   ).manual_seed(_UpperCAmelCase   )

       snake_case_			= floats_tensor((1, 3, 32, 32)						, rng=random.Random(_UpperCAmelCase   )   ).to(_UpperCAmelCase   )

       snake_case_			= {
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''image''': image,
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''output_type''': '''numpy''',
       }

       return inputs



 def 		UpperCamelCase__	(      self   ):
       self._test_save_load_optional_components()



 @unittest.skipIf(
     torch_device != '''cuda''' or not is_xformers_available()						, reason='''XFormers attention is only available with CUDA and `xformers` installed'''						, )
 def 		UpperCamelCase__	(      self   ):
       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3   )



 @unittest.skipIf(torch_device != '''cuda'''						, reason='''float16 requires CUDA'''   )
 def 		UpperCamelCase__	(      self   ):
       # Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
       super().test_save_load_floataa(expected_max_diff=1E-1   )



 @unittest.skipIf(torch_device != '''cuda'''						, reason='''float16 requires CUDA'''   )
 def 		UpperCamelCase__	(      self   ):
       super().test_floataa_inference(expected_max_diff=1E-1   )



 def 		UpperCamelCase__	(      self   ):
       self._test_attention_slicing_forward_pass(expected_max_diff=1E-2   )



 def 		UpperCamelCase__	(      self   ):
       self._test_save_load_local()





 def 		UpperCamelCase__	(      self   ):
       self._test_inference_batch_single_identical(
           expected_max_diff=1E-2						, )