




'''simple docstring'''
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

A							:      str 	=       logging.get_logger(__name__)
def 				snake_case_  (   ):



    """simple docstring"""




    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        SCREAMING_SNAKE_CASE(				__A     ):

       def __init__(    self					,		lowerCamelCase__     )    ->					Tuple:

           """simple docstring"""



           set_hf_deepspeed_config(self     )
           dep_version_check("""accelerate"""     )
           dep_version_check("""deepspeed"""     )
           super().__init__(lowerCamelCase__     )






class        SCREAMING_SNAKE_CASE(				__A     ):

       def __init__(    self					,		lowerCamelCase__     )    ->					Optional[int]:

           """simple docstring"""



           super().__init__(lowerCamelCase__     )
           __lowercase							=						None
           __lowercase							=						[]

       def     snake_case__			(    self     )    ->					List[str]:

           """simple docstring"""



           if self._dtype is None:
               raise ValueError("""trainer_config_process() wasn't called yet to tell dtype"""     )
           return self._dtype

       def     snake_case__			(    self					,		lowerCamelCase__     )    ->					Optional[int]:

           """simple docstring"""



           __lowercase							=						self.get_value(lowerCamelCase__     )
           if val is None:
               return False
           else:
               return val == "auto"

       def     snake_case__			(    self					,		lowerCamelCase__					,		lowerCamelCase__					,		lowerCamelCase__=None					,		lowerCamelCase__=True     )    ->					Any:

           """simple docstring"""



           __lowercase				,__lowercase							=						self.find_config_node(lowerCamelCase__     )
           if config is None:
               return

           if config.get(lowerCamelCase__     ) == "auto":
               __lowercase							=						hf_val
               return

           if not must_match:
               return

           __lowercase							=						config.get(lowerCamelCase__     )
           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}'     )

       snake_case_ :     Union[str, Any]    =					partialmethod(__A   ,      must_match=__A     )

       def     snake_case__			(    self					,		lowerCamelCase__     )    ->					Optional[int]:

           """simple docstring"""



           __lowercase							=						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"""					,		lowerCamelCase__					,		"""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:
               __lowercase							=						"""apex""" if args.fpaa_backend == """apex""" else """amp"""
           else:
               __lowercase							=						None

           if args.save_on_each_node:
               # deepspeed uses shared storage by default. Let's override this setting if save_on_each_node == True
               __lowercase							=						self.config.get("""checkpoint"""					,		{}     )
               __lowercase							=						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"""     ):
               __lowercase							=						torch.bfloataa
           elif self.is_false("""fp16.enabled"""     ):
               __lowercase							=						torch.floataa
           else:
               __lowercase							=						torch.floataa




       def     snake_case__			(    self					,		lowerCamelCase__					,		lowerCamelCase__					,		lowerCamelCase__     )    ->					List[str]:

           """simple docstring"""



           __lowercase							=						[
               """zero_optimization.reduce_bucket_size""",
               """zero_optimization.stage3_prefetch_bucket_size""",
               """zero_optimization.stage3_param_persistence_threshold""",
           ]
           __lowercase							=						[x for x in hidden_size_based_keys if self.is_auto(lowerCamelCase__     )]

           if len(lowerCamelCase__     ) > 0:
               if hasattr(model.config					,		"""hidden_size"""     ):
                   __lowercase							=						model.config.hidden_size
               elif hasattr(model.config					,		"""hidden_sizes"""     ):
                   # if there are many hidden sizes pick the largest one
                   __lowercase							=						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"""					,		lowerCamelCase__					,		"""num_training_steps (calculated)"""     )
           self.fill_match("""scheduler.params.warmup_num_steps"""					,		args.get_warmup_steps(lowerCamelCase__     )					,		"""warmup_steps"""     )

           if len(self.mismatches     ) > 0:
               __lowercase							=						"""\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\'.'     )

A							:      Tuple 	=       None
def 				snake_case_  (   a__		:	Tuple			):



    """simple docstring"""




    global _hf_deepspeed_config_weak_ref
    # will go away automatically when HfDeepSpeedConfig is destroyed (when TrainingArguments is destroyed)
    __lowercase							=						weakref.ref(a__			)
def 				snake_case_  (   ):



    """simple docstring"""




    global _hf_deepspeed_config_weak_ref
    __lowercase							=						None
def 				snake_case_  (   ):



    """simple docstring"""




    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 				snake_case_  (   ):



    """simple docstring"""




    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 				snake_case_  (   a__		:	Tuple    ,a__		:	Optional[int]    ,a__		:	Tuple    ,a__		:	Optional[Any]    ,a__		:	Optional[int]			):



    """simple docstring"""




    from accelerate.utils import DummyOptim, DummyScheduler

    __lowercase							=						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)

    __lowercase							=						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."""			)
        __lowercase							=						DummyOptim(params=a__			)
    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.
        __lowercase							=						trainer.create_optimizer()
        # To use other optimizers requires voiding warranty with: `zero_allow_untested_optimizer`
        __lowercase							=						True

    __lowercase							=						None
    if "scheduler" in config:
        __lowercase							=						DummyScheduler(a__			)
    else:
        if isinstance(a__    ,a__			):
            raise ValueError(
                """Found `optimizer` configured in the DeepSpeed config, but no `scheduler`. """
                """Please configure a scheduler in the DeepSpeed config."""			)
        __lowercase							=						trainer.create_scheduler(num_training_steps=a__    ,optimizer=a__			)

    return optimizer, lr_scheduler
def 				snake_case_  (   a__		:	Dict    ,a__		:	List[Any]    ,a__		:	List[str]=False			):



    """simple docstring"""




    from deepspeed.utils import logger as ds_logger

    __lowercase							=						trainer.model
    __lowercase							=						trainer.args

    __lowercase							=						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(a__    ,a__    ,a__			)

    # 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"""			)
        __lowercase				,__lowercase							=						None, None
        __lowercase							=						None
    else:
        __lowercase							=						None  # important for when deepspeed_init is used as re-init
        __lowercase							=						list(filter(lambda a__			: p.requires_grad    ,model.parameters()			)			)
        __lowercase				,__lowercase							=						deepspeed_optim_sched(
            a__    ,a__    ,a__    ,a__    ,a__			)

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

    return optimizer, lr_scheduler



def 				snake_case_  (   a__		:	Tuple    ,a__		:	Dict			):



    """simple docstring"""




    import glob

    __lowercase							=						sorted(glob.glob(f'{checkpoint_path}/global_step*'			)			)

    if len(a__			) > 0:
        logger.info(f'Attempting to resume from {checkpoint_path}'			)
        # this magically updates self.optimizer and self.lr_scheduler
        __lowercase				,__lowercase							=						deepspeed_engine.load_checkpoint(
            a__    ,load_optimizer_states=a__    ,load_lr_scheduler_states=a__			)
        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}'			)










'''simple docstring'''
import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    require_torch_gpu,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






class        SCREAMING_SNAKE_CASE(				__A   ,      __A   ,      unittest.TestCase     ):
       snake_case_ :     str    =					StableDiffusionParadigmsPipeline
       snake_case_ :     Optional[Any]    =					TEXT_TO_IMAGE_PARAMS
       snake_case_ :     Tuple    =					TEXT_TO_IMAGE_BATCH_PARAMS
       snake_case_ :     Optional[Any]    =					TEXT_TO_IMAGE_IMAGE_PARAMS
       snake_case_ :     str    =					TEXT_TO_IMAGE_IMAGE_PARAMS

       def     snake_case__			(    self     )    ->					int:

           """simple docstring"""



           torch.manual_seed(0     )
           __lowercase							=						UNetaDConditionModel(
               block_out_channels=(32, 64)					,		layers_per_block=2					,		sample_size=32					,		in_channels=4					,		out_channels=4					,		down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")					,		up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")					,		cross_attention_dim=32					,		attention_head_dim=(2, 4)					,		use_linear_projection=lowerCamelCase__					,		)
           __lowercase							=						DDIMParallelScheduler(
               beta_start=0.0_00_85					,		beta_end=0.0_12					,		beta_schedule="""scaled_linear"""					,		clip_sample=lowerCamelCase__					,		set_alpha_to_one=lowerCamelCase__					,		)
           torch.manual_seed(0     )
           __lowercase							=						AutoencoderKL(
               block_out_channels=[32, 64]					,		in_channels=3					,		out_channels=3					,		down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]					,		up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]					,		latent_channels=4					,		sample_size=128					,		)
           torch.manual_seed(0     )
           __lowercase							=						CLIPTextConfig(
               bos_token_id=0					,		eos_token_id=2					,		hidden_size=32					,		intermediate_size=37					,		layer_norm_eps=1E-0_5					,		num_attention_heads=4					,		num_hidden_layers=5					,		pad_token_id=1					,		vocab_size=1000					,		hidden_act="""gelu"""					,		projection_dim=512					,		)
           __lowercase							=						CLIPTextModel(lowerCamelCase__     )
           __lowercase							=						CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     )

           __lowercase							=						{
               """unet""": unet,
               """scheduler""": scheduler,
               """vae""": vae,
               """text_encoder""": text_encoder,
               """tokenizer""": tokenizer,
               """safety_checker""": None,
               """feature_extractor""": None,
           }
           return components

       def     snake_case__			(    self					,		lowerCamelCase__					,		lowerCamelCase__=0     )    ->					Any:

           """simple docstring"""



           if str(lowerCamelCase__     ).startswith("""mps"""     ):
               __lowercase							=						torch.manual_seed(lowerCamelCase__     )
           else:
               __lowercase							=						torch.Generator(device=lowerCamelCase__     ).manual_seed(lowerCamelCase__     )
           __lowercase							=						{
               """prompt""": """a photograph of an astronaut riding a horse""",
               """generator""": generator,
               """num_inference_steps""": 10,
               """guidance_scale""": 6.0,
               """output_type""": """numpy""",
               """parallel""": 3,
               """debug""": True,
           }
           return inputs

       def     snake_case__			(    self     )    ->					List[str]:

           """simple docstring"""



           __lowercase							=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
           __lowercase							=						self.get_dummy_components()
           __lowercase							=						StableDiffusionParadigmsPipeline(**lowerCamelCase__     )
           __lowercase							=						sd_pipe.to(lowerCamelCase__     )
           sd_pipe.set_progress_bar_config(disable=lowerCamelCase__     )

           __lowercase							=						self.get_dummy_inputs(lowerCamelCase__     )
           __lowercase							=						sd_pipe(**lowerCamelCase__     ).images
           __lowercase							=						image[0, -3:, -3:, -1]
           assert image.shape == (1, 64, 64, 3)

           __lowercase							=						np.array([0.47_73, 0.54_17, 0.47_23, 0.49_25, 0.56_31, 0.47_52, 0.52_40, 0.49_35, 0.50_23]     )

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

       def     snake_case__			(    self     )    ->					str:

           """simple docstring"""



           __lowercase							=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
           __lowercase							=						self.get_dummy_components()
           torch.manual_seed(0     )
           __lowercase							=						DDPMParallelScheduler()
           torch.manual_seed(0     )
           __lowercase							=						StableDiffusionParadigmsPipeline(**lowerCamelCase__     )
           __lowercase							=						sd_pipe.to(lowerCamelCase__     )
           sd_pipe.set_progress_bar_config(disable=lowerCamelCase__     )

           __lowercase							=						self.get_dummy_inputs(lowerCamelCase__     )
           __lowercase							=						sd_pipe(**lowerCamelCase__     ).images
           __lowercase							=						image[0, -3:, -3:, -1]
           assert image.shape == (1, 64, 64, 3)

           __lowercase							=						np.array([0.35_73, 0.44_20, 0.49_60, 0.47_99, 0.37_96, 0.38_79, 0.48_19, 0.43_65, 0.44_68]     )

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

       def     snake_case__			(    self     )    ->					Any:

           """simple docstring"""



           super().test_inference_batch_consistent(batch_sizes=[1, 2]     )

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

           """simple docstring"""



           super().test_inference_batch_single_identical(batch_size=2					,		expected_max_diff=3E-3     )




       def     snake_case__			(    self     )    ->					str:

           """simple docstring"""



           __lowercase							=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
           __lowercase							=						self.get_dummy_components()
           __lowercase							=						StableDiffusionParadigmsPipeline(**lowerCamelCase__     )
           __lowercase							=						sd_pipe.to(lowerCamelCase__     )
           sd_pipe.set_progress_bar_config(disable=lowerCamelCase__     )

           __lowercase							=						self.get_dummy_inputs(lowerCamelCase__     )
           __lowercase							=						"""french fries"""
           __lowercase							=						sd_pipe(**lowerCamelCase__					,		negative_prompt=lowerCamelCase__     )
           __lowercase							=						output.images
           __lowercase							=						image[0, -3:, -3:, -1]

           assert image.shape == (1, 64, 64, 3)

           __lowercase							=						np.array([0.47_71, 0.54_20, 0.46_83, 0.49_18, 0.56_36, 0.47_25, 0.52_30, 0.49_23, 0.50_15]     )

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






@slow
@require_torch_gpu
class        SCREAMING_SNAKE_CASE(				unittest.TestCase     ):

       def     snake_case__			(    self     )    ->					str:

           """simple docstring"""



           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()

       def     snake_case__			(    self					,		lowerCamelCase__=0     )    ->					int:

           """simple docstring"""



           __lowercase							=						torch.Generator(device=lowerCamelCase__     ).manual_seed(lowerCamelCase__     )
           __lowercase							=						{
               """prompt""": """a photograph of an astronaut riding a horse""",
               """generator""": generator,
               """num_inference_steps""": 10,
               """guidance_scale""": 7.5,
               """output_type""": """numpy""",
               """parallel""": 3,
               """debug""": True,
           }
           return inputs




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

           """simple docstring"""



           __lowercase							=						"""stabilityai/stable-diffusion-2-base"""
           __lowercase							=						DDIMParallelScheduler.from_pretrained(lowerCamelCase__					,		subfolder="""scheduler"""     )
           __lowercase							=						StableDiffusionParadigmsPipeline.from_pretrained(lowerCamelCase__					,		scheduler=lowerCamelCase__					,		safety_checker=lowerCamelCase__     )
           pipe.to(lowerCamelCase__     )
           pipe.set_progress_bar_config(disable=lowerCamelCase__     )
           pipe.enable_attention_slicing()

           __lowercase							=						self.get_inputs()
           __lowercase							=						pipe(**lowerCamelCase__     ).images
           __lowercase							=						image[0, -3:, -3:, -1].flatten()

           assert image.shape == (1, 512, 512, 3)

           __lowercase							=						np.array([0.96_22, 0.96_02, 0.97_48, 0.95_91, 0.96_30, 0.96_91, 0.96_61, 0.96_31, 0.97_41]     )

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




