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.
This commit is contained in:
2026-09-01 13:22:35 +08:00
parent 4d4b892d64
commit 2ff7b04f56
24 changed files with 767 additions and 32 deletions

View File

@@ -13,6 +13,7 @@ def infer_one(
image_path,
cancel_check=None,
preserve_alpha=False,
release_memory=True,
):
from PIL import Image
@@ -30,7 +31,12 @@ def infer_one(
parameters.get("max_input_size", DEFAULT_MAX_AI_INPUT_SIZE),
)
source_alpha = _source_alpha(source) if preserve_alpha else None
prediction = infer(model, source, cancel_check=cancel_check)
prediction = infer(
model,
source,
cancel_check=cancel_check,
release_memory=release_memory,
)
if source_alpha is None:
return prediction
alpha = source_alpha.resize(