





"""simple docstring"""





import inspect
import unittest
from typing import List

import numpy as np

from transformers import EfficientFormerConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
  import tensorflow as tf

  from transformers import (
      TFEfficientFormerForImageClassification,
      TFEfficientFormerForImageClassificationWithTeacher,
      TFEfficientFormerModel,
  )
  from transformers.models.efficientformer.modeling_tf_efficientformer import (
      TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
  )


if is_vision_available():
  from PIL import Image

  from transformers import EfficientFormerImageProcessor






class    _lowerCAmelCase  :
       """simple docstring"""




       def __init__(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase = 1_3       ,					__UpperCAmelCase = 6_4       ,					__UpperCAmelCase = 2       ,					__UpperCAmelCase = 3       ,					__UpperCAmelCase = 3       ,					__UpperCAmelCase = True       ,					__UpperCAmelCase = True       ,					__UpperCAmelCase = 1_2_8       ,					__UpperCAmelCase=[1_6, 3_2, 6_4, 1_2_8]       ,					__UpperCAmelCase = 7       ,					__UpperCAmelCase = 4       ,					__UpperCAmelCase = 3_7       ,					__UpperCAmelCase = "gelu"       ,					__UpperCAmelCase = 0.1       ,					__UpperCAmelCase = 0.1       ,					__UpperCAmelCase = 1_0       ,					__UpperCAmelCase = 0.02       ,					__UpperCAmelCase = 2       ,					__UpperCAmelCase = 1       ,					__UpperCAmelCase = 1_2_8       ,					__UpperCAmelCase = [2, 2, 2, 2]       ,					__UpperCAmelCase = 2       ,					__UpperCAmelCase = 2       ,					):



         '''simple docstring'''

         lowerCAmelCase__     :List[str]											=				parent
         lowerCAmelCase__     :str											=				batch_size
         lowerCAmelCase__     :int											=				image_size
         lowerCAmelCase__     :Optional[int]											=				patch_size
         lowerCAmelCase__     :Any											=				num_channels
         lowerCAmelCase__     :Tuple											=				is_training
         lowerCAmelCase__     :Tuple											=				use_labels
         lowerCAmelCase__     :Tuple											=				hidden_size
         lowerCAmelCase__     :Optional[Any]											=				num_hidden_layers
         lowerCAmelCase__     :Optional[Any]											=				num_attention_heads
         lowerCAmelCase__     :Optional[int]											=				intermediate_size
         lowerCAmelCase__     :str											=				hidden_act
         lowerCAmelCase__     :Tuple											=				hidden_dropout_prob
         lowerCAmelCase__     :str											=				attention_probs_dropout_prob
         lowerCAmelCase__     :Any											=				type_sequence_label_size
         lowerCAmelCase__     :Dict											=				initializer_range
         lowerCAmelCase__     :Optional[int]											=				encoder_stride
         lowerCAmelCase__     :Dict											=				num_attention_outputs
         lowerCAmelCase__     :List[Any]											=				embed_dim
         lowerCAmelCase__     :Optional[int]											=				embed_dim + 1
         lowerCAmelCase__     :Tuple											=				resolution
         lowerCAmelCase__     :List[Any]											=				depths
         lowerCAmelCase__     :List[Any]											=				hidden_sizes
         lowerCAmelCase__     :Optional[int]											=				dim
         lowerCAmelCase__     :List[str]											=				mlp_expansion_ratio






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Optional[Any]											=				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

         lowerCAmelCase__     :Any											=				None
         if self.use_labels:
           lowerCAmelCase__     :Optional[int]											=				ids_tensor([self.batch_size]       ,					self.type_sequence_label_size )

         lowerCAmelCase__     :List[Any]											=				self.get_config()

         return config, pixel_values, labels






       def 	snake_case		(       self ):



         '''simple docstring'''

         return EfficientFormerConfig(
             image_size=self.image_size       ,					patch_size=self.patch_size       ,					num_channels=self.num_channels       ,					hidden_size=self.hidden_size       ,					num_hidden_layers=self.num_hidden_layers       ,					num_attention_heads=self.num_attention_heads       ,					intermediate_size=self.intermediate_size       ,					hidden_act=self.hidden_act       ,					hidden_dropout_prob=self.hidden_dropout_prob       ,					attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,					is_decoder=__UpperCAmelCase       ,					initializer_range=self.initializer_range       ,					encoder_stride=self.encoder_stride       ,					resolution=self.resolution       ,					depths=self.depths       ,					hidden_sizes=self.hidden_sizes       ,					dim=self.dim       ,					mlp_expansion_ratio=self.mlp_expansion_ratio       ,					)






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase ):



         '''simple docstring'''

         lowerCAmelCase__     :Optional[Any]											=				TFEfficientFormerModel(config=__UpperCAmelCase )
         lowerCAmelCase__     :str											=				model(__UpperCAmelCase       ,					training=__UpperCAmelCase )
         self.parent.assertEqual(result.last_hidden_state.shape       ,					(self.batch_size, self.seq_length, self.hidden_size) )






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase ):



         '''simple docstring'''

         lowerCAmelCase__     :Optional[Any]											=				self.type_sequence_label_size
         lowerCAmelCase__     :List[str]											=				TFEfficientFormerForImageClassification(__UpperCAmelCase )
         lowerCAmelCase__     :Union[str, Any]											=				model(__UpperCAmelCase       ,					labels=__UpperCAmelCase       ,					training=__UpperCAmelCase )
         self.parent.assertEqual(result.logits.shape       ,					(self.batch_size, self.type_sequence_label_size) )

         # test greyscale images
         lowerCAmelCase__     :Tuple											=				1
         lowerCAmelCase__     :str											=				TFEfficientFormerForImageClassification(__UpperCAmelCase )

         lowerCAmelCase__     :Optional[Any]											=				floats_tensor([self.batch_size, 1, self.image_size, self.image_size] )
         lowerCAmelCase__     :Dict											=				model(__UpperCAmelCase       ,					labels=__UpperCAmelCase )
         self.parent.assertEqual(result.logits.shape       ,					(self.batch_size, self.type_sequence_label_size) )






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :str											=				self.prepare_config_and_inputs()
         lowerCAmelCase__				,							lowerCAmelCase__				,							lowerCAmelCase__     :str											=				config_and_inputs
         lowerCAmelCase__     :Any											=				{'pixel_values': pixel_values}
         return config, inputs_dict






@require_tf
class    _lowerCAmelCase  (	a    ,  a    ,  unittest.TestCase			):
       """simple docstring"""




       __magic_name__	:Dict						    =   (
           (
               TFEfficientFormerModel,
               TFEfficientFormerForImageClassificationWithTeacher,
               TFEfficientFormerForImageClassification,
           )
           if is_tf_available()
           else ()
       )
       __magic_name__	:Optional[int]						    =   (
           {
               """feature-extraction""": TFEfficientFormerModel,
               """image-classification""": (
                   TFEfficientFormerForImageClassification,
                   TFEfficientFormerForImageClassificationWithTeacher,
               ),
           }
           if is_tf_available()
           else {}
       )

       __magic_name__	:Union[str, Any]						    =   False

       __magic_name__	:List[str]						    =   False
       __magic_name__	:Union[str, Any]						    =   False
       __magic_name__	:Optional[int]						    =   False
       __magic_name__	:Any						    =   False






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Union[str, Any]											=				TFEfficientFormerModelTester(self )
         lowerCAmelCase__     :Any											=				ConfigTester(
             self       ,					config_class=__UpperCAmelCase       ,					has_text_modality=__UpperCAmelCase       ,					hidden_size=3_7 )






       def 	snake_case		(       self ):



         '''simple docstring'''

         self.config_tester.run_common_tests()






       @unittest.skip(reason='EfficientFormer does not use inputs_embeds' )
       def 	snake_case		(       self ):



         '''simple docstring'''

         pass






       @unittest.skip(reason='EfficientFormer does not support input and output embeddings' )
       def 	snake_case		(       self ):



         '''simple docstring'''

         pass






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__				,							lowerCAmelCase__     :str											=				self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
           lowerCAmelCase__     :Dict											=				model_class(__UpperCAmelCase )
           lowerCAmelCase__     :str											=				inspect.signature(model.call )
           # signature.parameters is an OrderedDict => so arg_names order is deterministic
           lowerCAmelCase__     :List[Any]											=				[*signature.parameters.keys()]

           lowerCAmelCase__     :Tuple											=				['pixel_values']
           self.assertListEqual(arg_names[:1]       ,					__UpperCAmelCase )






       def 	snake_case		(       self ):



         '''simple docstring'''

         def check_hidden_states_output(__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase ):
           lowerCAmelCase__     :Union[str, Any]											=				model_class(__UpperCAmelCase )

           lowerCAmelCase__     :Union[str, Any]											=				model(**self._prepare_for_class(__UpperCAmelCase       ,					__UpperCAmelCase )       ,					training=__UpperCAmelCase )
           lowerCAmelCase__     :List[Any]											=				outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

           lowerCAmelCase__     :Any											=				getattr(
               self.model_tester       ,					'expected_num_hidden_layers'       ,					self.model_tester.num_hidden_layers + 1 )
           self.assertEqual(len(__UpperCAmelCase )       ,					__UpperCAmelCase )

           if hasattr(self.model_tester       ,					'encoder_seq_length' ):
             lowerCAmelCase__     :Optional[Any]											=				self.model_tester.encoder_seq_length
             if hasattr(self.model_tester       ,					'chunk_length' ) and self.model_tester.chunk_length > 1:
               lowerCAmelCase__     :Optional[int]											=				seq_length * self.model_tester.chunk_length
           else:
             lowerCAmelCase__     :Any											=				self.model_tester.seq_length

           self.assertListEqual(
               list(hidden_states[-1].shape[-2:] )       ,					[seq_length, self.model_tester.hidden_size]       ,					)

           if config.is_encoder_decoder:
             lowerCAmelCase__     :List[Any]											=				outputs.decoder_hidden_states

             self.asseretIsInstance(__UpperCAmelCase       ,					(list, tuple) )
             self.assertEqual(len(__UpperCAmelCase )       ,					__UpperCAmelCase )
             lowerCAmelCase__     :List[str]											=				getattr(self.model_tester       ,					'seq_length'       ,					__UpperCAmelCase )
             lowerCAmelCase__     :int											=				getattr(self.model_tester       ,					'decoder_seq_length'       ,					__UpperCAmelCase )

             self.assertListEqual(
                 list(hidden_states[-1].shape[-2:] )       ,					[decoder_seq_length, self.model_tester.hidden_size]       ,					)

         lowerCAmelCase__				,							lowerCAmelCase__     :Tuple											=				self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
           lowerCAmelCase__     :Dict											=				True
           check_hidden_states_output(__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase )

           # check that output_hidden_states also work using config
           del inputs_dict["output_hidden_states"]
           lowerCAmelCase__     :int											=				True

           check_hidden_states_output(__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase )






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase=False ):



         '''simple docstring'''

         lowerCAmelCase__     :Any											=				super()._prepare_for_class(__UpperCAmelCase       ,					__UpperCAmelCase       ,					return_labels=__UpperCAmelCase )

         if return_labels:
           if model_class.__name__ == "TFEfficientFormerForImageClassificationWithTeacher":
             del inputs_dict["labels"]

         return inputs_dict






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Dict											=				self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__UpperCAmelCase )






       @unittest.skip(reason='EfficientFormer does not implement masked image modeling yet' )
       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Union[str, Any]											=				self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_image_modeling(*__UpperCAmelCase )






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Union[str, Any]											=				self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*__UpperCAmelCase )






       @slow
       def 	snake_case		(       self ):



         '''simple docstring'''

         for model_name in TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           lowerCAmelCase__     :List[str]											=				TFEfficientFormerModel.from_pretrained(__UpperCAmelCase )
           self.assertIsNotNone(__UpperCAmelCase )






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__				,							lowerCAmelCase__     :List[str]											=				self.model_tester.prepare_config_and_inputs_for_common()
         lowerCAmelCase__     :Any											=				True

         lowerCAmelCase__     :Tuple											=				getattr(self.model_tester       ,					'seq_length'       ,					__UpperCAmelCase )
         lowerCAmelCase__     :Tuple											=				getattr(self.model_tester       ,					'encoder_seq_length'       ,					__UpperCAmelCase )
         lowerCAmelCase__     :List[str]											=				getattr(self.model_tester       ,					'key_length'       ,					__UpperCAmelCase )
         lowerCAmelCase__     :List[str]											=				getattr(self.model_tester       ,					'chunk_length'       ,					__UpperCAmelCase )

         if chunk_length is not None and hasattr(self.model_tester       ,					'num_hashes' ):
           lowerCAmelCase__     :Tuple											=				encoder_seq_length * self.model_tester.num_hashes

         for model_class in self.all_model_classes:
           lowerCAmelCase__     :Tuple											=				True
           lowerCAmelCase__     :Optional[int]											=				False
           lowerCAmelCase__     :Optional[Any]											=				True
           lowerCAmelCase__     :str											=				model_class(__UpperCAmelCase )

           lowerCAmelCase__     :Optional[Any]											=				model(**self._prepare_for_class(__UpperCAmelCase       ,					__UpperCAmelCase )       ,					training=__UpperCAmelCase )
           lowerCAmelCase__     :Optional[Any]											=				outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
           self.assertEqual(len(__UpperCAmelCase )       ,					self.model_tester.num_attention_outputs )

           # check that output_attentions also work using config
           del inputs_dict["output_attentions"]
           lowerCAmelCase__     :Optional[Any]											=				True
           lowerCAmelCase__     :Union[str, Any]											=				model_class(__UpperCAmelCase )
           lowerCAmelCase__     :Any											=				model(**self._prepare_for_class(__UpperCAmelCase       ,					__UpperCAmelCase )       ,					training=__UpperCAmelCase )

           lowerCAmelCase__     :Dict											=				outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
           self.assertEqual(len(__UpperCAmelCase )       ,					self.model_tester.num_attention_outputs )

           if chunk_length is not None:
             self.assertListEqual(
                 list(attentions[0].shape[-4:] )       ,					[self.model_tester.num_attention_heads, encoder_seq_length, chunk_length, encoder_key_length]       ,					)
           else:
             self.assertListEqual(
                 list(attentions[0].shape[-3:] )       ,					[self.model_tester.num_attention_heads, encoder_seq_length, encoder_key_length]       ,					)






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__				,							lowerCAmelCase__     :Tuple											=				self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
           # Prepare our model
           lowerCAmelCase__     :int											=				model_class(__UpperCAmelCase )
           # These are maximally general inputs for the model, with multiple None dimensions
           # Hopefully this will catch any conditionals that fail for flexible shapes
           lowerCAmelCase__     :Optional[int]											=				{
               key: tf.keras.Input(shape=val.shape[1:]       ,					dtype=val.dtype       ,					name=__UpperCAmelCase )
               for key, val in model.input_signature.items()
               if key in model.dummy_inputs
           }
           lowerCAmelCase__     :str											=				model(__UpperCAmelCase )
           self.assertTrue(outputs_dict is not None )






def 					__A  ()   ->Optional[Any]:





  """simple docstring"""


  lowerCAmelCase__     :Union[str, Any]											=				Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'     )
  return image






@require_tf
@require_vision
class    _lowerCAmelCase  (	unittest.TestCase			):
       """simple docstring"""




       @cached_property
       def 	snake_case		(       self ):



         '''simple docstring'''

         return (
             EfficientFormerImageProcessor.from_pretrained('snap-research/efficientformer-l1-300' )
             if is_vision_available()
             else None
         )






       @slow
       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Optional[int]											=				TFEfficientFormerForImageClassification.from_pretrained('snap-research/efficientformer-l1-300' )
         lowerCAmelCase__     :List[Any]											=				self.default_image_processor
         lowerCAmelCase__     :List[str]											=				prepare_img()
         lowerCAmelCase__     :Optional[Any]											=				image_processor(images=__UpperCAmelCase       ,					return_tensors='tf' )
         # forward pass
         lowerCAmelCase__     :Union[str, Any]											=				model(**__UpperCAmelCase       ,					training=__UpperCAmelCase )
         # verify the logits
         lowerCAmelCase__     :List[Any]											=				tf.TensorShape((1, 1_0_0_0) )
         self.assertEqual(outputs.logits.shape       ,					__UpperCAmelCase )
         lowerCAmelCase__     :Optional[Any]											=				tf.constant([-0.05_55, 0.48_25, -0.08_52] )
         self.assertTrue(np.allclose(outputs.logits[0, :3]       ,					__UpperCAmelCase       ,					atol=1E-4 ) )






       @slow
       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Tuple											=				TFEfficientFormerForImageClassificationWithTeacher.from_pretrained(
             'snap-research/efficientformer-l1-300' )
         lowerCAmelCase__     :Optional[Any]											=				self.default_image_processor
         lowerCAmelCase__     :Optional[int]											=				prepare_img()
         lowerCAmelCase__     :str											=				image_processor(images=__UpperCAmelCase       ,					return_tensors='tf' )
         # forward pass
         lowerCAmelCase__     :int											=				model(**__UpperCAmelCase       ,					training=__UpperCAmelCase )
         # verify the logits
         lowerCAmelCase__     :Any											=				tf.TensorShape((1, 1_0_0_0) )
         self.assertEqual(outputs.logits.shape       ,					__UpperCAmelCase )
         lowerCAmelCase__     :Any											=				tf.constant([-0.13_12, 0.43_53, -1.04_99] )
         self.assertTrue(np.allclose(outputs.logits[0, :3]       ,					__UpperCAmelCase       ,					atol=1E-4 ) )







"""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


__A											=						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	,  )   ->Tuple:





  """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     )   ->List[Any]:





  """simple docstring"""


  lowerCAmelCase__     :List[Any]											=				torch.floataa if fpaa else torch.floataa
  if fpaa and torch.cuda.is_available():
    lowerCAmelCase__     :Tuple											=				'cuda'
  elif fpaa and not torch.cuda.is_available():
    raise ValueError('`float16` model export is only supported on GPUs with CUDA'     )
  else:
    lowerCAmelCase__     :List[Any]											=				'cpu'
  lowerCAmelCase__     :List[str]											=				StableDiffusionPipeline.from_pretrained(_SCREAMING_SNAKE_CASE	,  torch_dtype=_SCREAMING_SNAKE_CASE     ).to(_SCREAMING_SNAKE_CASE     )
  lowerCAmelCase__     :Dict											=				Path(_SCREAMING_SNAKE_CASE     )

  # TEXT ENCODER
  lowerCAmelCase__     :str											=				pipeline.text_encoder.config.max_position_embeddings
  lowerCAmelCase__     :Dict											=				pipeline.text_encoder.config.hidden_size
  lowerCAmelCase__     :List[Any]											=				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
  lowerCAmelCase__     :int											=				pipeline.unet.config.in_channels
  lowerCAmelCase__     :Optional[Any]											=				pipeline.unet.config.sample_size
  lowerCAmelCase__     :Dict											=				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	,  )
  lowerCAmelCase__     :List[Any]											=				str(unet_path.absolute().as_posix()     )
  lowerCAmelCase__     :int											=				os.path.dirname(_SCREAMING_SNAKE_CASE     )
  lowerCAmelCase__     :str											=				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
  lowerCAmelCase__     :int											=				pipeline.vae
  lowerCAmelCase__     :Optional[Any]											=				vae_encoder.config.in_channels
  lowerCAmelCase__     :int											=				vae_encoder.config.sample_size
  # need to get the raw tensor output (sample) from the encoder
  lowerCAmelCase__     :str											=				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
  lowerCAmelCase__     :int											=				pipeline.vae
  lowerCAmelCase__     :List[Any]											=				vae_decoder.config.latent_channels
  lowerCAmelCase__     :Optional[int]											=				vae_decoder.config.out_channels
  # forward only through the decoder part
  lowerCAmelCase__     :Any											=				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:
    lowerCAmelCase__     :Optional[int]											=				pipeline.safety_checker
    lowerCAmelCase__     :Optional[int]											=				safety_checker.config.vision_config.num_channels
    lowerCAmelCase__     :Any											=				safety_checker.config.vision_config.image_size
    lowerCAmelCase__     :List[str]											=				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
    lowerCAmelCase__     :Union[str, Any]											=				OnnxRuntimeModel.from_pretrained(output_path / 'safety_checker'     )
    lowerCAmelCase__     :Dict											=				pipeline.feature_extractor
  else:
    lowerCAmelCase__     :Tuple											=				None
    lowerCAmelCase__     :Optional[int]											=				None

  lowerCAmelCase__     :List[str]											=				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
  lowerCAmelCase__     :Dict											=				OnnxStableDiffusionPipeline.from_pretrained(_SCREAMING_SNAKE_CASE	,  provider='CPUExecutionProvider'     )
  print('ONNX pipeline is loadable'     )


if __name__ == "__main__":
  __A											=						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""")

  __A											=						parser.parse_args()

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