






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        UpperCAmelCase__		(						lowerCamelCase_       :       bool   ,  lowerCamelCase_       :       bool				):
			def run_func(lowerCamelCase_       :       str				):
						@wraps(lowerCamelCase_				)
						def run_in_eager_mode(*lowerCamelCase_       :       Union[str, Any]   ,  **lowerCamelCase_       :       List[str]				):
									return func(*lowerCamelCase_   ,  **lowerCamelCase_				)

						@wraps(lowerCamelCase_				)
						@tf.function(experimental_compile=lowerCamelCase_				)
						def run_in_graph_mode(*lowerCamelCase_       :       Dict   ,  **lowerCamelCase_       :       Optional[Any]				):
									return func(*lowerCamelCase_   ,  **lowerCamelCase_				)

						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__		(						lowerCamelCase_       :       int   ,  lowerCamelCase_       :       int   ,  lowerCamelCase_       :       int				):
			__a						:  List[str]									=     random.Random()
			__a						:  Optional[Any]									=     [rng.randint(0   ,  vocab_size - 1				) for i in range(batch_size * sequence_length				)]
			return tf.constant(lowerCamelCase_   ,  shape=(batch_size, sequence_length)   ,  dtype=tf.intaa				)






class 			_UpperCamelCase(       __lowerCamelCase		):
		__SCREAMING_SNAKE_CASE  :       TensorFlowBenchmarkArguments
		__SCREAMING_SNAKE_CASE  :       PretrainedConfig
		__SCREAMING_SNAKE_CASE  :       str							=							"TensorFlow"

		@property
		def    __lowerCAmelCase  (					self		:    int    ):



					'''simple docstring'''



					return tf.__version__

		def    __lowerCAmelCase  (					self		:    int    ,       SCREAMING_SNAKE_CASE__		:    str    ,       SCREAMING_SNAKE_CASE__		:    int    ,       SCREAMING_SNAKE_CASE__		:    int    ):



					'''simple docstring'''



					__a						:  int									=     self.args.strategy
					if strategy is None:
								raise ValueError('A device strategy has to be initialized before using TensorFlow.'    )
					__a						:  Optional[int]									=     self._prepare_inference_func(SCREAMING_SNAKE_CASE__    ,       SCREAMING_SNAKE_CASE__    ,       SCREAMING_SNAKE_CASE__    )
					return self._measure_speed(_inference    )

		def    __lowerCAmelCase  (					self		:    Optional[Any]    ,       SCREAMING_SNAKE_CASE__		:    str    ,       SCREAMING_SNAKE_CASE__		:    int    ,       SCREAMING_SNAKE_CASE__		:    int    ):



					'''simple docstring'''



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

		def    __lowerCAmelCase  (					self		:    Optional[Any]    ,       SCREAMING_SNAKE_CASE__		:    str    ,       SCREAMING_SNAKE_CASE__		:    int    ,       SCREAMING_SNAKE_CASE__		:    int    ):



					'''simple docstring'''



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

		def    __lowerCAmelCase  (					self		:    Union[str, Any]    ,       SCREAMING_SNAKE_CASE__		:    str    ,       SCREAMING_SNAKE_CASE__		:    int    ,       SCREAMING_SNAKE_CASE__		:    int    ):



					'''simple docstring'''



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

					__a						:  List[Any]									=     self._prepare_train_func(SCREAMING_SNAKE_CASE__    ,       SCREAMING_SNAKE_CASE__    ,       SCREAMING_SNAKE_CASE__    )
					return self._measure_memory(_train    )

		def    __lowerCAmelCase  (					self		:    Optional[int]    ,       SCREAMING_SNAKE_CASE__		:    str    ,       SCREAMING_SNAKE_CASE__		:    int    ,       SCREAMING_SNAKE_CASE__		:    int    ):



					'''simple docstring'''



					__a						:  str									=     self.config_dict[model_name]

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

					__a						:  Any									=     (
					    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:
											__a						:  Optional[Any]									=     'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
											__a						:  List[Any]									=     __import__('transformers'    ,       fromlist=[model_class]    )
											__a						:  Dict									=     getattr(SCREAMING_SNAKE_CASE__    ,       SCREAMING_SNAKE_CASE__    )
											__a						:  Any									=     model_cls(SCREAMING_SNAKE_CASE__    )
								except ImportError:
											raise ImportError(
											    f'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
											    ' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'    )
					else:
								__a						:  Optional[int]									=     TF_MODEL_MAPPING[config.__class__](SCREAMING_SNAKE_CASE__    )

					# encoder-decoder has vocab size saved differently
					__a						:  List[str]									=     config.vocab_size if hasattr(SCREAMING_SNAKE_CASE__    ,       'vocab_size'    ) else config.encoder.vocab_size
					__a						:  Optional[Any]									=     random_input_ids(SCREAMING_SNAKE_CASE__    ,       SCREAMING_SNAKE_CASE__    ,       SCREAMING_SNAKE_CASE__    )

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

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

					__a						:  Any									=     encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

					return _inference

		def    __lowerCAmelCase  (					self		:    List[Any]    ,       SCREAMING_SNAKE_CASE__		:    str    ,       SCREAMING_SNAKE_CASE__		:    int    ,       SCREAMING_SNAKE_CASE__		:    int    ):



					'''simple docstring'''



					__a						:  Optional[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.'    )

					__a						:  Optional[int]									=     (
					    hasattr(SCREAMING_SNAKE_CASE__    ,       'architectures'    )
					    and isinstance(config.architectures    ,       SCREAMING_SNAKE_CASE__    )
					    and len(config.architectures    ) > 0
					)
					if not self.args.only_pretrain_model and has_model_class_in_config:
								try:
											__a						:  List[str]									=     'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
											__a						:  Dict									=     __import__('transformers'    ,       fromlist=[model_class]    )
											__a						:  str									=     getattr(SCREAMING_SNAKE_CASE__    ,       SCREAMING_SNAKE_CASE__    )
											__a						:  Tuple									=     model_cls(SCREAMING_SNAKE_CASE__    )
								except ImportError:
											raise ImportError(
											    f'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
											    ' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'    )
					else:
								__a						:  Union[str, Any]									=     TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](SCREAMING_SNAKE_CASE__    )

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

					@run_with_tf_optimizations(self.args.eager_mode    ,       self.args.use_xla    )
					def encoder_decoder_train():
								__a						:  Union[str, Any]									=     model(SCREAMING_SNAKE_CASE__    ,       decoder_input_ids=SCREAMING_SNAKE_CASE__    ,       labels=SCREAMING_SNAKE_CASE__    ,       training=SCREAMING_SNAKE_CASE__    )[0]
								__a						:  List[str]									=     tf.gradients(SCREAMING_SNAKE_CASE__    ,       model.trainable_variables    )
								return gradients

					@run_with_tf_optimizations(self.args.eager_mode    ,       self.args.use_xla    )
					def encoder_train():
								__a						:  Tuple									=     model(SCREAMING_SNAKE_CASE__    ,       labels=SCREAMING_SNAKE_CASE__    ,       training=SCREAMING_SNAKE_CASE__    )[0]
								__a						:  Tuple									=     tf.gradients(SCREAMING_SNAKE_CASE__    ,       model.trainable_variables    )
								return gradients

					__a						:  Optional[int]									=     encoder_decoder_train if config.is_encoder_decoder else encoder_train

					return _train

		def    __lowerCAmelCase  (					self		:    Dict    ,       SCREAMING_SNAKE_CASE__		:    Dict    ):



					'''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(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
											__a						:  Optional[Any]									=     timeit.repeat(
											    SCREAMING_SNAKE_CASE__    ,       repeat=self.args.repeat    ,       number=1_0    ,       )

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



		def    __lowerCAmelCase  (					self		:    Any    ,       SCREAMING_SNAKE_CASE__		:    Callable[[], None]    ):



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

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionPanoramaPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

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


enable_full_determinism()

@skip_mps
class       _a  ( SCREAMING_SNAKE_CASE__     ,	SCREAMING_SNAKE_CASE__     ,	unittest.TestCase):
			__magic_name__     					=     StableDiffusionPanoramaPipeline
			__magic_name__     					=     TEXT_TO_IMAGE_PARAMS
			__magic_name__     					=     TEXT_TO_IMAGE_BATCH_PARAMS
			__magic_name__     					=     TEXT_TO_IMAGE_IMAGE_PARAMS
			__magic_name__     					=     TEXT_TO_IMAGE_IMAGE_PARAMS



			def 				__lowercase					(			self							:				Optional[int]		)							->					int:
									torch.manual_seed(0		)
									snake_case   :			Optional[Any]						    = UNetaDConditionModel(
									    block_out_channels=(32, 64)   ,	layers_per_block=1   ,	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   :			Tuple						    = DDIMScheduler()
									torch.manual_seed(0		)
									snake_case   :			List[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   ,	)
									torch.manual_seed(0		)
									snake_case   :			str						    = CLIPTextConfig(
									    bos_token_id=0   ,	eos_token_id=2   ,	hidden_size=32   ,	intermediate_size=37   ,	layer_norm_eps=1E-05   ,	num_attention_heads=4   ,	num_hidden_layers=5   ,	pad_token_id=1   ,	vocab_size=1000   ,	)
									snake_case   :			int						    = CLIPTextModel(_lowercase		)
									snake_case   :			Union[str, Any]						    = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

									snake_case   :			Any						    = {
									    "unet": unet,
									    "scheduler": scheduler,
									    "vae": vae,
									    "text_encoder": text_encoder,
									    "tokenizer": tokenizer,
									    "safety_checker": None,
									    "feature_extractor": None,
									}
									return components



			def 				__lowercase					(			self							:				List[str]   ,	_lowercase							:				int   ,	_lowercase							:				Dict=0		)							->					Optional[int]:
									snake_case   :			Union[str, Any]						    = torch.manual_seed(_lowercase		)
									snake_case   :			str						    = {
									    "prompt": "a photo of the dolomites",
									    "generator": generator,
									    # Setting height and width to None to prevent OOMs on CPU.
									    "height": None,
									    "width": None,
									    "num_inference_steps": 1,
									    "guidance_scale": 6.0,
									    "output_type": "numpy",
									}
									return inputs



			def 				__lowercase					(			self							:				str		)							->					List[str]:
									snake_case   :			Any						    = "cpu"  # ensure determinism for the device-dependent torch.Generator
									snake_case   :			Dict						    = self.get_dummy_components()
									snake_case   :			Any						    = StableDiffusionPanoramaPipeline(**_lowercase		)
									snake_case   :			List[Any]						    = sd_pipe.to(_lowercase		)
									sd_pipe.set_progress_bar_config(disable=_lowercase		)

									snake_case   :			Optional[Any]						    = self.get_dummy_inputs(_lowercase		)
									snake_case   :			Union[str, Any]						    = sd_pipe(**_lowercase		).images
									snake_case   :			List[Any]						    = image[0, -3:, -3:, -1]
									assert image.shape == (1, 64, 64, 3)

									snake_case   :			List[Any]						    = np.array([0.6186, 0.5374, 0.4915, 0.4135, 0.4114, 0.4563, 0.5128, 0.4977, 0.4757]		)

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



			def 				__lowercase					(			self							:				int		)							->					Union[str, Any]:
									super().test_inference_batch_consistent(batch_sizes=[1, 2]		)



			def 				__lowercase					(			self							:				Tuple		)							->					Tuple:
									super().test_inference_batch_single_identical(batch_size=2   ,	expected_max_diff=3.25E-3		)



			def 				__lowercase					(			self							:				Any		)							->					List[Any]:
									snake_case   :			Any						    = "cpu"  # ensure determinism for the device-dependent torch.Generator
									snake_case   :			Union[str, Any]						    = self.get_dummy_components()
									snake_case   :			Tuple						    = StableDiffusionPanoramaPipeline(**_lowercase		)
									snake_case   :			Tuple						    = sd_pipe.to(_lowercase		)
									sd_pipe.set_progress_bar_config(disable=_lowercase		)

									snake_case   :			List[str]						    = self.get_dummy_inputs(_lowercase		)
									snake_case   :			int						    = "french fries"
									snake_case   :			Union[str, Any]						    = sd_pipe(**_lowercase   ,	negative_prompt=_lowercase		)
									snake_case   :			str						    = output.images
									snake_case   :			Optional[Any]						    = image[0, -3:, -3:, -1]

									assert image.shape == (1, 64, 64, 3)

									snake_case   :			Union[str, Any]						    = np.array([0.6187, 0.5375, 0.4915, 0.4136, 0.4114, 0.4563, 0.5128, 0.4976, 0.4757]		)

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



			def 				__lowercase					(			self							:				str		)							->					Any:
									snake_case   :			Optional[Any]						    = "cpu"  # ensure determinism for the device-dependent torch.Generator
									snake_case   :			List[Any]						    = self.get_dummy_components()
									snake_case   :			List[Any]						    = StableDiffusionPanoramaPipeline(**_lowercase		)
									snake_case   :			int						    = sd_pipe.to(_lowercase		)
									sd_pipe.set_progress_bar_config(disable=_lowercase		)

									snake_case   :			Tuple						    = self.get_dummy_inputs(_lowercase		)
									snake_case   :			str						    = sd_pipe(**_lowercase   ,	view_batch_size=2		)
									snake_case   :			Optional[Any]						    = output.images
									snake_case   :			int						    = image[0, -3:, -3:, -1]

									assert image.shape == (1, 64, 64, 3)

									snake_case   :			str						    = np.array([0.6187, 0.5375, 0.4915, 0.4136, 0.4114, 0.4563, 0.5128, 0.4976, 0.4757]		)

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



			def 				__lowercase					(			self							:				int		)							->					Optional[Any]:
									snake_case   :			Dict						    = "cpu"  # ensure determinism for the device-dependent torch.Generator
									snake_case   :			List[str]						    = self.get_dummy_components()
									snake_case   :			Any						    = EulerAncestralDiscreteScheduler(
									    beta_start=0.00085   ,	beta_end=0.012   ,	beta_schedule="scaled_linear"		)
									snake_case   :			List[str]						    = StableDiffusionPanoramaPipeline(**_lowercase		)
									snake_case   :			List[str]						    = sd_pipe.to(_lowercase		)
									sd_pipe.set_progress_bar_config(disable=_lowercase		)

									snake_case   :			List[Any]						    = self.get_dummy_inputs(_lowercase		)
									snake_case   :			Optional[Any]						    = sd_pipe(**_lowercase		).images
									snake_case   :			int						    = image[0, -3:, -3:, -1]

									assert image.shape == (1, 64, 64, 3)

									snake_case   :			Tuple						    = np.array([0.4024, 0.6510, 0.4901, 0.5378, 0.5813, 0.5622, 0.4795, 0.4467, 0.4952]		)

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



			def 				__lowercase					(			self							:				Tuple		)							->					Union[str, Any]:
									snake_case   :			str						    = "cpu"  # ensure determinism for the device-dependent torch.Generator
									snake_case   :			str						    = self.get_dummy_components()
									snake_case   :			Optional[int]						    = PNDMScheduler(
									    beta_start=0.00085   ,	beta_end=0.012   ,	beta_schedule="scaled_linear"   ,	skip_prk_steps=_lowercase		)
									snake_case   :			Dict						    = StableDiffusionPanoramaPipeline(**_lowercase		)
									snake_case   :			Dict						    = sd_pipe.to(_lowercase		)
									sd_pipe.set_progress_bar_config(disable=_lowercase		)

									snake_case   :			Dict						    = self.get_dummy_inputs(_lowercase		)
									snake_case   :			Optional[int]						    = sd_pipe(**_lowercase		).images
									snake_case   :			str						    = image[0, -3:, -3:, -1]

									assert image.shape == (1, 64, 64, 3)

									snake_case   :			Optional[int]						    = np.array([0.6391, 0.6291, 0.4861, 0.5134, 0.5552, 0.4578, 0.5032, 0.5023, 0.4539]		)

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



@slow
@require_torch_gpu
class       _a  ( unittest.TestCase):



			def 				__lowercase					(			self							:				Union[str, Any]		)							->					str:
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()



			def 				__lowercase					(			self							:				Dict   ,	_lowercase							:				Any=0		)							->					Optional[Any]:
									snake_case   :			Any						    = torch.manual_seed(_lowercase		)
									snake_case   :			Any						    = {
									    "prompt": "a photo of the dolomites",
									    "generator": generator,
									    "num_inference_steps": 3,
									    "guidance_scale": 7.5,
									    "output_type": "numpy",
									}
									return inputs



			def 				__lowercase					(			self							:				Optional[Any]		)							->					Union[str, Any]:
									snake_case   :			List[Any]						    = "stabilityai/stable-diffusion-2-base"
									snake_case   :			Dict						    = DDIMScheduler.from_pretrained(_lowercase   ,	subfolder="scheduler"		)
									snake_case   :			Any						    = StableDiffusionPanoramaPipeline.from_pretrained(_lowercase   ,	scheduler=_lowercase   ,	safety_checker=_lowercase		)
									pipe.to(_lowercase		)
									pipe.set_progress_bar_config(disable=_lowercase		)
									pipe.enable_attention_slicing()

									snake_case   :			str						    = self.get_inputs()
									snake_case   :			List[Any]						    = pipe(**_lowercase		).images
									snake_case   :			Optional[int]						    = image[0, -3:, -3:, -1].flatten()

									assert image.shape == (1, 512, 2048, 3)

									snake_case   :			List[str]						    = np.array(
									    [
									        0.36968392,
									        0.27025372,
									        0.32446766,
									        0.28379387,
									        0.36363274,
									        0.30733347,
									        0.27100027,
									        0.27054125,
									        0.25536096,
									    ]		)

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



			def 				__lowercase					(			self							:				Tuple		)							->					List[str]:
									snake_case   :			List[str]						    = StableDiffusionPanoramaPipeline.from_pretrained(
									    "stabilityai/stable-diffusion-2-base"   ,	safety_checker=_lowercase		)
									snake_case   :			int						    = LMSDiscreteScheduler.from_config(pipe.scheduler.config		)
									pipe.to(_lowercase		)
									pipe.set_progress_bar_config(disable=_lowercase		)
									pipe.enable_attention_slicing()

									snake_case   :			Dict						    = self.get_inputs()
									snake_case   :			int						    = pipe(**_lowercase		).images
									snake_case   :			Optional[Any]						    = image[0, -3:, -3:, -1].flatten()

									assert image.shape == (1, 512, 2048, 3)

									snake_case   :			List[str]						    = np.array(
									    [
									        [
									            0.0,
									            0.0,
									            0.0,
									            0.0,
									            0.0,
									            0.0,
									            0.0,
									            0.0,
									            0.0,
									        ]
									    ]		)

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



			def 				__lowercase					(			self							:				str		)							->					Any:
									snake_case   :			Any						    = 0

									def callback_fn(_lowercase							:				int   ,	_lowercase							:				int   ,	_lowercase							:				torch.FloatTensor		) -> None:
															snake_case   :			List[Any]						    = True
															nonlocal number_of_steps
															number_of_steps += 1
															if step == 1:
																					snake_case   :			List[Any]						    = latents.detach().cpu().numpy()
																					assert latents.shape == (1, 4, 64, 256)
																					snake_case   :			int						    = latents[0, -3:, -3:, -1]

																					snake_case   :			Tuple						    = np.array(
																					    [
																					        0.18681869,
																					        0.33907816,
																					        0.5361276,
																					        0.14432865,
																					        -0.02856611,
																					        -0.73941123,
																					        0.23397987,
																					        0.47322682,
																					        -0.37823164,
																					    ]		)
																					assert np.abs(latents_slice.flatten() - expected_slice		).max() < 5E-2
															elif step == 2:
																					snake_case   :			Union[str, Any]						    = latents.detach().cpu().numpy()
																					assert latents.shape == (1, 4, 64, 256)
																					snake_case   :			Tuple						    = latents[0, -3:, -3:, -1]

																					snake_case   :			Tuple						    = np.array(
																					    [
																					        0.18539645,
																					        0.33987248,
																					        0.5378559,
																					        0.14437142,
																					        -0.02455261,
																					        -0.7338317,
																					        0.23990755,
																					        0.47356272,
																					        -0.3786505,
																					    ]		)

																					assert np.abs(latents_slice.flatten() - expected_slice		).max() < 5E-2

									snake_case   :			Dict						    = False

									snake_case   :			List[Any]						    = "stabilityai/stable-diffusion-2-base"
									snake_case   :			List[str]						    = DDIMScheduler.from_pretrained(_lowercase   ,	subfolder="scheduler"		)
									snake_case   :			Optional[Any]						    = StableDiffusionPanoramaPipeline.from_pretrained(_lowercase   ,	scheduler=_lowercase   ,	safety_checker=_lowercase		)
									snake_case   :			Any						    = pipe.to(_lowercase		)
									pipe.set_progress_bar_config(disable=_lowercase		)
									pipe.enable_attention_slicing()

									snake_case   :			Tuple						    = self.get_inputs()
									pipe(**_lowercase   ,	callback=_lowercase   ,	callback_steps=1		)
									assert callback_fn.has_been_called
									assert number_of_steps == 3



			def 				__lowercase					(			self							:				Any		)							->					Tuple:
									torch.cuda.empty_cache()
									torch.cuda.reset_max_memory_allocated()
									torch.cuda.reset_peak_memory_stats()

									snake_case   :			Optional[int]						    = "stabilityai/stable-diffusion-2-base"
									snake_case   :			Tuple						    = DDIMScheduler.from_pretrained(_lowercase   ,	subfolder="scheduler"		)
									snake_case   :			Any						    = StableDiffusionPanoramaPipeline.from_pretrained(_lowercase   ,	scheduler=_lowercase   ,	safety_checker=_lowercase		)
									snake_case   :			Tuple						    = pipe.to(_lowercase		)
									pipe.set_progress_bar_config(disable=_lowercase		)
									pipe.enable_attention_slicing(1		)
									pipe.enable_sequential_cpu_offload()

									snake_case   :			Any						    = self.get_inputs()
									snake_case   :			Optional[int]						    = pipe(**_lowercase		)

									snake_case   :			int						    = torch.cuda.max_memory_allocated()
									# make sure that less than 5.2 GB is allocated
									assert mem_bytes < 5.5 * 10**9



