Files
AnyImage/docs/internals/runtime.md
nan 2ff7b04f56 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.
2026-09-01 13:22:35 +08:00

96 lines
4.9 KiB
Markdown
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# Job Runtime 与 Server
AnyImage 使用 BlendJob 管理独立 Python Environment、Server 进程与单 Worker 队列。Blender 侧准备输入、提交普通字典参数并导入结果。
## 目录职责
```text
src/anyimage/
├── runtime.py # Environment 与 JobRuntime 配置
└── server/
├── app.py # 公开 Job 注册
├── model_catalog.py # 模型声明
├── model_download.py # 下载与校验
├── model_manager.py # 模型目录与 Session 生命周期
├── jobs/ # 公开协议与流程编排
├── media/ # 图片、视频、序列帧和尺寸处理
├── selection/ # 通用 BEN2 Selection 细化
├── geometry/ # Depth / Normal 文件产物与相机元数据
├── models/ # 模型适配与 ONNX 推理
└── outputs/ # Debug Prediction 序列化
```
`jobs/` 编排任务,领域目录提供媒体、Selection、Geometry、模型和输出能力。
## 生命周期
```mermaid
flowchart TD
Request["Operator request()"] --> Runtime["JobRuntime"]
Runtime --> Queue["单 Worker 队列"]
Queue --> Job["app.py 注册的 Job"]
Job --> Resource["ModelManager / Media"]
Resource --> Directory["jobs/<job-id>/"]
Directory --> Response["Operator response()"]
```
Job 在阶段边界调用 `context.check_cancelled()`,并把稳定产物写入当前 Job 目录。Blender 响应完成后,再由输入所有者清理临时 Packed Image 文件。
## 公开 Job
| Job | 主要用途 |
| --- | --- |
| `download-model`、`download-required-models` | 准备模型文件 |
| `remove-background` | 生成 BEN2 透明前景 |
| `upscale-image` | 4× 图片或帧序列 |
| `refine-image-selection` | 为 Image Box / Lasso 执行 Refine Selection |
| `generate-depth-plane-geometry` | 为 Depth Plane 复制颜色并生成 Z Depth 与相机元数据 |
| `generate-cutout-artifacts` | 按需执行 Refine Selection,并生成 Depth 或 Normal 文件产物 |
| `debug-*` | 保存模型适配层的直接输出 |
`generate-depth-plane-geometry` 对应 `jobs/depth_plane.py`;`generate-cutout-artifacts` 对应 `jobs/cutout.py`。共享的 MoGe-2 推理编排位于 `geometry/prediction_artifacts.py`,它把最终 Selection RGBA 合成到黑底后执行推理,并把单帧 `GeometryFrame` 交给 Depth、Normal 与 Metadata Writer。Depth EXR 与元数据写入位于 `geometry/depth_texture.py`,Normal 编码位于 `geometry/normal_texture.py`,Refine Selection 的通用实现位于 `selection/refine_selection.py`。
## 尺寸约束
`server/media/resolution.py` 提供 **Maximum AI Input Size**,默认 `2048 px`。它限制 Upscale 的可提交源图长边,以及 Depth、Normal Map、Refine Selection 的模型分析图长边;材质颜色图使用原始分辨率。
Blender 侧的 **Maximum Cutout Mesh Resolution** 统一限制所有 Cutout Shape 的网格采样密度。
`generate-depth-plane-geometry` 与 `generate-cutout-artifacts` 只接受一帧,因此 Depth Plane 与 Cutout 是静态图工作流。Remove Background 与 Upscale 继续通过媒体入口处理单图、视频和序列帧。
## 模型
Depth Plane、Cutout Depth 与 Normal Map 使用 MoGe-2,Remove Background 与 Refine Selection 使用 BEN2,Upscale 使用所选高清模型。MoGe-2 生产推理直接返回 `GeometryFrame`;DA3 调试推理保留模型原始 Prediction,并可按帧适配到相同结构。
`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`:
```text
~/.anyimage/
├── .venv/
├── models/
├── jobs/
├── environment.json
├── manifest.json
├── install.log
└── server.log
```
切换 Storage Root 等同于切换整套 Environment、模型与任务数据,不自动迁移旧目录。
各 Job 的稳定产物见 [Job 产物](outputs.md)。