


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


__UpperCAmelCase			:			List[Any]			      =					logging.get_logger(__name__)



class  UpperCAmelCase_						(				_a):
       '''simple docstring'''
       __UpperCamelCase   :  PyTorchBenchmarkArguments
       __UpperCamelCase   :  PretrainedConfig
       __UpperCamelCase   :  str				=       "PyTorch"






       @property
       def 			_lowercase (  self    ):




              """simple docstring"""




              return torch.__version__






       def 			_lowercase (  self			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ):




              """simple docstring"""




              UpperCamelCase     :						Any 						=		self._prepare_inference_func(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )
              return self._measure_speed(_inference    )






       def 			_lowercase (  self			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ):




              """simple docstring"""




              UpperCamelCase     :						Optional[Any] 						=		self._prepare_inference_func(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )
              return self._measure_memory(_inference    )






       def 			_lowercase (  self			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ):




              """simple docstring"""




              UpperCamelCase     :						Tuple 						=		self._prepare_train_func(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )
              return self._measure_speed(_train    )






       def 			_lowercase (  self			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ):




              """simple docstring"""




              UpperCamelCase     :						str 						=		self._prepare_train_func(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )
              return self._measure_memory(_train    )






       def 			_lowercase (  self			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ):




              """simple docstring"""




              UpperCamelCase     :						List[Any] 						=		self.config_dict[model_name]

              if self.args.torchscript:
                     UpperCamelCase     :						Dict 						=		True

              UpperCamelCase     :						int 						=		(
                  hasattr(__SCREAMING_SNAKE_CASE			,       '''architectures'''    )
                  and isinstance(config.architectures			,       __SCREAMING_SNAKE_CASE    )
                  and len(config.architectures    ) > 0
              )
              if not self.args.only_pretrain_model and has_model_class_in_config:
                     try:
                            UpperCamelCase     :						str 						=		config.architectures[0]
                            UpperCamelCase     :						Tuple 						=		__import__('''transformers'''			,       fromlist=[model_class]    )
                            UpperCamelCase     :						Dict 						=		getattr(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )
                            UpperCamelCase     :						Tuple 						=		model_cls(__SCREAMING_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     :						str 						=		MODEL_MAPPING[config.__class__](__SCREAMING_SNAKE_CASE    )

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

              # encoder-decoder has vocab size saved differently
              UpperCamelCase     :						Any 						=		config.vocab_size if hasattr(__SCREAMING_SNAKE_CASE			,       '''vocab_size'''    ) else config.encoder.vocab_size
              UpperCamelCase     :						int 						=		torch.randint(__SCREAMING_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     :						Optional[Any] 						=		torch.jit.trace(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )
              else:
                     UpperCamelCase     :						int 						=		model

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

              def encoder_forward():
                     with torch.no_grad():
                            UpperCamelCase     :						Tuple 						=		inference_model(__SCREAMING_SNAKE_CASE    )
                     return outputs

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






       def 			_lowercase (  self			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ):




              """simple docstring"""




              UpperCamelCase     :						int 						=		self.config_dict[model_name]

              UpperCamelCase     :						str 						=		(
                  hasattr(__SCREAMING_SNAKE_CASE			,       '''architectures'''    )
                  and isinstance(config.architectures			,       __SCREAMING_SNAKE_CASE    )
                  and len(config.architectures    ) > 0
              )
              if not self.args.only_pretrain_model and has_model_class_in_config:
                     try:
                            UpperCamelCase     :						int 						=		config.architectures[0]
                            UpperCamelCase     :						Union[str, Any] 						=		__import__('''transformers'''			,       fromlist=[model_class]    )
                            UpperCamelCase     :						int 						=		getattr(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )
                            UpperCamelCase     :						List[Any] 						=		model_cls(__SCREAMING_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     :						List[str] 						=		MODEL_WITH_LM_HEAD_MAPPING[config.__class__](__SCREAMING_SNAKE_CASE    )

              if self.args.torchscript:
                     raise NotImplementedError('''Training for torchscript is currently not implemented'''    )
              else:
                     UpperCamelCase     :						List[str] 						=		model

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

              # encoder-decoder has vocab size saved differently
              UpperCamelCase     :						Dict 						=		config.vocab_size if hasattr(__SCREAMING_SNAKE_CASE			,       '''vocab_size'''    ) else config.encoder.vocab_size
              UpperCamelCase     :						Optional[Any] 						=		torch.randint(__SCREAMING_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     :						Optional[int] 						=		train_model(__SCREAMING_SNAKE_CASE			,       labels=__SCREAMING_SNAKE_CASE    )[0]
                     loss.backward()
                     return loss

              def compute_loss_and_backprob_encoder_decoder():
                     UpperCamelCase     :						List[str] 						=		train_model(__SCREAMING_SNAKE_CASE			,       decoder_input_ids=__SCREAMING_SNAKE_CASE			,       labels=__SCREAMING_SNAKE_CASE    )[0]
                     loss.backward()
                     return loss

              UpperCamelCase     :						Optional[Any] 						=		(
                  compute_loss_and_backprob_encoder_decoder
                  if config.is_encoder_decoder
                  else compute_loss_and_backprob_encoder
              )
              return _train






       def 			_lowercase (  self			,       __SCREAMING_SNAKE_CASE    ):




              """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(
                                __SCREAMING_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     :						List[Any] 						=		timeit.repeat(
                         __SCREAMING_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(__SCREAMING_SNAKE_CASE    ) / 10.0
              except RuntimeError as e:
                     self.print_fn(f"""Doesn't fit on GPU. {e}"""    )
                     return "N/A"






       def 			_lowercase (  self			,       __SCREAMING_SNAKE_CASE    ):




              """simple docstring"""




              try:
                     if self.args.trace_memory_line_by_line:
                            UpperCamelCase     :						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.'''    )
                                   UpperCamelCase     :						Any 						=		'''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     :						Tuple 						=		nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx    )
                                   UpperCamelCase     :						List[str] 						=		nvml.nvmlDeviceGetMemoryInfo(__SCREAMING_SNAKE_CASE    )
                                   UpperCamelCase     :						Dict 						=		meminfo.used
                                   UpperCamelCase     :						List[Any] 						=		Memory(__SCREAMING_SNAKE_CASE    )
                                   # shutdown nvml
                                   nvml.nvmlShutdown()
                     else:
                            # cpu
                            UpperCamelCase     :						List[Any] 						=		measure_peak_memory_cpu(__SCREAMING_SNAKE_CASE    )
                            UpperCamelCase     :						Any 						=		Memory(__SCREAMING_SNAKE_CASE    ) if isinstance(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ) else memory_bytes

                     if self.args.trace_memory_line_by_line:
                            UpperCamelCase     :						Tuple 						=		stop_memory_tracing(__SCREAMING_SNAKE_CASE    )
                     else:
                            UpperCamelCase     :						List[Any] 						=		None

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





import tempfile
import unittest

import torch

from diffusers import UNetaDConditionModel
from diffusers.training_utils import EMAModel
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps, torch_device


enable_full_determinism()



class  UpperCAmelCase_						(				unittest.TestCase):
       '''simple docstring'''
       __UpperCamelCase   :  Tuple				=       "hf-internal-testing/tiny-stable-diffusion-pipe"
       __UpperCamelCase   :  Optional[Any]				=       1
       __UpperCamelCase   :  str				=       77
       __UpperCamelCase   :  Tuple				=       32
       __UpperCamelCase   :  List[str]				=       4
       __UpperCamelCase   :  Optional[int]				=       64
       __UpperCamelCase   :  Union[str, Any]				=       torch.manual_seed(0)






       def 			_lowercase (  self			,       __SCREAMING_SNAKE_CASE=0.9_999    ):




              """simple docstring"""




              UpperCamelCase     :						Optional[int] 						=		UNetaDConditionModel.from_pretrained(self.model_id			,       subfolder='''unet'''    )
              UpperCamelCase     :						Any 						=		unet.to(__SCREAMING_SNAKE_CASE    )
              UpperCamelCase     :						Dict 						=		EMAModel(unet.parameters()			,       decay=__SCREAMING_SNAKE_CASE			,       model_cls=__SCREAMING_SNAKE_CASE			,       model_config=unet.config    )
              return unet, ema_unet






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase     :						Tuple 						=		torch.randn(
                  self.batch_size			,       self.num_in_channels			,       self.latent_height			,       self.latent_width			,       generator=self.generator    ).to(__SCREAMING_SNAKE_CASE    )
              UpperCamelCase     :						Optional[int] 						=		torch.randint(0			,       1_000			,       size=(self.batch_size,)			,       generator=self.generator    ).to(__SCREAMING_SNAKE_CASE    )
              UpperCamelCase     :						List[Any] 						=		torch.randn(
                  self.batch_size			,       self.prompt_length			,       self.text_encoder_hidden_dim			,       generator=self.generator    ).to(__SCREAMING_SNAKE_CASE    )
              return noisy_latents, timesteps, encoder_hidden_states






       def 			_lowercase (  self			,       __SCREAMING_SNAKE_CASE    ):




              """simple docstring"""




              UpperCamelCase     :						Dict 						=		{}
              for k, param in unet.state_dict().items():
                     UpperCamelCase     :						int 						=		torch.randn_like(__SCREAMING_SNAKE_CASE    ) + (param * torch.randn_like(__SCREAMING_SNAKE_CASE    ))
                     updated_state_dict.update({k: updated_param}    )
              unet.load_state_dict(__SCREAMING_SNAKE_CASE    )
              return unet






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase						,			UpperCamelCase     :						Optional[Any] 						=		self.get_models()
              # Take the first (hypothetical) EMA step.
              ema_unet.step(unet.parameters()    )
              assert ema_unet.optimization_step == 1

              # Take two more.
              for _ in range(2    ):
                     ema_unet.step(unet.parameters()    )
              assert ema_unet.optimization_step == 3






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase						,			UpperCamelCase     :						List[Any] 						=		self.get_models()
              # Since the `unet` is not being updated (i.e., backprop'd)
              # there won't be any difference between the `params` of `unet`
              # and `ema_unet` even if we call `ema_unet.step(unet.parameters())`.
              ema_unet.step(unet.parameters()    )
              UpperCamelCase     :						int 						=		list(unet.parameters()    )
              for s_param, param in zip(ema_unet.shadow_params			,       __SCREAMING_SNAKE_CASE    ):
                     assert torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )

              # The above holds true even if we call `ema.step()` multiple times since
              # `unet` params are still not being updated.
              for _ in range(4    ):
                     ema_unet.step(unet.parameters()    )
              for s_param, param in zip(ema_unet.shadow_params			,       __SCREAMING_SNAKE_CASE    ):
                     assert torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase						,			UpperCamelCase     :						Any 						=		self.get_models()
              # Here we simulate the parameter updates for `unet`. Since there might
              # be some parameters which are initialized to zero we take extra care to
              # initialize their values to something non-zero before the multiplication.
              UpperCamelCase     :						Union[str, Any] 						=		self.simulate_backprop(__SCREAMING_SNAKE_CASE    )

              # Take the EMA step.
              ema_unet.step(unet_pseudo_updated_step_one.parameters()    )

              # Now the EMA'd parameters won't be equal to the original model parameters.
              UpperCamelCase     :						Any 						=		list(unet_pseudo_updated_step_one.parameters()    )
              for s_param, param in zip(ema_unet.shadow_params			,       __SCREAMING_SNAKE_CASE    ):
                     assert ~torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )

              # Ensure this is the case when we take multiple EMA steps.
              for _ in range(4    ):
                     ema_unet.step(unet.parameters()    )
              for s_param, param in zip(ema_unet.shadow_params			,       __SCREAMING_SNAKE_CASE    ):
                     assert ~torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase						,			UpperCamelCase     :						Optional[int] 						=		self.get_models()

              # First backprop + EMA
              UpperCamelCase     :						Dict 						=		self.simulate_backprop(__SCREAMING_SNAKE_CASE    )
              ema_unet.step(unet_step_one.parameters()    )
              UpperCamelCase     :						List[str] 						=		ema_unet.shadow_params

              # Second backprop + EMA
              UpperCamelCase     :						str 						=		self.simulate_backprop(__SCREAMING_SNAKE_CASE    )
              ema_unet.step(unet_step_two.parameters()    )
              UpperCamelCase     :						List[Any] 						=		ema_unet.shadow_params

              for step_one, step_two in zip(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ):
                     assert ~torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase						,			UpperCamelCase     :						int 						=		self.get_models(decay=0.0    )
              UpperCamelCase     :						Optional[int] 						=		self.simulate_backprop(__SCREAMING_SNAKE_CASE    )
              ema_unet.step(unet_step_one.parameters()    )
              UpperCamelCase     :						Union[str, Any] 						=		ema_unet.shadow_params

              UpperCamelCase     :						Any 						=		self.simulate_backprop(__SCREAMING_SNAKE_CASE    )
              ema_unet.step(unet_step_two.parameters()    )
              UpperCamelCase     :						int 						=		ema_unet.shadow_params

              for step_one, step_two in zip(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ):
                     assert torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )






       @skip_mps
       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase						,			UpperCamelCase     :						str 						=		self.get_models()
              UpperCamelCase						,			UpperCamelCase						,			UpperCamelCase     :						Union[str, Any] 						=		self.get_dummy_inputs()

              with tempfile.TemporaryDirectory() as tmpdir:
                     ema_unet.save_pretrained(__SCREAMING_SNAKE_CASE    )
                     UpperCamelCase     :						Optional[int] 						=		UNetaDConditionModel.from_pretrained(__SCREAMING_SNAKE_CASE			,       model_cls=__SCREAMING_SNAKE_CASE    )
                     UpperCamelCase     :						List[str] 						=		loaded_unet.to(unet.device    )

              # Since no EMA step has been performed the outputs should match.
              UpperCamelCase     :						Optional[Any] 						=		unet(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ).sample
              UpperCamelCase     :						List[str] 						=		loaded_unet(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ).sample

              assert torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       atol=1e-4    )

