





'''simple docstring'''


import argparse
from pathlib import Path

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

from diffusers import AutoencoderKL


lowerCAmelCase					:     int												=version.parse(version.parse(torch.__version__).base_version) < version.parse('''1.11''')
def 							UpperCAmelCase_  (						__lowerCamelCase					:					Optional[Any]	,__lowerCamelCase					:					tuple	,__lowerCamelCase					:					Path	,__lowerCamelCase					:					Dict	,__lowerCamelCase					:					Union[str, Any]	,__lowerCamelCase					:					int	,__lowerCamelCase					:					Optional[int]	,__lowerCamelCase					:					Optional[Any]=False	,):
   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					:					str	,__lowerCamelCase					:					str	,__lowerCamelCase					:					int	,__lowerCamelCase					:					bool = False ):
   lowercase_			:List[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_			:Optional[Any]    			= "cpu"
   lowercase_			:Dict    			= Path(__lowerCamelCase )

   # VAE DECODER
   lowercase_			:Dict    			= AutoencoderKL.from_pretrained(model_path + "/vae" )
   lowercase_			:str    			= vae_decoder.config.latent_channels
   # forward only through the decoder part
   lowercase_			:List[str]    			= vae_decoder.decode
   onnx_export(
       __lowerCamelCase	,model_args=(
           torch.randn(1	,__lowerCamelCase	,25	,25 ).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 vae_decoder


if __name__ == "__main__":
 lowerCAmelCase					:     Optional[Any]												=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					:     Any												=parser.parse_args()
 print(args.output_path)
 convert_models(args.model_path, args.output_path, args.opset, args.fpaa)
 print('''SD: Done: ONNX''')









'''simple docstring'''


import unittest

import numpy as np

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
 from transformers.models.bert.modeling_flax_bert import (
     FlaxBertForMaskedLM,
     FlaxBertForMultipleChoice,
     FlaxBertForNextSentencePrediction,
     FlaxBertForPreTraining,
     FlaxBertForQuestionAnswering,
     FlaxBertForSequenceClassification,
     FlaxBertForTokenClassification,
     FlaxBertModel,
 )
class       a_   (				unittest.TestCase							):





     def __init__( self		:					List[str]	,      lowercase		:					str	,      lowercase		:					Union[str, Any]=13	,      lowercase		:					int=7	,      lowercase		:					List[str]=True	,      lowercase		:					int=True	,      lowercase		:					str=True	,      lowercase		:					Any=True	,      lowercase		:					List[str]=99	,      lowercase		:					Union[str, Any]=32	,      lowercase		:					Optional[Any]=5	,      lowercase		:					Dict=4	,      lowercase		:					Dict=37	,      lowercase		:					Dict="gelu"	,      lowercase		:					Optional[int]=0.1	,      lowercase		:					str=0.1	,      lowercase		:					List[Any]=512	,      lowercase		:					str=16	,      lowercase		:					Dict=2	,      lowercase		:					Any=0.02	,      lowercase		:					Any=4	,      ):





        """simple docstring"""




        lowercase_			:List[str]    			= parent
        lowercase_			:Any    			= batch_size
        lowercase_			:Dict    			= seq_length
        lowercase_			:Union[str, Any]    			= is_training
        lowercase_			:Optional[int]    			= use_attention_mask
        lowercase_			:Any    			= use_token_type_ids
        lowercase_			:Union[str, Any]    			= use_labels
        lowercase_			:Dict    			= vocab_size
        lowercase_			:Tuple    			= hidden_size
        lowercase_			:Tuple    			= num_hidden_layers
        lowercase_			:Optional[int]    			= num_attention_heads
        lowercase_			:Optional[Any]    			= intermediate_size
        lowercase_			:str    			= hidden_act
        lowercase_			:Tuple    			= hidden_dropout_prob
        lowercase_			:Optional[Any]    			= attention_probs_dropout_prob
        lowercase_			:Tuple    			= max_position_embeddings
        lowercase_			:Any    			= type_vocab_size
        lowercase_			:int    			= type_sequence_label_size
        lowercase_			:Tuple    			= initializer_range
        lowercase_			:Optional[Any]    			= num_choices





     def 							lowercase__      ( self		:					List[Any]						):





        """simple docstring"""




        lowercase_			:int    			= ids_tensor([self.batch_size, self.seq_length]	,      self.vocab_size						)

        lowercase_			:Union[str, Any]    			= None
        if self.use_attention_mask:
           lowercase_			:Dict    			= random_attention_mask([self.batch_size, self.seq_length]						)

        lowercase_			:List[str]    			= None
        if self.use_token_type_ids:
           lowercase_			:str    			= ids_tensor([self.batch_size, self.seq_length]	,      self.type_vocab_size						)

        lowercase_			:Optional[Any]    			= BertConfig(
            vocab_size=self.vocab_size	,      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	,      max_position_embeddings=self.max_position_embeddings	,      type_vocab_size=self.type_vocab_size	,      is_decoder=lowercase	,      initializer_range=self.initializer_range	,      )

        return config, input_ids, token_type_ids, attention_mask





     def 							lowercase__      ( self		:					Union[str, Any]						):





        """simple docstring"""




        lowercase_			:int    			= self.prepare_config_and_inputs()
        lowercase_			,							lowercase_			,							lowercase_			,							lowercase_			:Tuple    			= config_and_inputs
        lowercase_			:Tuple    			= {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
        return config, inputs_dict







     def 							lowercase__      ( self		:					List[Any]						):





        """simple docstring"""




        lowercase_			:Any    			= self.prepare_config_and_inputs()
        lowercase_			,							lowercase_			,							lowercase_			,							lowercase_			:Union[str, Any]    			= config_and_inputs

        lowercase_			:Dict    			= True
        lowercase_			:Optional[int]    			= floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
        lowercase_			:str    			= ids_tensor([self.batch_size, self.seq_length]	,      vocab_size=2						)

        return (
            config,
            input_ids,
            attention_mask,
            encoder_hidden_states,
            encoder_attention_mask,
        )







@require_flax
class       a_   (				_lowerCAmelCase	,						unittest.TestCase							):
     __A      =			True

     __A      =			(
         (
             FlaxBertModel,
             FlaxBertForPreTraining,
             FlaxBertForMaskedLM,
             FlaxBertForMultipleChoice,
             FlaxBertForQuestionAnswering,
             FlaxBertForNextSentencePrediction,
             FlaxBertForSequenceClassification,
             FlaxBertForTokenClassification,
             FlaxBertForQuestionAnswering,
         )
         if is_flax_available()
         else ()
     )





     def 							lowercase__      ( self		:					Any						):





        """simple docstring"""




        lowercase_			:Optional[Any]    			= FlaxBertModelTester(self						)







     @slow
     def 							lowercase__      ( self		:					List[str]						):





        """simple docstring"""




        lowercase_			:List[str]    			= FlaxBertModel.from_pretrained("bert-base-cased"						)
        lowercase_			:str    			= model(np.ones((1, 1)						)						)
        self.assertIsNotNone(lowercase						)


