From 79d3342c3dad33dd40a35c71c09d8b65c9f61ff5 Mon Sep 17 00:00:00 2001
From: Mert <101130780+mertalev@users.noreply.github.com>
Date: Sat, 3 Feb 2024 00:35:26 -0500
Subject: [PATCH] fix(ml): openvino not working with dynamic axes (#6871)

* convert to static

* add comment about gross code

* formatting

* fixed test

* fix typing

* cleanup

* formatting

* Revert "formatting"

This reverts commit 073965c47e7e31f1255fd461cd34ee19917d78bb.

* Revert "cleanup"

This reverts commit bb56bd3297303332b25bdc3b112cfc278ee593ba.

* formatting
---
 machine-learning/app/models/base.py      | 48 +++++++++++++++++++++++-
 machine-learning/app/models/constants.py |  9 ++---
 2 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/machine-learning/app/models/base.py b/machine-learning/app/models/base.py
index e5ebb828f3..1c15a9d41a 100644
--- a/machine-learning/app/models/base.py
+++ b/machine-learning/app/models/base.py
@@ -6,12 +6,15 @@ from pathlib import Path
 from shutil import rmtree
 from typing import Any
 
+import onnx
 import onnxruntime as ort
 from huggingface_hub import snapshot_download
+from onnx.shape_inference import infer_shapes
+from onnx.tools.update_model_dims import update_inputs_outputs_dims
 from typing_extensions import Buffer
 
 import ann.ann
-from app.models.constants import SUPPORTED_PROVIDERS
+from app.models.constants import STATIC_INPUT_PROVIDERS, SUPPORTED_PROVIDERS
 
 from ..config import get_cache_dir, get_hf_model_name, log, settings
 from ..schemas import ModelRuntime, ModelType
@@ -114,6 +117,13 @@ class InferenceModel(ABC):
             )
             model_path = onnx_path
 
+        if any(provider in STATIC_INPUT_PROVIDERS for provider in self.providers):
+            static_path = model_path.parent / "static_1" / "model.onnx"
+            static_path.parent.mkdir(parents=True, exist_ok=True)
+            if not static_path.is_file():
+                self._convert_to_static(model_path, static_path)
+            model_path = static_path
+
         match model_path.suffix:
             case ".armnn":
                 session = AnnSession(model_path)
@@ -128,6 +138,42 @@ class InferenceModel(ABC):
                 raise ValueError(f"Unsupported model file type: {model_path.suffix}")
         return session
 
+    def _convert_to_static(self, source_path: Path, target_path: Path) -> None:
+        inferred = infer_shapes(onnx.load(source_path))
+        inputs = self._get_static_dims(inferred.graph.input)
+        outputs = self._get_static_dims(inferred.graph.output)
+
+        # check_model gets called in update_inputs_outputs_dims and doesn't work for large models
+        check_model = onnx.checker.check_model
+        try:
+
+            def check_model_stub(*args: Any, **kwargs: Any) -> None:
+                pass
+
+            onnx.checker.check_model = check_model_stub
+            updated_model = update_inputs_outputs_dims(inferred, inputs, outputs)
+        finally:
+            onnx.checker.check_model = check_model
+
+        onnx.save(
+            updated_model,
+            target_path,
+            save_as_external_data=True,
+            all_tensors_to_one_file=False,
+            size_threshold=1048576,
+        )
+
+    def _get_static_dims(self, graph_io: Any, dim_size: int = 1) -> dict[str, list[int]]:
+        return {
+            field.name: [
+                d.dim_value if d.HasField("dim_value") else dim_size
+                for shape in field.type.ListFields()
+                if (dim := shape[1].shape.dim)
+                for d in dim
+            ]
+            for field in graph_io
+        }
+
     @property
     def model_type(self) -> ModelType:
         return self._model_type
diff --git a/machine-learning/app/models/constants.py b/machine-learning/app/models/constants.py
index edfe68aa70..18965d2b1d 100644
--- a/machine-learning/app/models/constants.py
+++ b/machine-learning/app/models/constants.py
@@ -51,11 +51,10 @@ _INSIGHTFACE_MODELS = {
 }
 
 
-SUPPORTED_PROVIDERS = [
-    "CUDAExecutionProvider",
-    "OpenVINOExecutionProvider",
-    "CPUExecutionProvider",
-]
+SUPPORTED_PROVIDERS = ["CUDAExecutionProvider", "OpenVINOExecutionProvider", "CPUExecutionProvider"]
+
+
+STATIC_INPUT_PROVIDERS = ["OpenVINOExecutionProvider"]
 
 
 def is_openclip(model_name: str) -> bool: