


"""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 		snake_case_(       __SCREAMING_SNAKE_CASE							,					unittest.TestCase			):
					__UpperCamelCase  =					'''hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'''






					def      lowerCamelCase__     (	self  :					int					,				UpperCamelCase_  :					Optional[int]=0				):
											lowerCAmelCase					:	Any       =   np.random.RandomState(_a				)
											lowerCAmelCase					:	List[str]       =   {
											    """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  :					Union[str, Any]				):
											lowerCAmelCase					:	Union[str, Any]       =   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,				provider='''CPUExecutionProvider'''				)
											pipe.set_progress_bar_config(disable=_a				)

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

											assert image.shape == (1, 1_2_8, 1_2_8, 3)
											lowerCAmelCase					:	List[Any]       =   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  :					Any				):
											lowerCAmelCase					:	Union[str, Any]       =   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,				provider='''CPUExecutionProvider'''				)
											lowerCAmelCase					:	Any       =   PNDMScheduler.from_config(pipe.scheduler.config					,				skip_prk_steps=_a				)
											pipe.set_progress_bar_config(disable=_a				)

											lowerCAmelCase					:	Dict       =   self.get_dummy_inputs()
											lowerCAmelCase					:	Union[str, Any]       =   pipe(**_a				).images
											lowerCAmelCase					:	int       =   image[0, -3:, -3:, -1]

											assert image.shape == (1, 1_2_8, 1_2_8, 3)
											lowerCAmelCase					:	int       =   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  :					Dict				):
											lowerCAmelCase					:	Any       =   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,				provider='''CPUExecutionProvider'''				)
											lowerCAmelCase					:	Dict       =   LMSDiscreteScheduler.from_config(pipe.scheduler.config				)
											pipe.set_progress_bar_config(disable=_a				)

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

											assert image.shape == (1, 1_2_8, 1_2_8, 3)
											lowerCAmelCase					:	int       =   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  :					Union[str, Any]				):
											lowerCAmelCase					:	List[Any]       =   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,				provider='''CPUExecutionProvider'''				)
											lowerCAmelCase					:	List[str]       =   EulerDiscreteScheduler.from_config(pipe.scheduler.config				)
											pipe.set_progress_bar_config(disable=_a				)

											lowerCAmelCase					:	Dict       =   self.get_dummy_inputs()
											lowerCAmelCase					:	Union[str, Any]       =   pipe(**_a				).images
											lowerCAmelCase					:	int       =   image[0, -3:, -3:, -1]

											assert image.shape == (1, 1_2_8, 1_2_8, 3)
											lowerCAmelCase					:	Dict       =   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]				):
											lowerCAmelCase					:	List[str]       =   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,				provider='''CPUExecutionProvider'''				)
											lowerCAmelCase					:	Tuple       =   EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config				)
											pipe.set_progress_bar_config(disable=_a				)

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

											assert image.shape == (1, 1_2_8, 1_2_8, 3)
											lowerCAmelCase					:	Tuple       =   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  :					List[Any]				):
											lowerCAmelCase					:	Any       =   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,				provider='''CPUExecutionProvider'''				)
											lowerCAmelCase					:	Optional[int]       =   DPMSolverMultistepScheduler.from_config(pipe.scheduler.config				)
											pipe.set_progress_bar_config(disable=_a				)

											lowerCAmelCase					:	Tuple       =   self.get_dummy_inputs()
											lowerCAmelCase					:	int       =   pipe(**_a				).images
											lowerCAmelCase					:	List[str]       =   image[0, -3:, -3:, -1]

											assert image.shape == (1, 1_2_8, 1_2_8, 3)
											lowerCAmelCase					:	Union[str, Any]       =   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]				):
											lowerCAmelCase					:	Any       =   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,				provider='''CPUExecutionProvider'''				)
											pipe.set_progress_bar_config(disable=_a				)

											lowerCAmelCase					:	str       =   self.get_dummy_inputs()
											lowerCAmelCase					:	List[str]       =   3 * [inputs["""prompt"""]]

											# forward
											lowerCAmelCase					:	Optional[Any]       =   pipe(**_a				)
											lowerCAmelCase					:	Any       =   output.images[0, -3:, -3:, -1]

											lowerCAmelCase					:	str       =   self.get_dummy_inputs()
											lowerCAmelCase					:	Tuple       =   3 * [inputs.pop('''prompt'''				)]

											lowerCAmelCase					:	List[Any]       =   pipe.tokenizer(
											    _a					,				padding='''max_length'''					,				max_length=pipe.tokenizer.model_max_length					,				truncation=_a					,				return_tensors='''np'''					,				)
											lowerCAmelCase					:	Union[str, Any]       =   text_inputs["""input_ids"""]

											lowerCAmelCase					:	Tuple       =   pipe.text_encoder(input_ids=text_inputs.astype(np.intaa				)				)[0]

											lowerCAmelCase					:	Optional[int]       =   prompt_embeds

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

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






					def      lowerCamelCase__     (	self  :					Dict				):
											lowerCAmelCase					:	int       =   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,				provider='''CPUExecutionProvider'''				)
											pipe.set_progress_bar_config(disable=_a				)

											lowerCAmelCase					:	Tuple       =   self.get_dummy_inputs()
											lowerCAmelCase					:	Union[str, Any]       =   3 * ["""this is a negative prompt"""]
											lowerCAmelCase					:	Any       =   negative_prompt
											lowerCAmelCase					:	Optional[Any]       =   3 * [inputs["""prompt"""]]

											# forward
											lowerCAmelCase					:	Tuple       =   pipe(**_a				)
											lowerCAmelCase					:	Any       =   output.images[0, -3:, -3:, -1]

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

											lowerCAmelCase					:	List[Any]       =   []
											for p in [prompt, negative_prompt]:
																	lowerCAmelCase					:	int       =   pipe.tokenizer(
																	    _a					,				padding='''max_length'''					,				max_length=pipe.tokenizer.model_max_length					,				truncation=_a					,				return_tensors='''np'''					,				)
																	lowerCAmelCase					:	Tuple       =   text_inputs["""input_ids"""]

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

											lowerCAmelCase					:	Union[str, Any]       =   embeds

											# forward
											lowerCAmelCase					:	List[Any]       =   pipe(**_a				)
											lowerCAmelCase					:	str       =   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 		snake_case_(       unittest.TestCase			):






					@property
					def      lowerCamelCase__     (	self  :					Tuple				):
											return (
											    "CUDAExecutionProvider",
											    {
											        "gpu_mem_limit": "15000000000",  # 15GB
											        "arena_extend_strategy": "kSameAsRequested",
											    },
											)






					@property
					def      lowerCamelCase__     (	self  :					Any				):
											lowerCAmelCase					:	Union[str, Any]       =   ort.SessionOptions()
											lowerCAmelCase					:	int       =   False
											return options






					def      lowerCamelCase__     (	self  :					Any				):
											lowerCAmelCase					:	Tuple       =   OnnxStableDiffusionPipeline.from_pretrained(
											    '''CompVis/stable-diffusion-v1-4'''					,				revision='''onnx'''					,				safety_checker=_a					,				feature_extractor=_a					,				provider=self.gpu_provider					,				sess_options=self.gpu_options					,				)
											sd_pipe.set_progress_bar_config(disable=_a				)

											lowerCAmelCase					:	Optional[int]       =   """A painting of a squirrel eating a burger"""
											np.random.seed(0				)
											lowerCAmelCase					:	Dict       =   sd_pipe([prompt]					,				guidance_scale=6.0					,				num_inference_steps=1_0					,				output_type='''np'''				)
											lowerCAmelCase					:	Tuple       =   output.images

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

											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											lowerCAmelCase					:	Tuple       =   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  :					Optional[int]				):
											lowerCAmelCase					:	Dict       =   DDIMScheduler.from_pretrained(
											    '''runwayml/stable-diffusion-v1-5'''					,				subfolder='''scheduler'''					,				revision='''onnx'''				)
											lowerCAmelCase					:	int       =   OnnxStableDiffusionPipeline.from_pretrained(
											    '''runwayml/stable-diffusion-v1-5'''					,				revision='''onnx'''					,				scheduler=_a					,				safety_checker=_a					,				feature_extractor=_a					,				provider=self.gpu_provider					,				sess_options=self.gpu_options					,				)
											sd_pipe.set_progress_bar_config(disable=_a				)

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

											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											lowerCAmelCase					:	str       =   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  :					Any				):
											lowerCAmelCase					:	str       =   LMSDiscreteScheduler.from_pretrained(
											    '''runwayml/stable-diffusion-v1-5'''					,				subfolder='''scheduler'''					,				revision='''onnx'''				)
											lowerCAmelCase					:	List[str]       =   OnnxStableDiffusionPipeline.from_pretrained(
											    '''runwayml/stable-diffusion-v1-5'''					,				revision='''onnx'''					,				scheduler=_a					,				safety_checker=_a					,				feature_extractor=_a					,				provider=self.gpu_provider					,				sess_options=self.gpu_options					,				)
											sd_pipe.set_progress_bar_config(disable=_a				)

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

											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											lowerCAmelCase					:	str       =   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  :					Any				):
											lowerCAmelCase					:	int       =   0

											def test_callback_fn(UpperCamelCase_  :					int					,				UpperCamelCase_  :					int					,				UpperCamelCase_  :					np.ndarray				) -> None:
																	lowerCAmelCase					:	Any       =   True
																	nonlocal number_of_steps
																	number_of_steps += 1
																	if step == 0:
																							assert latents.shape == (1, 4, 6_4, 6_4)
																							lowerCAmelCase					:	Union[str, Any]       =   latents[0, -3:, -3:, -1]
																							lowerCAmelCase					:	str       =   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, 6_4, 6_4)
																							lowerCAmelCase					:	Any       =   latents[0, -3:, -3:, -1]
																							lowerCAmelCase					:	int       =   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

											lowerCAmelCase					:	int       =   False

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

											lowerCAmelCase					:	List[str]       =   """Andromeda galaxy in a bottle"""

											lowerCAmelCase					:	List[Any]       =   np.random.RandomState(0				)
											pipe(
											    prompt=_a					,				num_inference_steps=5					,				guidance_scale=7.5					,				generator=_a					,				callback=_a					,				callback_steps=1					,				)
											assert test_callback_fn.has_been_called
											assert number_of_steps == 6






					def      lowerCamelCase__     (	self  :					Dict				):
											lowerCAmelCase					:	Optional[Any]       =   OnnxStableDiffusionPipeline.from_pretrained(
											    '''runwayml/stable-diffusion-v1-5'''					,				revision='''onnx'''					,				safety_checker=_a					,				feature_extractor=_a					,				provider=self.gpu_provider					,				sess_options=self.gpu_options					,				)
											assert isinstance(_a					,				_a				)
											assert pipe.safety_checker is None

											lowerCAmelCase					:	List[str]       =   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(_a				)
																	lowerCAmelCase					:	Tuple       =   OnnxStableDiffusionPipeline.from_pretrained(_a				)

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






"""simple docstring"""





from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


snake_case__		:     Tuple    				=				logging.get_logger(__name__)

snake_case__		:     int    				=				{
    '''facebook/levit-128S''': '''https://huggingface.co/facebook/levit-128S/resolve/main/config.json''',
    # See all LeViT models at https://huggingface.co/models?filter=levit
}



class 		snake_case_(       a__			):
     __UpperCamelCase  =					'''levit'''






     def __init__(	self  :					str					,				UpperCamelCase_  :					Union[str, Any]=2_2_4					,				UpperCamelCase_  :					Union[str, Any]=3					,				UpperCamelCase_  :					Union[str, Any]=3					,				UpperCamelCase_  :					int=2					,				UpperCamelCase_  :					Union[str, Any]=1					,				UpperCamelCase_  :					Tuple=1_6					,				UpperCamelCase_  :					Dict=[1_2_8, 2_5_6, 3_8_4]					,				UpperCamelCase_  :					Optional[Any]=[4, 8, 1_2]					,				UpperCamelCase_  :					Dict=[4, 4, 4]					,				UpperCamelCase_  :					Any=[1_6, 1_6, 1_6]					,				UpperCamelCase_  :					str=0					,				UpperCamelCase_  :					int=[2, 2, 2]					,				UpperCamelCase_  :					Optional[Any]=[2, 2, 2]					,				UpperCamelCase_  :					str=0.02					,				**UpperCamelCase_  :					List[str]					,				):
           super().__init__(**UpperCamelCase_				)
           lowerCAmelCase					:	Tuple       =   image_size
           lowerCAmelCase					:	int       =   num_channels
           lowerCAmelCase					:	Optional[int]       =   kernel_size
           lowerCAmelCase					:	Dict       =   stride
           lowerCAmelCase					:	List[Any]       =   padding
           lowerCAmelCase					:	Dict       =   hidden_sizes
           lowerCAmelCase					:	List[str]       =   num_attention_heads
           lowerCAmelCase					:	Tuple       =   depths
           lowerCAmelCase					:	Dict       =   key_dim
           lowerCAmelCase					:	Union[str, Any]       =   drop_path_rate
           lowerCAmelCase					:	List[Any]       =   patch_size
           lowerCAmelCase					:	Tuple       =   attention_ratio
           lowerCAmelCase					:	Optional[int]       =   mlp_ratio
           lowerCAmelCase					:	Union[str, Any]       =   initializer_range
           lowerCAmelCase					:	List[str]       =   [
               ['''Subsample''', key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
               ['''Subsample''', key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
           ]



class 		snake_case_(       a__			):
     __UpperCamelCase  =					version.parse('''1.11'''			)






     @property
     def      lowerCamelCase__     (	self  :					Tuple				):
           return OrderedDict(
               [
                   ('''pixel_values''', {0: '''batch''', 1: '''num_channels''', 2: '''height''', 3: '''width'''}),
               ]				)






     @property
     def      lowerCamelCase__     (	self  :					Optional[Any]				):
           return 1E-4


