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_server_runtime.py

@@ -46,10 +46,12 @@ class ServerRuntimeTest(unittest.TestCase):
         importlib.import_module("server.jobs.cutout")
         from blendjob import JobServer
         from server.model_manager import ModelManager, ModelSessionCache
+        from server.models import onnx_runtime
 
         cls.JobServer = JobServer
 
         cls.ModelManager, cls.ModelSessionCache = ModelManager, ModelSessionCache
+        cls.onnx_runtime = onnx_runtime
         cls.server = cls.server_module.server
         cls.app_module = sys.modules["server.app"]
 
@@ -224,6 +226,88 @@ class ServerRuntimeTest(unittest.TestCase):
         self.assertEqual(loads, [("model", "cpu")])
         self.assertEqual(cache.snapshot()["upscale"], "model")
 
+    def test_inference_cleanup_keeps_every_cached_session(self):
+        class Session:
+            def get_providers(self):
+                return ["CUDAExecutionProvider"]
+
+            def run(self, _names, _feeds, run_options=None):
+                self.run_options = run_options
+                return [object()]
+
+        class RunOptions:
+            def add_run_config_entry(self, _key, _value):
+                return None
+
+        runtime = ModuleType("onnxruntime")
+        runtime.RunOptions = RunOptions
+        sessions = {
+            "ben2": Session(),
+            "da3": Session(),
+            "moge2": Session(),
+            "upscale": Session(),
+        }
+        ben2 = ModuleType("server.models.onnx_ben2")
+        ben2.create_session = lambda _directory, _device: sessions["ben2"]
+        da3 = ModuleType("server.models.onnx_da3")
+        da3.create_session = lambda _directory, _device: sessions["da3"]
+        moge2 = ModuleType("server.models.onnx_moge2")
+        moge2.create_session = lambda _directory, _device: sessions["moge2"]
+        upscale = ModuleType("server.models.onnx_upscale")
+        upscale.create_session = lambda _directory, _device: sessions["upscale"]
+        cache = self.create_cache()
+
+        def assert_reused(first, getter, release_name):
+            with (
+                patch.dict(sys.modules, {"onnxruntime": runtime}),
+                patch.object(
+                    cache,
+                    release_name,
+                    wraps=getattr(cache, release_name),
+                ) as release,
+            ):
+                self.onnx_runtime.run_session(first, ["output"], {"input": 1})
+                second = getter()
+            self.assertIs(first, second)
+            release.assert_not_called()
+
+        with patch.dict(
+            sys.modules,
+            {
+                "server.models.onnx_ben2": ben2,
+                "server.models.onnx_da3": da3,
+                "server.models.onnx_moge2": moge2,
+                "server.models.onnx_upscale": upscale,
+            },
+        ):
+            first_ben2, _load_ms = cache.get_ben2("ben2", "cuda")
+            assert_reused(
+                first_ben2,
+                lambda: cache.get_ben2("ben2", "cuda")[0],
+                "release_ben2",
+            )
+
+            first_da3 = cache.get_da3("da3", "cuda")
+            assert_reused(
+                first_da3,
+                lambda: cache.get_da3("da3", "cuda"),
+                "release_geometry",
+            )
+
+            first_moge2 = cache.get_moge2("moge2", "cuda")
+            assert_reused(
+                first_moge2,
+                lambda: cache.get_moge2("moge2", "cuda"),
+                "release_geometry",
+            )
+
+            first_upscale = cache.get_upscale("upscale", "cuda")
+            assert_reused(
+                first_upscale,
+                lambda: cache.get_upscale("upscale", "cuda"),
+                "release_upscale",
+            )
+
     def test_dedicated_server_runs_only_one_job_at_a_time(self):
         server = self.JobServer("Test", storage_root=self.storage_root)
         first_started = threading.Event()