


import math
from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import PIL
import torch
from transformers import CLIPTextModelWithProjection, CLIPTokenizer

from ...models import PriorTransformer
from ...pipelines import DiffusionPipeline
from ...schedulers import HeunDiscreteScheduler
from ...utils import (
    BaseOutput,
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)
from .renderer import ShapERenderer


_A							     =       logging.get_logger(__name__)  # pylint: disable=invalid-name

_A							     =       "\n    Examples:\n        ```py\n        >>> import torch\n        >>> from diffusers import DiffusionPipeline\n        >>> from diffusers.utils import export_to_gif\n\n        >>> device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n        >>> repo = \"openai/shap-e\"\n        >>> pipe = DiffusionPipeline.from_pretrained(repo, torch_dtype=torch.float16)\n        >>> pipe = pipe.to(device)\n\n        >>> guidance_scale = 15.0\n        >>> prompt = \"a shark\"\n\n        >>> images = pipe(\n        ...     prompt,\n        ...     guidance_scale=guidance_scale,\n        ...     num_inference_steps=64,\n        ...     frame_size=256,\n        ... ).images\n\n        >>> gif_path = export_to_gif(images[0], \"shark_3d.gif\")\n        ```\n"

@dataclass
class 			lowerCamelCase_					(    _UpperCamelCase    ):
							_lowerCamelCase				: Union[List[List[PIL.Image.Image]], List[List[np.ndarray]]]







class 			lowerCamelCase_					(    _UpperCamelCase    ):





							def __init__(							self			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE			,       ):
								super().__init__()

								self.register_modules(
								    prior=_SCREAMING_SNAKE_CASE			,       text_encoder=_SCREAMING_SNAKE_CASE			,       tokenizer=_SCREAMING_SNAKE_CASE			,       scheduler=_SCREAMING_SNAKE_CASE			,       renderer=_SCREAMING_SNAKE_CASE			,       )





							def   __magic_name__      (							self			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE    ):





								if latents is None:
									a_     			=   randn_tensor(_SCREAMING_SNAKE_CASE			,       generator=_SCREAMING_SNAKE_CASE			,       device=_SCREAMING_SNAKE_CASE			,       dtype=_SCREAMING_SNAKE_CASE    )
								else:
									if latents.shape != shape:
										raise ValueError(f"""Unexpected latents shape, got {latents.shape}, expected {shape}"""    )
									a_     			=   latents.to(_SCREAMING_SNAKE_CASE    )

								a_     			=   latents * scheduler.init_noise_sigma
								return latents





							def   __magic_name__      (							self			,       _SCREAMING_SNAKE_CASE=0    ):





								if is_accelerate_available():
									from accelerate import cpu_offload
								else:
									raise ImportError("""Please install accelerate via `pip install accelerate`"""    )

								a_     			=   torch.device(f"""cuda:{gpu_id}"""    )

								a_     			=   [self.text_encoder, self.prior]
								for cpu_offloaded_model in models:
									if cpu_offloaded_model is not None:
										cpu_offload(_SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE    )





							def   __magic_name__      (							self			,       _SCREAMING_SNAKE_CASE=0    ):





								if is_accelerate_available() and is_accelerate_version(""">="""			,       """0.17.0.dev0"""    ):
									from accelerate import cpu_offload_with_hook
								else:
									raise ImportError("""`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher."""    )

								a_     			=   torch.device(f"""cuda:{gpu_id}"""    )

								if self.device.type != "cpu":
									self.to("""cpu"""			,       silence_dtype_warnings=_SCREAMING_SNAKE_CASE    )
									torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

								a_     			=   None
								for cpu_offloaded_model in [self.text_encoder, self.prior, self.renderer]:
									a_			, a_     			=   cpu_offload_with_hook(_SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE			,       prev_module_hook=_SCREAMING_SNAKE_CASE    )

								if self.safety_checker is not None:
									a_			, a_     			=   cpu_offload_with_hook(self.safety_checker			,       _SCREAMING_SNAKE_CASE			,       prev_module_hook=_SCREAMING_SNAKE_CASE    )

								# We'll offload the last model manually.
								a_     			=   hook





							@property
							def   __magic_name__      (							self    ):





								if self.device != torch.device("""meta"""    ) or not hasattr(self.text_encoder			,       """_hf_hook"""    ):
									return self.device
								for module in self.text_encoder.modules():
									if (
									    hasattr(_SCREAMING_SNAKE_CASE			,       """_hf_hook"""    )
									    and hasattr(module._hf_hook			,       """execution_device"""    )
									    and module._hf_hook.execution_device is not None
									):
										return torch.device(module._hf_hook.execution_device    )
								return self.device





							def   __magic_name__      (							self			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE			,       ):
								len(_SCREAMING_SNAKE_CASE    ) if isinstance(_SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE    ) else 1

								# YiYi Notes: set pad_token_id to be 0, not sure why I can't set in the config file
								a_     			=   0
								# get prompt text embeddings
								a_     			=   self.tokenizer(
								    _SCREAMING_SNAKE_CASE			,       padding="""max_length"""			,       max_length=self.tokenizer.model_max_length			,       truncation=_SCREAMING_SNAKE_CASE			,       return_tensors="""pt"""			,       )
								a_     			=   text_inputs.input_ids
								a_     			=   self.tokenizer(_SCREAMING_SNAKE_CASE			,       padding="""longest"""			,       return_tensors="""pt"""    ).input_ids

								if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(_SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE    ):
									a_     			=   self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]    )
									logger.warning(
									    """The following part of your input was truncated because CLIP can only handle sequences up to"""
									    f""" {self.tokenizer.model_max_length} tokens: {removed_text}"""    )

								a_     			=   self.text_encoder(text_input_ids.to(_SCREAMING_SNAKE_CASE    )    )
								a_     			=   text_encoder_output.text_embeds

								a_     			=   prompt_embeds.repeat_interleave(_SCREAMING_SNAKE_CASE			,       dim=0    )
								# in Shap-E it normalize the prompt_embeds and then later rescale it
								a_     			=   prompt_embeds / torch.linalg.norm(_SCREAMING_SNAKE_CASE			,       dim=-1			,       keepdim=_SCREAMING_SNAKE_CASE    )

								if do_classifier_free_guidance:
									a_     			=   torch.zeros_like(_SCREAMING_SNAKE_CASE    )

									# For classifier free guidance, we need to do two forward passes.
									# Here we concatenate the unconditional and text embeddings into a single batch
									# to avoid doing two forward passes
									a_     			=   torch.cat([negative_prompt_embeds, prompt_embeds]    )

								# Rescale the features to have unit variance
								a_     			=   math.sqrt(prompt_embeds.shape[1]    ) * prompt_embeds

								return prompt_embeds





							@torch.no_grad()
							@replace_example_docstring(_SCREAMING_SNAKE_CASE    )
							def __call__(							self			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE = 1			,       _SCREAMING_SNAKE_CASE = 25			,       _SCREAMING_SNAKE_CASE = None			,       _SCREAMING_SNAKE_CASE = None			,       _SCREAMING_SNAKE_CASE = 4.0			,       _SCREAMING_SNAKE_CASE = 64			,       _SCREAMING_SNAKE_CASE = "pil"			,       _SCREAMING_SNAKE_CASE = True			,       ):





								if isinstance(_SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE    ):
									a_     			=   1
								elif isinstance(_SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE    ):
									a_     			=   len(_SCREAMING_SNAKE_CASE    )
								else:
									raise ValueError(f"""`prompt` has to be of type `str` or `list` but is {type(_SCREAMING_SNAKE_CASE    )}"""    )

								a_     			=   self._execution_device

								a_     			=   batch_size * num_images_per_prompt

								a_     			=   guidance_scale > 1.0
								a_     			=   self._encode_prompt(_SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE    )

								# prior

								self.scheduler.set_timesteps(_SCREAMING_SNAKE_CASE			,       device=_SCREAMING_SNAKE_CASE    )
								a_     			=   self.scheduler.timesteps

								a_     			=   self.prior.config.num_embeddings
								a_     			=   self.prior.config.embedding_dim

								a_     			=   self.prepare_latents(
								    (batch_size, num_embeddings * embedding_dim)			,       prompt_embeds.dtype			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE			,       self.scheduler			,       )

								# YiYi notes: for testing only to match ldm, we can directly create a latents with desired shape: batch_size, num_embeddings, embedding_dim
								a_     			=   latents.reshape(latents.shape[0]			,       _SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE    )

								for i, t in enumerate(self.progress_bar(_SCREAMING_SNAKE_CASE    )    ):
									# expand the latents if we are doing classifier free guidance
									a_     			=   torch.cat([latents] * 2    ) if do_classifier_free_guidance else latents
									a_     			=   self.scheduler.scale_model_input(_SCREAMING_SNAKE_CASE			,       _SCREAMING_SNAKE_CASE    )

									a_     			=   self.prior(
									    _SCREAMING_SNAKE_CASE			,       timestep=_SCREAMING_SNAKE_CASE			,       proj_embedding=_SCREAMING_SNAKE_CASE			,       ).predicted_image_embedding

									# remove the variance
									a_			, a_     			=   noise_pred.split(
									    scaled_model_input.shape[2]			,       dim=2    )  # batch_size, num_embeddings, embedding_dim

									if do_classifier_free_guidance is not None:
										a_			, a_     			=   noise_pred.chunk(2    )
										a_     			=   noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)

									a_     			=   self.scheduler.step(
									    _SCREAMING_SNAKE_CASE			,       timestep=_SCREAMING_SNAKE_CASE			,       sample=_SCREAMING_SNAKE_CASE			,       ).prev_sample

								if output_type == "latent":
									return ShapEPipelineOutput(images=_SCREAMING_SNAKE_CASE    )

								a_     			=   []
								for i, latent in enumerate(_SCREAMING_SNAKE_CASE    ):
									a_     			=   self.renderer.decode(
									    latent[None, :]			,       _SCREAMING_SNAKE_CASE			,       size=_SCREAMING_SNAKE_CASE			,       ray_batch_size=4096			,       n_coarse_samples=64			,       n_fine_samples=128			,       )
									images.append(_SCREAMING_SNAKE_CASE    )

								a_     			=   torch.stack(_SCREAMING_SNAKE_CASE    )

								if output_type not in ["np", "pil"]:
									raise ValueError(f"""Only the output types `pil` and `np` are supported not output_type={output_type}"""    )

								a_     			=   images.cpu().numpy()

								if output_type == "pil":
									a_     			=   [self.numpy_to_pil(_SCREAMING_SNAKE_CASE    ) for image in images]

								# Offload last model to CPU
								if hasattr(self			,       """final_offload_hook"""    ) and self.final_offload_hook is not None:
									self.final_offload_hook.offload()

								if not return_dict:
									return (images,)

								return ShapEPipelineOutput(images=_SCREAMING_SNAKE_CASE    )







"""simple docstring"""



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


_snake_case							:   str	  =       logging.get_logger(__name__)




class 							__lowercase(							_UpperCamelCase	):
						lowerCAmelCase_		:				PyTorchBenchmarkArguments
						lowerCAmelCase_		:				PretrainedConfig
						lowerCAmelCase_		:				str=							"PyTorch"






						@property
						def 		lowerCAmelCase (       self)      ->	List[Any]:




											"""simple docstring"""





											return torch.__version__






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	float:




											"""simple docstring"""





											lowerCamelCase__					=self._prepare_inference_func(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
											return self._measure_speed(_inference)






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	[Memory, Optional[MemorySummary]]:




											"""simple docstring"""





											lowerCamelCase__					=self._prepare_inference_func(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
											return self._measure_memory(_inference)






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	float:




											"""simple docstring"""





											lowerCamelCase__					=self._prepare_train_func(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
											return self._measure_speed(_train)






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	[Memory, Optional[MemorySummary]]:




											"""simple docstring"""





											lowerCamelCase__					=self._prepare_train_func(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
											return self._measure_memory(_train)






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	Callable[[], None]:




											"""simple docstring"""





											lowerCamelCase__					=self.config_dict[model_name]

											if self.args.torchscript:
																lowerCamelCase__					=True

											lowerCamelCase__					=(
											    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:
																					lowerCamelCase__					=config.architectures[0]
																					lowerCamelCase__					=__import__('transformers'     ,    fromlist=[model_class])
																					lowerCamelCase__					=getattr(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
																					lowerCamelCase__					=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:
																lowerCamelCase__					=MODEL_MAPPING[config.__class__](_SCREAMING_SNAKE_CASE)

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

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

											def encoder_decoder_forward():
																with torch.no_grad():
																					lowerCamelCase__					=inference_model(_SCREAMING_SNAKE_CASE     ,    decoder_input_ids=_SCREAMING_SNAKE_CASE)
																return outputs

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

											lowerCamelCase__					=encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
											return _forward






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	Callable[[], None]:




											"""simple docstring"""





											lowerCamelCase__					=self.config_dict[model_name]

											lowerCamelCase__					=(
											    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:
																					lowerCamelCase__					=config.architectures[0]
																					lowerCamelCase__					=__import__('transformers'     ,    fromlist=[model_class])
																					lowerCamelCase__					=getattr(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
																					lowerCamelCase__					=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:
																lowerCamelCase__					=MODEL_WITH_LM_HEAD_MAPPING[config.__class__](_SCREAMING_SNAKE_CASE)

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

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

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

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

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






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE)      ->	float:




											"""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
																lowerCamelCase__					=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 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE)      ->	[Memory, MemorySummary]:




											"""simple docstring"""





											try:
																if self.args.trace_memory_line_by_line:
																					lowerCamelCase__					=start_memory_tracing('transformers')

																if self.args.is_tpu:
																					# tpu
																					raise NotImplementedError(
																					    'Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking with'
																					    ' `--no-memory` or `args.memory=False`')
																elif self.args.is_gpu:
																					if not is_pyanvml_available():
																										logger.warning(
																										    'py3nvml not installed, we won\'t log GPU memory usage. '
																										    'Install py3nvml (pip install py3nvml) to log information about GPU.')
																										lowerCamelCase__					='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__					=nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx)
																										lowerCamelCase__					=nvml.nvmlDeviceGetMemoryInfo(_SCREAMING_SNAKE_CASE)
																										lowerCamelCase__					=meminfo.used
																										lowerCamelCase__					=Memory(_SCREAMING_SNAKE_CASE)
																										# shutdown nvml
																										nvml.nvmlShutdown()
																else:
																					# cpu
																					lowerCamelCase__					=measure_peak_memory_cpu(_SCREAMING_SNAKE_CASE)
																					lowerCamelCase__					=Memory(_SCREAMING_SNAKE_CASE) if isinstance(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE) else memory_bytes

																if self.args.trace_memory_line_by_line:
																					lowerCamelCase__					=stop_memory_tracing(_SCREAMING_SNAKE_CASE)
																else:
																					lowerCamelCase__					=None

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