


'''simple docstring'''






import random
import timeit
from functools import wraps
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_tf_auto import TF_MODEL_MAPPING, TF_MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_tf_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_tf_available():
  import tensorflow as tf
  from tensorflow.python.framework.errors_impl import ResourceExhaustedError

  from .benchmark_args_tf import TensorFlowBenchmarkArguments

if is_pyanvml_available():
  import pyanvml.pyanvml as nvml

__A :  str    							=				logging.get_logger(__name__)







def     UpperCamelCase_      (   A__ :       Dict							,   A__ :       str ):
  '''simple docstring'''





  def run_func(A__ :       Union[str, Any] ):
    @wraps(__a )
    def run_in_eager_mode(*A__ :       Union[str, Any]							,   **A__ :       List[Any] ):
      return func(*__a							,   **__a )

    @wraps(__a )
    @tf.function(experimental_compile=__a )
    def run_in_graph_mode(*A__ :       List[str]							,   **A__ :       Tuple ):
      return func(*__a							,   **__a )

    if do_eager_mode is True:
      if use_xla is not False:
        raise ValueError(
            """Cannot run model in XLA, if `args.eager_mode` is set to `True`. Please set `args.eager_mode=False`.""" )
      return run_in_eager_mode
    else:
      return run_in_graph_mode

  return run_func







def     UpperCamelCase_      (   A__ :       str							,   A__ :       Optional[int]							,   A__ :       Any ):
  '''simple docstring'''





  lowerCAmelCase_      :				Dict      =       random.Random()
  lowerCAmelCase_      :				int      =       [rng.randint(0							,   vocab_size - 1 ) for i in range(batch_size * sequence_length )]
  return tf.constant(__a							,   shape=(batch_size, sequence_length)							,   dtype=tf.intaa )






class        __snake_case	(       snake_case_):







    """simple docstring"""



    lowercase       			= 42
    lowercase       			= 42
    lowercase       			= "TensorFlow"

    @property
    def __lowercase       (       self   :  List[Any]		)       ->  int:
      return tf.__version__

    def __lowercase       (       self   :  Optional[Any]   ,						lowerCamelCase   :  str   ,						lowerCamelCase   :  int   ,						lowerCamelCase   :  int		)       ->  float:
      lowerCAmelCase_      :				Dict      =       self.args.strategy
      if strategy is None:
        raise ValueError("""A device strategy has to be initialized before using TensorFlow."""		)
      lowerCAmelCase_      :				Union[str, Any]      =       self._prepare_inference_func(_A   ,						_A   ,						_A		)
      return self._measure_speed(_inference		)

    def __lowercase       (       self   :  Tuple   ,						lowerCamelCase   :  str   ,						lowerCamelCase   :  int   ,						lowerCamelCase   :  int		)       ->  float:
      lowerCAmelCase_      :				List[str]      =       self.args.strategy
      if strategy is None:
        raise ValueError("""A device strategy has to be initialized before using TensorFlow."""		)
      lowerCAmelCase_      :				Optional[int]      =       self._prepare_train_func(_A   ,						_A   ,						_A		)
      return self._measure_speed(_train		)

    def __lowercase       (       self   :  Optional[int]   ,						lowerCamelCase   :  str   ,						lowerCamelCase   :  int   ,						lowerCamelCase   :  int		)       ->  [Memory, Optional[MemorySummary]]:


      if self.args.is_gpu:
        tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]   ,						_A		)
      lowerCAmelCase_      :				Any      =       self.args.strategy
      if strategy is None:
        raise ValueError("""A device strategy has to be initialized before using TensorFlow."""		)
      lowerCAmelCase_      :				Any      =       self._prepare_inference_func(_A   ,						_A   ,						_A		)
      return self._measure_memory(_inference		)

    def __lowercase       (       self   :  Dict   ,						lowerCamelCase   :  str   ,						lowerCamelCase   :  int   ,						lowerCamelCase   :  int		)       ->  [Memory, Optional[MemorySummary]]:


      if self.args.is_gpu:
        tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]   ,						_A		)
      lowerCAmelCase_      :				List[Any]      =       self.args.strategy
      if strategy is None:
        raise ValueError("""A device strategy has to be initialized before using TensorFlow."""		)

      lowerCAmelCase_      :				Union[str, Any]      =       self._prepare_train_func(_A   ,						_A   ,						_A		)
      return self._measure_memory(_train		)

    def __lowercase       (       self   :  str   ,						lowerCamelCase   :  str   ,						lowerCamelCase   :  int   ,						lowerCamelCase   :  int		)       ->  Callable[[], None]:
      lowerCAmelCase_      :				Tuple      =       self.config_dict[model_name]

      if self.args.fpaa:
        raise NotImplementedError("""Mixed precision is currently not supported."""		)

      lowerCAmelCase_      :				Tuple      =       (
          hasattr(_A   ,						"""architectures"""		)
          and isinstance(config.architectures   ,						_A		)
          and len(config.architectures		) > 0
      )
      if not self.args.only_pretrain_model and has_model_class_in_config:
        try:
          lowerCAmelCase_      :				Dict      =       'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
          lowerCAmelCase_      :				Union[str, Any]      =       __import__("""transformers"""   ,						fromlist=[model_class]		)
          lowerCAmelCase_      :				Optional[Any]      =       getattr(_A   ,						_A		)
          lowerCAmelCase_      :				Dict      =       model_cls(_A		)
        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_      :				int      =       TF_MODEL_MAPPING[config.__class__](_A		)

      # encoder-decoder has vocab size saved differently
      lowerCAmelCase_      :				Union[str, Any]      =       config.vocab_size if hasattr(_A   ,						"""vocab_size"""		) else config.encoder.vocab_size
      lowerCAmelCase_      :				List[str]      =       random_input_ids(_A   ,						_A   ,						_A		)

      @run_with_tf_optimizations(self.args.eager_mode   ,						self.args.use_xla		)
      def encoder_decoder_forward():
        return model(_A   ,						decoder_input_ids=_A   ,						training=_A		)

      @run_with_tf_optimizations(self.args.eager_mode   ,						self.args.use_xla		)
      def encoder_forward():
        return model(_A   ,						training=_A		)

      lowerCAmelCase_      :				List[Any]      =       encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

      return _inference

    def __lowercase       (       self   :  Tuple   ,						lowerCamelCase   :  str   ,						lowerCamelCase   :  int   ,						lowerCamelCase   :  int		)       ->  Callable[[], None]:
      lowerCAmelCase_      :				Any      =       self.config_dict[model_name]

      if self.args.eager_mode is not False:
        raise ValueError("""Training cannot be done in eager mode. Please make sure that `args.eager_mode = False`."""		)

      if self.args.fpaa:
        raise NotImplementedError("""Mixed precision is currently not supported."""		)

      lowerCAmelCase_      :				Union[str, Any]      =       (
          hasattr(_A   ,						"""architectures"""		)
          and isinstance(config.architectures   ,						_A		)
          and len(config.architectures		) > 0
      )
      if not self.args.only_pretrain_model and has_model_class_in_config:
        try:
          lowerCAmelCase_      :				Optional[Any]      =       'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
          lowerCAmelCase_      :				Dict      =       __import__("""transformers"""   ,						fromlist=[model_class]		)
          lowerCAmelCase_      :				List[Any]      =       getattr(_A   ,						_A		)
          lowerCAmelCase_      :				Union[str, Any]      =       model_cls(_A		)
        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_      :				str      =       TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](_A		)

      # encoder-decoder has vocab size saved differently
      lowerCAmelCase_      :				Tuple      =       config.vocab_size if hasattr(_A   ,						"""vocab_size"""		) else config.encoder.vocab_size
      lowerCAmelCase_      :				int      =       random_input_ids(_A   ,						_A   ,						_A		)

      @run_with_tf_optimizations(self.args.eager_mode   ,						self.args.use_xla		)
      def encoder_decoder_train():
        lowerCAmelCase_      :				Tuple      =       model(_A   ,						decoder_input_ids=_A   ,						labels=_A   ,						training=_A		)[0]
        lowerCAmelCase_      :				int      =       tf.gradients(_A   ,						model.trainable_variables		)
        return gradients

      @run_with_tf_optimizations(self.args.eager_mode   ,						self.args.use_xla		)
      def encoder_train():
        lowerCAmelCase_      :				Any      =       model(_A   ,						labels=_A   ,						training=_A		)[0]
        lowerCAmelCase_      :				Dict      =       tf.gradients(_A   ,						model.trainable_variables		)
        return gradients

      lowerCAmelCase_      :				int      =       encoder_decoder_train if config.is_encoder_decoder else encoder_train

      return _train

    def __lowercase       (       self   :  int   ,						lowerCamelCase   :  Union[str, Any]		)       ->  float:


      with self.args.strategy.scope():
        try:
          if self.args.is_tpu or self.args.use_xla:
            # run additional 10 times to stabilize compilation for tpu
            logger.info("""Do inference on TPU. Running model 5 times to stabilize compilation"""		)
            timeit.repeat(_A   ,						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_      :				Tuple      =       timeit.repeat(
              _A   ,						repeat=self.args.repeat   ,						number=10   ,						)

          return min(_A		) / 10.0
        except ResourceExhaustedError as e:
          self.print_fn(F'Doesn\'t fit on GPU. {e}'		)





    def __lowercase       (       self   :  Dict   ,						lowerCamelCase   :  Callable[[], None]		)       ->  [Memory, MemorySummary]:
      logger.info(
          """Note that TensorFlow allocates more memory than """
          """it might need to speed up computation. """
          """The memory reported here corresponds to the memory """
          """reported by `nvidia-smi`, which can vary depending """
          """on total available memory on the GPU that is used."""		)
      with self.args.strategy.scope():
        try:
          if self.args.trace_memory_line_by_line:
            if not self.args.eager_mode:
              raise ValueError(
                  """`args.eager_mode` is set to `False`. Make sure to run model in eager mode to measure memory"""
                  """ consumption line by line."""		)
            lowerCAmelCase_      :				Any      =       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 `args.memory=False`"""		)
          elif self.args.is_gpu:
            # 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_      :				str      =       '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_      :				List[Any]      =       nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx		)
              lowerCAmelCase_      :				Dict      =       nvml.nvmlDeviceGetMemoryInfo(_A		)
              lowerCAmelCase_      :				Tuple      =       meminfo.used
              lowerCAmelCase_      :				Optional[int]      =       Memory(_A		)
              # shutdown nvml
              nvml.nvmlShutdown()
          else:
            # cpu
            if self.args.trace_memory_line_by_line:
              logger.info(
                  """When enabling line by line tracing, the max peak memory for CPU is inaccurate in"""
                  """ TensorFlow."""		)
              lowerCAmelCase_      :				List[str]      =       None
            else:
              lowerCAmelCase_      :				Tuple      =       measure_peak_memory_cpu(_A		)
              lowerCAmelCase_      :				Optional[Any]      =       Memory(_A		) if isinstance(_A   ,						_A		) else memory_bytes
          if self.args.trace_memory_line_by_line:
            lowerCAmelCase_      :				List[Any]      =       stop_memory_tracing(_A		)
            if memory is None:
              lowerCAmelCase_      :				int      =       summary.total
          else:
            lowerCAmelCase_      :				List[str]      =       None

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










'''simple docstring'''






import unittest

from diffusers.pipelines.pipeline_utils import is_safetensors_compatible






class        __snake_case	(       unittest.TestCase):







    """simple docstring"""



    def __lowercase       (       self   :  Tuple		)       ->  Dict:
      lowerCAmelCase_      :				str      =       [
          """safety_checker/pytorch_model.bin""",
          """safety_checker/model.safetensors""",
          """vae/diffusion_pytorch_model.bin""",
          """vae/diffusion_pytorch_model.safetensors""",
          """text_encoder/pytorch_model.bin""",
          """text_encoder/model.safetensors""",
          """unet/diffusion_pytorch_model.bin""",
          """unet/diffusion_pytorch_model.safetensors""",
      ]
      self.assertTrue(is_safetensors_compatible(lowerCamelCase		)		)

    def __lowercase       (       self   :  List[Any]		)       ->  int:
      lowerCAmelCase_      :				Tuple      =       [
          """unet/diffusion_pytorch_model.bin""",
          """unet/diffusion_pytorch_model.safetensors""",
      ]
      self.assertTrue(is_safetensors_compatible(lowerCamelCase		)		)

    def __lowercase       (       self   :  Optional[int]		)       ->  Optional[Any]:
      lowerCAmelCase_      :				int      =       [
          """safety_checker/pytorch_model.bin""",
          """safety_checker/model.safetensors""",
          """vae/diffusion_pytorch_model.bin""",
          """vae/diffusion_pytorch_model.safetensors""",
          """text_encoder/pytorch_model.bin""",
          """text_encoder/model.safetensors""",
          """unet/diffusion_pytorch_model.bin""",
          # Removed: 'unet/diffusion_pytorch_model.safetensors',
      ]
      self.assertFalse(is_safetensors_compatible(lowerCamelCase		)		)

    def __lowercase       (       self   :  int		)       ->  List[Any]:
      lowerCAmelCase_      :				Dict      =       [
          """text_encoder/pytorch_model.bin""",
          """text_encoder/model.safetensors""",
      ]
      self.assertTrue(is_safetensors_compatible(lowerCamelCase		)		)

    def __lowercase       (       self   :  str		)       ->  List[str]:
      lowerCAmelCase_      :				Union[str, Any]      =       [
          """safety_checker/pytorch_model.bin""",
          """safety_checker/model.safetensors""",
          """vae/diffusion_pytorch_model.bin""",
          """vae/diffusion_pytorch_model.safetensors""",
          """text_encoder/pytorch_model.bin""",
          # Removed: 'text_encoder/model.safetensors',
          """unet/diffusion_pytorch_model.bin""",
          """unet/diffusion_pytorch_model.safetensors""",
      ]
      self.assertFalse(is_safetensors_compatible(lowerCamelCase		)		)

    def __lowercase       (       self   :  List[Any]		)       ->  Tuple:
      lowerCAmelCase_      :				Union[str, Any]      =       [
          """safety_checker/pytorch_model.fp16.bin""",
          """safety_checker/model.fp16.safetensors""",
          """vae/diffusion_pytorch_model.fp16.bin""",
          """vae/diffusion_pytorch_model.fp16.safetensors""",
          """text_encoder/pytorch_model.fp16.bin""",
          """text_encoder/model.fp16.safetensors""",
          """unet/diffusion_pytorch_model.fp16.bin""",
          """unet/diffusion_pytorch_model.fp16.safetensors""",
      ]
      lowerCAmelCase_      :				Union[str, Any]      =       """fp16"""
      self.assertTrue(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)

    def __lowercase       (       self   :  Optional[Any]		)       ->  List[str]:
      lowerCAmelCase_      :				str      =       [
          """unet/diffusion_pytorch_model.fp16.bin""",
          """unet/diffusion_pytorch_model.fp16.safetensors""",
      ]
      lowerCAmelCase_      :				Optional[int]      =       """fp16"""
      self.assertTrue(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)

    def __lowercase       (       self   :  Tuple		)       ->  Optional[Any]:
      # pass variant but use the non-variant filenames
      lowerCAmelCase_      :				Dict      =       [
          """unet/diffusion_pytorch_model.bin""",
          """unet/diffusion_pytorch_model.safetensors""",
      ]
      lowerCAmelCase_      :				str      =       """fp16"""
      self.assertTrue(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)

    def __lowercase       (       self   :  Optional[int]		)       ->  List[str]:
      lowerCAmelCase_      :				str      =       [
          """safety_checker/pytorch_model.fp16.bin""",
          """safety_checker/model.fp16.safetensors""",
          """vae/diffusion_pytorch_model.fp16.bin""",
          """vae/diffusion_pytorch_model.fp16.safetensors""",
          """text_encoder/pytorch_model.fp16.bin""",
          """text_encoder/model.fp16.safetensors""",
          """unet/diffusion_pytorch_model.fp16.bin""",
          # Removed: 'unet/diffusion_pytorch_model.fp16.safetensors',
      ]
      lowerCAmelCase_      :				List[str]      =       """fp16"""
      self.assertFalse(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)

    def __lowercase       (       self   :  Union[str, Any]		)       ->  Optional[int]:
      lowerCAmelCase_      :				str      =       [
          """text_encoder/pytorch_model.fp16.bin""",
          """text_encoder/model.fp16.safetensors""",
      ]
      lowerCAmelCase_      :				Optional[Any]      =       """fp16"""
      self.assertTrue(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)

    def __lowercase       (       self   :  List[Any]		)       ->  List[Any]:
      # pass variant but use the non-variant filenames
      lowerCAmelCase_      :				Dict      =       [
          """text_encoder/pytorch_model.bin""",
          """text_encoder/model.safetensors""",
      ]
      lowerCAmelCase_      :				Any      =       """fp16"""
      self.assertTrue(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)





    def __lowercase       (       self   :  Dict		)       ->  Any:
      lowerCAmelCase_      :				Optional[int]      =       [
          """safety_checker/pytorch_model.fp16.bin""",
          """safety_checker/model.fp16.safetensors""",
          """vae/diffusion_pytorch_model.fp16.bin""",
          """vae/diffusion_pytorch_model.fp16.safetensors""",
          """text_encoder/pytorch_model.fp16.bin""",
          # 'text_encoder/model.fp16.safetensors',
          """unet/diffusion_pytorch_model.fp16.bin""",
          """unet/diffusion_pytorch_model.fp16.safetensors""",
      ]
      lowerCAmelCase_      :				int      =       """fp16"""
      self.assertFalse(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)



