




# This model implementation is heavily based on:

import gc
import random
import tempfile
import unittest

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

from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DDIMScheduler,
    StableDiffusionControlNetInpaintPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel
from diffusers.utils import floats_tensor, load_image, load_numpy, randn_tensor, slow, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
    TEXT_TO_IMAGE_IMAGE_PARAMS,
)
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
)


enable_full_determinism()






class 					__lowerCAmelCase				( UpperCamelCase_      ,	UpperCamelCase_      ,	UpperCamelCase_      ,	unittest.TestCase				):
			snake_case_   :				List[Any]				  =  StableDiffusionControlNetInpaintPipeline
			snake_case_   :				List[str]				  =  TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
			snake_case_   :				Tuple				  =  TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
			snake_case_   :				List[str]				  =  frozenset({"control_image"}				)  # skip `image` and `mask` for now, only test for control_image
			snake_case_   :				List[str]				  =  TEXT_TO_IMAGE_IMAGE_PARAMS







			def UpperCamelCase    ( self					:    str   ):




									"""simple docstring"""


									torch.manual_seed(0   )
									_UpperCAmelCase									=  UNetaDConditionModel(
									    block_out_channels=(32, 64)					,					layers_per_block=2					,					sample_size=32					,					in_channels=9					,					out_channels=4					,					down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")					,					up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")					,					cross_attention_dim=32					,					)
									torch.manual_seed(0   )
									_UpperCAmelCase									=  ControlNetModel(
									    block_out_channels=(32, 64)					,					layers_per_block=2					,					in_channels=4					,					down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")					,					cross_attention_dim=32					,					conditioning_embedding_out_channels=(16, 32)					,					)
									torch.manual_seed(0   )
									_UpperCAmelCase									=  DDIMScheduler(
									    beta_start=0.00_085					,					beta_end=0.012					,					beta_schedule="scaled_linear"					,					clip_sample=__SCREAMING_SNAKE_CASE					,					set_alpha_to_one=__SCREAMING_SNAKE_CASE					,					)
									torch.manual_seed(0   )
									_UpperCAmelCase									=  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   )
									_UpperCAmelCase									=  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=1_000					,					)
									_UpperCAmelCase									=  CLIPTextModel(__SCREAMING_SNAKE_CASE   )
									_UpperCAmelCase									=  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"   )

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







			def UpperCamelCase    ( self					:    Any					,					snake_case__					:    Optional[int]					,					snake_case__					:    Dict=0   ):




									"""simple docstring"""


									if str(__SCREAMING_SNAKE_CASE   ).startswith("mps"   ):
															_UpperCAmelCase									=  torch.manual_seed(__SCREAMING_SNAKE_CASE   )
									else:
															_UpperCAmelCase									=  torch.Generator(device=__SCREAMING_SNAKE_CASE   ).manual_seed(__SCREAMING_SNAKE_CASE   )

									_UpperCAmelCase									=  2
									_UpperCAmelCase									=  randn_tensor(
									    (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)					,					generator=__SCREAMING_SNAKE_CASE					,					device=torch.device(__SCREAMING_SNAKE_CASE   )					,					)
									_UpperCAmelCase									=  floats_tensor((1, 3, 32, 32)					,					rng=random.Random(__SCREAMING_SNAKE_CASE   )   ).to(__SCREAMING_SNAKE_CASE   )
									_UpperCAmelCase									=  init_image.cpu().permute(0					,					2					,					3					,					1   )[0]

									_UpperCAmelCase									=  Image.fromarray(np.uinta(__SCREAMING_SNAKE_CASE   )   ).convert("RGB"   ).resize((64, 64)   )
									_UpperCAmelCase									=  Image.fromarray(np.uinta(init_image + 4   )   ).convert("RGB"   ).resize((64, 64)   )

									_UpperCAmelCase									=  {
									    "prompt": "A painting of a squirrel eating a burger",
									    "generator": generator,
									    "num_inference_steps": 2,
									    "guidance_scale": 6.0,
									    "output_type": "numpy",
									    "image": image,
									    "mask_image": mask_image,
									    "control_image": control_image,
									}

									return inputs







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




									"""simple docstring"""


									return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3   )







			@unittest.skipIf(
			    torch_device != "cuda" or not is_xformers_available()					,					reason="XFormers attention is only available with CUDA and `xformers` installed"					,					)
			def UpperCamelCase    ( self					:    Tuple   ):




									"""simple docstring"""


									self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3   )







			def UpperCamelCase    ( self					:    Any   ):




									"""simple docstring"""


									self._test_inference_batch_single_identical(expected_max_diff=2e-3   )






class 					__lowerCAmelCase				( UpperCamelCase_				):
			snake_case_   :				Dict				  =  StableDiffusionControlNetInpaintPipeline
			snake_case_   :				Dict				  =  TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
			snake_case_   :				Any				  =  TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
			snake_case_   :				int				  =  frozenset([]				)







			def UpperCamelCase    ( self					:    Optional[int]   ):




									"""simple docstring"""


									torch.manual_seed(0   )
									_UpperCAmelCase									=  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					,					)
									torch.manual_seed(0   )
									_UpperCAmelCase									=  ControlNetModel(
									    block_out_channels=(32, 64)					,					layers_per_block=2					,					in_channels=4					,					down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")					,					cross_attention_dim=32					,					conditioning_embedding_out_channels=(16, 32)					,					)
									torch.manual_seed(0   )
									_UpperCAmelCase									=  DDIMScheduler(
									    beta_start=0.00_085					,					beta_end=0.012					,					beta_schedule="scaled_linear"					,					clip_sample=__SCREAMING_SNAKE_CASE					,					set_alpha_to_one=__SCREAMING_SNAKE_CASE					,					)
									torch.manual_seed(0   )
									_UpperCAmelCase									=  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   )
									_UpperCAmelCase									=  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=1_000					,					)
									_UpperCAmelCase									=  CLIPTextModel(__SCREAMING_SNAKE_CASE   )
									_UpperCAmelCase									=  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"   )

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






class 					__lowerCAmelCase				( UpperCamelCase_      ,	UpperCamelCase_      ,	unittest.TestCase				):
			snake_case_   :				Optional[int]				  =  StableDiffusionControlNetInpaintPipeline
			snake_case_   :				Optional[int]				  =  TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
			snake_case_   :				Any				  =  TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS







			def UpperCamelCase    ( self					:    Optional[int]   ):




									"""simple docstring"""


									torch.manual_seed(0   )
									_UpperCAmelCase									=  UNetaDConditionModel(
									    block_out_channels=(32, 64)					,					layers_per_block=2					,					sample_size=32					,					in_channels=9					,					out_channels=4					,					down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")					,					up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")					,					cross_attention_dim=32					,					)
									torch.manual_seed(0   )

									def init_weights(snake_case__					:    Union[str, Any]   ):
															if isinstance(__SCREAMING_SNAKE_CASE					,					torch.nn.Convad   ):
																					torch.nn.init.normal(m.weight   )
																					m.bias.data.fill_(1.0   )

									_UpperCAmelCase									=  ControlNetModel(
									    block_out_channels=(32, 64)					,					layers_per_block=2					,					in_channels=4					,					down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")					,					cross_attention_dim=32					,					conditioning_embedding_out_channels=(16, 32)					,					)
									controlneta.controlnet_down_blocks.apply(__SCREAMING_SNAKE_CASE   )

									torch.manual_seed(0   )
									_UpperCAmelCase									=  ControlNetModel(
									    block_out_channels=(32, 64)					,					layers_per_block=2					,					in_channels=4					,					down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")					,					cross_attention_dim=32					,					conditioning_embedding_out_channels=(16, 32)					,					)
									controlneta.controlnet_down_blocks.apply(__SCREAMING_SNAKE_CASE   )

									torch.manual_seed(0   )
									_UpperCAmelCase									=  DDIMScheduler(
									    beta_start=0.00_085					,					beta_end=0.012					,					beta_schedule="scaled_linear"					,					clip_sample=__SCREAMING_SNAKE_CASE					,					set_alpha_to_one=__SCREAMING_SNAKE_CASE					,					)
									torch.manual_seed(0   )
									_UpperCAmelCase									=  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   )
									_UpperCAmelCase									=  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=1_000					,					)
									_UpperCAmelCase									=  CLIPTextModel(__SCREAMING_SNAKE_CASE   )
									_UpperCAmelCase									=  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"   )

									_UpperCAmelCase									=  MultiControlNetModel([controlneta, controlneta]   )

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







			def UpperCamelCase    ( self					:    List[Any]					,					snake_case__					:    List[Any]					,					snake_case__					:    Any=0   ):




									"""simple docstring"""


									if str(__SCREAMING_SNAKE_CASE   ).startswith("mps"   ):
															_UpperCAmelCase									=  torch.manual_seed(__SCREAMING_SNAKE_CASE   )
									else:
															_UpperCAmelCase									=  torch.Generator(device=__SCREAMING_SNAKE_CASE   ).manual_seed(__SCREAMING_SNAKE_CASE   )

									_UpperCAmelCase									=  2

									_UpperCAmelCase									=  [
									    randn_tensor(
									        (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)					,					generator=__SCREAMING_SNAKE_CASE					,					device=torch.device(__SCREAMING_SNAKE_CASE   )					,					),
									    randn_tensor(
									        (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)					,					generator=__SCREAMING_SNAKE_CASE					,					device=torch.device(__SCREAMING_SNAKE_CASE   )					,					),
									]
									_UpperCAmelCase									=  floats_tensor((1, 3, 32, 32)					,					rng=random.Random(__SCREAMING_SNAKE_CASE   )   ).to(__SCREAMING_SNAKE_CASE   )
									_UpperCAmelCase									=  init_image.cpu().permute(0					,					2					,					3					,					1   )[0]

									_UpperCAmelCase									=  Image.fromarray(np.uinta(__SCREAMING_SNAKE_CASE   )   ).convert("RGB"   ).resize((64, 64)   )
									_UpperCAmelCase									=  Image.fromarray(np.uinta(init_image + 4   )   ).convert("RGB"   ).resize((64, 64)   )

									_UpperCAmelCase									=  {
									    "prompt": "A painting of a squirrel eating a burger",
									    "generator": generator,
									    "num_inference_steps": 2,
									    "guidance_scale": 6.0,
									    "output_type": "numpy",
									    "image": image,
									    "mask_image": mask_image,
									    "control_image": control_image,
									}

									return inputs







			def UpperCamelCase    ( self					:    Dict   ):




									"""simple docstring"""


									_UpperCAmelCase									=  self.get_dummy_components()
									_UpperCAmelCase									=  self.pipeline_class(**__SCREAMING_SNAKE_CASE   )
									pipe.to(__SCREAMING_SNAKE_CASE   )

									_UpperCAmelCase									=  10.0
									_UpperCAmelCase									=  4

									_UpperCAmelCase									=  self.get_dummy_inputs(__SCREAMING_SNAKE_CASE   )
									_UpperCAmelCase									=  steps
									_UpperCAmelCase									=  scale
									_UpperCAmelCase									=  pipe(**__SCREAMING_SNAKE_CASE   )[0]

									_UpperCAmelCase									=  self.get_dummy_inputs(__SCREAMING_SNAKE_CASE   )
									_UpperCAmelCase									=  steps
									_UpperCAmelCase									=  scale
									_UpperCAmelCase									=  pipe(**__SCREAMING_SNAKE_CASE					,					control_guidance_start=0.1					,					control_guidance_end=0.2   )[0]

									_UpperCAmelCase									=  self.get_dummy_inputs(__SCREAMING_SNAKE_CASE   )
									_UpperCAmelCase									=  steps
									_UpperCAmelCase									=  scale
									_UpperCAmelCase									=  pipe(**__SCREAMING_SNAKE_CASE					,					control_guidance_start=[0.1, 0.3]					,					control_guidance_end=[0.2, 0.7]   )[0]

									_UpperCAmelCase									=  self.get_dummy_inputs(__SCREAMING_SNAKE_CASE   )
									_UpperCAmelCase									=  steps
									_UpperCAmelCase									=  scale
									_UpperCAmelCase									=  pipe(**__SCREAMING_SNAKE_CASE					,					control_guidance_start=0.4					,					control_guidance_end=[0.5, 0.8]   )[0]

									# make sure that all outputs are different
									assert np.sum(np.abs(output_a - output_a   )   ) > 1e-3
									assert np.sum(np.abs(output_a - output_a   )   ) > 1e-3
									assert np.sum(np.abs(output_a - output_a   )   ) > 1e-3







			def UpperCamelCase    ( self					:    List[str]   ):




									"""simple docstring"""


									return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3   )







			@unittest.skipIf(
			    torch_device != "cuda" or not is_xformers_available()					,					reason="XFormers attention is only available with CUDA and `xformers` installed"					,					)
			def UpperCamelCase    ( self					:    Any   ):




									"""simple docstring"""


									self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3   )







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




									"""simple docstring"""


									self._test_inference_batch_single_identical(expected_max_diff=2e-3   )







			def UpperCamelCase    ( self					:    List[Any]   ):




									"""simple docstring"""


									_UpperCAmelCase									=  self.get_dummy_components()
									_UpperCAmelCase									=  self.pipeline_class(**__SCREAMING_SNAKE_CASE   )
									pipe.to(__SCREAMING_SNAKE_CASE   )
									pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE   )
									with tempfile.TemporaryDirectory() as tmpdir:
															try:
																					# save_pretrained is not implemented for Multi-ControlNet
																					pipe.save_pretrained(__SCREAMING_SNAKE_CASE   )
															except NotImplementedError:
																					pass






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







			def UpperCamelCase    ( self					:    List[str]   ):




									"""simple docstring"""


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







			def UpperCamelCase    ( self					:    str   ):




									"""simple docstring"""


									_UpperCAmelCase									=  ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny"   )

									_UpperCAmelCase									=  StableDiffusionControlNetInpaintPipeline.from_pretrained(
									    "runwayml/stable-diffusion-inpainting"					,					safety_checker=__SCREAMING_SNAKE_CASE					,					controlnet=__SCREAMING_SNAKE_CASE   )
									pipe.enable_model_cpu_offload()
									pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE   )

									_UpperCAmelCase									=  torch.Generator(device="cpu"   ).manual_seed(0   )
									_UpperCAmelCase									=  load_image(
									    "https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png"   ).resize((512, 512)   )

									_UpperCAmelCase									=  load_image(
									    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"
									    "/stable_diffusion_inpaint/input_bench_mask.png"   ).resize((512, 512)   )

									_UpperCAmelCase									=  "pitch black hole"

									_UpperCAmelCase									=  load_image(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png"   ).resize((512, 512)   )

									_UpperCAmelCase									=  pipe(
									    __SCREAMING_SNAKE_CASE					,					image=__SCREAMING_SNAKE_CASE					,					mask_image=__SCREAMING_SNAKE_CASE					,					control_image=__SCREAMING_SNAKE_CASE					,					generator=__SCREAMING_SNAKE_CASE					,					output_type="np"					,					num_inference_steps=3					,					)

									_UpperCAmelCase									=  output.images[0]

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

									_UpperCAmelCase									=  load_numpy(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/inpaint.npy"   )

									assert np.abs(expected_image - image   ).max() < 9e-2







			def UpperCamelCase    ( self					:    str   ):




									"""simple docstring"""


									_UpperCAmelCase									=  ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_inpaint"   )

									_UpperCAmelCase									=  StableDiffusionControlNetInpaintPipeline.from_pretrained(
									    "runwayml/stable-diffusion-v1-5"					,					safety_checker=__SCREAMING_SNAKE_CASE					,					controlnet=__SCREAMING_SNAKE_CASE   )
									_UpperCAmelCase									=  DDIMScheduler.from_config(pipe.scheduler.config   )
									pipe.enable_model_cpu_offload()
									pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE   )

									_UpperCAmelCase									=  torch.Generator(device="cpu"   ).manual_seed(33   )

									_UpperCAmelCase									=  load_image(
									    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy.png"   )
									_UpperCAmelCase									=  init_image.resize((512, 512)   )

									_UpperCAmelCase									=  load_image(
									    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy_mask.png"   )
									_UpperCAmelCase									=  mask_image.resize((512, 512)   )

									_UpperCAmelCase									=  "a handsome man with ray-ban sunglasses"

									def make_inpaint_condition(snake_case__					:    Dict					,					snake_case__					:    Any   ):
															_UpperCAmelCase									=  np.array(image.convert("RGB"   )   ).astype(np.floataa   ) / 255.0
															_UpperCAmelCase									=  np.array(image_mask.convert("L"   )   ).astype(np.floataa   ) / 255.0

															assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"
															_UpperCAmelCase									=  -1.0  # set as masked pixel
															_UpperCAmelCase									=  np.expand_dims(__SCREAMING_SNAKE_CASE					,					0   ).transpose(0					,					3					,					1					,					2   )
															_UpperCAmelCase									=  torch.from_numpy(__SCREAMING_SNAKE_CASE   )
															return image

									_UpperCAmelCase									=  make_inpaint_condition(__SCREAMING_SNAKE_CASE					,					__SCREAMING_SNAKE_CASE   )

									_UpperCAmelCase									=  pipe(
									    __SCREAMING_SNAKE_CASE					,					image=__SCREAMING_SNAKE_CASE					,					mask_image=__SCREAMING_SNAKE_CASE					,					control_image=__SCREAMING_SNAKE_CASE					,					guidance_scale=9.0					,					eta=1.0					,					generator=__SCREAMING_SNAKE_CASE					,					num_inference_steps=20					,					output_type="np"					,					)
									_UpperCAmelCase									=  output.images[0]

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

									_UpperCAmelCase									=  load_numpy(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/boy_ray_ban.npy"   )

									assert np.abs(expected_image - image   ).max() < 9e-2










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


lowercase_			:       int						       =				logging.get_logger(__name__)






class 					__lowerCAmelCase				( UpperCAmelCase__				):
			snake_case_   :				PyTorchBenchmarkArguments
			snake_case_   :				PretrainedConfig
			snake_case_   :				str				  =  "PyTorch"







			@property
			def UpperCamelCase    ( self					:    List[Any]   ):




									"""simple docstring"""


									return torch.__version__







			def UpperCamelCase    ( self					:    str					,					snake_case__					:    str					,					snake_case__					:    int					,					snake_case__					:    int   ):




									"""simple docstring"""


									_UpperCAmelCase									=  self._prepare_inference_func(snake_case__					,					snake_case__					,					snake_case__   )
									return self._measure_speed(_inference   )







			def UpperCamelCase    ( self					:    List[Any]					,					snake_case__					:    str					,					snake_case__					:    int					,					snake_case__					:    int   ):




									"""simple docstring"""


									_UpperCAmelCase									=  self._prepare_inference_func(snake_case__					,					snake_case__					,					snake_case__   )
									return self._measure_memory(_inference   )







			def UpperCamelCase    ( self					:    Any					,					snake_case__					:    str					,					snake_case__					:    int					,					snake_case__					:    int   ):




									"""simple docstring"""


									_UpperCAmelCase									=  self._prepare_train_func(snake_case__					,					snake_case__					,					snake_case__   )
									return self._measure_speed(_train   )







			def UpperCamelCase    ( self					:    List[str]					,					snake_case__					:    str					,					snake_case__					:    int					,					snake_case__					:    int   ):




									"""simple docstring"""


									_UpperCAmelCase									=  self._prepare_train_func(snake_case__					,					snake_case__					,					snake_case__   )
									return self._measure_memory(_train   )







			def UpperCamelCase    ( self					:    Dict					,					snake_case__					:    str					,					snake_case__					:    int					,					snake_case__					:    int   ):




									"""simple docstring"""


									_UpperCAmelCase									=  self.config_dict[model_name]

									if self.args.torchscript:
															_UpperCAmelCase									=  True

									_UpperCAmelCase									=  (
									    hasattr(snake_case__					,					"architectures"   )
									    and isinstance(config.architectures					,					snake_case__   )
									    and len(config.architectures   ) > 0
									)
									if not self.args.only_pretrain_model and has_model_class_in_config:
															try:
																					_UpperCAmelCase									=  config.architectures[0]
																					_UpperCAmelCase									=  __import__("transformers"					,					fromlist=[model_class]   )
																					_UpperCAmelCase									=  getattr(snake_case__					,					snake_case__   )
																					_UpperCAmelCase									=  model_cls(snake_case__   )
															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:
															_UpperCAmelCase									=  MODEL_MAPPING[config.__class__](snake_case__   )

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

									# encoder-decoder has vocab size saved differently
									_UpperCAmelCase									=  config.vocab_size if hasattr(snake_case__					,					"vocab_size"   ) else config.encoder.vocab_size
									_UpperCAmelCase									=  torch.randint(snake_case__					,					(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():
																					_UpperCAmelCase									=  torch.jit.trace(snake_case__					,					snake_case__   )
									else:
															_UpperCAmelCase									=  model

									def encoder_decoder_forward():
															with torch.no_grad():
																					_UpperCAmelCase									=  inference_model(snake_case__					,					decoder_input_ids=snake_case__   )
															return outputs

									def encoder_forward():
															with torch.no_grad():
																					_UpperCAmelCase									=  inference_model(snake_case__   )
															return outputs

									_UpperCAmelCase									=  encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
									return _forward







			def UpperCamelCase    ( self					:    str					,					snake_case__					:    str					,					snake_case__					:    int					,					snake_case__					:    int   ):




									"""simple docstring"""


									_UpperCAmelCase									=  self.config_dict[model_name]

									_UpperCAmelCase									=  (
									    hasattr(snake_case__					,					"architectures"   )
									    and isinstance(config.architectures					,					snake_case__   )
									    and len(config.architectures   ) > 0
									)
									if not self.args.only_pretrain_model and has_model_class_in_config:
															try:
																					_UpperCAmelCase									=  config.architectures[0]
																					_UpperCAmelCase									=  __import__("transformers"					,					fromlist=[model_class]   )
																					_UpperCAmelCase									=  getattr(snake_case__					,					snake_case__   )
																					_UpperCAmelCase									=  model_cls(snake_case__   )
															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:
															_UpperCAmelCase									=  MODEL_WITH_LM_HEAD_MAPPING[config.__class__](snake_case__   )

									if self.args.torchscript:
															raise NotImplementedError("Training for torchscript is currently not implemented"   )
									else:
															_UpperCAmelCase									=  model

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

									# encoder-decoder has vocab size saved differently
									_UpperCAmelCase									=  config.vocab_size if hasattr(snake_case__					,					"vocab_size"   ) else config.encoder.vocab_size
									_UpperCAmelCase									=  torch.randint(snake_case__					,					(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():
															_UpperCAmelCase									=  train_model(snake_case__					,					labels=snake_case__   )[0]
															loss.backward()
															return loss

									def compute_loss_and_backprob_encoder_decoder():
															_UpperCAmelCase									=  train_model(snake_case__					,					decoder_input_ids=snake_case__					,					labels=snake_case__   )[0]
															loss.backward()
															return loss

									_UpperCAmelCase									=  (
									    compute_loss_and_backprob_encoder_decoder
									    if config.is_encoder_decoder
									    else compute_loss_and_backprob_encoder
									)
									return _train







			def UpperCamelCase    ( self					:    Any					,					snake_case__					:    Tuple   ):




									"""simple docstring"""


									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(
																					    snake_case__					,					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
															_UpperCAmelCase									=  timeit.repeat(
															    snake_case__					,					repeat=self.args.repeat					,					number=10					,					)

															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(snake_case__   ) / 10.0
									except RuntimeError as e:
															self.print_fn(F"""Doesn't fit on GPU. {e}"""   )
															return "N/A"







			def UpperCamelCase    ( self					:    Tuple					,					snake_case__					:    Callable[[], None]   ):




									"""simple docstring"""


									try:
															if self.args.trace_memory_line_by_line:
																					_UpperCAmelCase									=  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."   )
																											_UpperCAmelCase									=  "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()
																											_UpperCAmelCase									=  nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx   )
																											_UpperCAmelCase									=  nvml.nvmlDeviceGetMemoryInfo(snake_case__   )
																											_UpperCAmelCase									=  meminfo.used
																											_UpperCAmelCase									=  Memory(snake_case__   )
																											# shutdown nvml
																											nvml.nvmlShutdown()
															else:
																					# cpu
																					_UpperCAmelCase									=  measure_peak_memory_cpu(snake_case__   )
																					_UpperCAmelCase									=  Memory(snake_case__   ) if isinstance(snake_case__					,					snake_case__   ) else memory_bytes

															if self.args.trace_memory_line_by_line:
																					_UpperCAmelCase									=  stop_memory_tracing(snake_case__   )
															else:
																					_UpperCAmelCase									=  None

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




