





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

_SCREAMING_SNAKE_CASE 		=						logging.get_logger(__name__)


def 							SCREAMING_SNAKE_CASE__  (      __a				,	__a				):
	def run_func(__a				):
		@wraps(__a				)
		def run_in_eager_mode(*__a				,	**__a				):
			return func(*__a				,	**__a				)

		@wraps(__a				)
		@tf.function(experimental_compile=__a				)
		def run_in_graph_mode(*__a				,	**__a				):
			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 							SCREAMING_SNAKE_CASE__  (      __a				,	__a				,	__a				):
	snake_case_						:     Dict     =		random.Random()
	snake_case_						:     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 SCREAMING_SNAKE_CASE_  (   snake_case_   ):
						__magic_name__:     TensorFlowBenchmarkArguments
						__magic_name__:     PretrainedConfig
						__magic_name__:     str	 =    "TensorFlow"






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







							"""simple docstring"""


							return tf.__version__






						def UpperCAmelCase_			( self : Optional[Any]						,				_A : str						,				_A : int						,				_A : int		)							->      float:







							"""simple docstring"""


							snake_case_						:     Dict     =		self.args.strategy
							if strategy is None:
								raise ValueError('A device strategy has to be initialized before using TensorFlow.'		)
							snake_case_						:     Union[str, Any]     =		self._prepare_inference_func(_A						,				_A						,				_A		)
							return self._measure_speed(_inference		)






						def UpperCAmelCase_			( self : Tuple						,				_A : str						,				_A : int						,				_A : int		)							->      float:







							"""simple docstring"""


							snake_case_						:     List[str]     =		self.args.strategy
							if strategy is None:
								raise ValueError('A device strategy has to be initialized before using TensorFlow.'		)
							snake_case_						:     Optional[int]     =		self._prepare_train_func(_A						,				_A						,				_A		)
							return self._measure_speed(_train		)






						def UpperCAmelCase_			( self : Optional[int]						,				_A : str						,				_A : int						,				_A : int		)							->      [Memory, Optional[MemorySummary]]:







							"""simple docstring"""


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






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







							"""simple docstring"""


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

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






						def UpperCAmelCase_			( self : str						,				_A : str						,				_A : int						,				_A : int		)							->      Callable[[], None]:







							"""simple docstring"""


							snake_case_						:     Tuple     =		self.config_dict[model_name]

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

							snake_case_						:     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:
									snake_case_						:     Dict     =		'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
									snake_case_						:     Union[str, Any]     =		__import__('transformers'						,				fromlist=[model_class]		)
									snake_case_						:     Optional[Any]     =		getattr(_A						,				_A		)
									snake_case_						:     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:
								snake_case_						:     int     =		TF_MODEL_MAPPING[config.__class__](_A		)

							# encoder-decoder has vocab size saved differently
							snake_case_						:     Union[str, Any]     =		config.vocab_size if hasattr(_A						,				'vocab_size'		) else config.encoder.vocab_size
							snake_case_						:     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		)

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

							return _inference






						def UpperCAmelCase_			( self : Tuple						,				_A : str						,				_A : int						,				_A : int		)							->      Callable[[], None]:







							"""simple docstring"""


							snake_case_						:     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.'		)

							snake_case_						:     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:
									snake_case_						:     Optional[Any]     =		'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
									snake_case_						:     Dict     =		__import__('transformers'						,				fromlist=[model_class]		)
									snake_case_						:     List[Any]     =		getattr(_A						,				_A		)
									snake_case_						:     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:
								snake_case_						:     str     =		TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](_A		)

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

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

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

							snake_case_						:     int     =		encoder_decoder_train if config.is_encoder_decoder else encoder_train

							return _train






						def UpperCAmelCase_			( self : int						,				_A : Union[str, Any]		)							->      float:







							"""simple docstring"""


							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
									snake_case_						:     Tuple     =		timeit.repeat(
									    _A						,				repeat=self.args.repeat						,				number=10						,				)

									return min(_A		) / 1_0.0
								except ResourceExhaustedError as e:
									self.print_fn(F"""Doesn't fit on GPU. {e}"""		)






						def UpperCAmelCase_			( self : Dict						,				_A : Callable[[], None]		)							->      [Memory, MemorySummary]:







							"""simple docstring"""


							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.'		)
										snake_case_						:     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.'		)
											snake_case_						:     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()
											snake_case_						:     List[Any]     =		nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx		)
											snake_case_						:     Dict     =		nvml.nvmlDeviceGetMemoryInfo(_A		)
											snake_case_						:     Tuple     =		meminfo.used
											snake_case_						:     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.'		)
											snake_case_						:     List[str]     =		None
										else:
											snake_case_						:     Tuple     =		measure_peak_memory_cpu(_A		)
											snake_case_						:     Optional[Any]     =		Memory(_A		) if isinstance(_A						,				_A		) else memory_bytes
									if self.args.trace_memory_line_by_line:
										snake_case_						:     List[Any]     =		stop_memory_tracing(_A		)
										if memory is None:
											snake_case_						:     int     =		summary.total
									else:
										snake_case_						:     List[str]     =		None

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








import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, XLMRobertaTokenizer

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
class SCREAMING_SNAKE_CASE_  (   snake_case_    ,			snake_case_    ,			snake_case_    ,			unittest.TestCase   ):
						__magic_name__:     int	 =    AltDiffusionPipeline
						__magic_name__:     Any	 =    TEXT_TO_IMAGE_PARAMS
						__magic_name__:     Any	 =    TEXT_TO_IMAGE_BATCH_PARAMS
						__magic_name__:     Any	 =    TEXT_TO_IMAGE_IMAGE_PARAMS
						__magic_name__:     Union[str, Any]	 =    TEXT_TO_IMAGE_IMAGE_PARAMS






						def UpperCAmelCase_			( self : List[Any]		)							->      int:







							"""simple docstring"""


							torch.manual_seed(0		)
							snake_case_						:     List[str]     =		UNetaDConditionModel(
							    block_out_channels=(32, 64)						,				layers_per_block=2						,				sample_size=32						,				in_channels=4						,				out_channels=4						,				down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')						,				up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')						,				cross_attention_dim=32						,				)
							snake_case_						:     Optional[Any]     =		DDIMScheduler(
							    beta_start=0.0_0_0_8_5						,				beta_end=0.0_1_2						,				beta_schedule='scaled_linear'						,				clip_sample=_A						,				set_alpha_to_one=_A						,				)
							torch.manual_seed(0		)
							snake_case_						:     Union[str, Any]     =		AutoencoderKL(
							    block_out_channels=[32, 64]						,				in_channels=3						,				out_channels=3						,				down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']						,				up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']						,				latent_channels=4						,				)

							# TODO: address the non-deterministic text encoder (fails for save-load tests)
							# torch.manual_seed(0)
							# text_encoder_config = RobertaSeriesConfig(
							#     hidden_size=32,
							#     project_dim=32,
							#     intermediate_size=37,
							#     layer_norm_eps=1e-05,
							#     num_attention_heads=4,
							#     num_hidden_layers=5,
							#     vocab_size=5002,
							# )
							# text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

							torch.manual_seed(0		)
							snake_case_						:     Optional[Any]     =		CLIPTextConfig(
							    bos_token_id=0						,				eos_token_id=2						,				hidden_size=32						,				projection_dim=32						,				intermediate_size=37						,				layer_norm_eps=1E-05						,				num_attention_heads=4						,				num_hidden_layers=5						,				pad_token_id=1						,				vocab_size=5002						,				)
							snake_case_						:     Any     =		CLIPTextModel(_A		)

							snake_case_						:     Any     =		XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'		)
							snake_case_						:     Dict     =		77

							snake_case_						:     List[Any]     =		{
							    'unet': unet,
							    'scheduler': scheduler,
							    'vae': vae,
							    'text_encoder': text_encoder,
							    'tokenizer': tokenizer,
							    'safety_checker': None,
							    'feature_extractor': None,
							}
							return components






						def UpperCAmelCase_			( self : int						,				_A : Optional[int]						,				_A : int=0		)							->      Dict:







							"""simple docstring"""


							if str(_A		).startswith('mps'		):
								snake_case_						:     Union[str, Any]     =		torch.manual_seed(_A		)
							else:
								snake_case_						:     Union[str, Any]     =		torch.Generator(device=_A		).manual_seed(_A		)
							snake_case_						:     Any     =		{
							    'prompt': 'A painting of a squirrel eating a burger',
							    'generator': generator,
							    'num_inference_steps': 2,
							    'guidance_scale': 6.0,
							    'output_type': 'numpy',
							}
							return inputs






						def UpperCAmelCase_			( self : Optional[Any]		)							->      Any:







							"""simple docstring"""


							super().test_attention_slicing_forward_pass(expected_max_diff=3E-3		)






						def UpperCAmelCase_			( self : List[Any]		)							->      Dict:







							"""simple docstring"""


							super().test_inference_batch_single_identical(expected_max_diff=3E-3		)






						def UpperCAmelCase_			( self : Dict		)							->      Any:







							"""simple docstring"""


							snake_case_						:     Tuple     =		'cpu'  # ensure determinism for the device-dependent torch.Generator

							snake_case_						:     Any     =		self.get_dummy_components()
							torch.manual_seed(0		)
							snake_case_						:     Any     =		RobertaSeriesConfig(
							    hidden_size=32						,				project_dim=32						,				intermediate_size=37						,				layer_norm_eps=1E-05						,				num_attention_heads=4						,				num_hidden_layers=5						,				vocab_size=5002						,				)
							# TODO: remove after fixing the non-deterministic text encoder
							snake_case_						:     Optional[Any]     =		RobertaSeriesModelWithTransformation(_A		)
							snake_case_						:     Optional[Any]     =		text_encoder

							snake_case_						:     Optional[Any]     =		AltDiffusionPipeline(**_A		)
							snake_case_						:     List[Any]     =		alt_pipe.to(_A		)
							alt_pipe.set_progress_bar_config(disable=_A		)

							snake_case_						:     Optional[Any]     =		self.get_dummy_inputs(_A		)
							snake_case_						:     int     =		'A photo of an astronaut'
							snake_case_						:     Tuple     =		alt_pipe(**_A		)
							snake_case_						:     Any     =		output.images
							snake_case_						:     Optional[int]     =		image[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)
							snake_case_						:     Any     =		np.array(
							    [0.5_7_4_8_1_6_2, 0.6_0_4_4_7_1_4_5, 0.4_8_8_2_1_2_1_7, 0.5_0_1_0_0_6_3_6, 0.5_4_3_1_1_8_5, 0.4_5_7_6_3_6_8_3, 0.4_9_6_5_7_6_9_6, 0.4_8_1_3_2_7_3_3, 0.4_7_5_7_3_0_9_3]		)

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






						def UpperCAmelCase_			( self : Optional[int]		)							->      Optional[int]:







							"""simple docstring"""


							snake_case_						:     Optional[Any]     =		'cpu'  # ensure determinism for the device-dependent torch.Generator

							snake_case_						:     Any     =		self.get_dummy_components()
							snake_case_						:     List[str]     =		PNDMScheduler(skip_prk_steps=_A		)
							torch.manual_seed(0		)
							snake_case_						:     Optional[int]     =		RobertaSeriesConfig(
							    hidden_size=32						,				project_dim=32						,				intermediate_size=37						,				layer_norm_eps=1E-05						,				num_attention_heads=4						,				num_hidden_layers=5						,				vocab_size=5002						,				)
							# TODO: remove after fixing the non-deterministic text encoder
							snake_case_						:     Tuple     =		RobertaSeriesModelWithTransformation(_A		)
							snake_case_						:     Any     =		text_encoder
							snake_case_						:     Tuple     =		AltDiffusionPipeline(**_A		)
							snake_case_						:     Dict     =		alt_pipe.to(_A		)
							alt_pipe.set_progress_bar_config(disable=_A		)

							snake_case_						:     Dict     =		self.get_dummy_inputs(_A		)
							snake_case_						:     Tuple     =		alt_pipe(**_A		)
							snake_case_						:     int     =		output.images
							snake_case_						:     List[Any]     =		image[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)
							snake_case_						:     Optional[int]     =		np.array(
							    [0.5_1_6_0_5_0_9_3, 0.5_7_0_7_2_4_1, 0.4_7_3_6_5_5_0_7, 0.5_0_5_7_8_8_8_6, 0.5_6_3_3_8_7_7, 0.4_6_4_2_5_0_3, 0.5_1_8_2_0_8_1, 0.4_8_7_6_3_4_8_4, 0.4_9_0_8_4_2_3_7]		)

							assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2
@slow
@require_torch_gpu
class SCREAMING_SNAKE_CASE_  (   unittest.TestCase   ):






						def UpperCAmelCase_			( self : List[Any]		)							->      Tuple:







							"""simple docstring"""


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






						def UpperCAmelCase_			( self : int		)							->      List[str]:







							"""simple docstring"""


							snake_case_						:     Optional[int]     =		AltDiffusionPipeline.from_pretrained('BAAI/AltDiffusion'						,				safety_checker=_A		)
							snake_case_						:     Optional[int]     =		alt_pipe.to(_A		)
							alt_pipe.set_progress_bar_config(disable=_A		)

							snake_case_						:     str     =		'A painting of a squirrel eating a burger'
							snake_case_						:     Tuple     =		torch.manual_seed(0		)
							snake_case_						:     str     =		alt_pipe([prompt]						,				generator=_A						,				guidance_scale=6.0						,				num_inference_steps=20						,				output_type='np'		)

							snake_case_						:     Any     =		output.images

							snake_case_						:     List[Any]     =		image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							snake_case_						:     Union[str, Any]     =		np.array([0.1_0_1_0, 0.0_8_0_0, 0.0_7_9_4, 0.0_8_8_5, 0.0_8_4_3, 0.0_7_6_2, 0.0_7_6_9, 0.0_7_2_9, 0.0_5_8_6]		)

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






						def UpperCAmelCase_			( self : Optional[int]		)							->      Union[str, Any]:







							"""simple docstring"""


							snake_case_						:     Optional[Any]     =		DDIMScheduler.from_pretrained('BAAI/AltDiffusion'						,				subfolder='scheduler'		)

							snake_case_						:     Union[str, Any]     =		AltDiffusionPipeline.from_pretrained('BAAI/AltDiffusion'						,				scheduler=_A						,				safety_checker=_A		)
							snake_case_						:     List[str]     =		alt_pipe.to(_A		)
							alt_pipe.set_progress_bar_config(disable=_A		)

							snake_case_						:     List[Any]     =		'A painting of a squirrel eating a burger'
							snake_case_						:     int     =		torch.manual_seed(0		)

							snake_case_						:     List[Any]     =		alt_pipe([prompt]						,				generator=_A						,				num_inference_steps=2						,				output_type='numpy'		)
							snake_case_						:     Any     =		output.images

							snake_case_						:     Dict     =		image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							snake_case_						:     List[Any]     =		np.array([0.4_0_1_9, 0.4_0_5_2, 0.3_8_1_0, 0.4_1_1_9, 0.3_9_1_6, 0.3_9_8_2, 0.4_6_5_1, 0.4_1_9_5, 0.5_3_2_3]		)

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

