
import timeit
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_auto import MODEL_MAPPING, MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_torch_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_torch_available():
		import torch

		from .benchmark_args import PyTorchBenchmarkArguments


if is_pyanvml_available():
		import pyanvml.pyanvml as nvml


__SCREAMING_SNAKE_CASE			:   Tuple         =     logging.get_logger(__name__)


class 		A   (						snake_case    ):
		_lowerCAmelCase		=	42
		_lowerCAmelCase		=	42
		_lowerCAmelCase		=	"PyTorch"





		@property
		def       _A							(							self    :	Optional[Any]				):
							return torch.__version__





		def       _A							(							self    :	List[str]						,  lowerCamelCase__    :	str						,  lowerCamelCase__    :	int						,  lowerCamelCase__    :	int				):
							lowerCAmelCase		:					str  	=     self._prepare_inference_func(lowerCamelCase__						,  lowerCamelCase__						,  lowerCamelCase__				)
							return self._measure_speed(_inference				)





		def       _A							(							self    :	Any						,  lowerCamelCase__    :	str						,  lowerCamelCase__    :	int						,  lowerCamelCase__    :	int				):
							lowerCAmelCase		:					Union[str, Any]  	=     self._prepare_inference_func(lowerCamelCase__						,  lowerCamelCase__						,  lowerCamelCase__				)
							return self._measure_memory(_inference				)





		def       _A							(							self    :	Any						,  lowerCamelCase__    :	str						,  lowerCamelCase__    :	int						,  lowerCamelCase__    :	int				):
							lowerCAmelCase		:					Dict  	=     self._prepare_train_func(lowerCamelCase__						,  lowerCamelCase__						,  lowerCamelCase__				)
							return self._measure_speed(_train				)





		def       _A							(							self    :	Dict						,  lowerCamelCase__    :	str						,  lowerCamelCase__    :	int						,  lowerCamelCase__    :	int				):
							lowerCAmelCase		:					Optional[int]  	=     self._prepare_train_func(lowerCamelCase__						,  lowerCamelCase__						,  lowerCamelCase__				)
							return self._measure_memory(_train				)





		def       _A							(							self    :	List[Any]						,  lowerCamelCase__    :	str						,  lowerCamelCase__    :	int						,  lowerCamelCase__    :	int				):
							lowerCAmelCase		:					List[Any]  	=     self.config_dict[model_name]

							if self.args.torchscript:
												lowerCAmelCase		:					Optional[Any]  	=     True

							lowerCAmelCase		:					str  	=     (
							    hasattr(lowerCamelCase__						,  '''architectures'''				)
							    and isinstance(config.architectures						,  lowerCamelCase__				)
							    and len(config.architectures				) > 0
							)
							if not self.args.only_pretrain_model and has_model_class_in_config:
												try:
																	lowerCAmelCase		:					str  	=     config.architectures[0]
																	lowerCAmelCase		:					Dict  	=     __import__('''transformers'''						,  fromlist=[model_class]				)
																	lowerCAmelCase		:					Optional[int]  	=     getattr(lowerCamelCase__						,  lowerCamelCase__				)
																	lowerCAmelCase		:					List[str]  	=     model_cls(lowerCamelCase__				)
												except ImportError:
																	raise ImportError(
																	    f'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
																	    ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'''				)
							else:
												lowerCAmelCase		:					List[str]  	=     MODEL_MAPPING[config.__class__](lowerCamelCase__				)

							model.eval()
							model.to(self.args.device				)

							# encoder-decoder has vocab size saved differently
							lowerCAmelCase		:					Union[str, Any]  	=     config.vocab_size if hasattr(lowerCamelCase__						,  '''vocab_size'''				) else config.encoder.vocab_size
							lowerCAmelCase		:					str  	=     torch.randint(lowerCamelCase__						,  (batch_size, sequence_length)						,  dtype=torch.long						,  device=self.args.device				)

							if self.args.fpaa:
												logger.info('''Running training in Mixed Precision...'''				)
												if not self.args.is_gpu:
																	raise ValueError('''Mixed precision is possible only for GPU.'''				)
												# amp seems to have memory leaks so that memory usage
												# is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
												model.half()

							if self.args.torchscript:
												with torch.no_grad():
																	lowerCAmelCase		:					Union[str, Any]  	=     torch.jit.trace(lowerCamelCase__						,  lowerCamelCase__				)
							else:
												lowerCAmelCase		:					Optional[Any]  	=     model

							def encoder_decoder_forward():
												with torch.no_grad():
																	lowerCAmelCase		:					Union[str, Any]  	=     inference_model(lowerCamelCase__						,  decoder_input_ids=lowerCamelCase__				)
												return outputs

							def encoder_forward():
												with torch.no_grad():
																	lowerCAmelCase		:					Dict  	=     inference_model(lowerCamelCase__				)
												return outputs

							lowerCAmelCase		:					int  	=     encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
							return _forward





		def       _A							(							self    :	Union[str, Any]						,  lowerCamelCase__    :	str						,  lowerCamelCase__    :	int						,  lowerCamelCase__    :	int				):
							lowerCAmelCase		:					Tuple  	=     self.config_dict[model_name]

							lowerCAmelCase		:					str  	=     (
							    hasattr(lowerCamelCase__						,  '''architectures'''				)
							    and isinstance(config.architectures						,  lowerCamelCase__				)
							    and len(config.architectures				) > 0
							)
							if not self.args.only_pretrain_model and has_model_class_in_config:
												try:
																	lowerCAmelCase		:					List[Any]  	=     config.architectures[0]
																	lowerCAmelCase		:					str  	=     __import__('''transformers'''						,  fromlist=[model_class]				)
																	lowerCAmelCase		:					Union[str, Any]  	=     getattr(lowerCamelCase__						,  lowerCamelCase__				)
																	lowerCAmelCase		:					List[Any]  	=     model_cls(lowerCamelCase__				)
												except ImportError:
																	raise ImportError(
																	    f'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
																	    ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'''				)
							else:
												lowerCAmelCase		:					Any  	=     MODEL_WITH_LM_HEAD_MAPPING[config.__class__](lowerCamelCase__				)

							if self.args.torchscript:
												raise NotImplementedError('''Training for torchscript is currently not implemented'''				)
							else:
												lowerCAmelCase		:					Dict  	=     model

							model.train()
							model.to(self.args.device				)

							# encoder-decoder has vocab size saved differently
							lowerCAmelCase		:					int  	=     config.vocab_size if hasattr(lowerCamelCase__						,  '''vocab_size'''				) else config.encoder.vocab_size
							lowerCAmelCase		:					Optional[Any]  	=     torch.randint(lowerCamelCase__						,  (batch_size, sequence_length)						,  dtype=torch.long						,  device=self.args.device				)

							if self.args.fpaa:
												logger.info('''Running training in Mixed Precision...'''				)
												if not self.args.is_gpu:
																	raise ValueError('''Mixed precision is possible only for GPU.'''				)

												# amp seems to have memory leaks so that memory usage
												# is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
												model.half()

							def compute_loss_and_backprob_encoder():
												lowerCAmelCase		:					Union[str, Any]  	=     train_model(lowerCamelCase__						,  labels=lowerCamelCase__				)[0]
												loss.backward()
												return loss

							def compute_loss_and_backprob_encoder_decoder():
												lowerCAmelCase		:					Optional[int]  	=     train_model(lowerCamelCase__						,  decoder_input_ids=lowerCamelCase__						,  labels=lowerCamelCase__				)[0]
												loss.backward()
												return loss

							lowerCAmelCase		:					Any  	=     (
							    compute_loss_and_backprob_encoder_decoder
							    if config.is_encoder_decoder
							    else compute_loss_and_backprob_encoder
							)
							return _train





		def       _A							(							self    :	int						,  lowerCamelCase__    :	Optional[Any]				):
							try:
												if self.args.is_tpu or self.args.torchscript:
																	# run additional 10 times to stabilize compilation for tpu and torchscript
																	logger.info('''Do inference on TPU or torchscript. Running model 5 times to stabilize compilation'''				)
																	timeit.repeat(
																	    lowerCamelCase__						,  repeat=1						,  number=5						,  )

												# as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
												lowerCAmelCase		:					str  	=     timeit.repeat(
												    lowerCamelCase__						,  repeat=self.args.repeat						,  number=1_0						,  )

												if self.args.is_tpu and self.args.torch_xla_tpu_print_metrics:
																	import torch_xla.debug.metrics as met

																	self.print_fn(met.metrics_report()				)

												return min(lowerCamelCase__				) / 1_0.0
							except RuntimeError as e:
												self.print_fn(f'''Doesn\'t fit on GPU. {e}'''				)
												return "N/A"





		def       _A							(							self    :	List[Any]						,  lowerCamelCase__    :	Callable[[], None]				):
							try:
												if self.args.trace_memory_line_by_line:
																	lowerCAmelCase		:					Optional[int]  	=     start_memory_tracing('''transformers'''				)

												if self.args.is_tpu:
																	# tpu
																	raise NotImplementedError(
																	    '''Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking with'''
																	    ''' `--no-memory` or `args.memory=False`'''				)
												elif self.args.is_gpu:
																	if not is_pyanvml_available():
																						logger.warning(
																						    '''py3nvml not installed, we won\'t log GPU memory usage. '''
																						    '''Install py3nvml (pip install py3nvml) to log information about GPU.'''				)
																						lowerCAmelCase		:					Tuple  	=     '''N/A'''
																	else:
																						logger.info(
																						    '''Measuring total GPU usage on GPU device. Make sure to not have additional processes running'''
																						    ''' on the same GPU.'''				)
																						# init nvml
																						nvml.nvmlInit()
																						func()
																						lowerCAmelCase		:					str  	=     nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx				)
																						lowerCAmelCase		:					List[str]  	=     nvml.nvmlDeviceGetMemoryInfo(lowerCamelCase__				)
																						lowerCAmelCase		:					List[str]  	=     meminfo.used
																						lowerCAmelCase		:					int  	=     Memory(lowerCamelCase__				)
																						# shutdown nvml
																						nvml.nvmlShutdown()
												else:
																	# cpu
																	lowerCAmelCase		:					List[str]  	=     measure_peak_memory_cpu(lowerCamelCase__				)
																	lowerCAmelCase		:					str  	=     Memory(lowerCamelCase__				) if isinstance(lowerCamelCase__						,  lowerCamelCase__				) else memory_bytes

												if self.args.trace_memory_line_by_line:
																	lowerCAmelCase		:					int  	=     stop_memory_tracing(lowerCamelCase__				)
												else:
																	lowerCAmelCase		:					Any  	=     None

												return memory, summary
							except RuntimeError as e:
												self.print_fn(f'''Doesn\'t fit on GPU. {e}'''				)
												return "N/A", None


import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPImageProcessor, CLIPVisionConfig

from diffusers import AutoencoderKL, PaintByExamplePipeline, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.paint_by_example import PaintByExampleImageEncoder
from diffusers.utils import floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import IMAGE_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, IMAGE_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


class 		__magic_name__   (						snake_case,				unittest.TestCase    ):
		_lowerCAmelCase		=	PaintByExamplePipeline
		_lowerCAmelCase		=	IMAGE_GUIDED_IMAGE_INPAINTING_PARAMS
		_lowerCAmelCase		=	IMAGE_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
		_lowerCAmelCase		=	frozenset([]    )  # TO_DO: update the image_prams once refactored VaeImageProcessor.preprocess





		def       _A							(							self    :	Dict				):
							torch.manual_seed(0				)
							lowerCAmelCase		:					Optional[int]  	=     UNetaDConditionModel(
							    block_out_channels=(3_2, 6_4)						,  layers_per_block=2						,  sample_size=3_2						,  in_channels=9						,  out_channels=4						,  down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')						,  up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')						,  cross_attention_dim=3_2						,  )
							lowerCAmelCase		:					List[str]  	=     PNDMScheduler(skip_prk_steps=lowerCamelCase__				)
							torch.manual_seed(0				)
							lowerCAmelCase		:					List[str]  	=     AutoencoderKL(
							    block_out_channels=[3_2, 6_4]						,  in_channels=3						,  out_channels=3						,  down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']						,  up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']						,  latent_channels=4						,  )
							torch.manual_seed(0				)
							lowerCAmelCase		:					Union[str, Any]  	=     CLIPVisionConfig(
							    hidden_size=3_2						,  projection_dim=3_2						,  intermediate_size=3_7						,  layer_norm_eps=1E-05						,  num_attention_heads=4						,  num_hidden_layers=5						,  image_size=3_2						,  patch_size=4						,  )
							lowerCAmelCase		:					List[Any]  	=     PaintByExampleImageEncoder(lowerCamelCase__						,  proj_size=3_2				)
							lowerCAmelCase		:					Union[str, Any]  	=     CLIPImageProcessor(crop_size=3_2						,  size=3_2				)

							lowerCAmelCase		:					str  	=     {
							    '''unet''': unet,
							    '''scheduler''': scheduler,
							    '''vae''': vae,
							    '''image_encoder''': image_encoder,
							    '''safety_checker''': None,
							    '''feature_extractor''': feature_extractor,
							}
							return components





		def       _A							(							self    :	str						,  lowerCamelCase__    :	List[Any]				):
							lowerCAmelCase		:					List[Any]  	=     np.array(image.convert('''RGB'''				)				)
							lowerCAmelCase		:					int  	=     image[None].transpose(0						,  3						,  1						,  2				)
							lowerCAmelCase		:					Any  	=     torch.from_numpy(lowerCamelCase__				).to(dtype=torch.floataa				) / 1_2_7.5 - 1.0
							return image





		def       _A							(							self    :	str						,  lowerCamelCase__    :	Union[str, Any]="cpu"						,  lowerCamelCase__    :	Any=0				):
							# TODO: use tensor inputs instead of PIL, this is here just to leave the old expected_slices untouched
							lowerCAmelCase		:					str  	=     floats_tensor((1, 3, 3_2, 3_2)						,  rng=random.Random(lowerCamelCase__				)				).to(lowerCamelCase__				)
							lowerCAmelCase		:					Any  	=     image.cpu().permute(0						,  2						,  3						,  1				)[0]
							lowerCAmelCase		:					Dict  	=     Image.fromarray(np.uinta(lowerCamelCase__				)				).convert('''RGB'''				).resize((6_4, 6_4)				)
							lowerCAmelCase		:					str  	=     Image.fromarray(np.uinta(image + 4				)				).convert('''RGB'''				).resize((6_4, 6_4)				)
							lowerCAmelCase		:					Optional[int]  	=     Image.fromarray(np.uinta(lowerCamelCase__				)				).convert('''RGB'''				).resize((3_2, 3_2)				)

							if str(lowerCamelCase__				).startswith('''mps'''				):
												lowerCAmelCase		:					Dict  	=     torch.manual_seed(lowerCamelCase__				)
							else:
												lowerCAmelCase		:					str  	=     torch.Generator(device=lowerCamelCase__				).manual_seed(lowerCamelCase__				)
							lowerCAmelCase		:					Union[str, Any]  	=     {
							    '''example_image''': example_image,
							    '''image''': init_image,
							    '''mask_image''': mask_image,
							    '''generator''': generator,
							    '''num_inference_steps''': 2,
							    '''guidance_scale''': 6.0,
							    '''output_type''': '''numpy''',
							}
							return inputs





		def       _A							(							self    :	List[Any]				):
							lowerCAmelCase		:					Tuple  	=     self.get_dummy_components()

							# make sure here that pndm scheduler skips prk
							lowerCAmelCase		:					Union[str, Any]  	=     PaintByExamplePipeline(**lowerCamelCase__				)
							lowerCAmelCase		:					List[Any]  	=     pipe.to('''cpu'''				)
							pipe.set_progress_bar_config(disable=lowerCamelCase__				)

							lowerCAmelCase		:					Tuple  	=     self.get_dummy_inputs()
							lowerCAmelCase		:					List[str]  	=     pipe(**lowerCamelCase__				)
							lowerCAmelCase		:					Optional[Any]  	=     output.images

							lowerCAmelCase		:					Optional[int]  	=     image[0, -3:, -3:, -1]

							assert image.shape == (1, 6_4, 6_4, 3)
							lowerCAmelCase		:					List[Any]  	=     np.array([0.4_6_8_6, 0.5_6_8_7, 0.4_0_0_7, 0.5_2_1_8, 0.5_7_4_1, 0.4_4_8_2, 0.4_9_4_0, 0.4_6_2_9, 0.4_5_0_3]				)

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





		def       _A							(							self    :	List[str]				):
							lowerCAmelCase		:					List[str]  	=     '''cpu'''
							lowerCAmelCase		:					Union[str, Any]  	=     self.get_dummy_inputs()
							inputs.pop('''mask_image'''				)
							lowerCAmelCase		:					Union[str, Any]  	=     self.convert_to_pt(inputs.pop('''image'''				)				)
							lowerCAmelCase		:					Dict  	=     image.clamp(0						,  1				) / 2

							# make sure here that pndm scheduler skips prk
							lowerCAmelCase		:					str  	=     PaintByExamplePipeline(**self.get_dummy_components()				)
							lowerCAmelCase		:					Tuple  	=     pipe.to(lowerCamelCase__				)
							pipe.set_progress_bar_config(disable=lowerCamelCase__				)

							lowerCAmelCase		:					Optional[int]  	=     pipe(image=lowerCamelCase__						,  mask_image=mask_image[:, 0]						,  **lowerCamelCase__				)
							lowerCAmelCase		:					Optional[int]  	=     output.images

							lowerCAmelCase		:					int  	=     image.cpu().permute(0						,  2						,  3						,  1				)[0]
							lowerCAmelCase		:					int  	=     mask_image.cpu().permute(0						,  2						,  3						,  1				)[0]

							lowerCAmelCase		:					List[str]  	=     Image.fromarray(np.uinta(lowerCamelCase__				)				).convert('''RGB'''				)
							lowerCAmelCase		:					str  	=     Image.fromarray(np.uinta(lowerCamelCase__				)				).convert('''RGB'''				)

							lowerCAmelCase		:					List[Any]  	=     pipe(**self.get_dummy_inputs()				)
							lowerCAmelCase		:					Any  	=     output.images

							assert out_a.shape == (1, 6_4, 6_4, 3)
							assert np.abs(out_a.flatten() - out_a.flatten()				).max() < 5E-2





		def       _A							(							self    :	Optional[Any]				):
							super().test_inference_batch_single_identical(expected_max_diff=3E-3				)



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





		def       _A							(							self    :	Tuple				):
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()





		def       _A							(							self    :	Optional[int]				):
							# make sure here that pndm scheduler skips prk
							lowerCAmelCase		:					Union[str, Any]  	=     load_image(
							    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
							    '''/paint_by_example/dog_in_bucket.png'''				)
							lowerCAmelCase		:					Any  	=     load_image(
							    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
							    '''/paint_by_example/mask.png'''				)
							lowerCAmelCase		:					Any  	=     load_image(
							    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
							    '''/paint_by_example/panda.jpg'''				)

							lowerCAmelCase		:					Union[str, Any]  	=     PaintByExamplePipeline.from_pretrained('''Fantasy-Studio/Paint-by-Example'''				)
							lowerCAmelCase		:					str  	=     pipe.to(lowerCamelCase__				)
							pipe.set_progress_bar_config(disable=lowerCamelCase__				)

							lowerCAmelCase		:					int  	=     torch.manual_seed(3_2_1				)
							lowerCAmelCase		:					List[str]  	=     pipe(
							    image=lowerCamelCase__						,  mask_image=lowerCamelCase__						,  example_image=lowerCamelCase__						,  generator=lowerCamelCase__						,  guidance_scale=5.0						,  num_inference_steps=5_0						,  output_type='''np'''						,  )

							lowerCAmelCase		:					Any  	=     output.images

							lowerCAmelCase		:					Tuple  	=     image[0, -3:, -3:, -1]

							assert image.shape == (1, 5_1_2, 5_1_2, 3)
							lowerCAmelCase		:					Dict  	=     np.array([0.4_8_3_4, 0.4_8_1_1, 0.4_8_7_4, 0.5_1_2_2, 0.5_0_8_1, 0.5_1_4_4, 0.5_2_9_1, 0.5_2_9_0, 0.5_3_7_4]				)

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