


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
#
########################################################################


_snake_case				=			1_6
_snake_case				=			3_2







def 	A (		A__			):
     return int(x / 2**2_0			)




class  A_ :







 """simple docstring"""





 def __enter__(      self       :			Any					)     ->      Any:
      gc.collect()
      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
      UpperCAmelCase    :					int			=      torch.cuda.memory_allocated()
      UpperCAmelCase    :					int			=      psutil.Process()

      UpperCAmelCase    :					Optional[Any]			=      self.cpu_mem_used()
      UpperCAmelCase    :					Tuple			=      True
      UpperCAmelCase    :					Optional[int]			=      threading.Thread(target=self.peak_monitor_func					)
      UpperCAmelCase    :					List[Any]			=      True
      peak_monitor_thread.start()
      return self
 def 	_UpperCamelCase     (      self       :			Optional[Any]					)     ->      Union[str, Any]:
      return self.process.memory_info().rss
 def 	_UpperCamelCase     (      self       :			Any					)     ->      int:
      UpperCAmelCase    :					str			=      -1

      while True:
           UpperCAmelCase    :					List[Any]			=      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							,     *_a       :			Dict					)     ->      Dict:
      UpperCAmelCase    :					List[str]			=      False

      gc.collect()
      torch.cuda.empty_cache()
      UpperCAmelCase    :					int			=      torch.cuda.memory_allocated()
      UpperCAmelCase    :					str			=      torch.cuda.max_memory_allocated()
      UpperCAmelCase    :					int			=      bamb(self.end - self.begin					)
      UpperCAmelCase    :					Tuple			=      bamb(self.peak - self.begin					)

      UpperCAmelCase    :					Optional[int]			=      self.cpu_mem_used()
      UpperCAmelCase    :					List[Any]			=      bamb(self.cpu_end - self.cpu_begin					)
      UpperCAmelCase    :					Any			=      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

      _snake_case				=			mocked_dataloaders  # noqa: F811







def 	A (		A__		,		A__			):





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

     # New Code #
     # Pass the advanced FSDP settings not part of the accelerate config by creating fsdp_plugin
     UpperCAmelCase    :					List[Any]			=      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:
          UpperCAmelCase    :					int			=      Accelerator(
              cpu=args.cpu		,		mixed_precision=args.mixed_precision		,		log_with="wandb"		,		project_dir=args.logging_dir		,		fsdp_plugin=__UpperCamelCase		,		)
     else:
          UpperCAmelCase    :					Union[str, Any]			=      Accelerator(fsdp_plugin=__UpperCamelCase			)
     accelerator.print(accelerator.distributed_type			)

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

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

     UpperCAmelCase    :					str			=      AutoTokenizer.from_pretrained(args.model_name_or_path			)
     UpperCAmelCase    :					Optional[Any]			=      load_dataset("glue"		,		"mrpc"			)
     UpperCAmelCase    :					Tuple			=      evaluate.load("glue"		,		"mrpc"			)

     def tokenize_function(A__			):
          # max_length=None => use the model max length (it's actually the default)
          UpperCAmelCase    :					str			=      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():
          UpperCAmelCase    :					Union[str, Any]			=      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
     UpperCAmelCase    :					Optional[Any]			=      tokenized_datasets.rename_column("label"		,		"labels"			)

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

     def collate_fn(A__			):
          # On TPU it's best to pad everything to the same length or training will be very slow.
          UpperCAmelCase    :					List[Any]			=      1_2_8 if accelerator.distributed_type == DistributedType.TPU else None
          # When using mixed precision we want round multiples of 8/16
          if accelerator.mixed_precision == "fp8":
               UpperCAmelCase    :					Optional[int]			=      1_6
          elif accelerator.mixed_precision != "no":
               UpperCAmelCase    :					Union[str, Any]			=      8
          else:
               UpperCAmelCase    :					List[str]			=      None

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

     # Instantiate dataloaders.
     UpperCAmelCase    :					Optional[Any]			=      DataLoader(
         tokenized_datasets["train"]		,		shuffle=__UpperCamelCase		,		collate_fn=__UpperCamelCase		,		batch_size=__UpperCamelCase			)
     UpperCAmelCase    :					int			=      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)
     UpperCAmelCase    :					Tuple			=      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
     UpperCAmelCase    :					List[Any]			=      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
     UpperCAmelCase    :					Optional[int]			=      torch.optim.AdamW(params=model.parameters()		,		lr=__UpperCamelCase		,		weight_decay=2E-4			)

     # Instantiate scheduler
     UpperCAmelCase    :					Union[str, Any]			=      get_linear_schedule_with_warmup(
         optimizer=__UpperCamelCase		,		num_warmup_steps=1_0		,		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.
     UpperCAmelCase	,			UpperCAmelCase	,			UpperCAmelCase	,			UpperCAmelCase    :					Tuple			=      accelerator.prepare(
         __UpperCamelCase		,		__UpperCamelCase		,		__UpperCamelCase		,		__UpperCamelCase			)

     UpperCAmelCase    :					Tuple			=      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			)
               UpperCAmelCase    :					str			=      os.path.basename(args.resume_from_checkpoint			)
          else:
               # Get the most recent checkpoint
               UpperCAmelCase    :					List[str]			=      [f.name for f in os.scandir(os.getcwd()			) if f.is_dir()]
               dirs.sort(key=os.path.getctime			)
               UpperCAmelCase    :					Optional[int]			=      dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
          # Extract `epoch_{i}` or `step_{i}`
          UpperCAmelCase    :					List[Any]			=      os.path.splitext(__UpperCamelCase			)[0]

          if "epoch" in training_difference:
               num_epochs -= int(training_difference.replace("epoch_"		,		""			)			)
               UpperCAmelCase    :					List[str]			=      None
          else:
               UpperCAmelCase    :					Optional[int]			=      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
               UpperCAmelCase    :					Union[str, Any]			=      (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:
                    UpperCAmelCase    :					int			=      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			)
                    UpperCAmelCase    :					List[str]			=      model(**__UpperCamelCase			)
                    UpperCAmelCase    :					int			=      outputs.loss
                    UpperCAmelCase    :					int			=      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			):
                         UpperCAmelCase    :					List[Any]			=      f'step_{overall_step}'
                         if overall_step % checkpointing_steps == 0:
                              if args.output_dir is not None:
                                   UpperCAmelCase    :					str			=      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():
                         UpperCAmelCase    :					Union[str, Any]			=      model(**__UpperCamelCase			)
                    UpperCAmelCase    :					Any			=      outputs.logits.argmax(dim=-1			)
                    UpperCAmelCase	,			UpperCAmelCase    :					Any			=      accelerator.gather_for_metrics((predictions, batch["labels"])			)
                    metric.add_batch(
                        predictions=__UpperCamelCase		,		references=__UpperCamelCase		,		)

               UpperCAmelCase    :					str			=      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":
                    UpperCAmelCase    :					Optional[int]			=      f'epoch_{epoch}'
                    if args.output_dir is not None:
                         UpperCAmelCase    :					Optional[int]			=      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 	A (		):
     UpperCAmelCase    :					Any			=      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		,		)
     UpperCAmelCase    :					List[Any]			=      parser.parse_args()
     UpperCAmelCase    :					List[str]			=      {"lr": 2E-5, "num_epochs": 3, "seed": 4_2, "batch_size": 1_6}
     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
