


"""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   (lowerCamelCase__  ,						unittest.TestCase):





							'''simple docstring'''

							__lowercase:						Optional[Any]												=			'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'



							def 						lowerCAmelCase				( self							:					Optional[Any]						,				UpperCAmelCase_							:					Tuple=0				)    ->Optional[Any]:



														"""simple docstring"""



														snake_case_							=						np.random.RandomState(lowercase__				)
														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[Any]				)    ->List[Any]:



														"""simple docstring"""



														snake_case_							=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint						,				provider="""CPUExecutionProvider"""				)
														pipe.set_progress_bar_config(disable=lowercase__				)

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

														assert image.shape == (1, 128, 128, 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							:					Any				)    ->List[Any]:



														"""simple docstring"""



														snake_case_							=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint						,				provider="""CPUExecutionProvider"""				)
														snake_case_							=						PNDMScheduler.from_config(pipe.scheduler.config						,				skip_prk_steps=lowercase__				)
														pipe.set_progress_bar_config(disable=lowercase__				)

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

														assert image.shape == (1, 128, 128, 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				)    ->List[str]:



														"""simple docstring"""



														snake_case_							=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint						,				provider="""CPUExecutionProvider"""				)
														snake_case_							=						LMSDiscreteScheduler.from_config(pipe.scheduler.config				)
														pipe.set_progress_bar_config(disable=lowercase__				)

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

														assert image.shape == (1, 128, 128, 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[int]				)    ->str:



														"""simple docstring"""



														snake_case_							=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint						,				provider="""CPUExecutionProvider"""				)
														snake_case_							=						EulerDiscreteScheduler.from_config(pipe.scheduler.config				)
														pipe.set_progress_bar_config(disable=lowercase__				)

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

														assert image.shape == (1, 128, 128, 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							:					Tuple				)    ->List[Any]:



														"""simple docstring"""



														snake_case_							=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint						,				provider="""CPUExecutionProvider"""				)
														snake_case_							=						EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config				)
														pipe.set_progress_bar_config(disable=lowercase__				)

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

														assert image.shape == (1, 128, 128, 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							:					List[Any]				)    ->str:



														"""simple docstring"""



														snake_case_							=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint						,				provider="""CPUExecutionProvider"""				)
														snake_case_							=						DPMSolverMultistepScheduler.from_config(pipe.scheduler.config				)
														pipe.set_progress_bar_config(disable=lowercase__				)

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

														assert image.shape == (1, 128, 128, 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							:					Optional[int]				)    ->int:



														"""simple docstring"""



														snake_case_							=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint						,				provider="""CPUExecutionProvider"""				)
														pipe.set_progress_bar_config(disable=lowercase__				)

														snake_case_							=						self.get_dummy_inputs()
														snake_case_							=						3 * [inputs["""prompt"""]]

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

														snake_case_							=						self.get_dummy_inputs()
														snake_case_							=						3 * [inputs.pop("""prompt"""				)]

														snake_case_							=						pipe.tokenizer(
														    lowercase__						,				padding="""max_length"""						,				max_length=pipe.tokenizer.model_max_length						,				truncation=lowercase__						,				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(**lowercase__				)
														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:



														"""simple docstring"""



														snake_case_							=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint						,				provider="""CPUExecutionProvider"""				)
														pipe.set_progress_bar_config(disable=lowercase__				)

														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(**lowercase__				)
														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(
																					    lowercase__						,				padding="""max_length"""						,				max_length=pipe.tokenizer.model_max_length						,				truncation=lowercase__						,				return_tensors="""np"""						,				)
																					snake_case_							=						text_inputs["""input_ids"""]

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

														snake_case_				,  snake_case_							=						embeds

														# forward
														snake_case_							=						pipe(**lowercase__				)
														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    __A   (unittest.TestCase):





							'''simple docstring'''

							@property
							def 						lowerCAmelCase				( self							:					Dict				)    ->List[str]:



														"""simple docstring"""



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



							@property
							def 						lowerCAmelCase				( self							:					int				)    ->List[Any]:



														"""simple docstring"""



														snake_case_							=						ort.SessionOptions()
														snake_case_							=						False
														return options



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



														"""simple docstring"""



														snake_case_							=						OnnxStableDiffusionPipeline.from_pretrained(
														    """CompVis/stable-diffusion-v1-4"""						,				revision="""onnx"""						,				safety_checker=lowercase__						,				feature_extractor=lowercase__						,				provider=self.gpu_provider						,				sess_options=self.gpu_options						,				)
														sd_pipe.set_progress_bar_config(disable=lowercase__				)

														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, 512, 512, 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							:					Tuple				)    ->int:



														"""simple docstring"""



														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=lowercase__						,				safety_checker=lowercase__						,				feature_extractor=lowercase__						,				provider=self.gpu_provider						,				sess_options=self.gpu_options						,				)
														sd_pipe.set_progress_bar_config(disable=lowercase__				)

														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=lowercase__						,				output_type="""np"""				)
														snake_case_							=						output.images
														snake_case_							=						image[0, -3:, -3:, -1]

														assert image.shape == (1, 512, 512, 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							:					List[Any]				)    ->List[str]:



														"""simple docstring"""



														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=lowercase__						,				safety_checker=lowercase__						,				feature_extractor=lowercase__						,				provider=self.gpu_provider						,				sess_options=self.gpu_options						,				)
														sd_pipe.set_progress_bar_config(disable=lowercase__				)

														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=lowercase__						,				output_type="""np"""				)
														snake_case_							=						output.images
														snake_case_							=						image[0, -3:, -3:, -1]

														assert image.shape == (1, 512, 512, 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							:					Optional[int]				)    ->Optional[Any]:



														"""simple docstring"""



														snake_case_							=						0

														def test_callback_fn(UpperCAmelCase_							:					int						,				UpperCAmelCase_							:					int						,				UpperCAmelCase_							:					np.ndarray				) -> 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=lowercase__						,				feature_extractor=lowercase__						,				provider=self.gpu_provider						,				sess_options=self.gpu_options						,				)
														pipe.set_progress_bar_config(disable=lowercase__				)

														snake_case_							=						"""Andromeda galaxy in a bottle"""

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



							def 						lowerCAmelCase				( self							:					Optional[Any]				)    ->Dict:



														"""simple docstring"""



														snake_case_							=						OnnxStableDiffusionPipeline.from_pretrained(
														    """runwayml/stable-diffusion-v1-5"""						,				revision="""onnx"""						,				safety_checker=lowercase__						,				feature_extractor=lowercase__						,				provider=self.gpu_provider						,				sess_options=self.gpu_options						,				)
														assert isinstance(lowercase__						,				lowercase__				)
														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(lowercase__				)
																					snake_case_							=						OnnxStableDiffusionPipeline.from_pretrained(lowercase__				)

														# 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'''






import argparse
from pathlib import Path

import torch
from packaging import version
from torch.onnx import export

from diffusers import AutoencoderKL


lowerCAmelCase__											=	version.parse(version.parse(torch.__version__).base_version) < version.parse('''1.11''')


def 		_A     (    A__     ,				A__     ,				A__     ,				A__     ,				A__     ,				A__     ,				A__     ,				A__=False     ,				):






 """simple docstring"""



 output_path.parent.mkdir(parents=A__     ,				exist_ok=A__      )
 # PyTorch deprecated the `enable_onnx_checker` and `use_external_data_format` arguments in v1.11,
 # so we check the torch version for backwards compatibility
 if is_torch_less_than_1_11:
  export(
      A__     ,				A__     ,				f=output_path.as_posix()     ,				input_names=A__     ,				output_names=A__     ,				dynamic_axes=A__     ,				do_constant_folding=A__     ,				use_external_data_format=A__     ,				enable_onnx_checker=A__     ,				opset_version=A__     ,				)
 else:
  export(
      A__     ,				A__     ,				f=output_path.as_posix()     ,				input_names=A__     ,				output_names=A__     ,				dynamic_axes=A__     ,				do_constant_folding=A__     ,				opset_version=A__     ,				)


@torch.no_grad()
def 		_A     (    A__     ,				A__     ,				A__     ,				A__ = False      ):






 """simple docstring"""



 __lowercase							  = torch.floataa if fpaa else torch.floataa
 if fpaa and torch.cuda.is_available():
  __lowercase							  = '''cuda'''
 elif fpaa and not torch.cuda.is_available():
  raise ValueError('''`float16` model export is only supported on GPUs with CUDA'''      )
 else:
  __lowercase							  = '''cpu'''
 __lowercase							  = Path(A__      )

 # VAE DECODER
 __lowercase							  = AutoencoderKL.from_pretrained(model_path + '''/vae'''      )
 __lowercase							  = vae_decoder.config.latent_channels
 # forward only through the decoder part
 __lowercase							  = vae_decoder.decode
 onnx_export(
     A__     ,				model_args=(
         torch.randn(1     ,				A__     ,				25     ,				25      ).to(device=A__     ,				dtype=A__      ),
         False,
     )     ,				output_path=output_path / '''vae_decoder''' / '''model.onnx'''     ,				ordered_input_names=['''latent_sample''', '''return_dict''']     ,				output_names=['''sample''']     ,				dynamic_axes={
         '''latent_sample''': {0: '''batch''', 1: '''channels''', 2: '''height''', 3: '''width'''},
     }     ,				opset=A__     ,				)
 del vae_decoder


if __name__ == "__main__":
 lowerCAmelCase__											=	argparse.ArgumentParser()

 parser.add_argument(
     '''--model_path''',
     type=str,
     required=True,
     help='''Path to the `diffusers` checkpoint to convert (either a local directory or on the Hub).''',
 )

 parser.add_argument('''--output_path''', type=str, required=True, help='''Path to the output model.''')
 parser.add_argument(
     '''--opset''',
     default=14,
     type=int,
     help='''The version of the ONNX operator set to use.''',
 )
 parser.add_argument('''--fp16''', action='''store_true''', default=False, help='''Export the models in `float16` mode''')

 lowerCAmelCase__											=	parser.parse_args()
 print(args.output_path)
 convert_models(args.model_path, args.output_path, args.opset, args.fpaa)
 print('''SD: Done: ONNX''')

