



'''simple docstring'''

import tempfile
import unittest

import numpy as np

from diffusers import (
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionPipeline,
    PNDMScheduler,
)
from diffusers.utils.testing_utils import is_onnx_available, nightly, require_onnxruntime, require_torch_gpu

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
		import onnxruntime as ort



class 		lowerCamelCase__      (						__UpperCAmelCase	,   unittest.TestCase     ):




	'''simple docstring'''



	A_						    =						"hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"




	def 						__UpperCAmelCase     (  self				:	List[str]     ,  UpperCamelCase_				:	Optional[Any]=0					)  ->       List[str]:



		'''simple docstring'''





		_lowercase					:       int       =		np.random.RandomState(_lowerCamelCase					)
		_lowercase					:       Any       =		{
		    'prompt': 'A painting of a squirrel eating a burger',
		    'generator': generator,
		    'num_inference_steps': 2,
		    'guidance_scale': 7.5,
		    'output_type': 'numpy',
		}
		return inputs




	def 						__UpperCAmelCase     (  self				:	int					)  ->       str:



		'''simple docstring'''





		_lowercase					:       Union[str, Any]       =		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,  provider='CPUExecutionProvider'					)
		pipe.set_progress_bar_config(disable=_lowerCamelCase					)

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

		assert image.shape == (1, 128, 128, 3)
		_lowercase					:       List[Any]       =		np.array([0.6_50_72, 0.5_84_92, 0.4_82_19, 0.5_55_21, 0.5_31_80, 0.5_59_39, 0.5_06_97, 0.3_98_00, 0.4_64_55]					)

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




	def 						__UpperCAmelCase     (  self				:	Any					)  ->       Tuple:



		'''simple docstring'''





		_lowercase					:       Dict       =		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,  provider='CPUExecutionProvider'					)
		_lowercase					:       Tuple       =		PNDMScheduler.from_config(pipe.scheduler.config     ,  skip_prk_steps=_lowerCamelCase					)
		pipe.set_progress_bar_config(disable=_lowerCamelCase					)

		_lowercase					:       int       =		self.get_dummy_inputs()
		_lowercase					:       str       =		pipe(**_lowerCamelCase					).images
		_lowercase					:       Optional[Any]       =		image[0, -3:, -3:, -1]

		assert image.shape == (1, 128, 128, 3)
		_lowercase					:       List[Any]       =		np.array([0.6_58_63, 0.5_94_25, 0.4_93_26, 0.5_63_13, 0.5_38_75, 0.5_66_27, 0.5_10_65, 0.3_97_77, 0.4_63_30]					)

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




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



		'''simple docstring'''





		_lowercase					:       Union[str, Any]       =		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,  provider='CPUExecutionProvider'					)
		_lowercase					:       Optional[Any]       =		LMSDiscreteScheduler.from_config(pipe.scheduler.config					)
		pipe.set_progress_bar_config(disable=_lowerCamelCase					)

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

		assert image.shape == (1, 128, 128, 3)
		_lowercase					:       str       =		np.array([0.5_37_55, 0.6_07_86, 0.4_74_02, 0.4_94_88, 0.5_18_69, 0.4_98_19, 0.4_79_85, 0.3_89_57, 0.4_42_79]					)

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




	def 						__UpperCAmelCase     (  self				:	Any					)  ->       Optional[int]:



		'''simple docstring'''





		_lowercase					:       Tuple       =		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,  provider='CPUExecutionProvider'					)
		_lowercase					:       str       =		EulerDiscreteScheduler.from_config(pipe.scheduler.config					)
		pipe.set_progress_bar_config(disable=_lowerCamelCase					)

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

		assert image.shape == (1, 128, 128, 3)
		_lowercase					:       Optional[Any]       =		np.array([0.5_37_55, 0.6_07_86, 0.4_74_02, 0.4_94_88, 0.5_18_69, 0.4_98_19, 0.4_79_85, 0.3_89_57, 0.4_42_79]					)

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




	def 						__UpperCAmelCase     (  self				:	Union[str, Any]					)  ->       Dict:



		'''simple docstring'''





		_lowercase					:       int       =		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,  provider='CPUExecutionProvider'					)
		_lowercase					:       List[Any]       =		EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config					)
		pipe.set_progress_bar_config(disable=_lowerCamelCase					)

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

		assert image.shape == (1, 128, 128, 3)
		_lowercase					:       Dict       =		np.array([0.5_38_17, 0.6_08_12, 0.4_73_84, 0.4_95_30, 0.5_18_94, 0.4_98_14, 0.4_79_84, 0.3_89_58, 0.4_42_71]					)

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




	def 						__UpperCAmelCase     (  self				:	str					)  ->       Union[str, Any]:



		'''simple docstring'''





		_lowercase					:       Optional[int]       =		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,  provider='CPUExecutionProvider'					)
		_lowercase					:       List[Any]       =		DPMSolverMultistepScheduler.from_config(pipe.scheduler.config					)
		pipe.set_progress_bar_config(disable=_lowerCamelCase					)

		_lowercase					:       Tuple       =		self.get_dummy_inputs()
		_lowercase					:       Dict       =		pipe(**_lowerCamelCase					).images
		_lowercase					:       Dict       =		image[0, -3:, -3:, -1]

		assert image.shape == (1, 128, 128, 3)
		_lowercase					:       Optional[Any]       =		np.array([0.5_38_95, 0.6_08_08, 0.4_79_33, 0.4_96_08, 0.5_18_86, 0.4_99_50, 0.4_80_53, 0.3_89_57, 0.4_42_00]					)

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




	def 						__UpperCAmelCase     (  self				:	List[str]					)  ->       List[Any]:



		'''simple docstring'''





		_lowercase					:       int       =		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,  provider='CPUExecutionProvider'					)
		pipe.set_progress_bar_config(disable=_lowerCamelCase					)

		_lowercase					:       Optional[Any]       =		self.get_dummy_inputs()
		_lowercase					:       Tuple       =		3 * [inputs['prompt']]

		# forward
		_lowercase					:       Optional[int]       =		pipe(**_lowerCamelCase					)
		_lowercase					:       Optional[int]       =		output.images[0, -3:, -3:, -1]

		_lowercase					:       List[Any]       =		self.get_dummy_inputs()
		_lowercase					:       Optional[int]       =		3 * [inputs.pop('prompt'					)]

		_lowercase					:       str       =		pipe.tokenizer(
		    _lowerCamelCase     ,  padding='max_length'     ,  max_length=pipe.tokenizer.model_max_length     ,  truncation=_lowerCamelCase     ,  return_tensors='np'     ,  )
		_lowercase					:       str       =		text_inputs['input_ids']

		_lowercase					:       Optional[Any]       =		pipe.text_encoder(input_ids=text_inputs.astype(np.intaa					)					)[0]

		_lowercase					:       List[str]       =		prompt_embeds

		# forward
		_lowercase					:       Dict       =		pipe(**_lowerCamelCase					)
		_lowercase					:       Tuple       =		output.images[0, -3:, -3:, -1]

		assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()					).max() < 1E-4




	def 						__UpperCAmelCase     (  self				:	Optional[int]					)  ->       str:



		'''simple docstring'''





		_lowercase					:       Tuple       =		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,  provider='CPUExecutionProvider'					)
		pipe.set_progress_bar_config(disable=_lowerCamelCase					)

		_lowercase					:       Union[str, Any]       =		self.get_dummy_inputs()
		_lowercase					:       Tuple       =		3 * ['this is a negative prompt']
		_lowercase					:       str       =		negative_prompt
		_lowercase					:       Dict       =		3 * [inputs['prompt']]

		# forward
		_lowercase					:       Dict       =		pipe(**_lowerCamelCase					)
		_lowercase					:       Dict       =		output.images[0, -3:, -3:, -1]

		_lowercase					:       List[str]       =		self.get_dummy_inputs()
		_lowercase					:       Union[str, Any]       =		3 * [inputs.pop('prompt'					)]

		_lowercase					:       int       =		[]
		for p in [prompt, negative_prompt]:
			_lowercase					:       str       =		pipe.tokenizer(
			    _lowerCamelCase     ,  padding='max_length'     ,  max_length=pipe.tokenizer.model_max_length     ,  truncation=_lowerCamelCase     ,  return_tensors='np'     ,  )
			_lowercase					:       List[str]       =		text_inputs['input_ids']

			embeds.append(pipe.text_encoder(input_ids=text_inputs.astype(np.intaa					)					)[0]					)

		_lowercase  ,    _lowercase					:       str       =		embeds

		# forward
		_lowercase					:       str       =		pipe(**_lowerCamelCase					)
		_lowercase					:       Union[str, Any]       =		output.images[0, -3:, -3:, -1]

		assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()					).max() < 1E-4






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




	'''simple docstring'''



	@property
	def 						__UpperCAmelCase     (  self				:	Dict					)  ->       Optional[Any]:



		'''simple docstring'''





		return (
		    "CUDAExecutionProvider",
		    {
		        "gpu_mem_limit": "15000000000",  # 15GB
		        "arena_extend_strategy": "kSameAsRequested",
		    },
		)




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



		'''simple docstring'''





		_lowercase					:       int       =		ort.SessionOptions()
		_lowercase					:       List[str]       =		False
		return options




	def 						__UpperCAmelCase     (  self				:	List[Any]					)  ->       List[str]:



		'''simple docstring'''





		_lowercase					:       List[str]       =		OnnxStableDiffusionPipeline.from_pretrained(
		    'CompVis/stable-diffusion-v1-4'     ,  revision='onnx'     ,  safety_checker=_lowerCamelCase     ,  feature_extractor=_lowerCamelCase     ,  provider=self.gpu_provider     ,  sess_options=self.gpu_options     ,  )
		sd_pipe.set_progress_bar_config(disable=_lowerCamelCase					)

		_lowercase					:       Tuple       =		'A painting of a squirrel eating a burger'
		np.random.seed(0					)
		_lowercase					:       str       =		sd_pipe([prompt]     ,  guidance_scale=6.0     ,  num_inference_steps=10     ,  output_type='np'					)
		_lowercase					:       Union[str, Any]       =		output.images

		_lowercase					:       Any       =		image[0, -3:, -3:, -1]

		assert image.shape == (1, 512, 512, 3)
		_lowercase					:       int       =		np.array([0.04_52, 0.03_90, 0.00_87, 0.03_50, 0.06_17, 0.03_64, 0.05_44, 0.05_23, 0.07_20]					)

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




	def 						__UpperCAmelCase     (  self				:	List[str]					)  ->       List[str]:



		'''simple docstring'''





		_lowercase					:       Union[str, Any]       =		DDIMScheduler.from_pretrained(
		    'runwayml/stable-diffusion-v1-5'     ,  subfolder='scheduler'     ,  revision='onnx'					)
		_lowercase					:       List[str]       =		OnnxStableDiffusionPipeline.from_pretrained(
		    'runwayml/stable-diffusion-v1-5'     ,  revision='onnx'     ,  scheduler=_lowerCamelCase     ,  safety_checker=_lowerCamelCase     ,  feature_extractor=_lowerCamelCase     ,  provider=self.gpu_provider     ,  sess_options=self.gpu_options     ,  )
		sd_pipe.set_progress_bar_config(disable=_lowerCamelCase					)

		_lowercase					:       Optional[int]       =		'open neural network exchange'
		_lowercase					:       Union[str, Any]       =		np.random.RandomState(0					)
		_lowercase					:       Tuple       =		sd_pipe([prompt]     ,  guidance_scale=7.5     ,  num_inference_steps=10     ,  generator=_lowerCamelCase     ,  output_type='np'					)
		_lowercase					:       List[str]       =		output.images
		_lowercase					:       str       =		image[0, -3:, -3:, -1]

		assert image.shape == (1, 512, 512, 3)
		_lowercase					:       Tuple       =		np.array([0.28_67, 0.19_74, 0.14_81, 0.72_94, 0.72_51, 0.66_67, 0.41_94, 0.56_42, 0.64_86]					)

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




	def 						__UpperCAmelCase     (  self				:	Any					)  ->       int:



		'''simple docstring'''





		_lowercase					:       Union[str, Any]       =		LMSDiscreteScheduler.from_pretrained(
		    'runwayml/stable-diffusion-v1-5'     ,  subfolder='scheduler'     ,  revision='onnx'					)
		_lowercase					:       Dict       =		OnnxStableDiffusionPipeline.from_pretrained(
		    'runwayml/stable-diffusion-v1-5'     ,  revision='onnx'     ,  scheduler=_lowerCamelCase     ,  safety_checker=_lowerCamelCase     ,  feature_extractor=_lowerCamelCase     ,  provider=self.gpu_provider     ,  sess_options=self.gpu_options     ,  )
		sd_pipe.set_progress_bar_config(disable=_lowerCamelCase					)

		_lowercase					:       Optional[int]       =		'open neural network exchange'
		_lowercase					:       List[str]       =		np.random.RandomState(0					)
		_lowercase					:       Optional[int]       =		sd_pipe([prompt]     ,  guidance_scale=7.5     ,  num_inference_steps=10     ,  generator=_lowerCamelCase     ,  output_type='np'					)
		_lowercase					:       str       =		output.images
		_lowercase					:       Any       =		image[0, -3:, -3:, -1]

		assert image.shape == (1, 512, 512, 3)
		_lowercase					:       Optional[Any]       =		np.array([0.23_06, 0.19_59, 0.15_93, 0.65_49, 0.63_94, 0.54_08, 0.50_65, 0.60_10, 0.61_61]					)

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




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



		'''simple docstring'''





		_lowercase					:       str       =		0

		def test_callback_fn(UpperCamelCase_				:	int     ,  UpperCamelCase_				:	int     ,  UpperCamelCase_				:	np.ndarray					) -> None:
			_lowercase					:       Union[str, Any]       =		True
			nonlocal number_of_steps
			number_of_steps += 1
			if step == 0:
				assert latents.shape == (1, 4, 64, 64)
				_lowercase					:       Tuple       =		latents[0, -3:, -3:, -1]
				_lowercase					:       str       =		np.array(
				    [-0.67_72, -0.38_35, -1.24_56, 0.19_05, -1.09_74, 0.69_67, -1.93_53, 0.01_78, 1.01_67]					)

				assert np.abs(latents_slice.flatten() - expected_slice					).max() < 1E-3
			elif step == 5:
				assert latents.shape == (1, 4, 64, 64)
				_lowercase					:       Tuple       =		latents[0, -3:, -3:, -1]
				_lowercase					:       List[Any]       =		np.array(
				    [-0.33_51, 0.22_41, -0.18_37, -0.23_25, -0.65_77, 0.33_93, -0.02_41, 0.58_99, 1.38_75]					)

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

		_lowercase					:       Dict       =		False

		_lowercase					:       Optional[int]       =		OnnxStableDiffusionPipeline.from_pretrained(
		    'runwayml/stable-diffusion-v1-5'     ,  revision='onnx'     ,  safety_checker=_lowerCamelCase     ,  feature_extractor=_lowerCamelCase     ,  provider=self.gpu_provider     ,  sess_options=self.gpu_options     ,  )
		pipe.set_progress_bar_config(disable=_lowerCamelCase					)

		_lowercase					:       Any       =		'Andromeda galaxy in a bottle'

		_lowercase					:       Dict       =		np.random.RandomState(0					)
		pipe(
		    prompt=_lowerCamelCase     ,  num_inference_steps=5     ,  guidance_scale=7.5     ,  generator=_lowerCamelCase     ,  callback=_lowerCamelCase     ,  callback_steps=1     ,  )
		assert test_callback_fn.has_been_called
		assert number_of_steps == 6




	def 						__UpperCAmelCase     (  self				:	str					)  ->       str:



		'''simple docstring'''





		_lowercase					:       Dict       =		OnnxStableDiffusionPipeline.from_pretrained(
		    'runwayml/stable-diffusion-v1-5'     ,  revision='onnx'     ,  safety_checker=_lowerCamelCase     ,  feature_extractor=_lowerCamelCase     ,  provider=self.gpu_provider     ,  sess_options=self.gpu_options     ,  )
		assert isinstance(_lowerCamelCase     ,  _lowerCamelCase					)
		assert pipe.safety_checker is None

		_lowercase					:       Any       =		pipe('example prompt'     ,  num_inference_steps=2					).images[0]
		assert image is not None

		# check that there's no error when saving a pipeline with one of the models being None
		with tempfile.TemporaryDirectory() as tmpdirname:
			pipe.save_pretrained(_lowerCamelCase					)
			_lowercase					:       Optional[int]       =		OnnxStableDiffusionPipeline.from_pretrained(_lowerCamelCase					)

		# sanity check that the pipeline still works
		assert pipe.safety_checker is None
		_lowercase					:       Optional[Any]       =		pipe('example prompt'     ,  num_inference_steps=2					).images[0]
		assert image is not None





'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


_A				:    Union[str, Any]     						={'''configuration_reformer''': ['''REFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ReformerConfig''']}

try:
		if not is_sentencepiece_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_A				:    Dict     						=['''ReformerTokenizer''']

try:
		if not is_tokenizers_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_A				:    Union[str, Any]     						=['''ReformerTokenizerFast''']

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_A				:    str     						=[
		    '''REFORMER_PRETRAINED_MODEL_ARCHIVE_LIST''',
		    '''ReformerAttention''',
		    '''ReformerForMaskedLM''',
		    '''ReformerForQuestionAnswering''',
		    '''ReformerForSequenceClassification''',
		    '''ReformerLayer''',
		    '''ReformerModel''',
		    '''ReformerModelWithLMHead''',
		    '''ReformerPreTrainedModel''',
		]


if TYPE_CHECKING:
		from .configuration_reformer import REFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, ReformerConfig

		try:
				if not is_sentencepiece_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .tokenization_reformer import ReformerTokenizer

		try:
				if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .tokenization_reformer_fast import ReformerTokenizerFast

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_reformer import (
				    REFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
				    ReformerAttention,
				    ReformerForMaskedLM,
				    ReformerForQuestionAnswering,
				    ReformerForSequenceClassification,
				    ReformerLayer,
				    ReformerModel,
				    ReformerModelWithLMHead,
				    ReformerPreTrainedModel,
				)

else:
		import sys

		_A				:    Union[str, Any]     						=_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)
