feat: Implement unified memory management for ONNX model inference

- Introduced a new memory management policy for ONNX Runtime to release temporary GPU memory after inference completion.
- Established a shared inference entry point for all ONNX models, differentiating between long-lived sessions and temporary memory arenas.
- Updated CUDA sessions to utilize a shrinkable memory arena, reclaiming GPU memory post-inference.
- Ensured consistent memory release semantics across BEN2, DA3, MoGe-2, and Upscale models while reusing existing sessions.
- Added lifecycle tests to validate temporary memory release and session reuse.
- Modified model adapters to integrate the new memory management features, ensuring backward compatibility with existing inference behavior.
- Enhanced unit tests to cover new memory release scenarios and verify correct behavior across different execution providers.

tests/test_model_onnx_runtime.py

@@ -0,0 +1,135 @@
+import sys
+import unittest
+from pathlib import Path
+from types import ModuleType
+from unittest.mock import patch
+
+
+SOURCE_DIR = (
+    Path(__file__).resolve().parents[1]
+    / "src"
+    / "anyimage"
+    / "server"
+    / "models"
+)
+sys.path.insert(0, str(SOURCE_DIR))
+
+import onnx_runtime
+
+
+class FakeSession:
+    def __init__(self, providers):
+        self.providers = providers
+        self.calls = []
+
+    def get_providers(self):
+        return self.providers
+
+    def run(self, output_names, input_feed, run_options=None):
+        self.calls.append((output_names, input_feed, run_options))
+        return ["prediction"]
+
+
+class OnnxRuntimeTest(unittest.TestCase):
+    def create_runtime(self, available):
+        calls = []
+        runtime = ModuleType("onnxruntime")
+        runtime.get_available_providers = lambda: available
+        runtime.InferenceSession = lambda path, providers: calls.append(
+            (path, providers)
+        )
+        return runtime, calls
+
+    def test_cuda_session_uses_shrinkable_arena(self):
+        runtime, calls = self.create_runtime(
+            ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        )
+
+        with patch.dict(sys.modules, {"onnxruntime": runtime}):
+            onnx_runtime.create_session("model.onnx", "cuda")
+
+        self.assertEqual(
+            calls[0],
+            (
+                "model.onnx",
+                [
+                    (
+                        "CUDAExecutionProvider",
+                        {"arena_extend_strategy": "kSameAsRequested"},
+                    )
+                ],
+            ),
+        )
+
+    def test_coreml_and_cpu_session_options_are_preserved(self):
+        runtime, calls = self.create_runtime(
+            ["CoreMLExecutionProvider", "CPUExecutionProvider"]
+        )
+
+        with patch.dict(sys.modules, {"onnxruntime": runtime}):
+            onnx_runtime.create_session("model.onnx", "auto")
+
+        self.assertEqual(
+            calls[0][1],
+            [
+                (
+                    "CoreMLExecutionProvider",
+                    dict(onnx_runtime.COREML_PROVIDER_OPTIONS),
+                ),
+                "CPUExecutionProvider",
+            ],
+        )
+
+    def test_completing_cuda_run_shrinks_memory_arena(self):
+        configs = []
+
+        class FakeRunOptions:
+            def add_run_config_entry(self, key, value):
+                configs.append((key, value))
+
+        runtime = ModuleType("onnxruntime")
+        runtime.RunOptions = FakeRunOptions
+        session = FakeSession(["CUDAExecutionProvider", "CPUExecutionProvider"])
+
+        with patch.dict(sys.modules, {"onnxruntime": runtime}):
+            result = onnx_runtime.run_session(session, ["output"], {"input": 1})
+
+        self.assertEqual(result, ["prediction"])
+        self.assertIsInstance(session.calls[0][2], FakeRunOptions)
+        self.assertEqual(
+            configs,
+            [("memory.enable_memory_arena_shrinkage", "gpu:0")],
+        )
+
+    def test_intermediate_and_non_cuda_runs_use_plain_session_run(self):
+        cuda = FakeSession(["CUDAExecutionProvider"])
+        directml = FakeSession(["DmlExecutionProvider"])
+
+        onnx_runtime.run_session(
+            cuda,
+            ["output"],
+            {"input": 1},
+            release_memory=False,
+        )
+        result = onnx_runtime.run_session(
+            directml,
+            ["output"],
+            {"input": 2},
+        )
+
+        self.assertIsNone(cuda.calls[0][2])
+        self.assertIsNone(directml.calls[0][2])
+        self.assertEqual(result, ["prediction"])
+
+    def test_installed_runtime_accepts_memory_arena_shrinkage_option(self):
+        import onnxruntime
+
+        run_options = onnxruntime.RunOptions()
+        run_options.add_run_config_entry(
+            "memory.enable_memory_arena_shrinkage",
+            "gpu:0",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()