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:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user