






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      __SCREAMING_SNAKE_CASE(		UpperCAmelCase__			,     unittest.TestCase		):
							_UpperCAmelCase       			=				"""hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"""
							def lowerCAmelCase_       (   self:					List[str]				,		UpperCamelCase:					str=0  )						->						int:
											snake_case__													=   np.random.RandomState(snake_case_  )
											snake_case__													=   {
											    '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 lowerCAmelCase_       (   self:					Optional[int]  )						->						Optional[Any]:
											snake_case__													=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint				,		provider='CPUExecutionProvider'  )
											pipe.set_progress_bar_config(disable=snake_case_  )

											snake_case__													=   self.get_dummy_inputs()
											snake_case__													=   pipe(**snake_case_  ).images
											snake_case__													=   image[0, -3:, -3:, -1]

											assert image.shape == (1, 1_28, 1_28, 3)
											snake_case__													=   np.array([0.65_072, 0.58_492, 0.48_219, 0.55_521, 0.53_180, 0.55_939, 0.50_697, 0.39_800, 0.46_455]  )

											assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2
							def lowerCAmelCase_       (   self:					Dict  )						->						Tuple:
											snake_case__													=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint				,		provider='CPUExecutionProvider'  )
											snake_case__													=   PNDMScheduler.from_config(pipe.scheduler.config				,		skip_prk_steps=snake_case_  )
											pipe.set_progress_bar_config(disable=snake_case_  )

											snake_case__													=   self.get_dummy_inputs()
											snake_case__													=   pipe(**snake_case_  ).images
											snake_case__													=   image[0, -3:, -3:, -1]

											assert image.shape == (1, 1_28, 1_28, 3)
											snake_case__													=   np.array([0.65_863, 0.59_425, 0.49_326, 0.56_313, 0.53_875, 0.56_627, 0.51_065, 0.39_777, 0.46_330]  )

											assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2
							def lowerCAmelCase_       (   self:					Any  )						->						Dict:
											snake_case__													=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint				,		provider='CPUExecutionProvider'  )
											snake_case__													=   LMSDiscreteScheduler.from_config(pipe.scheduler.config  )
											pipe.set_progress_bar_config(disable=snake_case_  )

											snake_case__													=   self.get_dummy_inputs()
											snake_case__													=   pipe(**snake_case_  ).images
											snake_case__													=   image[0, -3:, -3:, -1]

											assert image.shape == (1, 1_28, 1_28, 3)
											snake_case__													=   np.array([0.53_755, 0.60_786, 0.47_402, 0.49_488, 0.51_869, 0.49_819, 0.47_985, 0.38_957, 0.44_279]  )

											assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2
							def lowerCAmelCase_       (   self:					Optional[Any]  )						->						Dict:
											snake_case__													=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint				,		provider='CPUExecutionProvider'  )
											snake_case__													=   EulerDiscreteScheduler.from_config(pipe.scheduler.config  )
											pipe.set_progress_bar_config(disable=snake_case_  )

											snake_case__													=   self.get_dummy_inputs()
											snake_case__													=   pipe(**snake_case_  ).images
											snake_case__													=   image[0, -3:, -3:, -1]

											assert image.shape == (1, 1_28, 1_28, 3)
											snake_case__													=   np.array([0.53_755, 0.60_786, 0.47_402, 0.49_488, 0.51_869, 0.49_819, 0.47_985, 0.38_957, 0.44_279]  )

											assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2
							def lowerCAmelCase_       (   self:					str  )						->						Union[str, Any]:
											snake_case__													=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint				,		provider='CPUExecutionProvider'  )
											snake_case__													=   EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config  )
											pipe.set_progress_bar_config(disable=snake_case_  )

											snake_case__													=   self.get_dummy_inputs()
											snake_case__													=   pipe(**snake_case_  ).images
											snake_case__													=   image[0, -3:, -3:, -1]

											assert image.shape == (1, 1_28, 1_28, 3)
											snake_case__													=   np.array([0.53_817, 0.60_812, 0.47_384, 0.49_530, 0.51_894, 0.49_814, 0.47_984, 0.38_958, 0.44_271]  )

											assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2
							def lowerCAmelCase_       (   self:					Dict  )						->						Tuple:
											snake_case__													=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint				,		provider='CPUExecutionProvider'  )
											snake_case__													=   DPMSolverMultistepScheduler.from_config(pipe.scheduler.config  )
											pipe.set_progress_bar_config(disable=snake_case_  )

											snake_case__													=   self.get_dummy_inputs()
											snake_case__													=   pipe(**snake_case_  ).images
											snake_case__													=   image[0, -3:, -3:, -1]

											assert image.shape == (1, 1_28, 1_28, 3)
											snake_case__													=   np.array([0.53_895, 0.60_808, 0.47_933, 0.49_608, 0.51_886, 0.49_950, 0.48_053, 0.38_957, 0.44_200]  )

											assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2
							def lowerCAmelCase_       (   self:					Union[str, Any]  )						->						List[str]:
											snake_case__													=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint				,		provider='CPUExecutionProvider'  )
											pipe.set_progress_bar_config(disable=snake_case_  )

											snake_case__													=   self.get_dummy_inputs()
											snake_case__													=   3 * [inputs['prompt']]

											# forward
											snake_case__													=   pipe(**snake_case_  )
											snake_case__													=   output.images[0, -3:, -3:, -1]

											snake_case__													=   self.get_dummy_inputs()
											snake_case__													=   3 * [inputs.pop('prompt'  )]

											snake_case__													=   pipe.tokenizer(
											    snake_case_				,		padding='max_length'				,		max_length=pipe.tokenizer.model_max_length				,		truncation=snake_case_				,		return_tensors='np'				,		)
											snake_case__													=   text_inputs['input_ids']

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

											snake_case__													=   prompt_embeds

											# forward
											snake_case__													=   pipe(**snake_case_  )
											snake_case__													=   output.images[0, -3:, -3:, -1]

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


							def lowerCAmelCase_       (   self:					Optional[Any]  )						->						Any:
											snake_case__													=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint				,		provider='CPUExecutionProvider'  )
											pipe.set_progress_bar_config(disable=snake_case_  )

											snake_case__													=   self.get_dummy_inputs()
											snake_case__													=   3 * ['this is a negative prompt']
											snake_case__													=   negative_prompt
											snake_case__													=   3 * [inputs['prompt']]

											# forward
											snake_case__													=   pipe(**snake_case_  )
											snake_case__													=   output.images[0, -3:, -3:, -1]

											snake_case__													=   self.get_dummy_inputs()
											snake_case__													=   3 * [inputs.pop('prompt'  )]

											snake_case__													=   []
											for p in [prompt, negative_prompt]:
															snake_case__													=   pipe.tokenizer(
															    snake_case_				,		padding='max_length'				,		max_length=pipe.tokenizer.model_max_length				,		truncation=snake_case_				,		return_tensors='np'				,		)
															snake_case__													=   text_inputs['input_ids']

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

											snake_case__													=   embeds

											# forward
											snake_case__													=   pipe(**snake_case_  )
											snake_case__													=   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      __SCREAMING_SNAKE_CASE(		unittest.TestCase		):
							@property
							def lowerCAmelCase_       (   self:					Any  )						->						Optional[int]:
											return (
											    "CUDAExecutionProvider",
											    {
											        "gpu_mem_limit": "15000000000",  # 15GB
											        "arena_extend_strategy": "kSameAsRequested",
											    },
											)
							@property
							def lowerCAmelCase_       (   self:					Optional[int]  )						->						Any:
											snake_case__													=   ort.SessionOptions()
											snake_case__													=   False
											return options
							def lowerCAmelCase_       (   self:					Union[str, Any]  )						->						List[Any]:
											snake_case__													=   OnnxStableDiffusionPipeline.from_pretrained(
											    'CompVis/stable-diffusion-v1-4'				,		revision='onnx'				,		safety_checker=snake_case_				,		feature_extractor=snake_case_				,		provider=self.gpu_provider				,		sess_options=self.gpu_options				,		)
											sd_pipe.set_progress_bar_config(disable=snake_case_  )

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

											snake_case__													=   image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_12, 5_12, 3)
											snake_case__													=   np.array([0.0_452, 0.0_390, 0.0_087, 0.0_350, 0.0_617, 0.0_364, 0.0_544, 0.0_523, 0.0_720]  )

											assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-3
							def lowerCAmelCase_       (   self:					int  )						->						str:
											snake_case__													=   DDIMScheduler.from_pretrained(
											    'runwayml/stable-diffusion-v1-5'				,		subfolder='scheduler'				,		revision='onnx'  )
											snake_case__													=   OnnxStableDiffusionPipeline.from_pretrained(
											    'runwayml/stable-diffusion-v1-5'				,		revision='onnx'				,		scheduler=snake_case_				,		safety_checker=snake_case_				,		feature_extractor=snake_case_				,		provider=self.gpu_provider				,		sess_options=self.gpu_options				,		)
											sd_pipe.set_progress_bar_config(disable=snake_case_  )

											snake_case__													=   'open neural network exchange'
											snake_case__													=   np.random.RandomState(0  )
											snake_case__													=   sd_pipe([prompt]				,		guidance_scale=7.5				,		num_inference_steps=10				,		generator=snake_case_				,		output_type='np'  )
											snake_case__													=   output.images
											snake_case__													=   image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_12, 5_12, 3)
											snake_case__													=   np.array([0.2_867, 0.1_974, 0.1_481, 0.7_294, 0.7_251, 0.6_667, 0.4_194, 0.5_642, 0.6_486]  )

											assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-3
							def lowerCAmelCase_       (   self:					Optional[int]  )						->						List[str]:
											snake_case__													=   LMSDiscreteScheduler.from_pretrained(
											    'runwayml/stable-diffusion-v1-5'				,		subfolder='scheduler'				,		revision='onnx'  )
											snake_case__													=   OnnxStableDiffusionPipeline.from_pretrained(
											    'runwayml/stable-diffusion-v1-5'				,		revision='onnx'				,		scheduler=snake_case_				,		safety_checker=snake_case_				,		feature_extractor=snake_case_				,		provider=self.gpu_provider				,		sess_options=self.gpu_options				,		)
											sd_pipe.set_progress_bar_config(disable=snake_case_  )

											snake_case__													=   'open neural network exchange'
											snake_case__													=   np.random.RandomState(0  )
											snake_case__													=   sd_pipe([prompt]				,		guidance_scale=7.5				,		num_inference_steps=10				,		generator=snake_case_				,		output_type='np'  )
											snake_case__													=   output.images
											snake_case__													=   image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_12, 5_12, 3)
											snake_case__													=   np.array([0.2_306, 0.1_959, 0.1_593, 0.6_549, 0.6_394, 0.5_408, 0.5_065, 0.6_010, 0.6_161]  )

											assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-3
							def lowerCAmelCase_       (   self:					int  )						->						List[str]:
											snake_case__													=   0

											def test_callback_fn(UpperCamelCase:					Optional[Any]				,		UpperCamelCase:					Any				,		UpperCamelCase:					List[str]  ) -> None:
															snake_case__													=   True
															nonlocal number_of_steps
															number_of_steps += 1
															if step == 0:
																			assert latents.shape == (1, 4, 64, 64)
																			snake_case__													=   latents[0, -3:, -3:, -1]
																			snake_case__													=   np.array(
																			    [-0.6_772, -0.3_835, -1.2_456, 0.1_905, -1.0_974, 0.6_967, -1.9_353, 0.0_178, 1.0_167]  )

																			assert np.abs(latents_slice.flatten() - expected_slice  ).max() < 1e-3
															elif step == 5:
																			assert latents.shape == (1, 4, 64, 64)
																			snake_case__													=   latents[0, -3:, -3:, -1]
																			snake_case__													=   np.array(
																			    [-0.3_351, 0.2_241, -0.1_837, -0.2_325, -0.6_577, 0.3_393, -0.0_241, 0.5_899, 1.3_875]  )

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

											snake_case__													=   False

											snake_case__													=   OnnxStableDiffusionPipeline.from_pretrained(
											    'runwayml/stable-diffusion-v1-5'				,		revision='onnx'				,		safety_checker=snake_case_				,		feature_extractor=snake_case_				,		provider=self.gpu_provider				,		sess_options=self.gpu_options				,		)
											pipe.set_progress_bar_config(disable=snake_case_  )

											snake_case__													=   'Andromeda galaxy in a bottle'

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


							def lowerCAmelCase_       (   self:					str  )						->						Any:
											snake_case__													=   OnnxStableDiffusionPipeline.from_pretrained(
											    'runwayml/stable-diffusion-v1-5'				,		revision='onnx'				,		safety_checker=snake_case_				,		feature_extractor=snake_case_				,		provider=self.gpu_provider				,		sess_options=self.gpu_options				,		)
											assert isinstance(snake_case_				,		snake_case_  )
											assert pipe.safety_checker is None

											snake_case__													=   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(snake_case_  )
															snake_case__													=   OnnxStableDiffusionPipeline.from_pretrained(snake_case_  )

											# sanity check that the pipeline still works
											assert pipe.safety_checker is None
											snake_case__													=   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


lowerCamelCase_     :						int						=					{}


try:
   if not is_sentencepiece_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   lowerCamelCase_     :						Tuple						=					['MLukeTokenizer']

if TYPE_CHECKING:
   try:
      if not is_sentencepiece_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_mluke import MLukeTokenizer


else:
   import sys

   lowerCamelCase_     :						List[str]						=					_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)