






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




							def 		lowerCamelCase_    (							self:    str      ,							UpperCamelCase__:    List[str]=0					):
														lowerCamelCase__ :     List[Any]	  =	np.random.RandomState(UpperCamelCase__					)
														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:    Any					):
														lowerCamelCase__ :     Dict	  =	OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint      ,							provider="""CPUExecutionProvider"""					)
														pipe.set_progress_bar_config(disable=UpperCamelCase__					)

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

														assert image.shape == (1, 128, 128, 3)
														lowerCamelCase__ :     Union[str, 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:    str					):
														lowerCamelCase__ :     str	  =	OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint      ,							provider="""CPUExecutionProvider"""					)
														lowerCamelCase__ :     Optional[int]	  =	PNDMScheduler.from_config(pipe.scheduler.config      ,							skip_prk_steps=UpperCamelCase__					)
														pipe.set_progress_bar_config(disable=UpperCamelCase__					)

														lowerCamelCase__ :     Optional[Any]	  =	self.get_dummy_inputs()
														lowerCamelCase__ :     List[Any]	  =	pipe(**UpperCamelCase__					).images
														lowerCamelCase__ :     Any	  =	image[0, -3:, -3:, -1]

														assert image.shape == (1, 128, 128, 3)
														lowerCamelCase__ :     Optional[Any]	  =	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:    Tuple					):
														lowerCamelCase__ :     List[str]	  =	OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint      ,							provider="""CPUExecutionProvider"""					)
														lowerCamelCase__ :     Tuple	  =	LMSDiscreteScheduler.from_config(pipe.scheduler.config					)
														pipe.set_progress_bar_config(disable=UpperCamelCase__					)

														lowerCamelCase__ :     Tuple	  =	self.get_dummy_inputs()
														lowerCamelCase__ :     str	  =	pipe(**UpperCamelCase__					).images
														lowerCamelCase__ :     Tuple	  =	image[0, -3:, -3:, -1]

														assert image.shape == (1, 128, 128, 3)
														lowerCamelCase__ :     str	  =	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:    Dict					):
														lowerCamelCase__ :     Dict	  =	OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint      ,							provider="""CPUExecutionProvider"""					)
														lowerCamelCase__ :     Optional[int]	  =	EulerDiscreteScheduler.from_config(pipe.scheduler.config					)
														pipe.set_progress_bar_config(disable=UpperCamelCase__					)

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

														assert image.shape == (1, 128, 128, 3)
														lowerCamelCase__ :     Union[str, Any]	  =	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:    List[str]					):
														lowerCamelCase__ :     Union[str, Any]	  =	OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint      ,							provider="""CPUExecutionProvider"""					)
														lowerCamelCase__ :     str	  =	EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config					)
														pipe.set_progress_bar_config(disable=UpperCamelCase__					)

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

														assert image.shape == (1, 128, 128, 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:    str					):
														lowerCamelCase__ :     str	  =	OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint      ,							provider="""CPUExecutionProvider"""					)
														lowerCamelCase__ :     List[str]	  =	DPMSolverMultistepScheduler.from_config(pipe.scheduler.config					)
														pipe.set_progress_bar_config(disable=UpperCamelCase__					)

														lowerCamelCase__ :     Tuple	  =	self.get_dummy_inputs()
														lowerCamelCase__ :     int	  =	pipe(**UpperCamelCase__					).images
														lowerCamelCase__ :     Any	  =	image[0, -3:, -3:, -1]

														assert image.shape == (1, 128, 128, 3)
														lowerCamelCase__ :     Dict	  =	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:    int					):
														lowerCamelCase__ :     Any	  =	OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint      ,							provider="""CPUExecutionProvider"""					)
														pipe.set_progress_bar_config(disable=UpperCamelCase__					)

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

														# forward
														lowerCamelCase__ :     List[str]	  =	pipe(**UpperCamelCase__					)
														lowerCamelCase__ :     Tuple	  =	output.images[0, -3:, -3:, -1]

														lowerCamelCase__ :     Tuple	  =	self.get_dummy_inputs()
														lowerCamelCase__ :     List[Any]	  =	3 * [inputs.pop("""prompt"""					)]

														lowerCamelCase__ :     List[str]	  =	pipe.tokenizer(
														    UpperCamelCase__      ,							padding="""max_length"""      ,							max_length=pipe.tokenizer.model_max_length      ,							truncation=UpperCamelCase__      ,							return_tensors="""np"""      ,							)
														lowerCamelCase__ :     Tuple	  =	text_inputs['''input_ids''']

														lowerCamelCase__ :     int	  =	pipe.text_encoder(input_ids=text_inputs.astype(np.intaa					)					)[0]

														lowerCamelCase__ :     List[Any]	  =	prompt_embeds

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

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




							def 		lowerCamelCase_    (							self:    Tuple					):
														lowerCamelCase__ :     Optional[Any]	  =	OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint      ,							provider="""CPUExecutionProvider"""					)
														pipe.set_progress_bar_config(disable=UpperCamelCase__					)

														lowerCamelCase__ :     Optional[int]	  =	self.get_dummy_inputs()
														lowerCamelCase__ :     Any	  =	3 * ['''this is a negative prompt''']
														lowerCamelCase__ :     Any	  =	negative_prompt
														lowerCamelCase__ :     Tuple	  =	3 * [inputs['''prompt''']]

														# forward
														lowerCamelCase__ :     List[str]	  =	pipe(**UpperCamelCase__					)
														lowerCamelCase__ :     Any	  =	output.images[0, -3:, -3:, -1]

														lowerCamelCase__ :     Tuple	  =	self.get_dummy_inputs()
														lowerCamelCase__ :     Optional[Any]	  =	3 * [inputs.pop("""prompt"""					)]

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

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

														lowerCamelCase__ :     List[str]	  =	embeds

														# forward
														lowerCamelCase__ :     List[Any]	  =	pipe(**UpperCamelCase__					)
														lowerCamelCase__ :     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 _lowercase     (    unittest.TestCase		):




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




							@property
							def 		lowerCamelCase_    (							self:    List[str]					):
														lowerCamelCase__ :     Any	  =	ort.SessionOptions()
														lowerCamelCase__ :     Optional[Any]	  =	False
														return options




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

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

														lowerCamelCase__ :     Optional[int]	  =	image[0, -3:, -3:, -1]

														assert image.shape == (1, 512, 512, 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[Any]					):
														lowerCamelCase__ :     Optional[Any]	  =	DDIMScheduler.from_pretrained(
														    """runwayml/stable-diffusion-v1-5"""      ,							subfolder="""scheduler"""      ,							revision="""onnx"""					)
														lowerCamelCase__ :     Optional[int]	  =	OnnxStableDiffusionPipeline.from_pretrained(
														    """runwayml/stable-diffusion-v1-5"""      ,							revision="""onnx"""      ,							scheduler=UpperCamelCase__      ,							safety_checker=UpperCamelCase__      ,							feature_extractor=UpperCamelCase__      ,							provider=self.gpu_provider      ,							sess_options=self.gpu_options      ,							)
														sd_pipe.set_progress_bar_config(disable=UpperCamelCase__					)

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

														assert image.shape == (1, 512, 512, 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:    Optional[Any]					):
														lowerCamelCase__ :     Tuple	  =	LMSDiscreteScheduler.from_pretrained(
														    """runwayml/stable-diffusion-v1-5"""      ,							subfolder="""scheduler"""      ,							revision="""onnx"""					)
														lowerCamelCase__ :     str	  =	OnnxStableDiffusionPipeline.from_pretrained(
														    """runwayml/stable-diffusion-v1-5"""      ,							revision="""onnx"""      ,							scheduler=UpperCamelCase__      ,							safety_checker=UpperCamelCase__      ,							feature_extractor=UpperCamelCase__      ,							provider=self.gpu_provider      ,							sess_options=self.gpu_options      ,							)
														sd_pipe.set_progress_bar_config(disable=UpperCamelCase__					)

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

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

														def test_callback_fn(UpperCamelCase__:    int      ,							UpperCamelCase__:    int      ,							UpperCamelCase__:    np.ndarray					) -> None:
																					lowerCamelCase__ :     List[str]	  =	True
																					nonlocal number_of_steps
																					number_of_steps += 1
																					if step == 0:
																												assert latents.shape == (1, 4, 64, 64)
																												lowerCamelCase__ :     str	  =	latents[0, -3:, -3:, -1]
																												lowerCamelCase__ :     Any	  =	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__ :     List[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__ :     List[str]	  =	False

														lowerCamelCase__ :     Any	  =	OnnxStableDiffusionPipeline.from_pretrained(
														    """runwayml/stable-diffusion-v1-5"""      ,							revision="""onnx"""      ,							safety_checker=UpperCamelCase__      ,							feature_extractor=UpperCamelCase__      ,							provider=self.gpu_provider      ,							sess_options=self.gpu_options      ,							)
														pipe.set_progress_bar_config(disable=UpperCamelCase__					)

														lowerCamelCase__ :     Tuple	  =	'''Andromeda galaxy in a bottle'''

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




							def 		lowerCamelCase_    (							self:    Any					):
														lowerCamelCase__ :     str	  =	OnnxStableDiffusionPipeline.from_pretrained(
														    """runwayml/stable-diffusion-v1-5"""      ,							revision="""onnx"""      ,							safety_checker=UpperCamelCase__      ,							feature_extractor=UpperCamelCase__      ,							provider=self.gpu_provider      ,							sess_options=self.gpu_options      ,							)
														assert isinstance(UpperCamelCase__      ,							UpperCamelCase__					)
														assert pipe.safety_checker is None

														lowerCamelCase__ :     List[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(UpperCamelCase__					)
																					lowerCamelCase__ :     List[Any]	  =	OnnxStableDiffusionPipeline.from_pretrained(UpperCamelCase__					)

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









'''simple docstring'''







import argparse

from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection

from diffusers import UnCLIPImageVariationPipeline, UnCLIPPipeline


if __name__ == "__main__":
			_A     :      List[Any]   					=argparse.ArgumentParser()

			parser.add_argument('''--dump_path''', default=None, type=str, required=True, help='''Path to the output model.''')

			parser.add_argument(
			    '''--txt2img_unclip''',
			    default='''kakaobrain/karlo-v1-alpha''',
			    type=str,
			    required=False,
			    help='''The pretrained txt2img unclip.''',
			)

			_A     :      Dict   					=parser.parse_args()

			_A     :      List[str]   					=UnCLIPPipeline.from_pretrained(args.txtaimg_unclip)

			_A     :      Any   					=CLIPImageProcessor()
			_A     :      Union[str, Any]   					=CLIPVisionModelWithProjection.from_pretrained('''openai/clip-vit-large-patch14''')

			_A     :      Union[str, Any]   					=UnCLIPImageVariationPipeline(
			    decoder=txtaimg.decoder,
			    text_encoder=txtaimg.text_encoder,
			    tokenizer=txtaimg.tokenizer,
			    text_proj=txtaimg.text_proj,
			    feature_extractor=feature_extractor,
			    image_encoder=image_encoder,
			    super_res_first=txtaimg.super_res_first,
			    super_res_last=txtaimg.super_res_last,
			    decoder_scheduler=txtaimg.decoder_scheduler,
			    super_res_scheduler=txtaimg.super_res_scheduler,
			)

			imgaimg.save_pretrained(args.dump_path)

