onnx

export.py

@@ -5,6 +5,8 @@ from transformers.models.clip.modeling_clip import _get_vector_norm
 import coremltools as ct
 import torch
 import numpy as np
+import platform
+import sys
 
 processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
@@ -94,6 +96,42 @@ def convert_onnx():
 # convert_onnx()
 
 
+def infer_onnx():
+    import onnxruntime as ort
+
+    providers: list[str] = []
+    if sys.platform == "darwin":
+        providers.append("CoreMLExecutionProvider")
+
+    if ("linux" in sys.platform or "win" in sys.platform) and (
+        platform.machine() == "x86_64" or platform.machine() == "AMD64"
+    ):
+        providers.append(("CUDAExecutionProvider", {"device_id": 0}))
+
+    providers.append("CPUExecutionProvider")
+
+    vision_session = ort.InferenceSession("vision.onnx", providers=providers)
+    text_session = ort.InferenceSession("text.onnx", providers=providers)
+
+    vision_inputs = {vision_session.get_inputs()[0].name: np_inputs['pixel_values']}
+    text_inputs = {
+        text_session.get_inputs()[0].name: np_inputs['input_ids'],
+        text_session.get_inputs()[1].name: np_inputs['attention_mask']
+    }
+
+    vision_predictions = vision_session.run(None, vision_inputs)
+    text_predictions = text_session.run(None, text_inputs)
+
+    image_embeds = vision_predictions[0]
+    text_embeds = text_predictions[0]
+
+    logits_per_text = text_embeds @ image_embeds.T
+
+    print("similarity:", logits_per_text.item())
+
+infer_onnx()
+
+
 def convert_openvino():
     import openvino as ov
     ov_vision_model = ov.convert_model(traced_vision_model, example_input=inputs.data['pixel_values'])
@@ -122,7 +160,7 @@ def infer_openvino():
 
     print("similarity:", logits_per_text.item())
 
-infer_openvino()
+# infer_openvino()
 
 # convert_coreml()