





"""simple docstring"""





import argparse
import os
import shutil
from pathlib import Path

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

from diffusers import OnnxRuntimeModel, OnnxStableDiffusionPipeline, StableDiffusionPipeline


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


def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE=False  ,	):




		"""simple docstring"""


		output_path.parent.mkdir(parents=_SCREAMING_SNAKE_CASE  ,	exist_ok=_SCREAMING_SNAKE_CASE    )
		# 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(
				    _SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	f=output_path.as_posix()  ,	input_names=_SCREAMING_SNAKE_CASE  ,	output_names=_SCREAMING_SNAKE_CASE  ,	dynamic_axes=_SCREAMING_SNAKE_CASE  ,	do_constant_folding=_SCREAMING_SNAKE_CASE  ,	use_external_data_format=_SCREAMING_SNAKE_CASE  ,	enable_onnx_checker=_SCREAMING_SNAKE_CASE  ,	opset_version=_SCREAMING_SNAKE_CASE  ,	)
		else:
				export(
				    _SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	f=output_path.as_posix()  ,	input_names=_SCREAMING_SNAKE_CASE  ,	output_names=_SCREAMING_SNAKE_CASE  ,	dynamic_axes=_SCREAMING_SNAKE_CASE  ,	do_constant_folding=_SCREAMING_SNAKE_CASE  ,	opset_version=_SCREAMING_SNAKE_CASE  ,	)


@torch.no_grad()
def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE = False    ):




		"""simple docstring"""


		UpperCamelCase   =	torch.floataa if fpaa else torch.floataa
		if fpaa and torch.cuda.is_available():
				UpperCamelCase   =	"cuda"
		elif fpaa and not torch.cuda.is_available():
				raise ValueError("`float16` model export is only supported on GPUs with CUDA"    )
		else:
				UpperCamelCase   =	"cpu"
		UpperCamelCase   =	StableDiffusionPipeline.from_pretrained(_SCREAMING_SNAKE_CASE  ,	torch_dtype=_SCREAMING_SNAKE_CASE    ).to(_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	Path(_SCREAMING_SNAKE_CASE    )

		# TEXT ENCODER
		UpperCamelCase   =	pipeline.text_encoder.config.max_position_embeddings
		UpperCamelCase   =	pipeline.text_encoder.config.hidden_size
		UpperCamelCase   =	pipeline.tokenizer(
		    "A sample prompt"  ,	padding="max_length"  ,	max_length=pipeline.tokenizer.model_max_length  ,	truncation=_SCREAMING_SNAKE_CASE  ,	return_tensors="pt"  ,	)
		onnx_export(
		    pipeline.text_encoder  ,	model_args=(text_input.input_ids.to(device=_SCREAMING_SNAKE_CASE  ,	dtype=torch.intaa    ))  ,	output_path=output_path / "text_encoder" / "model.onnx"  ,	ordered_input_names=["input_ids"]  ,	output_names=["last_hidden_state", "pooler_output"]  ,	dynamic_axes={
		        "input_ids": {0: "batch", 1: "sequence"},
		    }  ,	opset=_SCREAMING_SNAKE_CASE  ,	)
		del pipeline.text_encoder

		# UNET
		UpperCamelCase   =	pipeline.unet.config.in_channels
		UpperCamelCase   =	pipeline.unet.config.sample_size
		UpperCamelCase   =	output_path / "unet" / "model.onnx"
		onnx_export(
		    pipeline.unet  ,	model_args=(
		        torch.randn(2  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ).to(device=_SCREAMING_SNAKE_CASE  ,	dtype=_SCREAMING_SNAKE_CASE    ),
		        torch.randn(2    ).to(device=_SCREAMING_SNAKE_CASE  ,	dtype=_SCREAMING_SNAKE_CASE    ),
		        torch.randn(2  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ).to(device=_SCREAMING_SNAKE_CASE  ,	dtype=_SCREAMING_SNAKE_CASE    ),
		        False,
		    )  ,	output_path=_SCREAMING_SNAKE_CASE  ,	ordered_input_names=["sample", "timestep", "encoder_hidden_states", "return_dict"]  ,	output_names=["out_sample"]  ,	dynamic_axes={
		        "sample": {0: "batch", 1: "channels", 2: "height", 3: "width"},
		        "timestep": {0: "batch"},
		        "encoder_hidden_states": {0: "batch", 1: "sequence"},
		    }  ,	opset=_SCREAMING_SNAKE_CASE  ,	use_external_data_format=_SCREAMING_SNAKE_CASE  ,	)
		UpperCamelCase   =	str(unet_path.absolute().as_posix()    )
		UpperCamelCase   =	os.path.dirname(_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	onnx.load(_SCREAMING_SNAKE_CASE    )
		# clean up existing tensor files
		shutil.rmtree(_SCREAMING_SNAKE_CASE    )
		os.mkdir(_SCREAMING_SNAKE_CASE    )
		# collate external tensor files into one
		onnx.save_model(
		    _SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	save_as_external_data=_SCREAMING_SNAKE_CASE  ,	all_tensors_to_one_file=_SCREAMING_SNAKE_CASE  ,	location="weights.pb"  ,	convert_attribute=_SCREAMING_SNAKE_CASE  ,	)
		del pipeline.unet

		# VAE ENCODER
		UpperCamelCase   =	pipeline.vae
		UpperCamelCase   =	vae_encoder.config.in_channels
		UpperCamelCase   =	vae_encoder.config.sample_size
		# need to get the raw tensor output (sample) from the encoder
		UpperCamelCase   =	lambda _SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    : vae_encoder.encode(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    )[0].sample()
		onnx_export(
		    _SCREAMING_SNAKE_CASE  ,	model_args=(
		        torch.randn(1  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ).to(device=_SCREAMING_SNAKE_CASE  ,	dtype=_SCREAMING_SNAKE_CASE    ),
		        False,
		    )  ,	output_path=output_path / "vae_encoder" / "model.onnx"  ,	ordered_input_names=["sample", "return_dict"]  ,	output_names=["latent_sample"]  ,	dynamic_axes={
		        "sample": {0: "batch", 1: "channels", 2: "height", 3: "width"},
		    }  ,	opset=_SCREAMING_SNAKE_CASE  ,	)

		# VAE DECODER
		UpperCamelCase   =	pipeline.vae
		UpperCamelCase   =	vae_decoder.config.latent_channels
		UpperCamelCase   =	vae_decoder.config.out_channels
		# forward only through the decoder part
		UpperCamelCase   =	vae_encoder.decode
		onnx_export(
		    _SCREAMING_SNAKE_CASE  ,	model_args=(
		        torch.randn(1  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ).to(device=_SCREAMING_SNAKE_CASE  ,	dtype=_SCREAMING_SNAKE_CASE    ),
		        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=_SCREAMING_SNAKE_CASE  ,	)
		del pipeline.vae

		# SAFETY CHECKER
		if pipeline.safety_checker is not None:
				UpperCamelCase   =	pipeline.safety_checker
				UpperCamelCase   =	safety_checker.config.vision_config.num_channels
				UpperCamelCase   =	safety_checker.config.vision_config.image_size
				UpperCamelCase   =	safety_checker.forward_onnx
				onnx_export(
				    pipeline.safety_checker  ,	model_args=(
				        torch.randn(
				            1  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	).to(device=_SCREAMING_SNAKE_CASE  ,	dtype=_SCREAMING_SNAKE_CASE    ),
				        torch.randn(1  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ).to(device=_SCREAMING_SNAKE_CASE  ,	dtype=_SCREAMING_SNAKE_CASE    ),
				    )  ,	output_path=output_path / "safety_checker" / "model.onnx"  ,	ordered_input_names=["clip_input", "images"]  ,	output_names=["out_images", "has_nsfw_concepts"]  ,	dynamic_axes={
				        "clip_input": {0: "batch", 1: "channels", 2: "height", 3: "width"},
				        "images": {0: "batch", 1: "height", 2: "width", 3: "channels"},
				    }  ,	opset=_SCREAMING_SNAKE_CASE  ,	)
				del pipeline.safety_checker
				UpperCamelCase   =	OnnxRuntimeModel.from_pretrained(output_path / "safety_checker"    )
				UpperCamelCase   =	pipeline.feature_extractor
		else:
				UpperCamelCase   =	None
				UpperCamelCase   =	None

		UpperCamelCase   =	OnnxStableDiffusionPipeline(
		    vae_encoder=OnnxRuntimeModel.from_pretrained(output_path / "vae_encoder"    )  ,	vae_decoder=OnnxRuntimeModel.from_pretrained(output_path / "vae_decoder"    )  ,	text_encoder=OnnxRuntimeModel.from_pretrained(output_path / "text_encoder"    )  ,	tokenizer=pipeline.tokenizer  ,	unet=OnnxRuntimeModel.from_pretrained(output_path / "unet"    )  ,	scheduler=pipeline.scheduler  ,	safety_checker=_SCREAMING_SNAKE_CASE  ,	feature_extractor=_SCREAMING_SNAKE_CASE  ,	requires_safety_checker=safety_checker is not None  ,	)

		onnx_pipeline.save_pretrained(_SCREAMING_SNAKE_CASE    )
		print("ONNX pipeline saved to"  ,	_SCREAMING_SNAKE_CASE    )

		del pipeline
		del onnx_pipeline
		UpperCamelCase   =	OnnxStableDiffusionPipeline.from_pretrained(_SCREAMING_SNAKE_CASE  ,	provider="CPUExecutionProvider"    )
		print("ONNX pipeline is loadable"    )


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()

					convert_models(args.model_path, args.output_path, args.opset, args.fpaa)













"""simple docstring"""





import importlib
import os
import sys


# This is required to make the module import works (when the python process is running from the root of the repo)
sys.path.append('''.''')


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	test_file.split(os.path.sep    )
		if components[0:2] != ["tests", "models"]:
				raise ValueError(
				    "`test_file` should start with `tests/models/` (with `/` being the OS specific path separator). Got "
				    F"{test_file} instead."    )
		UpperCamelCase   =	components[-1]
		if not test_fn.endswith("py"    ):
				raise ValueError(F"`test_file` should be a python file. Got {test_fn} instead."    )
		if not test_fn.startswith("test_modeling_"    ):
				raise ValueError(
				    F"`test_file` should point to a file name of the form `test_modeling_*.py`. Got {test_fn} instead."    )

		UpperCamelCase   =	components[:-1] + [test_fn.replace(".py"  ,	""    )]
		UpperCamelCase   =	".".join(_SCREAMING_SNAKE_CASE    )

		return test_module_path


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	get_module_path(_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	importlib.import_module(_SCREAMING_SNAKE_CASE    )

		return test_module


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	[]
		UpperCamelCase   =	get_test_module(_SCREAMING_SNAKE_CASE    )
		for attr in dir(_SCREAMING_SNAKE_CASE    ):
				if attr.endswith("ModelTester"    ):
						tester_classes.append(getattr(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    )    )

    # sort with class names
		return sorted(_SCREAMING_SNAKE_CASE  ,	key=lambda _SCREAMING_SNAKE_CASE    : x.__name__    )


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	[]
		UpperCamelCase   =	get_test_module(_SCREAMING_SNAKE_CASE    )
		for attr in dir(_SCREAMING_SNAKE_CASE    ):
				UpperCamelCase   =	getattr(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    )
				# (TF/Flax)ModelTesterMixin is also an attribute in specific model test module. Let's exclude them by checking
				# `all_model_classes` is not empty (which also excludes other special classes).
				UpperCamelCase   =	getattr(_SCREAMING_SNAKE_CASE  ,	"all_model_classes"  ,	[]    )
				if len(_SCREAMING_SNAKE_CASE    ) > 0:
						test_classes.append(_SCREAMING_SNAKE_CASE    )

    # sort with class names
		return sorted(_SCREAMING_SNAKE_CASE  ,	key=lambda _SCREAMING_SNAKE_CASE    : x.__name__    )


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	get_test_classes(_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	set()
		for test_class in test_classes:
				model_classes.update(test_class.all_model_classes    )

		# sort with class names
		return sorted(_SCREAMING_SNAKE_CASE  ,	key=lambda _SCREAMING_SNAKE_CASE    : x.__name__    )


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	test_class()
		if hasattr(_SCREAMING_SNAKE_CASE  ,	"setUp"    ):
				test.setUp()

		UpperCamelCase   =	None
		if hasattr(_SCREAMING_SNAKE_CASE  ,	"model_tester"    ):
				# `(TF/Flax)ModelTesterMixin` has this attribute default to `None`. Let's skip this case.
				if test.model_tester is not None:
						UpperCamelCase   =	test.model_tester.__class__

		return model_tester


def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	get_test_classes(_SCREAMING_SNAKE_CASE    )

		UpperCamelCase   =	[]
		for test_class in test_classes:
				if model_class in test_class.all_model_classes:
						target_test_classes.append(_SCREAMING_SNAKE_CASE    )

    # sort with class names
		return sorted(_SCREAMING_SNAKE_CASE  ,	key=lambda _SCREAMING_SNAKE_CASE    : x.__name__    )


def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	get_test_classes_for_model(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    )

		UpperCamelCase   =	[]
		for test_class in test_classes:
				UpperCamelCase   =	get_model_tester_from_test_class(_SCREAMING_SNAKE_CASE    )
				if tester_class is not None:
						tester_classes.append(_SCREAMING_SNAKE_CASE    )

    # sort with class names
		return sorted(_SCREAMING_SNAKE_CASE  ,	key=lambda _SCREAMING_SNAKE_CASE    : x.__name__    )


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	get_test_classes(_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	{test_class: get_model_tester_from_test_class(_SCREAMING_SNAKE_CASE    ) for test_class in test_classes}
		return test_tester_mapping


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	get_model_classes(_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	{
		    model_class: get_test_classes_for_model(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ) for model_class in model_classes
		}
		return model_test_mapping


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	get_model_classes(_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	{
		    model_class: get_tester_classes_for_model(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ) for model_class in model_classes
		}
		return model_to_tester_mapping






def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		if isinstance(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):
				return o
		elif isinstance(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):
				return o.__name__
		elif isinstance(_SCREAMING_SNAKE_CASE  ,	(list, tuple)    ):
				return [to_json(_SCREAMING_SNAKE_CASE    ) for x in o]
		elif isinstance(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):
				return {to_json(_SCREAMING_SNAKE_CASE    ): to_json(_SCREAMING_SNAKE_CASE    ) for k, v in o.items()}
		else:
				return o






