import argparse
import gc
import os
import threading

import evaluate
import psutil
import torch
from datasets import load_dataset
from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig
from torch.utils.data import DataLoader
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    get_linear_schedule_with_warmup,
    set_seed,
)

from accelerate import Accelerator, DistributedType, FullyShardedDataParallelPlugin


########################################################################
# This is a fully working simple example to use Accelerate
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#   - FSDP
#
# This example also demonstrates the checkpointing and sharding capabilities
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


__lowerCamelCase     :					Tuple				  =						16
__lowerCamelCase     :					int				  =						32



def       __SCREAMING_SNAKE_CASE    (   __UpperCamelCase       :			str      )				->   int:







  """simple docstring"""





  return int(x / 2**20      )







class __snake_case							:




   def __enter__(       self     :							Any ):



     """simple docstring"""






     gc.collect()
     torch.cuda.empty_cache()
     torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
     SCREAMING_SNAKE_CASE__     				= torch.cuda.memory_allocated()
     SCREAMING_SNAKE_CASE__     				= psutil.Process()

     SCREAMING_SNAKE_CASE__     				= self.cpu_mem_used()
     SCREAMING_SNAKE_CASE__     				= True
     SCREAMING_SNAKE_CASE__     				= threading.Thread(target=self.peak_monitor_func )
     SCREAMING_SNAKE_CASE__     				= True
     peak_monitor_thread.start()
     return self




   def       __a		(       self     :							Optional[Any] ):



     """simple docstring"""






     return self.process.memory_info().rss




   def       __a		(       self     :							Any ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= -1

     while True:
       SCREAMING_SNAKE_CASE__     				= max(self.cpu_mem_used()			,      self.cpu_peak )

       # can't sleep or will not catch the peak right (this comment is here on purpose)
       # time.sleep(0.001) # 1msec

       if not self.peak_monitoring:
         break




   def __exit__(       self     :							Tuple			,      *_lowercase     :							Dict ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= False

     gc.collect()
     torch.cuda.empty_cache()
     SCREAMING_SNAKE_CASE__     				= torch.cuda.memory_allocated()
     SCREAMING_SNAKE_CASE__     				= torch.cuda.max_memory_allocated()
     SCREAMING_SNAKE_CASE__     				= bamb(self.end - self.begin )
     SCREAMING_SNAKE_CASE__     				= bamb(self.peak - self.begin )

     SCREAMING_SNAKE_CASE__     				= self.cpu_mem_used()
     SCREAMING_SNAKE_CASE__     				= bamb(self.cpu_end - self.cpu_begin )
     SCREAMING_SNAKE_CASE__     				= bamb(self.cpu_peak - self.cpu_begin )
     # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")


# For testing only
if os.environ.get('''TESTING_MOCKED_DATALOADERS''', None) == "1":
   from accelerate.test_utils.training import mocked_dataloaders

   __lowerCamelCase     :					Tuple				  =						mocked_dataloaders  # noqa: F811



def       __SCREAMING_SNAKE_CASE    (   __UpperCamelCase       :			Dict					,	__UpperCamelCase       :			Any      )				->   Dict:







  """simple docstring"""





  if os.environ.get("""TESTING_MOCKED_DATALOADERS"""					,	__UpperCamelCase      ) == "1":
    SCREAMING_SNAKE_CASE__     				= 2

  # New Code #
  # Pass the advanced FSDP settings not part of the accelerate config by creating fsdp_plugin
  SCREAMING_SNAKE_CASE__     				= FullyShardedDataParallelPlugin(
      state_dict_config=FullStateDictConfig(offload_to_cpu=__UpperCamelCase					,	ranka_only=__UpperCamelCase      )					,	optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=__UpperCamelCase					,	ranka_only=__UpperCamelCase      )					,	)

  # Initialize accelerator
  if args.with_tracking:
    SCREAMING_SNAKE_CASE__     				= Accelerator(
        cpu=args.cpu					,	mixed_precision=args.mixed_precision					,	log_with="""wandb"""					,	project_dir=args.logging_dir					,	fsdp_plugin=__UpperCamelCase					,	)
  else:
    SCREAMING_SNAKE_CASE__     				= Accelerator(fsdp_plugin=__UpperCamelCase      )
  accelerator.print(accelerator.distributed_type      )

  if hasattr(args.checkpointing_steps					,	"""isdigit"""      ):
    if args.checkpointing_steps == "epoch":
      SCREAMING_SNAKE_CASE__     				= args.checkpointing_steps
    elif args.checkpointing_steps.isdigit():
      SCREAMING_SNAKE_CASE__     				= int(args.checkpointing_steps      )
    else:
      raise ValueError(
          f"""Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."""      )
  else:
    SCREAMING_SNAKE_CASE__     				= None
  # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
  SCREAMING_SNAKE_CASE__     				= config["""lr"""]
  SCREAMING_SNAKE_CASE__     				= int(config["""num_epochs"""]      )
  SCREAMING_SNAKE_CASE__     				= int(config["""seed"""]      )
  SCREAMING_SNAKE_CASE__     				= int(config["""batch_size"""]      )

  # We need to initialize the trackers we use, and also store our configuration
  if args.with_tracking:
    SCREAMING_SNAKE_CASE__     				= vars(__UpperCamelCase      )
    accelerator.init_trackers("""fsdp_glue_no_trainer"""					,	__UpperCamelCase      )

  SCREAMING_SNAKE_CASE__     				= AutoTokenizer.from_pretrained(args.model_name_or_path      )
  SCREAMING_SNAKE_CASE__     				= load_dataset("""glue"""					,	"""mrpc"""      )
  SCREAMING_SNAKE_CASE__     				= evaluate.load("""glue"""					,	"""mrpc"""      )

  def tokenize_function(__UpperCamelCase       :			Optional[Any]      ):
    # max_length=None => use the model max length (it's actually the default)
    SCREAMING_SNAKE_CASE__     				= tokenizer(examples["""sentence1"""]					,	examples["""sentence2"""]					,	truncation=__UpperCamelCase					,	max_length=__UpperCamelCase      )
    return outputs

  # Apply the method we just defined to all the examples in all the splits of the dataset
  # starting with the main process first:
  with accelerator.main_process_first():
    SCREAMING_SNAKE_CASE__     				= datasets.map(
        __UpperCamelCase					,	batched=__UpperCamelCase					,	remove_columns=["""idx""", """sentence1""", """sentence2"""]					,	)

  # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
  # transformers library
  SCREAMING_SNAKE_CASE__     				= tokenized_datasets.rename_column("""label"""					,	"""labels"""      )

  # If the batch size is too big we use gradient accumulation
  SCREAMING_SNAKE_CASE__     				= 1
  if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
    SCREAMING_SNAKE_CASE__     				= batch_size // MAX_GPU_BATCH_SIZE
    SCREAMING_SNAKE_CASE__     				= MAX_GPU_BATCH_SIZE

  def collate_fn(__UpperCamelCase       :			str      ):
    # On TPU it's best to pad everything to the same length or training will be very slow.
    SCREAMING_SNAKE_CASE__     				= 1_28 if accelerator.distributed_type == DistributedType.TPU else None
    # When using mixed precision we want round multiples of 8/16
    if accelerator.mixed_precision == "fp8":
      SCREAMING_SNAKE_CASE__     				= 16
    elif accelerator.mixed_precision != "no":
      SCREAMING_SNAKE_CASE__     				= 8
    else:
      SCREAMING_SNAKE_CASE__     				= None

    return tokenizer.pad(
        __UpperCamelCase					,	padding="""longest"""					,	max_length=__UpperCamelCase					,	pad_to_multiple_of=__UpperCamelCase					,	return_tensors="""pt"""					,	)

  # Instantiate dataloaders.
  SCREAMING_SNAKE_CASE__     				= DataLoader(
      tokenized_datasets["""train"""]					,	shuffle=__UpperCamelCase					,	collate_fn=__UpperCamelCase					,	batch_size=__UpperCamelCase      )
  SCREAMING_SNAKE_CASE__     				= DataLoader(
      tokenized_datasets["""validation"""]					,	shuffle=__UpperCamelCase					,	collate_fn=__UpperCamelCase					,	batch_size=__UpperCamelCase      )

  set_seed(__UpperCamelCase      )

  # Instantiate the model (we build the model here so that the seed also control new weights initialization)
  SCREAMING_SNAKE_CASE__     				= AutoModelForSequenceClassification.from_pretrained(
      args.model_name_or_path					,	return_dict=__UpperCamelCase					,	low_cpu_mem_usage=__UpperCamelCase      )

  # New Code #
  # For FSDP feature, it is highly recommended and efficient to prepare the model before creating optimizer
  SCREAMING_SNAKE_CASE__     				= accelerator.prepare(__UpperCamelCase      )
  accelerator.print(__UpperCamelCase      )

  # Instantiate optimizer
  # New Code #
  # For FSDP feature, at present it doesn't support multiple parameter groups,
  # so we need to create a single parameter group for the whole model
  SCREAMING_SNAKE_CASE__     				= torch.optim.AdamW(params=model.parameters()					,	lr=__UpperCamelCase					,	weight_decay=2E-4      )

  # Instantiate scheduler
  SCREAMING_SNAKE_CASE__     				= get_linear_schedule_with_warmup(
      optimizer=__UpperCamelCase					,	num_warmup_steps=10					,	num_training_steps=(len(__UpperCamelCase      ) * num_epochs) // gradient_accumulation_steps					,	)

  # New Code #
  # For FSDP feature, prepare everything except the model as we have already prepared the model
  # before creating the optimizer
  # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
  # prepare method.
  SCREAMING_SNAKE_CASE__				,	SCREAMING_SNAKE_CASE__				,	SCREAMING_SNAKE_CASE__				,	SCREAMING_SNAKE_CASE__     				= accelerator.prepare(
      __UpperCamelCase					,	__UpperCamelCase					,	__UpperCamelCase					,	__UpperCamelCase      )

  SCREAMING_SNAKE_CASE__     				= 0

  # Potentially load in the weights and states from a previous save
  if args.resume_from_checkpoint:
    if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
      accelerator.print(f"""Resumed from checkpoint: {args.resume_from_checkpoint}"""      )
      accelerator.load_state(args.resume_from_checkpoint      )
      SCREAMING_SNAKE_CASE__     				= os.path.basename(args.resume_from_checkpoint      )
    else:
      # Get the most recent checkpoint
      SCREAMING_SNAKE_CASE__     				= [f.name for f in os.scandir(os.getcwd()      ) if f.is_dir()]
      dirs.sort(key=os.path.getctime      )
      SCREAMING_SNAKE_CASE__     				= dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
    # Extract `epoch_{i}` or `step_{i}`
    SCREAMING_SNAKE_CASE__     				= os.path.splitext(__UpperCamelCase      )[0]

    if "epoch" in training_difference:
      num_epochs -= int(training_difference.replace("""epoch_"""					,	""""""      )      )
      SCREAMING_SNAKE_CASE__     				= None
    else:
      SCREAMING_SNAKE_CASE__     				= int(training_difference.replace("""step_"""					,	""""""      )      )
      num_epochs -= resume_step // len(__UpperCamelCase      )
      # If resuming by step, we also need to know exactly how far into the DataLoader we went
      SCREAMING_SNAKE_CASE__     				= (num_epochs * len(__UpperCamelCase      )) - resume_step

    # Now we train the model
  for epoch in range(__UpperCamelCase      ):
    # New Code #
    # context manager to track the peak memory usage during the training epoch
    with TorchTracemalloc() as tracemalloc:
      model.train()
      if args.with_tracking:
        SCREAMING_SNAKE_CASE__     				= 0
      for step, batch in enumerate(__UpperCamelCase      ):
        # We need to skip steps until we reach the resumed step
        if args.resume_from_checkpoint and epoch == 0:
          if resume_step is not None and step < resume_step:
            pass
                # We could avoid this line since we set the accelerator with `device_placement=True`.
        batch.to(accelerator.device      )
        SCREAMING_SNAKE_CASE__     				= model(**__UpperCamelCase      )
        SCREAMING_SNAKE_CASE__     				= outputs.loss
        SCREAMING_SNAKE_CASE__     				= loss / gradient_accumulation_steps
        # We keep track of the loss at each epoch
        if args.with_tracking:
          total_loss += loss.detach().float()
        accelerator.backward(__UpperCamelCase      )
        if step % gradient_accumulation_steps == 0:
          optimizer.step()
          lr_scheduler.step()
          optimizer.zero_grad()
          # accelerator.print(lr_scheduler.get_lr())

        overall_step += 1

        if isinstance(__UpperCamelCase					,	__UpperCamelCase      ):
          SCREAMING_SNAKE_CASE__     				= f"""step_{overall_step}"""
          if overall_step % checkpointing_steps == 0:
            if args.output_dir is not None:
              SCREAMING_SNAKE_CASE__     				= os.path.join(args.output_dir					,	__UpperCamelCase      )
            accelerator.save_state(__UpperCamelCase      )
        # New Code #
        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
    accelerator.print("""Memory before entering the train : {}""".format(bamb(tracemalloc.begin      )      )      )
    accelerator.print("""Memory consumed at the end of the train (end-begin): {}""".format(tracemalloc.used      )      )
    accelerator.print("""Peak Memory consumed during the train (max-begin): {}""".format(tracemalloc.peaked      )      )
    accelerator.print(
        """Total Peak Memory consumed during the train (max): {}""".format(
            tracemalloc.peaked + bamb(tracemalloc.begin      )      )      )
    # Logging the peak memory usage of the GPU to the tracker
    if args.with_tracking:
      accelerator.log(
          {
              """train_total_peak_memory""": tracemalloc.peaked + bamb(tracemalloc.begin      ),
          }					,	step=__UpperCamelCase					,	)

    # New Code #
    # context manager to track the peak memory usage during the evaluation
    with TorchTracemalloc() as tracemalloc:
      model.eval()
      for step, batch in enumerate(__UpperCamelCase      ):
        # We could avoid this line since we set the accelerator with `device_placement=True`.
        batch.to(accelerator.device      )
        with torch.no_grad():
          SCREAMING_SNAKE_CASE__     				= model(**__UpperCamelCase      )
        SCREAMING_SNAKE_CASE__     				= outputs.logits.argmax(dim=-1      )
        SCREAMING_SNAKE_CASE__				,	SCREAMING_SNAKE_CASE__     				= accelerator.gather_for_metrics((predictions, batch["""labels"""])      )
        metric.add_batch(
            predictions=__UpperCamelCase					,	references=__UpperCamelCase					,	)

      SCREAMING_SNAKE_CASE__     				= metric.compute()
      # Use accelerator.print to print only on the main process.
      accelerator.print(f"""epoch {epoch}:"""					,	__UpperCamelCase      )
      if args.with_tracking:
        accelerator.log(
            {
                """accuracy""": eval_metric["""accuracy"""],
                """f1""": eval_metric["""f1"""],
                """train_loss""": total_loss.item() / len(__UpperCamelCase      ),
            }					,	step=__UpperCamelCase					,	)

      if checkpointing_steps == "epoch":
        SCREAMING_SNAKE_CASE__     				= f"""epoch_{epoch}"""
        if args.output_dir is not None:
          SCREAMING_SNAKE_CASE__     				= os.path.join(args.output_dir					,	__UpperCamelCase      )
        accelerator.save_state(__UpperCamelCase      )
        # New Code #
        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
    accelerator.print("""Memory before entering the eval : {}""".format(bamb(tracemalloc.begin      )      )      )
    accelerator.print("""Memory consumed at the end of the eval (end-begin): {}""".format(tracemalloc.used      )      )
    accelerator.print("""Peak Memory consumed during the eval (max-begin): {}""".format(tracemalloc.peaked      )      )
    accelerator.print(
        """Total Peak Memory consumed during the eval (max): {}""".format(tracemalloc.peaked + bamb(tracemalloc.begin      )      )      )
    # Logging the peak memory usage of the GPU to the tracker
    if args.with_tracking:
      accelerator.log(
          {
              """eval_total_peak_memory""": tracemalloc.peaked + bamb(tracemalloc.begin      ),
          }					,	step=__UpperCamelCase					,	)

  if args.with_tracking:
    accelerator.end_training()



def       __SCREAMING_SNAKE_CASE    (   )				->   Tuple:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__     				= argparse.ArgumentParser(description="""Simple example of training script."""      )
  parser.add_argument(
      """--mixed_precision"""					,	type=__UpperCamelCase					,	default=__UpperCamelCase					,	choices=["""no""", """fp16""", """bf16""", """fp8"""]					,	help="""Whether to use mixed precision. Choose"""
      """between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."""
      """and an Nvidia Ampere GPU."""					,	)
  parser.add_argument("""--cpu"""					,	action="""store_true"""					,	help="""If passed, will train on the CPU."""      )
  parser.add_argument(
      """--checkpointing_steps"""					,	type=__UpperCamelCase					,	default=__UpperCamelCase					,	help="""Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch."""					,	)
  parser.add_argument(
      """--resume_from_checkpoint"""					,	type=__UpperCamelCase					,	default=__UpperCamelCase					,	help="""If the training should continue from a checkpoint folder."""					,	)
  parser.add_argument(
      """--with_tracking"""					,	action="""store_true"""					,	help="""Whether to load in all available experiment trackers from the environment and use them for logging."""					,	)
  parser.add_argument(
      """--output_dir"""					,	type=__UpperCamelCase					,	default="""."""					,	help="""Optional save directory where all checkpoint folders will be stored. Default is the current working directory."""					,	)
  parser.add_argument(
      """--logging_dir"""					,	type=__UpperCamelCase					,	default="""logs"""					,	help="""Location on where to store experiment tracking logs`"""					,	)
  parser.add_argument(
      """--model_name_or_path"""					,	type=__UpperCamelCase					,	help="""Path to pretrained model or model identifier from huggingface.co/models."""					,	required=__UpperCamelCase					,	)
  SCREAMING_SNAKE_CASE__     				= parser.parse_args()
  SCREAMING_SNAKE_CASE__     				= {"""lr""": 2E-5, """num_epochs""": 3, """seed""": 42, """batch_size""": 16}
  training_function(__UpperCamelCase					,	__UpperCamelCase      )


if __name__ == "__main__":
   main()

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 __snake_case							(		lowerCamelCase_  ,						lowerCamelCase_  ,						lowerCamelCase_  ,						unittest.TestCase						):
   lowerCAmelCase_								=				StableDiffusionModelEditingPipeline
   lowerCAmelCase_								=				TEXT_TO_IMAGE_PARAMS
   lowerCAmelCase_								=				TEXT_TO_IMAGE_BATCH_PARAMS
   lowerCAmelCase_								=				TEXT_TO_IMAGE_IMAGE_PARAMS
   lowerCAmelCase_								=				TEXT_TO_IMAGE_IMAGE_PARAMS




   def       __a		(       self     :							List[str] ):



     """simple docstring"""






     torch.manual_seed(0 )
     SCREAMING_SNAKE_CASE__     				= 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__     				= DDIMScheduler()
     torch.manual_seed(0 )
     SCREAMING_SNAKE_CASE__     				= 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__     				= 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=10_00			,      )
     SCREAMING_SNAKE_CASE__     				= CLIPTextModel(_lowercase )
     SCREAMING_SNAKE_CASE__     				= CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""" )

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




   def       __a		(       self     :							Optional[int]			,      _lowercase     :							Optional[int]			,      _lowercase     :							Union[str, Any]=0 ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= torch.manual_seed(_lowercase )
     SCREAMING_SNAKE_CASE__     				= {
         """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       __a		(       self     :							Dict ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= """cpu"""  # ensure determinism for the device-dependent torch.Generator
     SCREAMING_SNAKE_CASE__     				= self.get_dummy_components()
     SCREAMING_SNAKE_CASE__     				= StableDiffusionModelEditingPipeline(**_lowercase )
     SCREAMING_SNAKE_CASE__     				= sd_pipe.to(_lowercase )
     sd_pipe.set_progress_bar_config(disable=_lowercase )

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

     SCREAMING_SNAKE_CASE__     				= np.array([0.47_55, 0.51_32, 0.49_76, 0.39_04, 0.35_54, 0.47_65, 0.51_39, 0.51_58, 0.48_89] )

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




   def       __a		(       self     :							List[str] ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= """cpu"""  # ensure determinism for the device-dependent torch.Generator
     SCREAMING_SNAKE_CASE__     				= self.get_dummy_components()
     SCREAMING_SNAKE_CASE__     				= StableDiffusionModelEditingPipeline(**_lowercase )
     SCREAMING_SNAKE_CASE__     				= sd_pipe.to(_lowercase )
     sd_pipe.set_progress_bar_config(disable=_lowercase )

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

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

     SCREAMING_SNAKE_CASE__     				= np.array([0.49_92, 0.51_01, 0.50_04, 0.39_49, 0.36_04, 0.47_35, 0.52_16, 0.52_04, 0.49_13] )

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




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



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= """cpu"""  # ensure determinism for the device-dependent torch.Generator
     SCREAMING_SNAKE_CASE__     				= self.get_dummy_components()
     SCREAMING_SNAKE_CASE__     				= EulerAncestralDiscreteScheduler(
         beta_start=0.0_00_85			,      beta_end=0.0_12			,      beta_schedule="""scaled_linear""" )
     SCREAMING_SNAKE_CASE__     				= StableDiffusionModelEditingPipeline(**_lowercase )
     SCREAMING_SNAKE_CASE__     				= sd_pipe.to(_lowercase )
     sd_pipe.set_progress_bar_config(disable=_lowercase )

     SCREAMING_SNAKE_CASE__     				= self.get_dummy_inputs(_lowercase )
     SCREAMING_SNAKE_CASE__     				= sd_pipe(**_lowercase ).images
     SCREAMING_SNAKE_CASE__     				= image[0, -3:, -3:, -1]

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

     SCREAMING_SNAKE_CASE__     				= np.array([0.47_47, 0.53_72, 0.47_79, 0.49_82, 0.55_43, 0.48_16, 0.52_38, 0.49_04, 0.50_27] )

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




   def       __a		(       self     :							List[Any] ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= """cpu"""  # ensure determinism for the device-dependent torch.Generator
     SCREAMING_SNAKE_CASE__     				= self.get_dummy_components()
     SCREAMING_SNAKE_CASE__     				= PNDMScheduler()
     SCREAMING_SNAKE_CASE__     				= StableDiffusionModelEditingPipeline(**_lowercase )
     SCREAMING_SNAKE_CASE__     				= sd_pipe.to(_lowercase )
     sd_pipe.set_progress_bar_config(disable=_lowercase )

     SCREAMING_SNAKE_CASE__     				= self.get_dummy_inputs(_lowercase )
     # the pipeline does not expect pndm so test if it raises error.
     with self.assertRaises(_lowercase ):
       SCREAMING_SNAKE_CASE__     				= sd_pipe(**_lowercase ).images




   def       __a		(       self     :							List[str] ):



     """simple docstring"""






     super().test_inference_batch_single_identical(expected_max_diff=5E-3 )




   def       __a		(       self     :							Any ):



     """simple docstring"""






     super().test_attention_slicing_forward_pass(expected_max_diff=5E-3 )







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




   def       __a		(       self     :							str ):



     """simple docstring"""






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




   def       __a		(       self     :							List[Any]			,      _lowercase     :							Dict=0 ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= torch.manual_seed(_lowercase )
     SCREAMING_SNAKE_CASE__     				= {
         """prompt""": """A field of roses""",
         """generator""": generator,
         """num_inference_steps""": 3,
         """guidance_scale""": 7.5,
         """output_type""": """numpy""",
     }
     return inputs




   def       __a		(       self     :							Tuple ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= """CompVis/stable-diffusion-v1-4"""
     SCREAMING_SNAKE_CASE__     				= StableDiffusionModelEditingPipeline.from_pretrained(_lowercase			,      safety_checker=_lowercase )
     pipe.to(_lowercase )
     pipe.set_progress_bar_config(disable=_lowercase )
     pipe.enable_attention_slicing()

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

     assert image.shape == (1, 5_12, 5_12, 3)

     SCREAMING_SNAKE_CASE__     				= np.array(
         [0.6_74_94_96, 0.6_38_64_53, 0.51_44_32_67, 0.66_09_49_05, 0.61_92_12_15, 0.5_49_13_32, 0.5_74_44_17, 0.58_07_51_06, 0.5_17_46_58] )

     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__     				= pipe(**_lowercase ).images
     SCREAMING_SNAKE_CASE__     				= image[0, -3:, -3:, -1].flatten()

     assert image.shape == (1, 5_12, 5_12, 3)

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




   def       __a		(       self     :							Any ):



     """simple docstring"""






     torch.cuda.empty_cache()
     torch.cuda.reset_max_memory_allocated()
     torch.cuda.reset_peak_memory_stats()

     SCREAMING_SNAKE_CASE__     				= """CompVis/stable-diffusion-v1-4"""
     SCREAMING_SNAKE_CASE__     				= DDIMScheduler.from_pretrained(_lowercase			,      subfolder="""scheduler""" )
     SCREAMING_SNAKE_CASE__     				= StableDiffusionModelEditingPipeline.from_pretrained(
         _lowercase			,      scheduler=_lowercase			,      safety_checker=_lowercase )
     SCREAMING_SNAKE_CASE__     				= pipe.to(_lowercase )
     pipe.set_progress_bar_config(disable=_lowercase )
     pipe.enable_attention_slicing(1 )
     pipe.enable_sequential_cpu_offload()

     SCREAMING_SNAKE_CASE__     				= self.get_inputs()
     SCREAMING_SNAKE_CASE__     				= pipe(**_lowercase )

     SCREAMING_SNAKE_CASE__     				= torch.cuda.max_memory_allocated()
     # make sure that less than 4.4 GB is allocated
     assert mem_bytes < 4.4 * 10**9
