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

@@ -64,6 +64,17 @@ Depth Plane、Cutout Depth 与 Normal Map 使用 MoGe-2,Remove Background 与
`ModelManager` 按模型族管理下载目录与 ONNX Session,同一族只保留最近使用的模型和 Device,避免同类模型重复占用资源。
BEN2、DA3、MoGe-2 与 Upscale 共用 ONNX Runtime 推理入口。CUDA Session 使用可收缩的 memory arena;单次图片、完整帧集合或完整 Upscale tile 集合只在最后一次模型运行结束时收缩空闲 arena。Session 持有的模型权重与执行计划继续由 `ModelManager` 缓存,后续相同模型与 Device 直接复用。CPU、DirectML 与 CoreML 按各自 Provider 的普通运行方式执行。
```mermaid
flowchart TD
Request["模型请求"] --> Cache["ModelManager 获取缓存 Session"]
Cache --> Runs["执行单次或批量模型运行"]
Runs --> Final["最后一次模型运行"]
Final --> Shrink["CUDA 收缩空闲 memory arena"]
Shrink --> Reuse["Session 留在 ModelManager 中复用"]
```
## Environment 与 Storage Root
`runtime.py` 声明独立 Python 3.12 环境和固定依赖。首次安装会准备默认 MoGe-2、BEN2 与 Real-ESRGAN WDN 模型。Storage Root 默认是 `~/.anyimage`: