
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__				:      Tuple    		=     logging.get_logger(__name__)







def 				UpperCAmelCase_(    ):



    """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     a_       (			a__							):



       """simple docstring"""






       def __init__(		self				,       _lowerCamelCase   )					->List[Any]:
           # set global weakref object
           set_hf_deepspeed_config(self   )
           dep_version_check('''accelerate'''   )
           dep_version_check('''deepspeed'''   )
           super().__init__(_lowerCamelCase   )






class     a_       (			a__							):



       """simple docstring"""






       def __init__(		self				,       _lowerCamelCase   )					->Optional[Any]:
           super().__init__(_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				int				      = None
           SCREAMING_SNAKE_CASE				:				List[str]				      = []






       def 	__lowerCAmelCase     (		self   )					->Dict:
           if self._dtype is None:
               raise ValueError('''trainer_config_process() wasn\'t called yet to tell dtype'''   )
           return self._dtype






       def 	__lowerCAmelCase     (		self				,       _lowerCamelCase   )					->List[str]:
           SCREAMING_SNAKE_CASE				:				Dict				      = self.get_value(_lowerCamelCase   )
           if val is None:
               return False
           else:
               return val == "auto"






       def 	__lowerCAmelCase     (		self				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase=None				,       _lowerCamelCase=True   )					->str:
           SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE				:				List[Any]				      = self.find_config_node(_lowerCamelCase   )
           if config is None:
               return

           if config.get(_lowerCamelCase   ) == "auto":
               SCREAMING_SNAKE_CASE				:				List[str]				      = hf_val
               return

           if not must_match:
               return

           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = 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}"""   )

       __SCREAMING_SNAKE_CASE       :					Optional[Any]						=  partialmethod(a__				,    must_match=a__							)






       def 	__lowerCAmelCase     (		self				,       _lowerCamelCase   )					->List[Any]:
           SCREAMING_SNAKE_CASE				:				Any				      = 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:
               SCREAMING_SNAKE_CASE				:				Optional[Any]				      = '''apex''' if args.fpaa_backend == '''apex''' else '''amp'''
           else:
               SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = None

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






       def 	__lowerCAmelCase     (		self				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase   )					->Tuple:
           SCREAMING_SNAKE_CASE				:				int				      = [
               '''zero_optimization.reduce_bucket_size''',
               '''zero_optimization.stage3_prefetch_bucket_size''',
               '''zero_optimization.stage3_param_persistence_threshold''',
           ]
           SCREAMING_SNAKE_CASE				:				Any				      = [x for x in hidden_size_based_keys if self.is_auto(_lowerCamelCase   )]

           if len(_lowerCamelCase   ) > 0:
               if hasattr(model.config				,       '''hidden_size'''   ):
                   SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = model.config.hidden_size
               elif hasattr(model.config				,       '''hidden_sizes'''   ):
                   # if there are many hidden sizes pick the largest one
                   SCREAMING_SNAKE_CASE				:				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'''				,       _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:
               SCREAMING_SNAKE_CASE				:				Union[str, 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'."""   )





a__				:      str    		=     None







def 				UpperCAmelCase_(    a__   ):



    """simple docstring"""



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







def 				UpperCAmelCase_(    ):



    """simple docstring"""



    global _hf_deepspeed_config_weak_ref
    SCREAMING_SNAKE_CASE				:				str				      = None







def 				UpperCAmelCase_(    ):



    """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 				UpperCAmelCase_(    ):



    """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 				UpperCAmelCase_(    a__ ,						a__ ,						a__ ,						a__ ,						a__   ):



    """simple docstring"""



    from accelerate.utils import DummyOptim, DummyScheduler

    SCREAMING_SNAKE_CASE				:				List[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)

    SCREAMING_SNAKE_CASE				:				List[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.'''   )
        SCREAMING_SNAKE_CASE				:				str				      = 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.
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = trainer.create_optimizer()
        # To use other optimizers requires voiding warranty with: `zero_allow_untested_optimizer`
        SCREAMING_SNAKE_CASE				:				str				      = True

    SCREAMING_SNAKE_CASE				:				List[Any]				      = None
    if "scheduler" in config:
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = 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.'''   )
        SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = trainer.create_scheduler(num_training_steps=a__ ,						optimizer=a__   )

    return optimizer, lr_scheduler







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__=False   ):



    """simple docstring"""



    from deepspeed.utils import logger as ds_logger

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

    SCREAMING_SNAKE_CASE				:				List[str]				      = 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'''   )
        SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE				:				Tuple				      = None, None
        SCREAMING_SNAKE_CASE				:				int				      = None
    else:
        SCREAMING_SNAKE_CASE				:				int				      = None  # important for when deepspeed_init is used as re-init
        SCREAMING_SNAKE_CASE				:				str				      = list(filter(lambda a__   : p.requires_grad ,						model.parameters()   )   )
        SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE				:				Dict				      = deepspeed_optim_sched(
            a__ ,						a__ ,						a__ ,						a__ ,						a__   )

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

    return optimizer, lr_scheduler







def 				UpperCAmelCase_(    a__ ,						a__   ):



    """simple docstring"""



    import glob

    SCREAMING_SNAKE_CASE				:				Tuple				      = 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
        SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE				:				int				      = 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}"""   )


import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

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


enable_full_determinism()






@skip_mps
class     a_       (			a__				,    a__				,    a__				,    unittest.TestCase							):



       """simple docstring"""
       __SCREAMING_SNAKE_CASE       :					List[Any]						=  StableDiffusionModelEditingPipeline
       __SCREAMING_SNAKE_CASE       :					Optional[int]						=  TEXT_TO_IMAGE_PARAMS
       __SCREAMING_SNAKE_CASE       :					int						=  TEXT_TO_IMAGE_BATCH_PARAMS
       __SCREAMING_SNAKE_CASE       :					Any						=  TEXT_TO_IMAGE_IMAGE_PARAMS
       __SCREAMING_SNAKE_CASE       :					List[Any]						=  TEXT_TO_IMAGE_IMAGE_PARAMS






       def 	__lowerCAmelCase     (		self   )					->int:
           torch.manual_seed(0   )
           SCREAMING_SNAKE_CASE				:				Dict				      = 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				,       )
           SCREAMING_SNAKE_CASE				:				Dict				      = DDIMScheduler()
           torch.manual_seed(0   )
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = 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				,       )
           torch.manual_seed(0   )
           SCREAMING_SNAKE_CASE				:				List[str]				      = CLIPTextConfig(
               bos_token_id=0				,       eos_token_id=2				,       hidden_size=32				,       intermediate_size=37				,       layer_norm_eps=1e-05				,       num_attention_heads=4				,       num_hidden_layers=5				,       pad_token_id=1				,       vocab_size=1000				,       )
           SCREAMING_SNAKE_CASE				:				Dict				      = CLIPTextModel(_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Tuple				      = CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''   )

           SCREAMING_SNAKE_CASE				:				List[Any]				      = {
               '''unet''': unet,
               '''scheduler''': scheduler,
               '''vae''': vae,
               '''text_encoder''': text_encoder,
               '''tokenizer''': tokenizer,
               '''safety_checker''': None,
               '''feature_extractor''': None,
           }
           return components






       def 	__lowerCAmelCase     (		self				,       _lowerCamelCase				,       _lowerCamelCase=0   )					->Optional[int]:
           SCREAMING_SNAKE_CASE				:				Any				      = torch.manual_seed(_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = {
               '''prompt''': '''A field of roses''',
               '''generator''': generator,
               # Setting height and width to None to prevent OOMs on CPU.
               '''height''': None,
               '''width''': None,
               '''num_inference_steps''': 2,
               '''guidance_scale''': 6.0,
               '''output_type''': '''numpy''',
           }
           return inputs






       def 	__lowerCAmelCase     (		self   )					->Optional[int]:
           SCREAMING_SNAKE_CASE				:				Tuple				      = '''cpu'''  # ensure determinism for the device-dependent torch.Generator
           SCREAMING_SNAKE_CASE				:				Any				      = self.get_dummy_components()
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = StableDiffusionModelEditingPipeline(**_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				List[Any]				      = sd_pipe.to(_lowerCamelCase   )
           sd_pipe.set_progress_bar_config(disable=_lowerCamelCase   )

           SCREAMING_SNAKE_CASE				:				Tuple				      = self.get_dummy_inputs(_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Tuple				      = sd_pipe(**_lowerCamelCase   ).images
           SCREAMING_SNAKE_CASE				:				Any				      = image[0, -3:, -3:, -1]
           assert image.shape == (1, 64, 64, 3)

           SCREAMING_SNAKE_CASE				:				List[str]				      = np.array([0.4_7_5_5, 0.5_1_3_2, 0.4_9_7_6, 0.3_9_0_4, 0.3_5_5_4, 0.4_7_6_5, 0.5_1_3_9, 0.5_1_5_8, 0.4_8_8_9]   )

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






       def 	__lowerCAmelCase     (		self   )					->List[Any]:
           SCREAMING_SNAKE_CASE				:				Dict				      = '''cpu'''  # ensure determinism for the device-dependent torch.Generator
           SCREAMING_SNAKE_CASE				:				Tuple				      = self.get_dummy_components()
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = StableDiffusionModelEditingPipeline(**_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				str				      = sd_pipe.to(_lowerCamelCase   )
           sd_pipe.set_progress_bar_config(disable=_lowerCamelCase   )

           SCREAMING_SNAKE_CASE				:				Any				      = self.get_dummy_inputs(_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				List[str]				      = '''french fries'''
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = sd_pipe(**_lowerCamelCase				,       negative_prompt=_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Any				      = output.images
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = image[0, -3:, -3:, -1]

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

           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = np.array([0.4_9_9_2, 0.5_1_0_1, 0.5_0_0_4, 0.3_9_4_9, 0.3_6_0_4, 0.4_7_3_5, 0.5_2_1_6, 0.5_2_0_4, 0.4_9_1_3]   )

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






       def 	__lowerCAmelCase     (		self   )					->Union[str, Any]:
           SCREAMING_SNAKE_CASE				:				List[str]				      = '''cpu'''  # ensure determinism for the device-dependent torch.Generator
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = self.get_dummy_components()
           SCREAMING_SNAKE_CASE				:				List[Any]				      = EulerAncestralDiscreteScheduler(
               beta_start=0.0_0_0_8_5				,       beta_end=0.0_1_2				,       beta_schedule='''scaled_linear'''   )
           SCREAMING_SNAKE_CASE				:				List[Any]				      = StableDiffusionModelEditingPipeline(**_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = sd_pipe.to(_lowerCamelCase   )
           sd_pipe.set_progress_bar_config(disable=_lowerCamelCase   )

           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = self.get_dummy_inputs(_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = sd_pipe(**_lowerCamelCase   ).images
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = image[0, -3:, -3:, -1]

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

           SCREAMING_SNAKE_CASE				:				str				      = np.array([0.4_7_4_7, 0.5_3_7_2, 0.4_7_7_9, 0.4_9_8_2, 0.5_5_4_3, 0.4_8_1_6, 0.5_2_3_8, 0.4_9_0_4, 0.5_0_2_7]   )

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






       def 	__lowerCAmelCase     (		self   )					->Any:
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = '''cpu'''  # ensure determinism for the device-dependent torch.Generator
           SCREAMING_SNAKE_CASE				:				List[Any]				      = self.get_dummy_components()
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = PNDMScheduler()
           SCREAMING_SNAKE_CASE				:				Tuple				      = StableDiffusionModelEditingPipeline(**_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = sd_pipe.to(_lowerCamelCase   )
           sd_pipe.set_progress_bar_config(disable=_lowerCamelCase   )

           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = self.get_dummy_inputs(_lowerCamelCase   )
           # the pipeline does not expect pndm so test if it raises error.
           with self.assertRaises(_lowerCamelCase   ):
               SCREAMING_SNAKE_CASE				:				Dict				      = sd_pipe(**_lowerCamelCase   ).images






       def 	__lowerCAmelCase     (		self   )					->Dict:
           super().test_inference_batch_single_identical(expected_max_diff=5e-3   )






       def 	__lowerCAmelCase     (		self   )					->List[Any]:
           super().test_attention_slicing_forward_pass(expected_max_diff=5e-3   )






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



       """simple docstring"""






       def 	__lowerCAmelCase     (		self   )					->Dict:
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()






       def 	__lowerCAmelCase     (		self				,       _lowerCamelCase=0   )					->Dict:
           SCREAMING_SNAKE_CASE				:				str				      = torch.manual_seed(_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				List[Any]				      = {
               '''prompt''': '''A field of roses''',
               '''generator''': generator,
               '''num_inference_steps''': 3,
               '''guidance_scale''': 7.5,
               '''output_type''': '''numpy''',
           }
           return inputs






       def 	__lowerCAmelCase     (		self   )					->List[str]:
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = '''CompVis/stable-diffusion-v1-4'''
           SCREAMING_SNAKE_CASE				:				List[str]				      = StableDiffusionModelEditingPipeline.from_pretrained(_lowerCamelCase				,       safety_checker=_lowerCamelCase   )
           pipe.to(_lowerCamelCase   )
           pipe.set_progress_bar_config(disable=_lowerCamelCase   )
           pipe.enable_attention_slicing()

           SCREAMING_SNAKE_CASE				:				Any				      = self.get_inputs()
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = pipe(**_lowerCamelCase   ).images
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = image[0, -3:, -3:, -1].flatten()

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

           SCREAMING_SNAKE_CASE				:				str				      = np.array(
               [0.6_7_4_9_4_9_6, 0.6_3_8_6_4_5_3, 0.5_1_4_4_3_2_6_7, 0.6_6_0_9_4_9_0_5, 0.6_1_9_2_1_2_1_5, 0.5_4_9_1_3_3_2, 0.5_7_4_4_4_1_7, 0.5_8_0_7_5_1_0_6, 0.5_1_7_4_6_5_8]   )

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

           # make sure image changes after editing
           pipe.edit_model('''A pack of roses'''				,       '''A pack of blue roses'''   )

           SCREAMING_SNAKE_CASE				:				Any				      = pipe(**_lowerCamelCase   ).images
           SCREAMING_SNAKE_CASE				:				List[str]				      = image[0, -3:, -3:, -1].flatten()

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

           assert np.abs(expected_slice - image_slice   ).max() > 1e-1






       def 	__lowerCAmelCase     (		self   )					->str:
           torch.cuda.empty_cache()
           torch.cuda.reset_max_memory_allocated()
           torch.cuda.reset_peak_memory_stats()

           SCREAMING_SNAKE_CASE				:				Optional[int]				      = '''CompVis/stable-diffusion-v1-4'''
           SCREAMING_SNAKE_CASE				:				Any				      = DDIMScheduler.from_pretrained(_lowerCamelCase				,       subfolder='''scheduler'''   )
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = StableDiffusionModelEditingPipeline.from_pretrained(
               _lowerCamelCase				,       scheduler=_lowerCamelCase				,       safety_checker=_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Dict				      = pipe.to(_lowerCamelCase   )
           pipe.set_progress_bar_config(disable=_lowerCamelCase   )
           pipe.enable_attention_slicing(1   )
           pipe.enable_sequential_cpu_offload()

           SCREAMING_SNAKE_CASE				:				List[Any]				      = self.get_inputs()
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = pipe(**_lowerCamelCase   )

           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = torch.cuda.max_memory_allocated()
           # make sure that less than 4.4 GB is allocated
           assert mem_bytes < 4.4 * 10**9
