





import random
import unittest

import numpy as np

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionImgaImgPipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
       import onnxruntime as ort


class        __SCREAMING_SNAKE_CASE      (			snake_case__					,unittest.TestCase ):






       """simple docstring"""




       _UpperCAmelCase				:  Union[str, Any]			   =   """hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"""







       def 						__lowerCamelCase   (				self					:     Any		, A					:     str=0      )       ->Dict:
          lowerCamelCase__       :						Dict        =							floats_tensor((1, 3, 1_2_8, 1_2_8)		, rng=random.Random(_A      )      )
          lowerCamelCase__       :						Dict        =							np.random.RandomState(_A      )
          lowerCamelCase__       :						Tuple        =							{
              '''prompt''': '''A painting of a squirrel eating a burger''',
              '''image''': image,
              '''generator''': generator,
              '''num_inference_steps''': 3,
              '''strength''': 0.75,
              '''guidance_scale''': 7.5,
              '''output_type''': '''numpy''',
          }
          return inputs







       def 						__lowerCamelCase   (				self					:     Dict      )       ->Optional[Any]:
          lowerCamelCase__       :						List[Any]        =							OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint		, provider='''CPUExecutionProvider'''      )
          pipe.set_progress_bar_config(disable=_A      )

          lowerCamelCase__       :						Optional[int]        =							self.get_dummy_inputs()
          lowerCamelCase__       :						Optional[int]        =							pipe(**_A      ).images
          lowerCamelCase__       :						Union[str, Any]        =							image[0, -3:, -3:, -1].flatten()

          assert image.shape == (1, 1_2_8, 1_2_8, 3)
          lowerCamelCase__       :						List[Any]        =							np.array([0.6_96_43, 0.5_84_84, 0.5_03_14, 0.5_87_60, 0.5_53_68, 0.5_96_43, 0.5_15_29, 0.4_12_17, 0.4_90_87]      )
          assert np.abs(image_slice - expected_slice      ).max() < 1e-1







       def 						__lowerCamelCase   (				self					:     List[str]      )       ->int:
          lowerCamelCase__       :						Optional[int]        =							OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint		, provider='''CPUExecutionProvider'''      )
          lowerCamelCase__       :						Tuple        =							PNDMScheduler.from_config(pipe.scheduler.config		, skip_prk_steps=_A      )
          pipe.set_progress_bar_config(disable=_A      )

          lowerCamelCase__       :						Union[str, Any]        =							self.get_dummy_inputs()
          lowerCamelCase__       :						int        =							pipe(**_A      ).images
          lowerCamelCase__       :						List[str]        =							image[0, -3:, -3:, -1]

          assert image.shape == (1, 1_2_8, 1_2_8, 3)
          lowerCamelCase__       :						List[str]        =							np.array([0.6_17_37, 0.5_46_42, 0.5_31_83, 0.5_44_65, 0.5_27_42, 0.6_05_25, 0.4_99_69, 0.4_06_55, 0.4_81_54]      )

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







       def 						__lowerCamelCase   (				self					:     List[Any]      )       ->List[str]:
          lowerCamelCase__       :						Tuple        =							OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint		, provider='''CPUExecutionProvider'''      )
          lowerCamelCase__       :						Union[str, Any]        =							LMSDiscreteScheduler.from_config(pipe.scheduler.config      )
          pipe.set_progress_bar_config(disable=_A      )

          # warmup pass to apply optimizations
          lowerCamelCase__       :						int        =							pipe(**self.get_dummy_inputs()      )

          lowerCamelCase__       :						str        =							self.get_dummy_inputs()
          lowerCamelCase__       :						List[str]        =							pipe(**_A      ).images
          lowerCamelCase__       :						Optional[int]        =							image[0, -3:, -3:, -1]

          assert image.shape == (1, 1_2_8, 1_2_8, 3)
          lowerCamelCase__       :						List[Any]        =							np.array([0.5_27_61, 0.5_99_77, 0.4_90_33, 0.4_96_19, 0.5_42_82, 0.5_03_11, 0.4_76_00, 0.4_09_18, 0.4_52_03]      )

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







       def 						__lowerCamelCase   (				self					:     Dict      )       ->str:
          lowerCamelCase__       :						Any        =							OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint		, provider='''CPUExecutionProvider'''      )
          lowerCamelCase__       :						Optional[Any]        =							EulerDiscreteScheduler.from_config(pipe.scheduler.config      )
          pipe.set_progress_bar_config(disable=_A      )

          lowerCamelCase__       :						Tuple        =							self.get_dummy_inputs()
          lowerCamelCase__       :						Tuple        =							pipe(**_A      ).images
          lowerCamelCase__       :						str        =							image[0, -3:, -3:, -1]

          assert image.shape == (1, 1_2_8, 1_2_8, 3)
          lowerCamelCase__       :						Optional[Any]        =							np.array([0.5_29_11, 0.6_00_04, 0.4_92_29, 0.4_98_05, 0.5_45_02, 0.5_06_80, 0.4_77_77, 0.4_10_28, 0.4_53_04]      )

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







       def 						__lowerCamelCase   (				self					:     Optional[Any]      )       ->str:
          lowerCamelCase__       :						Tuple        =							OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint		, provider='''CPUExecutionProvider'''      )
          lowerCamelCase__       :						Tuple        =							EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config      )
          pipe.set_progress_bar_config(disable=_A      )

          lowerCamelCase__       :						Union[str, Any]        =							self.get_dummy_inputs()
          lowerCamelCase__       :						List[Any]        =							pipe(**_A      ).images
          lowerCamelCase__       :						Optional[Any]        =							image[0, -3:, -3:, -1]

          assert image.shape == (1, 1_2_8, 1_2_8, 3)
          lowerCamelCase__       :						str        =							np.array([0.5_29_11, 0.6_00_04, 0.4_92_29, 0.4_98_05, 0.5_45_02, 0.5_06_80, 0.4_77_77, 0.4_10_28, 0.4_53_04]      )

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







       def 						__lowerCamelCase   (				self					:     List[str]      )       ->Dict:
          lowerCamelCase__       :						Any        =							OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint		, provider='''CPUExecutionProvider'''      )
          lowerCamelCase__       :						Union[str, Any]        =							DPMSolverMultistepScheduler.from_config(pipe.scheduler.config      )
          pipe.set_progress_bar_config(disable=_A      )

          lowerCamelCase__       :						Tuple        =							self.get_dummy_inputs()
          lowerCamelCase__       :						Optional[int]        =							pipe(**_A      ).images
          lowerCamelCase__       :						Optional[Any]        =							image[0, -3:, -3:, -1]

          assert image.shape == (1, 1_2_8, 1_2_8, 3)
          lowerCamelCase__       :						Optional[Any]        =							np.array([0.6_53_31, 0.5_82_77, 0.4_82_04, 0.5_60_59, 0.5_36_65, 0.5_62_35, 0.5_09_69, 0.4_00_09, 0.4_65_52]      )

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


@nightly
@require_onnxruntime
@require_torch_gpu
class        __SCREAMING_SNAKE_CASE      (			unittest.TestCase ):






       """simple docstring"""




       @property
       def 						__lowerCamelCase   (				self					:     List[str]      )       ->str:
          return (
              "CUDAExecutionProvider",
              {
                  "gpu_mem_limit": "15000000000",  # 15GB
                  "arena_extend_strategy": "kSameAsRequested",
              },
          )







       @property
       def 						__lowerCamelCase   (				self					:     List[Any]      )       ->str:
          lowerCamelCase__       :						Union[str, Any]        =							ort.SessionOptions()
          lowerCamelCase__       :						Optional[Any]        =							False
          return options







       def 						__lowerCamelCase   (				self					:     List[str]      )       ->Any:
          lowerCamelCase__       :						Optional[int]        =							load_image(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
              '''/img2img/sketch-mountains-input.jpg'''      )
          lowerCamelCase__       :						int        =							init_image.resize((7_6_8, 5_1_2)      )
          # using the PNDM scheduler by default
          lowerCamelCase__       :						List[str]        =							OnnxStableDiffusionImgaImgPipeline.from_pretrained(
              '''CompVis/stable-diffusion-v1-4'''		, revision='''onnx'''		, safety_checker=_A		, feature_extractor=_A		, provider=self.gpu_provider		, sess_options=self.gpu_options		, )
          pipe.set_progress_bar_config(disable=_A      )

          lowerCamelCase__       :						Dict        =							'''A fantasy landscape, trending on artstation'''

          lowerCamelCase__       :						str        =							np.random.RandomState(0      )
          lowerCamelCase__       :						Optional[int]        =							pipe(
              prompt=_A		, image=_A		, strength=0.75		, guidance_scale=7.5		, num_inference_steps=1_0		, generator=_A		, output_type='''np'''		, )
          lowerCamelCase__       :						Tuple        =							output.images
          lowerCamelCase__       :						Any        =							images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

          assert images.shape == (1, 5_1_2, 7_6_8, 3)
          lowerCamelCase__       :						List[str]        =							np.array([0.49_09, 0.50_59, 0.53_72, 0.46_23, 0.48_76, 0.50_49, 0.48_20, 0.49_56, 0.50_19]      )
          # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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







       def 						__lowerCamelCase   (				self					:     int      )       ->Optional[Any]:
          lowerCamelCase__       :						Any        =							load_image(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
              '''/img2img/sketch-mountains-input.jpg'''      )
          lowerCamelCase__       :						Any        =							init_image.resize((7_6_8, 5_1_2)      )
          lowerCamelCase__       :						Dict        =							LMSDiscreteScheduler.from_pretrained(
              '''runwayml/stable-diffusion-v1-5'''		, subfolder='''scheduler'''		, revision='''onnx'''      )
          lowerCamelCase__       :						Dict        =							OnnxStableDiffusionImgaImgPipeline.from_pretrained(
              '''runwayml/stable-diffusion-v1-5'''		, revision='''onnx'''		, scheduler=_A		, safety_checker=_A		, feature_extractor=_A		, provider=self.gpu_provider		, sess_options=self.gpu_options		, )
          pipe.set_progress_bar_config(disable=_A      )

          lowerCamelCase__       :						Union[str, Any]        =							'''A fantasy landscape, trending on artstation'''

          lowerCamelCase__       :						List[str]        =							np.random.RandomState(0      )
          lowerCamelCase__       :						List[Any]        =							pipe(
              prompt=_A		, image=_A		, strength=0.75		, guidance_scale=7.5		, num_inference_steps=2_0		, generator=_A		, output_type='''np'''		, )
          lowerCamelCase__       :						List[Any]        =							output.images
          lowerCamelCase__       :						Any        =							images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

          assert images.shape == (1, 5_1_2, 7_6_8, 3)
          lowerCamelCase__       :						Union[str, Any]        =							np.array([0.80_43, 0.9_26, 0.95_81, 0.81_19, 0.89_54, 0.9_13, 0.72_09, 0.74_63, 0.74_31]      )
          # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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










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 :   Any  			=  logging.get_logger(__name__)



def 	_a						(   UpperCAmelCase      ,  UpperCAmelCase					)    ->    List[str]:







   """simple docstring"""
   def run_func(UpperCAmelCase					):
      @wraps(UpperCAmelCase					)
      def run_in_eager_mode(*UpperCAmelCase      ,  **UpperCAmelCase					):
         return func(*UpperCAmelCase      ,  **UpperCAmelCase					)

      @wraps(UpperCAmelCase					)
      @tf.function(experimental_compile=UpperCAmelCase					)
      def run_in_graph_mode(*UpperCAmelCase      ,  **UpperCAmelCase					):
         return func(*UpperCAmelCase      ,  **UpperCAmelCase					)

      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 	_a						(   UpperCAmelCase      ,  UpperCAmelCase      ,  UpperCAmelCase					)    ->    ["tf.Tensor"]:







   """simple docstring"""
   lowerCamelCase__       :						List[Any]        =							random.Random()
   lowerCamelCase__       :						str        =							[rng.randint(0      ,  vocab_size - 1					) for i in range(batch_size * sequence_length					)]
   return tf.constant(UpperCAmelCase      ,  shape=(batch_size, sequence_length)      ,  dtype=tf.intaa					)




class        __SCREAMING_SNAKE_CASE      (			lowerCAmelCase_ ):
       _UpperCAmelCase				:  TensorFlowBenchmarkArguments
       _UpperCAmelCase				:  PretrainedConfig
       _UpperCAmelCase				:  str			   =   "TensorFlow"







       @property
       def 						__lowerCamelCase   (				self					:     int      )       ->Optional[int]:
          return tf.__version__







       def 						__lowerCamelCase   (				self					:     Optional[int]		, A					:     str		, A					:     int		, A					:     int      )       ->float:
          # initialize GPU on separate process
          lowerCamelCase__       :						Dict        =							self.args.strategy
          if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''      )
          lowerCamelCase__       :						int        =							self._prepare_inference_func(A		, A		, A      )
          return self._measure_speed(_inference      )







       def 						__lowerCamelCase   (				self					:     str		, A					:     str		, A					:     int		, A					:     int      )       ->float:
          lowerCamelCase__       :						Optional[int]        =							self.args.strategy
          if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''      )
          lowerCamelCase__       :						List[Any]        =							self._prepare_train_func(A		, A		, A      )
          return self._measure_speed(_train      )







       def 						__lowerCamelCase   (				self					:     int		, A					:     str		, A					:     int		, A					:     int      )       ->[Memory, Optional[MemorySummary]]:
          # initialize GPU on separate process
          if self.args.is_gpu:
             tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]		, A      )
          lowerCamelCase__       :						int        =							self.args.strategy
          if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''      )
          lowerCamelCase__       :						str        =							self._prepare_inference_func(A		, A		, A      )
          return self._measure_memory(_inference      )







       def 						__lowerCamelCase   (				self					:     List[str]		, A					:     str		, A					:     int		, A					:     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__       :						str        =							self._prepare_train_func(A		, A		, A      )
          return self._measure_memory(_train      )







       def 						__lowerCamelCase   (				self					:     Dict		, A					:     str		, A					:     int		, A					:     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__       :						Any        =							'''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
                lowerCamelCase__       :						List[Any]        =							__import__('''transformers'''		, fromlist=[model_class]      )
                lowerCamelCase__       :						int        =							getattr(A		, A      )
                lowerCamelCase__       :						int        =							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__       :						Union[str, Any]        =							TF_MODEL_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__       :						Optional[Any]        =							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__       :						int        =							encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

          return _inference







       def 						__lowerCamelCase   (				self					:     List[str]		, A					:     str		, A					:     int		, A					:     int      )       ->Callable[[], None]:
          lowerCamelCase__       :						Tuple        =							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__       :						Optional[int]        =							(
              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__       :						Any        =							'''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
                lowerCamelCase__       :						List[str]        =							__import__('''transformers'''		, fromlist=[model_class]      )
                lowerCamelCase__       :						Optional[int]        =							getattr(A		, A      )
                lowerCamelCase__       :						Optional[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__       :						List[str]        =							TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](A      )

          # encoder-decoder has vocab size saved differently
          lowerCamelCase__       :						Optional[int]        =							config.vocab_size if hasattr(A		, '''vocab_size'''      ) else config.encoder.vocab_size
          lowerCamelCase__       :						Dict        =							random_input_ids(A		, A		, A      )

          @run_with_tf_optimizations(self.args.eager_mode		, self.args.use_xla      )
          def encoder_decoder_train():
             lowerCamelCase__       :						int        =							model(A		, decoder_input_ids=A		, labels=A		, training=A      )[0]
             lowerCamelCase__       :						List[Any]        =							tf.gradients(A		, model.trainable_variables      )
             return gradients

          @run_with_tf_optimizations(self.args.eager_mode		, self.args.use_xla      )
          def encoder_train():
             lowerCamelCase__       :						Optional[int]        =							model(A		, labels=A		, training=A      )[0]
             lowerCamelCase__       :						List[str]        =							tf.gradients(A		, model.trainable_variables      )
             return gradients

          lowerCamelCase__       :						Tuple        =							encoder_decoder_train if config.is_encoder_decoder else encoder_train

          return _train







       def 						__lowerCamelCase   (				self					:     Tuple		, A					:     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__       :						Optional[Any]        =							timeit.repeat(
                    A		, repeat=self.args.repeat		, number=1_0		, )

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







       def 						__lowerCamelCase   (				self					:     List[Any]		, A					:     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__       :						Union[str, 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__       :						Union[str, 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()
                      lowerCamelCase__       :						Any        =							nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx      )
                      lowerCamelCase__       :						Optional[int]        =							nvml.nvmlDeviceGetMemoryInfo(A      )
                      lowerCamelCase__       :						List[Any]        =							meminfo.used
                      lowerCamelCase__       :						Union[str, Any]        =							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__       :						Tuple        =							None
                   else:
                      lowerCamelCase__       :						Dict        =							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__       :						Union[str, Any]        =							stop_memory_tracing(A      )
                   if memory is None:
                      lowerCamelCase__       :						Dict        =							summary.total
                else:
                   lowerCamelCase__       :						Optional[int]        =							None

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



