
'''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					:			List[Any]			      =		version.parse(version.parse(torch.__version__).base_version) < version.parse("1.11")


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


@torch.no_grad()
def     UpperCAmelCase_  (						__lowerCAmelCase	,					__lowerCAmelCase	,					__lowerCAmelCase	,					__lowerCAmelCase = False	) ->						Dict:
    __lowercase		:			Any	    =       torch.floataa if fpaa else torch.floataa
    if fpaa and torch.cuda.is_available():
        __lowercase		:			List[str]	    =       '''cuda'''
    elif fpaa and not torch.cuda.is_available():
        raise ValueError('''`float16` model export is only supported on GPUs with CUDA'''	)
    else:
        __lowercase		:			Union[str, Any]	    =       '''cpu'''
    __lowercase		:			Optional[int]	    =       StableDiffusionPipeline.from_pretrained(__lowerCAmelCase	,					torch_dtype=__lowerCAmelCase	).to(__lowerCAmelCase	)
    __lowercase		:			Dict	    =       Path(__lowerCAmelCase	)

    # TEXT ENCODER
    __lowercase		:			str	    =       pipeline.text_encoder.config.max_position_embeddings
    __lowercase		:			Dict	    =       pipeline.text_encoder.config.hidden_size
    __lowercase		:			Dict	    =       pipeline.tokenizer(
        '''A sample prompt'''	,					padding='''max_length'''	,					max_length=pipeline.tokenizer.model_max_length	,					truncation=__lowerCAmelCase	,					return_tensors='''pt'''	,					)
    onnx_export(
        pipeline.text_encoder	,					model_args=(text_input.input_ids.to(device=__lowerCAmelCase	,					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=__lowerCAmelCase	,					)
    del pipeline.text_encoder

    # UNET
    __lowercase		:			List[str]	    =       pipeline.unet.config.in_channels
    __lowercase		:			Optional[Any]	    =       pipeline.unet.config.sample_size
    __lowercase		:			List[Any]	    =       output_path / '''unet''' / '''model.onnx'''
    onnx_export(
        pipeline.unet	,					model_args=(
            torch.randn(2	,					__lowerCAmelCase	,					__lowerCAmelCase	,					__lowerCAmelCase	).to(device=__lowerCAmelCase	,					dtype=__lowerCAmelCase	),
            torch.randn(2	).to(device=__lowerCAmelCase	,					dtype=__lowerCAmelCase	),
            torch.randn(2	,					__lowerCAmelCase	,					__lowerCAmelCase	).to(device=__lowerCAmelCase	,					dtype=__lowerCAmelCase	),
            False,
        )	,					output_path=__lowerCAmelCase	,					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=__lowerCAmelCase	,					use_external_data_format=__lowerCAmelCase	,					)
    __lowercase		:			Union[str, Any]	    =       str(unet_path.absolute().as_posix()	)
    __lowercase		:			Optional[int]	    =       os.path.dirname(__lowerCAmelCase	)
    __lowercase		:			Optional[Any]	    =       onnx.load(__lowerCAmelCase	)
    # clean up existing tensor files
    shutil.rmtree(__lowerCAmelCase	)
    os.mkdir(__lowerCAmelCase	)
    # collate external tensor files into one
    onnx.save_model(
        __lowerCAmelCase	,					__lowerCAmelCase	,					save_as_external_data=__lowerCAmelCase	,					all_tensors_to_one_file=__lowerCAmelCase	,					location='''weights.pb'''	,					convert_attribute=__lowerCAmelCase	,					)
    del pipeline.unet

    # VAE ENCODER
    __lowercase		:			Optional[int]	    =       pipeline.vae
    __lowercase		:			Dict	    =       vae_encoder.config.in_channels
    __lowercase		:			List[Any]	    =       vae_encoder.config.sample_size
    # need to get the raw tensor output (sample) from the encoder
    __lowercase		:			List[str]	    =       lambda __lowerCAmelCase	,					__lowerCAmelCase	: vae_encoder.encode(__lowerCAmelCase	,					__lowerCAmelCase	)[0].sample()
    onnx_export(
        __lowerCAmelCase	,					model_args=(
            torch.randn(1	,					__lowerCAmelCase	,					__lowerCAmelCase	,					__lowerCAmelCase	).to(device=__lowerCAmelCase	,					dtype=__lowerCAmelCase	),
            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=__lowerCAmelCase	,					)

    # VAE DECODER
    __lowercase		:			Union[str, Any]	    =       pipeline.vae
    __lowercase		:			Union[str, Any]	    =       vae_decoder.config.latent_channels
    __lowercase		:			Any	    =       vae_decoder.config.out_channels
    # forward only through the decoder part
    __lowercase		:			List[Any]	    =       vae_encoder.decode
    onnx_export(
        __lowerCAmelCase	,					model_args=(
            torch.randn(1	,					__lowerCAmelCase	,					__lowerCAmelCase	,					__lowerCAmelCase	).to(device=__lowerCAmelCase	,					dtype=__lowerCAmelCase	),
            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=__lowerCAmelCase	,					)
    del pipeline.vae

    # SAFETY CHECKER
    if pipeline.safety_checker is not None:
        __lowercase		:			Optional[int]	    =       pipeline.safety_checker
        __lowercase		:			Optional[Any]	    =       safety_checker.config.vision_config.num_channels
        __lowercase		:			Union[str, Any]	    =       safety_checker.config.vision_config.image_size
        __lowercase		:			str	    =       safety_checker.forward_onnx
        onnx_export(
            pipeline.safety_checker	,					model_args=(
                torch.randn(
                    1	,					__lowerCAmelCase	,					__lowerCAmelCase	,					__lowerCAmelCase	,					).to(device=__lowerCAmelCase	,					dtype=__lowerCAmelCase	),
                torch.randn(1	,					__lowerCAmelCase	,					__lowerCAmelCase	,					__lowerCAmelCase	).to(device=__lowerCAmelCase	,					dtype=__lowerCAmelCase	),
            )	,					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=__lowerCAmelCase	,					)
        del pipeline.safety_checker
        __lowercase		:			Union[str, Any]	    =       OnnxRuntimeModel.from_pretrained(output_path / '''safety_checker'''	)
        __lowercase		:			int	    =       pipeline.feature_extractor
    else:
        __lowercase		:			Optional[int]	    =       None
        __lowercase		:			Tuple	    =       None

    __lowercase		:			Optional[int]	    =       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=__lowerCAmelCase	,					feature_extractor=__lowerCAmelCase	,					requires_safety_checker=safety_checker is not None	,					)

    onnx_pipeline.save_pretrained(__lowerCAmelCase	)
    print('''ONNX pipeline saved to'''	,					__lowerCAmelCase	)

    del pipeline
    del onnx_pipeline
    __lowercase		:			List[Any]	    =       OnnxStableDiffusionPipeline.from_pretrained(__lowerCAmelCase	,					provider='''CPUExecutionProvider'''	)
    print('''ONNX pipeline is loadable'''	)


if __name__ == "__main__":
 __lowerCAmelCase					:			int			      =		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					:			Union[str, Any]			      =		parser.parse_args()

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




import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np


class 						__lowerCAmelCase	:



      """simple docstring"""







      def 		snake_case_			(      self		:       str				):
          torch.manual_seed(0				)
          __lowercase		:			Optional[int]	    =       TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''				)

          torch.manual_seed(0				)
          __lowercase		:			Optional[Any]	    =       AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5'''				)

          torch.manual_seed(0				)
          __lowercase		:			Tuple	    =       UNetaDConditionModel(
              sample_size=32				,       layers_per_block=1				,       block_out_channels=[32, 64]				,       down_block_types=[
                  '''ResnetDownsampleBlock2D''',
                  '''SimpleCrossAttnDownBlock2D''',
              ]				,       mid_block_type='''UNetMidBlock2DSimpleCrossAttn'''				,       up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D''']				,       in_channels=3				,       out_channels=6				,       cross_attention_dim=32				,       encoder_hid_dim=32				,       attention_head_dim=8				,       addition_embed_type='''text'''				,       addition_embed_type_num_heads=2				,       cross_attention_norm='''group_norm'''				,       resnet_time_scale_shift='''scale_shift'''				,       act_fn='''gelu'''				,       )
          unet.set_attn_processor(AttnAddedKVProcessor()				)  # For reproducibility tests

          torch.manual_seed(0				)
          __lowercase		:			List[str]	    =       DDPMScheduler(
              num_train_timesteps=1000				,       beta_schedule='''squaredcos_cap_v2'''				,       beta_start=0.00_01				,       beta_end=0.02				,       thresholding=_snake_case				,       dynamic_thresholding_ratio=0.95				,       sample_max_value=1.0				,       prediction_type='''epsilon'''				,       variance_type='''learned_range'''				,       )

          torch.manual_seed(0				)
          __lowercase		:			str	    =       IFWatermarker()

          return {
              "text_encoder": text_encoder,
              "tokenizer": tokenizer,
              "unet": unet,
              "scheduler": scheduler,
              "watermarker": watermarker,
              "safety_checker": None,
              "feature_extractor": None,
          }




      def 		snake_case_			(      self		:       Optional[int]				):
          torch.manual_seed(0				)
          __lowercase		:			Optional[int]	    =       TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''				)

          torch.manual_seed(0				)
          __lowercase		:			Tuple	    =       AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5'''				)

          torch.manual_seed(0				)
          __lowercase		:			Dict	    =       UNetaDConditionModel(
              sample_size=32				,       layers_per_block=[1, 2]				,       block_out_channels=[32, 64]				,       down_block_types=[
                  '''ResnetDownsampleBlock2D''',
                  '''SimpleCrossAttnDownBlock2D''',
              ]				,       mid_block_type='''UNetMidBlock2DSimpleCrossAttn'''				,       up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D''']				,       in_channels=6				,       out_channels=6				,       cross_attention_dim=32				,       encoder_hid_dim=32				,       attention_head_dim=8				,       addition_embed_type='''text'''				,       addition_embed_type_num_heads=2				,       cross_attention_norm='''group_norm'''				,       resnet_time_scale_shift='''scale_shift'''				,       act_fn='''gelu'''				,       class_embed_type='''timestep'''				,       mid_block_scale_factor=1.4_14				,       time_embedding_act_fn='''gelu'''				,       time_embedding_dim=32				,       )
          unet.set_attn_processor(AttnAddedKVProcessor()				)  # For reproducibility tests

          torch.manual_seed(0				)
          __lowercase		:			Dict	    =       DDPMScheduler(
              num_train_timesteps=1000				,       beta_schedule='''squaredcos_cap_v2'''				,       beta_start=0.00_01				,       beta_end=0.02				,       thresholding=_snake_case				,       dynamic_thresholding_ratio=0.95				,       sample_max_value=1.0				,       prediction_type='''epsilon'''				,       variance_type='''learned_range'''				,       )

          torch.manual_seed(0				)
          __lowercase		:			Optional[int]	    =       DDPMScheduler(
              num_train_timesteps=1000				,       beta_schedule='''squaredcos_cap_v2'''				,       beta_start=0.00_01				,       beta_end=0.02				,       )

          torch.manual_seed(0				)
          __lowercase		:			Any	    =       IFWatermarker()

          return {
              "text_encoder": text_encoder,
              "tokenizer": tokenizer,
              "unet": unet,
              "scheduler": scheduler,
              "image_noising_scheduler": image_noising_scheduler,
              "watermarker": watermarker,
              "safety_checker": None,
              "feature_extractor": None,
          }




      def 		snake_case_			(      self		:       Any				):
          __lowercase		:			Tuple	    =       self.get_dummy_components()
          __lowercase		:			Any	    =       self.pipeline_class(**_snake_case				)
          pipe.to(_snake_case				)
          pipe.set_progress_bar_config(disable=_snake_case				)

          __lowercase		:			List[str]	    =       self.get_dummy_inputs(_snake_case				)

          __lowercase		:			List[Any]	    =       inputs['''prompt''']
          __lowercase		:			int	    =       inputs['''generator''']
          __lowercase		:			Dict	    =       inputs['''num_inference_steps''']
          __lowercase		:			Optional[int]	    =       inputs['''output_type''']

          if "image" in inputs:
              __lowercase		:			Tuple	    =       inputs['''image''']
          else:
              __lowercase		:			Dict	    =       None

          if "mask_image" in inputs:
              __lowercase		:			Tuple	    =       inputs['''mask_image''']
          else:
              __lowercase		:			List[Any]	    =       None

          if "original_image" in inputs:
              __lowercase		:			List[Any]	    =       inputs['''original_image''']
          else:
              __lowercase		:			Optional[int]	    =       None

          __lowercase    ,      __lowercase		:			Optional[int]	    =       pipe.encode_prompt(_snake_case				)

          # inputs with prompt converted to embeddings
          __lowercase		:			int	    =       {
              '''prompt_embeds''': prompt_embeds,
              '''negative_prompt_embeds''': negative_prompt_embeds,
              '''generator''': generator,
              '''num_inference_steps''': num_inference_steps,
              '''output_type''': output_type,
          }

          if image is not None:
              __lowercase		:			Tuple	    =       image

          if mask_image is not None:
              __lowercase		:			List[str]	    =       mask_image

          if original_image is not None:
              __lowercase		:			List[Any]	    =       original_image

          # set all optional components to None
          for optional_component in pipe._optional_components:
              setattr(_snake_case				,       _snake_case				,       _snake_case				)

          __lowercase		:			List[str]	    =       pipe(**_snake_case				)[0]

          with tempfile.TemporaryDirectory() as tmpdir:
              pipe.save_pretrained(_snake_case				)
              __lowercase		:			List[Any]	    =       self.pipeline_class.from_pretrained(_snake_case				)
              pipe_loaded.to(_snake_case				)
              pipe_loaded.set_progress_bar_config(disable=_snake_case				)

          pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()				)  # For reproducibility tests

          for optional_component in pipe._optional_components:
              self.assertTrue(
                  getattr(_snake_case				,       _snake_case				) is None				,       F'`{optional_component}` did not stay set to None after loading.'				,       )

          __lowercase		:			int	    =       self.get_dummy_inputs(_snake_case				)

          __lowercase		:			Union[str, Any]	    =       inputs['''generator''']
          __lowercase		:			Any	    =       inputs['''num_inference_steps''']
          __lowercase		:			Tuple	    =       inputs['''output_type''']

          # inputs with prompt converted to embeddings
          __lowercase		:			List[str]	    =       {
              '''prompt_embeds''': prompt_embeds,
              '''negative_prompt_embeds''': negative_prompt_embeds,
              '''generator''': generator,
              '''num_inference_steps''': num_inference_steps,
              '''output_type''': output_type,
          }

          if image is not None:
              __lowercase		:			Dict	    =       image

          if mask_image is not None:
              __lowercase		:			Tuple	    =       mask_image

          if original_image is not None:
              __lowercase		:			Any	    =       original_image

          __lowercase		:			List[str]	    =       pipe_loaded(**_snake_case				)[0]

          __lowercase		:			Optional[int]	    =       np.abs(to_np(_snake_case				) - to_np(_snake_case				)				).max()
          self.assertLess(_snake_case				,       1E-4				)




      def 		snake_case_			(      self		:       Optional[int]				):
          __lowercase		:			Union[str, Any]	    =       self.get_dummy_components()
          __lowercase		:			Dict	    =       self.pipeline_class(**_snake_case				)
          pipe.to(_snake_case				)
          pipe.set_progress_bar_config(disable=_snake_case				)

          __lowercase		:			Dict	    =       self.get_dummy_inputs(_snake_case				)
          __lowercase		:			Any	    =       pipe(**_snake_case				)[0]

          with tempfile.TemporaryDirectory() as tmpdir:
              pipe.save_pretrained(_snake_case				)
              __lowercase		:			Dict	    =       self.pipeline_class.from_pretrained(_snake_case				)
              pipe_loaded.to(_snake_case				)
              pipe_loaded.set_progress_bar_config(disable=_snake_case				)

          pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()				)  # For reproducibility tests

          __lowercase		:			Optional[Any]	    =       self.get_dummy_inputs(_snake_case				)
          __lowercase		:			str	    =       pipe_loaded(**_snake_case				)[0]

          __lowercase		:			int	    =       np.abs(to_np(_snake_case				) - to_np(_snake_case				)				).max()
          self.assertLess(_snake_case				,       1E-4				)


