


'''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 						a__							(				a__							,							unittest.TestCase    ):





							'''simple docstring'''






							lowercase__     :  Dict      		= "hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_=0  )			->				Tuple:
														lowerCAmelCase__			  =      np.random.RandomState(lowerCamelCase_  )
														lowerCAmelCase__			  =      {
														    '''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     __SCREAMING_SNAKE_CASE				(						self  )			->				List[str]:
														lowerCAmelCase__			  =      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,       provider='''CPUExecutionProvider'''  )
														pipe.set_progress_bar_config(disable=lowerCamelCase_  )

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

														assert image.shape == (1, 1_28, 1_28, 3)
														lowerCAmelCase__			  =      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     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[int]:
														lowerCAmelCase__			  =      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,       provider='''CPUExecutionProvider'''  )
														lowerCAmelCase__			  =      PNDMScheduler.from_config(pipe.scheduler.config   ,       skip_prk_steps=lowerCamelCase_  )
														pipe.set_progress_bar_config(disable=lowerCamelCase_  )

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

														assert image.shape == (1, 1_28, 1_28, 3)
														lowerCAmelCase__			  =      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     __SCREAMING_SNAKE_CASE				(						self  )			->				List[Any]:
														lowerCAmelCase__			  =      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,       provider='''CPUExecutionProvider'''  )
														lowerCAmelCase__			  =      LMSDiscreteScheduler.from_config(pipe.scheduler.config  )
														pipe.set_progress_bar_config(disable=lowerCamelCase_  )

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

														assert image.shape == (1, 1_28, 1_28, 3)
														lowerCAmelCase__			  =      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     __SCREAMING_SNAKE_CASE				(						self  )			->				Dict:
														lowerCAmelCase__			  =      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,       provider='''CPUExecutionProvider'''  )
														lowerCAmelCase__			  =      EulerDiscreteScheduler.from_config(pipe.scheduler.config  )
														pipe.set_progress_bar_config(disable=lowerCamelCase_  )

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

														assert image.shape == (1, 1_28, 1_28, 3)
														lowerCAmelCase__			  =      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     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[int]:
														lowerCAmelCase__			  =      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,       provider='''CPUExecutionProvider'''  )
														lowerCAmelCase__			  =      EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config  )
														pipe.set_progress_bar_config(disable=lowerCamelCase_  )

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

														assert image.shape == (1, 1_28, 1_28, 3)
														lowerCAmelCase__			  =      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     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														lowerCAmelCase__			  =      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,       provider='''CPUExecutionProvider'''  )
														lowerCAmelCase__			  =      DPMSolverMultistepScheduler.from_config(pipe.scheduler.config  )
														pipe.set_progress_bar_config(disable=lowerCamelCase_  )

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

														assert image.shape == (1, 1_28, 1_28, 3)
														lowerCAmelCase__			  =      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     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														lowerCAmelCase__			  =      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,       provider='''CPUExecutionProvider'''  )
														pipe.set_progress_bar_config(disable=lowerCamelCase_  )

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

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

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

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

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

														lowerCAmelCase__			  =      prompt_embeds

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

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







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				str:
														lowerCAmelCase__			  =      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,       provider='''CPUExecutionProvider'''  )
														pipe.set_progress_bar_config(disable=lowerCamelCase_  )

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

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

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

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

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

														lowerCAmelCase__  ,   lowerCAmelCase__			  =      embeds

														# forward
														lowerCAmelCase__			  =      pipe(**lowerCamelCase_  )
														lowerCAmelCase__			  =      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 						a__							(				unittest.TestCase    ):





							'''simple docstring'''






							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				str:
														return (
														    "CUDAExecutionProvider",
														    {
														        "gpu_mem_limit": "15000000000",  # 15GB
														        "arena_extend_strategy": "kSameAsRequested",
														    },
														)







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[Any]:
														lowerCAmelCase__			  =      ort.SessionOptions()
														lowerCAmelCase__			  =      False
														return options







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[int]:
														# using the PNDM scheduler by default
														lowerCAmelCase__			  =      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_  )

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

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

														assert image.shape == (1, 5_12, 5_12, 3)
														lowerCAmelCase__			  =      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     __SCREAMING_SNAKE_CASE				(						self  )			->				Any:
														lowerCAmelCase__			  =      DDIMScheduler.from_pretrained(
														    '''runwayml/stable-diffusion-v1-5'''   ,       subfolder='''scheduler'''   ,       revision='''onnx'''  )
														lowerCAmelCase__			  =      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_  )

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

														assert image.shape == (1, 5_12, 5_12, 3)
														lowerCAmelCase__			  =      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     __SCREAMING_SNAKE_CASE				(						self  )			->				Dict:
														lowerCAmelCase__			  =      LMSDiscreteScheduler.from_pretrained(
														    '''runwayml/stable-diffusion-v1-5'''   ,       subfolder='''scheduler'''   ,       revision='''onnx'''  )
														lowerCAmelCase__			  =      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_  )

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

														assert image.shape == (1, 5_12, 5_12, 3)
														lowerCAmelCase__			  =      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     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														lowerCAmelCase__			  =      0

														def test_callback_fn(lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  ) -> None:
																					lowerCAmelCase__			  =      True
																					nonlocal number_of_steps
																					number_of_steps += 1
																					if step == 0:
																												assert latents.shape == (1, 4, 64, 64)
																												lowerCAmelCase__			  =      latents[0, -3:, -3:, -1]
																												lowerCAmelCase__			  =      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)
																												lowerCAmelCase__			  =      latents[0, -3:, -3:, -1]
																												lowerCAmelCase__			  =      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__			  =      False

														lowerCAmelCase__			  =      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_  )

														lowerCAmelCase__			  =      '''Andromeda galaxy in a bottle'''

														lowerCAmelCase__			  =      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     __SCREAMING_SNAKE_CASE				(						self  )			->				List[Any]:
														lowerCAmelCase__			  =      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

														lowerCAmelCase__			  =      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_  )
																					lowerCAmelCase__			  =      OnnxStableDiffusionPipeline.from_pretrained(lowerCamelCase_  )

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



'''simple docstring'''






# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

from ...utils.dataclasses import (
    ComputeEnvironment,
    DistributedType,
    DynamoBackend,
    PrecisionType,
    SageMakerDistributedType,
)
from ..menu import BulletMenu


__UpperCAmelCase       							=				[
    '''EAGER''',
    '''AOT_EAGER''',
    '''INDUCTOR''',
    '''NVFUSER''',
    '''AOT_NVFUSER''',
    '''AOT_CUDAGRAPHS''',
    '''OFI''',
    '''FX2TRT''',
    '''ONNXRT''',
    '''IPEX''',
]


def    _snake_case  (	A					,       A=None					,       A=None					,       A=None			)		->       Union[str, Any]:
							lowerCAmelCase__			  =      True
							while ask_again:
														lowerCAmelCase__			  =      input(A			)
														try:
																					if default is not None and len(A			) == 0:
																												return default
																					return convert_value(A			) if convert_value is not None else result
														except Exception:
																					if error_message is not None:
																												print(A			)


def    _snake_case  (	A					,       A=[]					,       A=None					,       A=0			)		->       List[Any]:
							lowerCAmelCase__			  =      BulletMenu(A					,       A			)
							lowerCAmelCase__			  =      menu.run(default_choice=A			)
							return convert_value(A			) if convert_value is not None else result


def    _snake_case  (	A			)		->       Tuple:
							lowerCAmelCase__			  =      int(A			)
							return ComputeEnvironment(['''LOCAL_MACHINE''', '''AMAZON_SAGEMAKER'''][value]			)


def    _snake_case  (	A			)		->       Union[str, Any]:
							lowerCAmelCase__			  =      int(A			)
							return DistributedType(['''NO''', '''MULTI_CPU''', '''MULTI_XPU''', '''MULTI_GPU''', '''MULTI_NPU''', '''TPU'''][value]			)


def    _snake_case  (	A			)		->       str:
							lowerCAmelCase__			  =      int(A			)
							return DynamoBackend(DYNAMO_BACKENDS[value]			).value


def    _snake_case  (	A			)		->       Tuple:
							lowerCAmelCase__			  =      int(A			)
							return PrecisionType(['''no''', '''fp16''', '''bf16''', '''fp8'''][value]			)


def    _snake_case  (	A			)		->       Union[str, Any]:
							lowerCAmelCase__			  =      int(A			)
							return SageMakerDistributedType(['''NO''', '''DATA_PARALLEL''', '''MODEL_PARALLEL'''][value]			)


def    _snake_case  (	A			)		->       List[str]:
							return {"yes": True, "no": False}[value.lower()]






class 						a__							(				argparse.RawDescriptionHelpFormatter    ):





							'''simple docstring'''






							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  )			->				Dict:
														lowerCAmelCase__			  =      super()._format_usage(lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  )
														lowerCAmelCase__			  =      usage.replace('''<command> [<args>] '''   ,       ''''''  )
														return usage