



"""simple docstring"""







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


A     =      1_6
A     =      3_2

def 	lowerCAmelCase__      (  lowerCamelCase__      )			->				Any:
							return int(x / 2**20      )






class   UpperCAmelCase__		:



		def __enter__(			self	:     Dict					)      ->    Union[str, Any]:






									'''simple docstring'''

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

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



		def        A_		(			self	:     Optional[int]					)      ->    Tuple:






									'''simple docstring'''

									return self.process.memory_info().rss



		def        A_		(			self	:     Union[str, Any]					)      ->    List[str]:






									'''simple docstring'''

									A    =       -1

									while True:
																A    =       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	:     List[Any]   ,				*snake_case	:     Tuple					)      ->    Optional[Any]:






									'''simple docstring'''

									A    =       False

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

									A    =       self.cpu_mem_used()
									A    =       bamb(self.cpu_end - self.cpu_begin					)
									A    =       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

							A     =      mocked_dataloaders  # noqa: F811

def 	lowerCAmelCase__      (  lowerCamelCase__					,   lowerCamelCase__      )			->				str:
							# For testing only
							if os.environ.get('TESTING_MOCKED_DATALOADERS'					,   UpperCamelCase__      ) == "1":
														A    =       2

							# New Code #
							# Pass the advanced FSDP settings not part of the accelerate config by creating fsdp_plugin
							A    =       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:
														A    =       Accelerator(
														    cpu=args.cpu					,   mixed_precision=args.mixed_precision					,   log_with='wandb'					,   project_dir=args.logging_dir					,   fsdp_plugin=UpperCamelCase__					,   )
							else:
														A    =       Accelerator(fsdp_plugin=UpperCamelCase__      )
							accelerator.print(accelerator.distributed_type      )

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

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

							A    =       AutoTokenizer.from_pretrained(args.model_name_or_path      )
							A    =       load_dataset('glue'					,   'mrpc'      )
							A    =       evaluate.load('glue'					,   'mrpc'      )

							def tokenize_function(lowerCamelCase__      ):
														# max_length=None => use the model max length (it's actually the default)
														A    =       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():
														A    =       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
							A    =       tokenized_datasets.rename_column('label'					,   'labels'      )

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

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

														return tokenizer.pad(
														    UpperCamelCase__					,   padding='longest'					,   max_length=UpperCamelCase__					,   pad_to_multiple_of=UpperCamelCase__					,   return_tensors='pt'					,   )

							# Instantiate dataloaders.
							A    =       DataLoader(
							    tokenized_datasets['train']					,   shuffle=UpperCamelCase__					,   collate_fn=UpperCamelCase__					,   batch_size=UpperCamelCase__      )
							A    =       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)
							A    =       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
							A    =       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
							A    =       torch.optim.AdamW(params=model.parameters()					,   lr=UpperCamelCase__					,   weight_decay=2E-4      )

							# Instantiate scheduler
							A    =       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.
							A    =       accelerator.prepare(
							    UpperCamelCase__					,   UpperCamelCase__					,   UpperCamelCase__					,   UpperCamelCase__      )

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

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

																					A    =       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":
																												A    =       f"""epoch_{epoch}"""
																												if args.output_dir is not None:
																																			A    =       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 	lowerCAmelCase__      (  )			->				List[Any]:
							A    =       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__					,   )
							A    =       parser.parse_args()
							A    =       {'''lr''': 2E-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
							training_function(UpperCamelCase__					,   UpperCamelCase__      )


if __name__ == "__main__":
							main()








'''simple docstring'''



import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLPipeline,
    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				(lowercase__,	lowercase__,	unittest.TestCase      ):
			A__								=  StableDiffusionXLPipeline
			A__								=  TEXT_TO_IMAGE_PARAMS
			A__								=  TEXT_TO_IMAGE_BATCH_PARAMS
			A__								=  TEXT_TO_IMAGE_IMAGE_PARAMS
			A__								=  TEXT_TO_IMAGE_IMAGE_PARAMS



			def 				lowerCAmelCase (     self		:						List[str]			)      ->      Tuple:
							"""simple docstring"""





							torch.manual_seed(0			)
							snake_case__      :   Any      =    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''')    ,			attention_head_dim=(2, 4)    ,			use_linear_projection=__UpperCamelCase    ,			addition_embed_type='''text_time'''    ,			addition_time_embed_dim=8    ,			transformer_layers_per_block=(1, 2)    ,			projection_class_embeddings_input_dim=80    ,			cross_attention_dim=64    ,			)
							snake_case__      :   Tuple      =    EulerDiscreteScheduler(
							    beta_start=0.00085    ,			beta_end=0.012    ,			steps_offset=1    ,			beta_schedule='''scaled_linear'''    ,			timestep_spacing='''leading'''    ,			)
							torch.manual_seed(0			)
							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    ,			sample_size=128    ,			)
							torch.manual_seed(0			)
							snake_case__      :   List[Any]      =    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    ,			hidden_act='''gelu'''    ,			projection_dim=32    ,			)
							snake_case__      :   List[str]      =    CLIPTextModel(__UpperCamelCase			)
							snake_case__      :   Dict      =    CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''    ,			local_files_only=__UpperCamelCase			)

							snake_case__      :   Optional[int]      =    CLIPTextModelWithProjection(__UpperCamelCase			)
							snake_case__      :   Optional[int]      =    CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''    ,			local_files_only=__UpperCamelCase			)

							snake_case__      :   List[str]      =    {
							    '''unet''': unet,
							    '''scheduler''': scheduler,
							    '''vae''': vae,
							    '''text_encoder''': text_encoder,
							    '''tokenizer''': tokenizer,
							    '''text_encoder_2''': text_encoder_a,
							    '''tokenizer_2''': tokenizer_a,
							    # "safety_checker": None,
							    # "feature_extractor": None,
							}
							return components



			def 				lowerCAmelCase (     self		:						Union[str, Any]    ,			__UpperCamelCase		:						int    ,			__UpperCamelCase		:						Union[str, Any]=0			)      ->      Optional[Any]:
							"""simple docstring"""





							if str(__UpperCamelCase			).startswith('''mps'''			):
											snake_case__      :   List[Any]      =    torch.manual_seed(__UpperCamelCase			)
							else:
											snake_case__      :   List[Any]      =    torch.Generator(device=__UpperCamelCase			).manual_seed(__UpperCamelCase			)
							snake_case__      :   Dict      =    {
							    '''prompt''': '''A painting of a squirrel eating a burger''',
							    '''generator''': generator,
							    '''num_inference_steps''': 2,
							    '''guidance_scale''': 5.0,
							    '''output_type''': '''numpy''',
							}
							return inputs



			def 				lowerCAmelCase (     self		:						List[Any]			)      ->      Tuple:
							"""simple docstring"""





							snake_case__      :   Any      =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							snake_case__      :   Union[str, Any]      =    self.get_dummy_components()
							snake_case__      :   List[Any]      =    StableDiffusionXLPipeline(**__UpperCamelCase			)
							snake_case__      :   Optional[int]      =    sd_pipe.to(__UpperCamelCase			)
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							snake_case__      :   Dict      =    self.get_dummy_inputs(__UpperCamelCase			)
							snake_case__      :   Union[str, Any]      =    sd_pipe(**__UpperCamelCase			).images
							snake_case__      :   List[Any]      =    image[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)
							snake_case__      :   Union[str, Any]      =    np.array([0.5873, 0.6128, 0.4797, 0.5122, 0.5674, 0.4639, 0.5227, 0.5149, 0.4747]			)

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



			def 				lowerCAmelCase (     self		:						Optional[Any]			)      ->      Optional[int]:
							"""simple docstring"""





							snake_case__      :   List[str]      =    self.get_dummy_components()
							snake_case__      :   Optional[Any]      =    StableDiffusionXLPipeline(**__UpperCamelCase			)
							snake_case__      :   str      =    sd_pipe.to(__UpperCamelCase			)
							snake_case__      :   Tuple      =    sd_pipe.to(__UpperCamelCase			)
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							# forward without prompt embeds
							snake_case__      :   str      =    self.get_dummy_inputs(__UpperCamelCase			)
							snake_case__      :   Any      =    3 * ['''this is a negative prompt''']
							snake_case__      :   Optional[int]      =    negative_prompt
							snake_case__      :   Tuple      =    3 * [inputs['''prompt''']]

							snake_case__      :   Tuple      =    sd_pipe(**__UpperCamelCase			)
							snake_case__      :   str      =    output.images[0, -3:, -3:, -1]

							# forward with prompt embeds
							snake_case__      :   Tuple      =    self.get_dummy_inputs(__UpperCamelCase			)
							snake_case__      :   List[Any]      =    3 * ['''this is a negative prompt''']
							snake_case__      :   List[Any]      =    3 * [inputs.pop('''prompt'''			)]

							(
							    (
							    snake_case__
							)  ,      (
							    snake_case__
							)  ,      (
							    snake_case__
							)  ,      (
							    snake_case__
							)  ,      
							)      :   List[Any]      =    sd_pipe.encode_prompt(__UpperCamelCase    ,			negative_prompt=__UpperCamelCase			)

							snake_case__      :   Dict      =    sd_pipe(
							    **__UpperCamelCase    ,			prompt_embeds=__UpperCamelCase    ,			negative_prompt_embeds=__UpperCamelCase    ,			pooled_prompt_embeds=__UpperCamelCase    ,			negative_pooled_prompt_embeds=__UpperCamelCase    ,			)
							snake_case__      :   Optional[Any]      =    output.images[0, -3:, -3:, -1]

							# make sure that it's equal
							assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()			).max() < 1e-4



			def 				lowerCAmelCase (     self		:						Dict			)      ->      Dict:
							"""simple docstring"""





							super().test_attention_slicing_forward_pass(expected_max_diff=3e-3			)







			def 				lowerCAmelCase (     self		:						int			)      ->      Dict:
							"""simple docstring"""





							super().test_inference_batch_single_identical(expected_max_diff=3e-3			)






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



			def 				lowerCAmelCase (     self		:						Union[str, Any]			)      ->      Union[str, Any]:
							"""simple docstring"""





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



			def 				lowerCAmelCase (     self		:						Tuple    ,			__UpperCamelCase		:						Optional[int]    ,			__UpperCamelCase		:						List[str]="cpu"    ,			__UpperCamelCase		:						Tuple=torch.floataa    ,			__UpperCamelCase		:						List[str]=0			)      ->      Dict:
							"""simple docstring"""





							snake_case__      :   Dict      =    torch.Generator(device=__UpperCamelCase			).manual_seed(__UpperCamelCase			)
							snake_case__      :   int      =    np.random.RandomState(__UpperCamelCase			).standard_normal((1, 4, 64, 64)			)
							snake_case__      :   List[Any]      =    torch.from_numpy(__UpperCamelCase			).to(device=__UpperCamelCase    ,			dtype=__UpperCamelCase			)
							snake_case__      :   int      =    {
							    '''prompt''': '''a photograph of an astronaut riding a horse''',
							    '''latents''': latents,
							    '''generator''': generator,
							    '''num_inference_steps''': 3,
							    '''guidance_scale''': 7.5,
							    '''output_type''': '''numpy''',
							}
							return inputs







			def 				lowerCAmelCase (     self		:						int			)      ->      Dict:
							"""simple docstring"""





							snake_case__      :   Dict      =    DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''			)
							pipe.to(__UpperCamelCase			)
							pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							snake_case__      :   Optional[Any]      =    self.get_inputs(__UpperCamelCase			)
							snake_case__      :   Any      =    pipe(**__UpperCamelCase			).images
							snake_case__      :   Optional[int]      =    image[0, -3:, -3:, -1].flatten()

							assert image.shape == (1, 512, 512, 3)
							snake_case__      :   Tuple      =    np.array([0.49493, 0.47896, 0.40798, 0.54214, 0.53212, 0.48202, 0.47656, 0.46329, 0.48506]			)
							assert np.abs(image_slice - expected_slice			).max() < 7e-3




