diff --git a/docs/internals/runtime.md b/docs/internals/runtime.md index e2f0ec3..25443aa 100644 --- a/docs/internals/runtime.md +++ b/docs/internals/runtime.md @@ -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`： diff --git a/openspec/changes/release-model-inference-memory/.openspec.yaml b/openspec/changes/release-model-inference-memory/.openspec.yaml new file mode 100644 index 0000000..b4b3ece --- /dev/null +++ b/openspec/changes/release-model-inference-memory/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-09-01 diff --git a/openspec/changes/release-model-inference-memory/design.md b/openspec/changes/release-model-inference-memory/design.md new file mode 100644 index 0000000..d16d27c --- /dev/null +++ b/openspec/changes/release-model-inference-memory/design.md @@ -0,0 +1,79 @@ +## Context + +See `proposal.md` - Why。当前 `ModelManager` 按 BEN2、Geometry 与 Upscale 三类缓存 ONNX `InferenceSession`；模型适配器各自直接调用 `session.run()`。Session 释放时的 `gc.collect()` 会连同模型一起丢弃，无法满足只清理单次推理内存的目标。 + +四类模型具有不同的运行粒度：BEN2 可逐帧执行，DA3 可 joint 或 per-frame 执行，MoGe-2 可依次处理多张图，Upscale 会把每帧拆成多个固定 tile。临时 arena 应在完整业务推理的最后一次 `session.run()` 后回收，而不是在内部每次运行后回收。 + +ONNX Runtime 的 arena shrinkage 通过最后一次 Run 的 `RunOptions` 发起。CUDA 默认的 `kNextPowerOfTwo` arena 不保证首个分配区域可收缩，因此需要为所有 CUDA 模型 Session 使用 `kSameAsRequested`；Session 与仍被模型使用的分配保持存活，只回收空闲区域。 + +## Goals / Non-Goals + +**Goals:** + +- 让所有 ONNX 模型通过一个共享运行边界选择 Provider 安全的临时内存释放行为。 +- 保持模型 Session、权重与执行计划缓存，下一次匹配请求不重新加载。 +- 把释放点放在单图、完整帧集合或完整 tile 集合的最后一次模型运行。 +- 保持 CPU、DirectML、CoreML 与现有 Provider fallback 行为不变。 + +**Non-Goals:** + +- 不释放 Blender 当前显示图片所占用的纹理显存，也不改变结果图片分辨率。 +- 不保证失败或取消发生在中间运行时仍能在不销毁 Session 的前提下立即收缩 arena；既有 OOM Session 失效策略保持不变。 +- 不引入 PyTorch CUDA allocator、外部 allocator 或新的运行依赖。 +- 不调整模型族缓存数量、切换模型时的淘汰规则或 `Clear Models` 行为。 + +## Decisions + +### 1. Session 创建与 Session 运行共享同一套内存策略 + +`server/models/onnx_runtime.py` 扩展为唯一策略入口： + +- `create_session()` 在选中 CUDA 时附加 `arena_extend_strategy=kSameAsRequested`；其他 Provider 保留现有配置。 +- 共享的模型运行函数接收“本次运行结束业务推理”的布尔值。仅当该值为真且 Session 实际使用 CUDA 时，创建 `RunOptions` 并设置 `memory.enable_memory_arena_shrinkage=gpu:0`。 +- 未请求释放或 Provider 不支持时，直接执行原 `session.run()`，不构造无效配置。 + +不把清理放进 `ModelManager`：Manager 只知道 Session 的缓存与淘汰，不知道一次模型调用是不是多帧或多 tile 的最后一次运行。也不在任务结束时销毁 Session，因为这会丢失模型缓存。 + +### 2. 低层模型运行显式接收完成边界 + +BEN2、DA3、MoGe-2 与 Upscale 的 ONNX 适配器不再直接调用 `session.run()`，统一调用共享运行函数，并接受是否在本次运行后释放临时内存的参数。单次调用默认把自身视为完整业务推理；批处理编排方对中间运行显式传入 false，仅对最后一次传入 true。 + +选择显式参数而不是全局状态或计数器，因为 Server 虽是单 Worker，模型入口仍同时服务生产 Job、Selection 与 Debug Job。调用方掌握准确的输入数量与 tile 数量，显式传递不会让取消、异常或嵌套调用破坏隐式计数。 + +### 3. 每种批处理在已有所有者处标记最后一次运行 + +- BEN2 的多帧入口只让最后一帧释放；单图 Remove Background 与 Refine Selection 直接释放。 +- DA3 joint 模式在唯一一次运行后释放；per-frame 模式只让最后一帧释放。 +- MoGe-2 的多图入口只让最后一张释放；单图 Artifact 与直接调试入口释放。 +- Upscale 的 tile 循环可确定最后一个 tile；上层多帧 Job 只让最后一帧的最后一个 tile 释放。生产与 Debug 的帧循环使用相同边界。 + +不选择每次 `session.run()` 后都收缩，因为视频、DA3 per-frame 与 Upscale tile 会在下一次运行立即重新分配同类空间，增加同步与 allocator 开销而没有为 Blender 提供中途使用显存的机会。 + +### 4. 清理能力按 Session 的实际 Provider 判定 + +判断依据是 Session 的实际 Provider，而不是用户请求的 Device。这样 CUDA 请求回退到 DirectML 或 CPU 时不会错误提交 `gpu:0` arena 配置；`auto` 实际选到 CUDA 时仍会启用清理。 + +本次只对 ONNX Runtime 明确支持的 CUDA GPU arena 使用 shrinkage。DirectML、CoreML、CPU 与其他 Provider 保持普通 Run，避免把 CUDA 配置推广为未经验证的跨 Provider 行为。 + +### 5. 测试同时固定清理时机与缓存身份 + +共享 Runtime 测试覆盖 CUDA Session 配置、最终 RunOptions、非 CUDA 普通 Run 与输出透传。各模型测试记录每次运行是否请求释放，分别固定 BEN2 多帧、DA3 两种模式、MoGe-2 多图、Upscale 多帧多 tile 的“仅最后一次”为真。 + +Server Runtime 测试在推理完成前后比较缓存 Session 身份和加载次数，证明 arena cleanup 不经过 `release_*()`。现有 Upscale OOM 测试继续固定失败时销毁 Session 的既有合同。 + +## Risks / Trade-offs + +- [arena shrinkage 增加一次任务结束时的同步与后续重新分配成本] → 只在完整业务推理末尾执行，不在帧或 tile 内部重复执行。 +- [CUDA Provider 配置在不同 ONNX Runtime 版本中存在差异] → 使用项目锁定版本支持的标准 Provider/Run 配置，并用真实 API 冒烟测试与模拟 Session 单元测试共同覆盖。 +- [实际显存不会回到模型加载前水平] → 明确只回收空闲推理 arena；模型权重、执行计划和 Blender 正在显示的纹理继续占用显存。 +- [调用方漏传最后一次标记会延迟释放] → 单次模型入口默认释放，所有已知批处理入口建立调用序列测试。 +- [中间运行取消或非 OOM 异常可能保留 arena] → 保留 Session 缓存优先；下一次成功业务推理会再次执行收缩，显式资源错误继续沿用既有 Session 淘汰路径。 + +## Migration Plan + +1. 建立共享 CUDA Session 配置和 Provider-aware Run 入口及其测试。 +2. 迁移四类 ONNX 适配器，保持输出与异常转换不变。 +3. 从模型批处理到生产与 Debug Job 逐层传递最后一次运行标记，并覆盖多帧、多图与多 tile 测试。 +4. 更新 Runtime 内部文档，运行全部测试。 + +该变更没有数据迁移。回退时同时回退共享 Runtime 配置、四类适配器参数和调用方标记；模型文件与已生成 Job 产物无需处理。 diff --git a/openspec/changes/release-model-inference-memory/proposal.md b/openspec/changes/release-model-inference-memory/proposal.md new file mode 100644 index 0000000..1e09a7d --- /dev/null +++ b/openspec/changes/release-model-inference-memory/proposal.md @@ -0,0 +1,27 @@ +## Why + +ONNX Runtime 会把模型权重与推理期间扩张的 GPU memory arena 一起保留在长期存活的 Session 中；一次高显存推理结束后，临时分配继续挤占 Blender 的 GPU 资源。模型 Session 需要继续缓存以避免重复加载，但 BEN2、DA3、MoGe-2 与 Upscale 的单次推理内存应在业务处理完成后统一释放。 + +## What Changes + +- 为所有 ONNX 模型建立统一的推理运行入口，区分长期缓存的 Session 与单次处理使用的临时 memory arena。 +- 让 CUDA Session 使用可收缩的 arena 配置，并在一次图片、帧批次或分块处理的最后一次模型运行后请求回收临时 GPU 分配。 +- BEN2、DA3、MoGe-2 与全部 Upscale 模型均使用相同的释放语义，同时继续复用 `ModelManager` 中的原 Session。 +- 对不支持 arena shrinkage 的 Execution Provider 保持原推理行为，不提交无效的 Provider 或 Run 配置。 +- 增加模型级与 Server 生命周期测试，验证临时内存释放、批次结束边界和 Session 复用。 + +## Capabilities + +### New Capabilities + +- `model-inference-memory`: 规定所有 ONNX 模型在复用 Session 的同时，于一次业务推理完成后释放 Provider 支持的临时推理内存。 + +### Modified Capabilities + +无。 + +## Impact + +- 影响 `server/models/onnx_runtime.py`、BEN2、DA3、MoGe-2、Upscale 的 ONNX 适配器及对应测试。 +- 影响多帧 BEN2 与分块 Upscale 的运行边界，需要只在最后一次 `session.run()` 请求 arena shrinkage，避免每帧或每块重复释放与重新分配。 +- `ModelManager` 的模型族缓存、Job 协议、模型文件、输出文件与外部依赖保持不变。 diff --git a/openspec/changes/release-model-inference-memory/specs/model-inference-memory/spec.md b/openspec/changes/release-model-inference-memory/specs/model-inference-memory/spec.md new file mode 100644 index 0000000..6cac80b --- /dev/null +++ b/openspec/changes/release-model-inference-memory/specs/model-inference-memory/spec.md @@ -0,0 +1,61 @@ +## Purpose + +规定所有 ONNX 模型在复用已加载 Session 的同时，于一次完整业务推理结束后释放 Execution Provider 支持回收的临时 GPU 内存，避免推理资源长期挤占 Blender。 + +## ADDED Requirements + +### Requirement: Cached model sessions survive inference cleanup + +系统 SHALL 在清理一次推理使用的临时内存后保留对应模型的已加载 Session，使相同模型、Device 与 Runtime 的下一次请求复用该 Session。 + +#### Scenario: Matching request reuses the loaded model + +- **WHEN** 一个模型操作成功完成临时推理内存清理，随后以相同模型、Device 与 Runtime 再次请求该模型 +- **THEN** 系统复用同一个已加载 Session，且不重新从模型文件创建 Session + +#### Scenario: All model families share the lifecycle + +- **WHEN** BEN2、DA3、MoGe-2 或任一 Upscale 模型在支持清理的 Provider 上完成推理 +- **THEN** 每个模型均遵循相同的 Session 保留与临时内存释放语义 + +### Requirement: Temporary GPU inference memory is released at operation completion + +系统 SHALL 在 Execution Provider 支持 memory arena shrinkage 时，于一次成功业务推理的最后一次模型运行请求释放不再使用的临时 GPU arena，同时保留 Session 持有的模型资源。 + +#### Scenario: Single-run inference releases temporary memory + +- **WHEN** 一次业务推理只执行一次模型运行并成功返回 +- **THEN** 该次模型运行请求在结束时收缩支持回收的 GPU memory arena + +#### Scenario: Multi-frame inference releases after the final frame + +- **WHEN** 一次业务推理依次处理多个输入帧 +- **THEN** 中间帧不触发 arena shrinkage，最后一帧的模型运行请求触发 arena shrinkage + +#### Scenario: Tiled inference releases after the final tile + +- **WHEN** Upscale 为一个或多个输入帧执行多次分块模型运行 +- **THEN** 中间分块与中间帧不触发 arena shrinkage，整个业务推理的最后一个分块触发 arena shrinkage + +### Requirement: Unsupported providers retain existing inference behavior + +系统 MUST 只向明确支持所选 memory arena shrinkage 配置的 Execution Provider 提交该配置；其他 Provider SHALL 按原有运行方式执行。 + +#### Scenario: Provider does not support GPU arena shrinkage + +- **WHEN** 模型 Session 使用 CPU、DirectML、CoreML 或其他未声明支持该 GPU arena shrinkage 配置的 Provider +- **THEN** 系统不提交对应的 Provider 或 Run 清理配置，且推理继续正常执行 + +### Requirement: Inference results and failure contracts remain stable + +临时内存释放 SHALL 不改变模型输入、模型输出、Job 产物、取消检查或既有资源错误处理。 + +#### Scenario: Cleanup-enabled inference returns the same model result + +- **WHEN** 对相同输入与模型分别执行启用和未启用临时内存释放的推理 +- **THEN** 两次推理返回相同结构、尺寸、dtype 与数值语义的模型结果 + +#### Scenario: Upscale provider resource failure + +- **WHEN** Upscale 的 Execution Provider 因 GPU 或系统内存不足而失败 +- **THEN** 系统继续释放失效的 Upscale Session，并返回既有的可读资源错误 diff --git a/openspec/changes/release-model-inference-memory/tasks.md b/openspec/changes/release-model-inference-memory/tasks.md new file mode 100644 index 0000000..c5381ac --- /dev/null +++ b/openspec/changes/release-model-inference-memory/tasks.md @@ -0,0 +1,24 @@ +## 1. Shared ONNX Runtime Memory Policy + +- [x] 1.1 Extend the shared ONNX Session creation path with CUDA `kSameAsRequested` arena configuration, and verify CUDA, CoreML and plain Provider configurations with runtime helper tests. +- [x] 1.2 Add the Provider-aware shared Session run entry that applies `memory.enable_memory_arena_shrinkage=gpu:0` only to a completing CUDA run, and verify CUDA cleanup, intermediate run, fallback Provider and output passthrough scenarios with unit tests. +- [x] 1.3 Add an ONNX Runtime 1.24 API smoke test for constructing the required `RunOptions` entry, and verify it runs without requiring a CUDA device. + +## 2. Model Adapter Integration + +- [x] 2.1 Migrate BEN2 ONNX inference to the shared run entry with an explicit completion flag, and verify preprocessing, alpha output and cleanup-enabled output remain unchanged in `test_model_ben2_onnx.py`. +- [x] 2.2 Migrate DA3 ONNX inference to the shared run entry with an explicit completion flag, and verify joint prediction fields, attention validation and completion forwarding in `test_model_da3_onnx.py`. +- [x] 2.3 Migrate MoGe-2 ONNX inference to the shared run entry with an explicit completion flag, and verify prediction fields, Provider selection and completion forwarding in `test_model_moge2_onnx.py`. +- [x] 2.4 Migrate Upscale tile inference to the shared run entry and mark only the final requested tile as completing, and verify multi-tile call order, output equality and existing resource-error conversion in `test_model_upscale_onnx.py`. + +## 3. Business Inference Boundaries + +- [x] 3.1 Propagate the completion flag through BEN2 single-image, multi-frame, Selection and Debug flows, and verify only the final frame requests cleanup while single-image flows still request it. +- [x] 3.2 Propagate the completion flag through DA3 joint/per-frame and MoGe-2 single/multi-image production and Debug flows, and verify each complete operation records exactly one cleanup request on its final model run. +- [x] 3.3 Propagate the completion flag through production and Debug Upscale frame loops so only the final frame's final tile requests cleanup, and verify single-frame, multi-frame and multi-tile sequences in Server Job tests. + +## 4. Lifecycle and Project Verification + +- [x] 4.1 Extend ModelManager lifecycle tests to prove BEN2, DA3, MoGe-2 and Upscale cleanup does not call their Session release paths and a matching subsequent request returns the same cached Session. +- [x] 4.2 Update `docs/internals/runtime.md` to describe cached Session ownership and end-of-operation temporary arena cleanup, and verify the model lifecycle section matches all four model families and Provider limitations. +- [x] 4.3 Run `uv run pytest` and verify the complete test suite passes without changes to Job protocols, model outputs or existing Upscale OOM behavior. diff --git a/src/anyimage/server/jobs/debug_predictions.py b/src/anyimage/server/jobs/debug_predictions.py index 067ecf4..fbdd4c3 100644 --- a/src/anyimage/server/jobs/debug_predictions.py +++ b/src/anyimage/server/jobs/debug_predictions.py @@ -34,6 +34,7 @@ def run_moge2(context, parameters): context.resource("model_manager"), parameters, image, + release_memory=index == len(images) - 1, ) suffix = "" if len(images) == 1 else f".{index + 1:04d}" write_moge2_prediction( @@ -60,6 +61,7 @@ def run_ben2(context, parameters): parameters["device"], model_cache=model_manager, cancel_check=context.check_cancelled, + release_memory=index == len(images) - 1, ) suffix = "" if len(images) == 1 else f".{index + 1:04d}" foreground.save(context.directory / f"foreground{suffix}.png") @@ -80,6 +82,7 @@ def run_upscale(context, parameters): parameters, image, cancel_check=context.check_cancelled, + release_memory=index == len(images) - 1, ) suffix = "" if len(images) == 1 else f".{index + 1:04d}" prediction.save( diff --git a/src/anyimage/server/jobs/upscale.py b/src/anyimage/server/jobs/upscale.py index 5461838..81fefb9 100644 --- a/src/anyimage/server/jobs/upscale.py +++ b/src/anyimage/server/jobs/upscale.py @@ -30,6 +30,7 @@ def run(context, parameters): image, cancel_check=context.check_cancelled, preserve_alpha=True, + release_memory=index == total - 1, ) prediction.save( context.directory / f"upscale.{index + 1:04d}.png", diff --git a/src/anyimage/server/models/ben2.py b/src/anyimage/server/models/ben2.py index 52a409a..b798bdc 100644 --- a/src/anyimage/server/models/ben2.py +++ b/src/anyimage/server/models/ben2.py @@ -11,6 +11,7 @@ def infer_foreground( model_cache=None, cancel_check=None, preserve_input_alpha=False, + release_memory=True, ): from PIL import Image, ImageChops @@ -39,7 +40,11 @@ def infer_foreground( source_rgba.getchannel("A"), ) inference_started = time.perf_counter() - foreground = infer(model, image) + foreground = infer( + model, + image, + release_memory=release_memory, + ) inference_ms = (time.perf_counter() - inference_started) * 1000.0 if preserve_input_alpha: foreground_alpha = foreground.convert("RGBA").getchannel("A") @@ -84,6 +89,7 @@ def infer_foreground_frames( model_cache=model_cache, cancel_check=cancel_check, preserve_input_alpha=preserve_input_alpha, + release_memory=index == len(images) - 1, ) output_path = output_directory / f"{output_prefix}.{index + 1:04d}.png" foreground.convert("RGBA").save( diff --git a/src/anyimage/server/models/da3.py b/src/anyimage/server/models/da3.py index 33414dc..6bf1581 100644 --- a/src/anyimage/server/models/da3.py +++ b/src/anyimage/server/models/da3.py @@ -46,9 +46,14 @@ def infer( "ref_view_strategy": parameters["ref_view_strategy"], } - def run(paths): + def run(paths, *, release_memory): try: - return infer_onnx(model, paths, **options) + return infer_onnx( + model, + paths, + **options, + release_memory=release_memory, + ) except UnicodeDecodeError as error: if provider_label(model) == "DirectML": raise RuntimeError( @@ -59,7 +64,10 @@ def infer( raise if parameters.get("inference_mode", "joint") != "per_frame": - return run([str(path) for path in images]) + return run( + [str(path) for path in images], + release_memory=True, + ) predictions = [] for index, image in enumerate(images): @@ -67,7 +75,12 @@ def infer( cancel_check() if progress_callback is not None: progress_callback(index, len(images)) - predictions.append(run([str(image)])) + predictions.append( + run( + [str(image)], + release_memory=index == len(images) - 1, + ) + ) return merge_frame_predictions(predictions) diff --git a/src/anyimage/server/models/moge2.py b/src/anyimage/server/models/moge2.py index 37a0834..9414141 100644 --- a/src/anyimage/server/models/moge2.py +++ b/src/anyimage/server/models/moge2.py @@ -51,14 +51,25 @@ def infer( progress_callback(index, len(images)) frames.append( _geometry_frame( - infer_one(model_manager, parameters, image_path), + infer_one( + model_manager, + parameters, + image_path, + release_memory=index == len(images) - 1, + ), include_points, ) ) return frames -def infer_one(model_manager, parameters, image_path): +def infer_one( + model_manager, + parameters, + image_path, + *, + release_memory=True, +): import numpy as np from PIL import Image @@ -75,4 +86,5 @@ def infer_one(model_manager, parameters, image_path): model, image, int(parameters["resolution_level"]), + release_memory=release_memory, ) diff --git a/src/anyimage/server/models/onnx_ben2.py b/src/anyimage/server/models/onnx_ben2.py index 3132e57..87fd427 100644 --- a/src/anyimage/server/models/onnx_ben2.py +++ b/src/anyimage/server/models/onnx_ben2.py @@ -4,8 +4,10 @@ from pathlib import Path try: from .onnx_runtime import create_session as create_runtime_session + from .onnx_runtime import run_session except ImportError: from onnx_runtime import create_session as create_runtime_session + from onnx_runtime import run_session ONNX_PATH = Path("onnx/model_fp16.onnx") @@ -32,13 +34,18 @@ def preprocess(image): return np.ascontiguousarray(array.transpose(2, 0, 1)[None]) -def infer(session, image): +def infer(session, image, *, release_memory=True): import numpy as np from PIL import Image input_name = session.get_inputs()[0].name output_name = session.get_outputs()[0].name - alpha = session.run([output_name], {input_name: preprocess(image)})[0] + alpha = run_session( + session, + [output_name], + {input_name: preprocess(image)}, + release_memory=release_memory, + )[0] alpha = np.asarray(alpha, dtype=np.float32).squeeze() minimum = float(alpha.min()) span = max(float(alpha.max()) - minimum, np.finfo(np.float32).eps) diff --git a/src/anyimage/server/models/onnx_da3.py b/src/anyimage/server/models/onnx_da3.py index 4c4a4b7..740c125 100644 --- a/src/anyimage/server/models/onnx_da3.py +++ b/src/anyimage/server/models/onnx_da3.py @@ -9,8 +9,10 @@ import numpy as np try: from .onnx_runtime import create_session as create_onnx_session + from .onnx_runtime import run_session except ImportError: from onnx_runtime import create_session as create_onnx_session + from onnx_runtime import run_session ONNX_PATH = "model.onnx" @@ -195,6 +197,8 @@ def infer( process_res=504, process_res_method="upper_bound_resize", ref_view_strategy="saddle_balanced", + *, + release_memory=True, ): inputs = session.get_inputs() model_input, processed_images = preprocess( @@ -216,12 +220,14 @@ def infer( dtype=np.int64, ) input_names = [value.name for value in inputs] - outputs = session.run( + outputs = run_session( + session, None, { input_names[0]: model_input, input_names[1]: strategy_code, }, + release_memory=release_memory, ) depth, confidence, pose_encoding = outputs[:3] height, width = depth.shape[-2:] diff --git a/src/anyimage/server/models/onnx_moge2.py b/src/anyimage/server/models/onnx_moge2.py index 4c31a2e..6c9e995 100644 --- a/src/anyimage/server/models/onnx_moge2.py +++ b/src/anyimage/server/models/onnx_moge2.py @@ -3,6 +3,7 @@ from pathlib import Path from .onnx_runtime import create_session as create_runtime_session +from .onnx_runtime import run_session from .onnx_runtime import select_providers as select_providers @@ -136,7 +137,7 @@ def postprocess(raw_output): } -def infer(session, image, resolution_level): +def infer(session, image, resolution_level, *, release_memory=True): import numpy as np level = min(max(int(resolution_level), 0), 9) @@ -149,7 +150,12 @@ def infer(session, image, resolution_level): output_names = {output.name for output in session.get_outputs()} scale_name = "metric_scale" if "metric_scale" in output_names else "scale" names = ("points", "normal", "mask", scale_name) - values = session.run(list(names), input_feed) + values = run_session( + session, + list(names), + input_feed, + release_memory=release_memory, + ) raw_output = dict(zip(names, values)) raw_output["metric_scale"] = raw_output.pop(scale_name) return postprocess(raw_output) diff --git a/src/anyimage/server/models/onnx_runtime.py b/src/anyimage/server/models/onnx_runtime.py index d699202..cb170d8 100644 --- a/src/anyimage/server/models/onnx_runtime.py +++ b/src/anyimage/server/models/onnx_runtime.py @@ -7,6 +7,10 @@ COREML_PROVIDER_OPTIONS = { "EnableOnSubgraphs": "0", } +CUDA_PROVIDER_OPTIONS = { + "arena_extend_strategy": "kSameAsRequested", +} + def select_providers(requested, available): available = set(available) @@ -47,13 +51,39 @@ def create_session(path, requested_device="auto"): requested_device, onnxruntime.get_available_providers(), ) - providers = [ - (provider, dict(COREML_PROVIDER_OPTIONS)) - if provider == "CoreMLExecutionProvider" - else provider - for provider in providers - ] - return onnxruntime.InferenceSession(str(path), providers=providers) + configured_providers = [] + for provider in providers: + if provider == "CUDAExecutionProvider": + provider = (provider, dict(CUDA_PROVIDER_OPTIONS)) + elif provider == "CoreMLExecutionProvider": + provider = (provider, dict(COREML_PROVIDER_OPTIONS)) + configured_providers.append(provider) + return onnxruntime.InferenceSession( + str(path), + providers=configured_providers, + ) + + +def run_session( + session, + output_names, + input_feed, + *, + release_memory=True, +): + if not release_memory or "CUDAExecutionProvider" not in ( + session.get_providers() + ): + return session.run(output_names, input_feed) + + import onnxruntime + + run_options = onnxruntime.RunOptions() + run_options.add_run_config_entry( + "memory.enable_memory_arena_shrinkage", + "gpu:0", + ) + return session.run(output_names, input_feed, run_options=run_options) def provider_label(session): diff --git a/src/anyimage/server/models/onnx_upscale.py b/src/anyimage/server/models/onnx_upscale.py index eaa4d15..0297cde 100644 --- a/src/anyimage/server/models/onnx_upscale.py +++ b/src/anyimage/server/models/onnx_upscale.py @@ -4,8 +4,10 @@ from pathlib import Path try: from .onnx_runtime import create_session as create_runtime_session + from .onnx_runtime import run_session except ImportError: from onnx_runtime import create_session as create_runtime_session + from onnx_runtime import run_session SCALE = 4 @@ -42,15 +44,17 @@ def preprocess(image): return np.ascontiguousarray(image_array.transpose(2, 0, 1)[None]) -def _run(session, input_tensor): +def _run(session, input_tensor, *, release_memory=True): import numpy as np input_name = session.get_inputs()[0].name output_name = session.get_outputs()[0].name try: - raw_prediction = session.run( + raw_prediction = run_session( + session, [output_name], {input_name: input_tensor}, + release_memory=release_memory, )[0] except (MemoryError, UnicodeDecodeError) as error: raise UpscaleResourceError( @@ -100,7 +104,7 @@ def _pad_source(source, extra_height, extra_width): ) -def infer(session, image, cancel_check=None): +def infer(session, image, cancel_check=None, *, release_memory=True): import numpy as np from PIL import Image @@ -122,7 +126,18 @@ def infer(session, image, cancel_check=None): padded[top : top + TILE_SIZE, left : left + TILE_SIZE], mode="RGB", ) - pixels = _prediction_pixels(_run(session, preprocess(tile))) + final_tile = ( + release_memory + and top + core_size >= height + and left + core_size >= width + ) + pixels = _prediction_pixels( + _run( + session, + preprocess(tile), + release_memory=final_tile, + ) + ) border = TILE_BORDER * SCALE output[ top * SCALE : (top + core_height) * SCALE, diff --git a/src/anyimage/server/models/upscale.py b/src/anyimage/server/models/upscale.py index 99882c2..124c38b 100644 --- a/src/anyimage/server/models/upscale.py +++ b/src/anyimage/server/models/upscale.py @@ -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( diff --git a/tests/test_model_ben2_onnx.py b/tests/test_model_ben2_onnx.py index 621d3a4..65e6696 100644 --- a/tests/test_model_ben2_onnx.py +++ b/tests/test_model_ben2_onnx.py @@ -2,6 +2,7 @@ import sys import unittest from pathlib import Path from types import SimpleNamespace +from unittest.mock import patch import numpy as np from PIL import Image @@ -20,6 +21,9 @@ import onnx_ben2 class FakeSession: + def get_providers(self): + return ["CPUExecutionProvider"] + def get_inputs(self): return [SimpleNamespace(name="input")] @@ -34,6 +38,25 @@ class FakeSession: class Ben2OnnxTest(unittest.TestCase): + def test_infer_forwards_memory_release_boundary(self): + session = FakeSession() + + with patch.object( + onnx_ben2, + "run_session", + side_effect=lambda model, names, feeds, **_options: model.run( + names, feeds + ), + ) as run: + result = onnx_ben2.infer( + session, + Image.new("RGB", (40, 20)), + release_memory=False, + ) + + self.assertEqual(result.size, (40, 20)) + self.assertFalse(run.call_args.kwargs["release_memory"]) + def test_preprocess_matches_model_layout(self): image = Image.new("RGB", (40, 20), (64, 128, 192)) result = onnx_ben2.preprocess(image) diff --git a/tests/test_model_da3_onnx.py b/tests/test_model_da3_onnx.py index 885487b..29cc522 100644 --- a/tests/test_model_da3_onnx.py +++ b/tests/test_model_da3_onnx.py @@ -3,6 +3,7 @@ import tempfile import unittest from pathlib import Path from types import SimpleNamespace +from unittest.mock import patch import numpy as np from PIL import Image @@ -21,6 +22,9 @@ import onnx_da3 class FakeSession: + def get_providers(self): + return ["CPUExecutionProvider"] + def get_inputs(self): return [SimpleNamespace(name="images"), SimpleNamespace(name="strategy")] @@ -36,6 +40,27 @@ class FakeSession: class Da3OnnxTest(unittest.TestCase): + def test_infer_forwards_memory_release_boundary(self): + with tempfile.TemporaryDirectory() as directory: + image_path = Path(directory) / "input.png" + Image.new("RGB", (80, 60)).save(image_path) + session = FakeSession() + with patch.object( + onnx_da3, + "run_session", + side_effect=lambda model, names, feeds, **_options: model.run( + names, feeds + ), + ) as run: + onnx_da3.infer( + session, + [image_path], + process_res=56, + release_memory=False, + ) + + self.assertFalse(run.call_args.kwargs["release_memory"]) + def test_target_size_uses_patch_multiple(self): self.assertEqual( onnx_da3.target_size((800, 600), 504, "upper_bound_resize"), diff --git a/tests/test_model_moge2_onnx.py b/tests/test_model_moge2_onnx.py index 6189f4c..2d43d81 100644 --- a/tests/test_model_moge2_onnx.py +++ b/tests/test_model_moge2_onnx.py @@ -56,6 +56,9 @@ class FakeSession: self.calls.append((names, inputs)) return [self.outputs[name] for name in names] + def get_providers(self): + return ["CPUExecutionProvider"] + def get_outputs(self): return [FakeInput(name) for name in self.outputs] @@ -145,6 +148,35 @@ class Moge2OnnxTest(unittest.TestCase): self.assertEqual(inputs["num_tokens"].dtype, np.int64) self.assertEqual(inputs["num_tokens"], 3600) + def test_infer_forwards_memory_release_boundary(self): + height, width = 8, 12 + raw = { + "points": np.zeros((1, height, width, 3), dtype=np.float32), + "normal": np.zeros((1, height, width, 3), dtype=np.float32), + "mask": np.ones((1, height, width), dtype=np.float32), + "scale": np.ones((1,), dtype=np.float32), + } + session = FakeSession(raw) + + with ( + patch.object(self.moge, "postprocess", return_value={}), + patch.object( + self.moge, + "run_session", + side_effect=lambda model, names, feeds, **_options: model.run( + names, feeds + ), + ) as run, + ): + self.moge.infer( + session, + np.zeros((height, width, 3), dtype=np.float32), + 0, + release_memory=False, + ) + + self.assertFalse(run.call_args.kwargs["release_memory"]) + def test_postprocess_recovers_camera_projection_and_metric_scale(self): height, width = 64, 80 focal = np.float32(0.9) diff --git a/tests/test_model_onnx_runtime.py b/tests/test_model_onnx_runtime.py new file mode 100644 index 0000000..de02999 --- /dev/null +++ b/tests/test_model_onnx_runtime.py @@ -0,0 +1,135 @@ +import sys +import unittest +from pathlib import Path +from types import ModuleType +from unittest.mock import patch + + +SOURCE_DIR = ( + Path(__file__).resolve().parents[1] + / "src" + / "anyimage" + / "server" + / "models" +) +sys.path.insert(0, str(SOURCE_DIR)) + +import onnx_runtime + + +class FakeSession: + def __init__(self, providers): + self.providers = providers + self.calls = [] + + def get_providers(self): + return self.providers + + def run(self, output_names, input_feed, run_options=None): + self.calls.append((output_names, input_feed, run_options)) + return ["prediction"] + + +class OnnxRuntimeTest(unittest.TestCase): + def create_runtime(self, available): + calls = [] + runtime = ModuleType("onnxruntime") + runtime.get_available_providers = lambda: available + runtime.InferenceSession = lambda path, providers: calls.append( + (path, providers) + ) + return runtime, calls + + def test_cuda_session_uses_shrinkable_arena(self): + runtime, calls = self.create_runtime( + ["CUDAExecutionProvider", "CPUExecutionProvider"] + ) + + with patch.dict(sys.modules, {"onnxruntime": runtime}): + onnx_runtime.create_session("model.onnx", "cuda") + + self.assertEqual( + calls[0], + ( + "model.onnx", + [ + ( + "CUDAExecutionProvider", + {"arena_extend_strategy": "kSameAsRequested"}, + ) + ], + ), + ) + + def test_coreml_and_cpu_session_options_are_preserved(self): + runtime, calls = self.create_runtime( + ["CoreMLExecutionProvider", "CPUExecutionProvider"] + ) + + with patch.dict(sys.modules, {"onnxruntime": runtime}): + onnx_runtime.create_session("model.onnx", "auto") + + self.assertEqual( + calls[0][1], + [ + ( + "CoreMLExecutionProvider", + dict(onnx_runtime.COREML_PROVIDER_OPTIONS), + ), + "CPUExecutionProvider", + ], + ) + + def test_completing_cuda_run_shrinks_memory_arena(self): + configs = [] + + class FakeRunOptions: + def add_run_config_entry(self, key, value): + configs.append((key, value)) + + runtime = ModuleType("onnxruntime") + runtime.RunOptions = FakeRunOptions + session = FakeSession(["CUDAExecutionProvider", "CPUExecutionProvider"]) + + with patch.dict(sys.modules, {"onnxruntime": runtime}): + result = onnx_runtime.run_session(session, ["output"], {"input": 1}) + + self.assertEqual(result, ["prediction"]) + self.assertIsInstance(session.calls[0][2], FakeRunOptions) + self.assertEqual( + configs, + [("memory.enable_memory_arena_shrinkage", "gpu:0")], + ) + + def test_intermediate_and_non_cuda_runs_use_plain_session_run(self): + cuda = FakeSession(["CUDAExecutionProvider"]) + directml = FakeSession(["DmlExecutionProvider"]) + + onnx_runtime.run_session( + cuda, + ["output"], + {"input": 1}, + release_memory=False, + ) + result = onnx_runtime.run_session( + directml, + ["output"], + {"input": 2}, + ) + + self.assertIsNone(cuda.calls[0][2]) + self.assertIsNone(directml.calls[0][2]) + self.assertEqual(result, ["prediction"]) + + def test_installed_runtime_accepts_memory_arena_shrinkage_option(self): + import onnxruntime + + run_options = onnxruntime.RunOptions() + run_options.add_run_config_entry( + "memory.enable_memory_arena_shrinkage", + "gpu:0", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_model_upscale_onnx.py b/tests/test_model_upscale_onnx.py index 17e9da8..7373b33 100644 --- a/tests/test_model_upscale_onnx.py +++ b/tests/test_model_upscale_onnx.py @@ -2,6 +2,7 @@ import sys import unittest from pathlib import Path from types import SimpleNamespace +from unittest.mock import patch import numpy as np from PIL import Image @@ -20,6 +21,9 @@ import onnx_upscale class FakeSession: + def get_providers(self): + return ["CPUExecutionProvider"] + def get_inputs(self): return [SimpleNamespace(name="input")] @@ -121,6 +125,40 @@ class UpscaleOnnxTest(unittest.TestCase): self.assertEqual(session.calls, 6) self.assertEqual(session.feeds["input"].shape, (1, 3, 256, 256)) + def test_only_final_tile_releases_memory(self): + releases = [] + session = FakeSession() + + def run(model, output_names, feeds, *, release_memory): + releases.append(release_memory) + return model.run(output_names, feeds) + + with patch.object(onnx_upscale, "run_session", side_effect=run): + result = onnx_upscale.infer( + session, + Image.new("RGB", (449, 225), (1, 2, 3)), + ) + + self.assertEqual(result.size, (1796, 900)) + self.assertEqual(releases, [False, False, False, False, False, True]) + + def test_non_final_image_does_not_release_its_final_tile(self): + releases = [] + session = FakeSession() + + def run(model, output_names, feeds, *, release_memory): + releases.append(release_memory) + return model.run(output_names, feeds) + + with patch.object(onnx_upscale, "run_session", side_effect=run): + onnx_upscale.infer( + session, + Image.new("RGB", (449, 225), (1, 2, 3)), + release_memory=False, + ) + + self.assertEqual(releases, [False] * 6) + def test_provider_unicode_error_is_presented_as_resource_error(self): class InvalidProviderSession(FakeSession): def run(self, _output_names, _feeds): diff --git a/tests/test_server_jobs.py b/tests/test_server_jobs.py index dda1e94..998ef35 100644 --- a/tests/test_server_jobs.py +++ b/tests/test_server_jobs.py @@ -504,8 +504,8 @@ class ServerJobsTest(unittest.TestCase): fake_torch.inference_mode = nullcontext fake_ben2 = ModuleType("server.models.onnx_ben2") fake_ben2.create_session = lambda _model_dir, _device: object() - fake_ben2.infer = lambda _model, image: ( - inference_options.append((image, {})) or FakeImage() + fake_ben2.infer = lambda _model, image, **options: ( + inference_options.append((image, options)) or FakeImage() ) fake_pil = ModuleType("PIL") fake_pil.Image = SimpleNamespace( @@ -537,7 +537,7 @@ class ServerJobsTest(unittest.TestCase): self.assertEqual( inference_options[0][1], - {}, + {"release_memory": True}, ) self.assertEqual( save_options[0], @@ -556,7 +556,7 @@ class ServerJobsTest(unittest.TestCase): model_result = Image.new("RGBA", (2, 1), (10, 20, 30, 128)) received = [] - def infer(_model, image): + def infer(_model, image, **_options): received.append(np.asarray(image.convert("RGB"), dtype=np.uint8)) return model_result @@ -644,8 +644,8 @@ class ServerJobsTest(unittest.TestCase): fake_torch.inference_mode = nullcontext fake_ben2 = ModuleType("server.models.onnx_ben2") fake_ben2.create_session = lambda _model_dir, _device: object() - fake_ben2.infer = lambda _model, image: ( - inference_options.append((image, {})) or FakeImage() + fake_ben2.infer = lambda _model, image, **options: ( + inference_options.append((image, options)) or FakeImage() ) fake_pil = ModuleType("PIL") fake_pil.Image = SimpleNamespace( @@ -689,8 +689,10 @@ class ServerJobsTest(unittest.TestCase): ) self.assertEqual(len(inference_options), 2) - for _image, options in inference_options: - self.assertEqual(options, {}) + self.assertEqual( + [options["release_memory"] for _image, options in inference_options], + [False, True], + ) def test_da3_debug_saves_only_direct_api_prediction(self): prediction = SimpleNamespace( @@ -754,6 +756,52 @@ class ServerJobsTest(unittest.TestCase): self.assertFalse(any(context.directory.glob("*.png"))) self.assertEqual(calls[0][0], [str(input_path)]) self.assertEqual(calls[0][1]["process_res"], 518) + self.assertTrue(calls[0][1]["release_memory"]) + + def test_da3_per_frame_releases_only_after_final_frame(self): + releases = [] + + def prediction(): + return SimpleNamespace( + depth=np.ones((1, 2, 2), dtype=np.float32), + conf=np.ones((1, 2, 2), dtype=np.float32), + extrinsics=np.eye(4, dtype=np.float32)[None, :3], + intrinsics=np.eye(3, dtype=np.float32)[None], + processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8), + ) + + fake_da3 = ModuleType("server.models.onnx_da3") + + def infer(_session, _images, **options): + releases.append(options["release_memory"]) + return prediction() + + fake_da3.infer = infer + fake_runtime = ModuleType("server.models.onnx_runtime") + fake_runtime.provider_label = lambda _session: "CPU" + parameters = { + "model": "DA3", + "device": "cpu", + "process_res": 56, + "process_res_method": "upper_bound_resize", + "ref_view_strategy": "middle", + "inference_mode": "per_frame", + } + + with patch.dict( + sys.modules, + { + "server.models.onnx_da3": fake_da3, + "server.models.onnx_runtime": fake_runtime, + }, + ): + self.da3_model.infer( + FakeModels(), + parameters, + [Path("a.png"), Path("b.png")], + ) + + self.assertEqual(releases, [False, True]) def test_moge2_debug_saves_each_direct_api_dictionary(self): prediction = { @@ -786,12 +834,42 @@ class ServerJobsTest(unittest.TestCase): }, ) self.assertEqual(infer_one.call_count, 2) + self.assertEqual( + [call.kwargs["release_memory"] for call in infer_one.call_args_list], + [False, True], + ) archives = sorted(context.directory.glob("prediction.*.npz")) self.assertEqual(len(archives), 2) with np.load(archives[0]) as result: self.assertEqual(set(result.files), set(prediction)) self.assertFalse(any(context.directory.glob("*.glb"))) + def test_moge2_batch_releases_only_after_final_image(self): + prediction = { + "points": np.ones((2, 2, 3), dtype=np.float32), + "depth": np.ones((2, 2), dtype=np.float32), + "normal": np.ones((2, 2, 3), dtype=np.float32), + "mask": np.ones((2, 2), dtype=bool), + "intrinsics": np.eye(3, dtype=np.float32), + } + with patch.object( + self.moge2_model, + "infer_one", + return_value=prediction, + ) as infer_one: + frames = self.moge2_model.infer( + FakeModels(), + {"model": "MOGE2_VITS_NORMAL", "device": "cpu"}, + [Path("a.png"), Path("b.png")], + include_points=True, + ) + + self.assertEqual(len(frames), 2) + self.assertEqual( + [call.kwargs["release_memory"] for call in infer_one.call_args_list], + [False, True], + ) + def test_ben2_debug_saves_direct_api_foreground(self): class Foreground: def save(self, path): @@ -817,6 +895,9 @@ class ServerJobsTest(unittest.TestCase): }, ) infer_foreground.assert_called_once() + self.assertTrue( + infer_foreground.call_args.kwargs["release_memory"] + ) self.assertEqual( (context.directory / "foreground.png").read_bytes(), b"foreground", @@ -848,6 +929,10 @@ class ServerJobsTest(unittest.TestCase): ) self.assertEqual(infer_one.call_count, 2) + self.assertEqual( + [call.kwargs["release_memory"] for call in infer_one.call_args_list], + [False, True], + ) predictions = sorted(context.directory.glob("prediction.*.png")) self.assertEqual(len(predictions), 2) with Image.open(predictions[0]) as prediction: @@ -880,6 +965,10 @@ class ServerJobsTest(unittest.TestCase): ) self.assertEqual(infer_one.call_count, 2) + self.assertEqual( + [call.kwargs["release_memory"] for call in infer_one.call_args_list], + [False, True], + ) self.assertEqual(result, {"upscale": "."}) colors = sorted(context.directory.glob("upscale.*.png")) self.assertEqual( diff --git a/tests/test_server_runtime.py b/tests/test_server_runtime.py index ba6f678..739e701 100644 --- a/tests/test_server_runtime.py +++ b/tests/test_server_runtime.py @@ -46,10 +46,12 @@ class ServerRuntimeTest(unittest.TestCase): importlib.import_module("server.jobs.cutout") from blendjob import JobServer from server.model_manager import ModelManager, ModelSessionCache + from server.models import onnx_runtime cls.JobServer = JobServer cls.ModelManager, cls.ModelSessionCache = ModelManager, ModelSessionCache + cls.onnx_runtime = onnx_runtime cls.server = cls.server_module.server cls.app_module = sys.modules["server.app"] @@ -224,6 +226,88 @@ class ServerRuntimeTest(unittest.TestCase): self.assertEqual(loads, [("model", "cpu")]) self.assertEqual(cache.snapshot()["upscale"], "model") + def test_inference_cleanup_keeps_every_cached_session(self): + class Session: + def get_providers(self): + return ["CUDAExecutionProvider"] + + def run(self, _names, _feeds, run_options=None): + self.run_options = run_options + return [object()] + + class RunOptions: + def add_run_config_entry(self, _key, _value): + return None + + runtime = ModuleType("onnxruntime") + runtime.RunOptions = RunOptions + sessions = { + "ben2": Session(), + "da3": Session(), + "moge2": Session(), + "upscale": Session(), + } + ben2 = ModuleType("server.models.onnx_ben2") + ben2.create_session = lambda _directory, _device: sessions["ben2"] + da3 = ModuleType("server.models.onnx_da3") + da3.create_session = lambda _directory, _device: sessions["da3"] + moge2 = ModuleType("server.models.onnx_moge2") + moge2.create_session = lambda _directory, _device: sessions["moge2"] + upscale = ModuleType("server.models.onnx_upscale") + upscale.create_session = lambda _directory, _device: sessions["upscale"] + cache = self.create_cache() + + def assert_reused(first, getter, release_name): + with ( + patch.dict(sys.modules, {"onnxruntime": runtime}), + patch.object( + cache, + release_name, + wraps=getattr(cache, release_name), + ) as release, + ): + self.onnx_runtime.run_session(first, ["output"], {"input": 1}) + second = getter() + self.assertIs(first, second) + release.assert_not_called() + + with patch.dict( + sys.modules, + { + "server.models.onnx_ben2": ben2, + "server.models.onnx_da3": da3, + "server.models.onnx_moge2": moge2, + "server.models.onnx_upscale": upscale, + }, + ): + first_ben2, _load_ms = cache.get_ben2("ben2", "cuda") + assert_reused( + first_ben2, + lambda: cache.get_ben2("ben2", "cuda")[0], + "release_ben2", + ) + + first_da3 = cache.get_da3("da3", "cuda") + assert_reused( + first_da3, + lambda: cache.get_da3("da3", "cuda"), + "release_geometry", + ) + + first_moge2 = cache.get_moge2("moge2", "cuda") + assert_reused( + first_moge2, + lambda: cache.get_moge2("moge2", "cuda"), + "release_geometry", + ) + + first_upscale = cache.get_upscale("upscale", "cuda") + assert_reused( + first_upscale, + lambda: cache.get_upscale("upscale", "cuda"), + "release_upscale", + ) + def test_dedicated_server_runs_only_one_job_at_a_time(self): server = self.JobServer("Test", storage_root=self.storage_root) first_started = threading.Event()