




from ..utils import is_flax_available, is_torch_available


if is_torch_available():
 from .autoencoder_kl import AutoencoderKL
 from .controlnet import ControlNetModel
 from .dual_transformer_ad import DualTransformeraDModel
 from .modeling_utils import ModelMixin
 from .prior_transformer import PriorTransformer
 from .ta_film_transformer import TaFilmDecoder
 from .transformer_ad import TransformeraDModel
 from .unet_ad import UNetaDModel
 from .unet_ad import UNetaDModel
 from .unet_ad_condition import UNetaDConditionModel
 from .unet_ad_condition import UNetaDConditionModel
 from .vq_model import VQModel

if is_flax_available():
 from .controlnet_flax import FlaxControlNetModel
 from .unet_ad_condition_flax import FlaxUNetaDConditionModel
 from .vae_flax import FlaxAutoencoderKL









import os
import time

import numpy as np
import onnxruntime as ort


lowerCamelCase    :		int			   =       "1"
lowerCamelCase    :		int			   =       "0"
lowerCamelCase    :		Union[str, Any]			   =       "1"

lowerCamelCase    :		List[Any]			   =       ort.SessionOptions()
lowerCamelCase    :		Optional[Any]			   =       ort.GraphOptimizationLevel.ORT_DISABLE_ALL
print("Create inference session...")
lowerCamelCase    :		Union[str, Any]			   =       ["TensorrtExecutionProvider", "CUDAExecutionProvider"]
lowerCamelCase    :		Tuple			   =       ort.InferenceSession("model.onnx", sess_options=sess_opt, providers=execution_provider)
lowerCamelCase    :		List[Any]			   =       ort.RunOptions()

lowerCamelCase    :		List[str]			   =       128
lowerCamelCase    :		List[Any]			   =       1
lowerCamelCase    :		Union[str, Any]			   =       np.ones((batch, sequence), dtype=np.intaa)
lowerCamelCase    :		Dict			   =       np.ones((batch, sequence), dtype=np.intaa)
lowerCamelCase    :		Optional[Any]			   =       np.ones((batch, sequence), dtype=np.intaa)

print("Warm up phase...")
sess.run(
    None,
    {
        sess.get_inputs()[0].name: input_ids,
        sess.get_inputs()[1].name: attention_mask,
        sess.get_inputs()[2].name: token_type_ids,
    },
    run_options=run_opt,
)

print("Start inference...")
lowerCamelCase    :		int			   =       time.time()
lowerCamelCase    :		Dict			   =       2_000
lowerCamelCase    :		Any			   =       {}
for iter in range(max_iters):
 lowerCamelCase    :		Union[str, Any]			   =       sess.run(
     None,
     {
         sess.get_inputs()[0].name: input_ids,
         sess.get_inputs()[1].name: attention_mask,
         sess.get_inputs()[2].name: token_type_ids,
     },
     run_options=run_opt,
 )
print("Average Inference Time = {:.3f} ms".format((time.time() - start_time) * 1_000 / max_iters))



