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