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.

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`：

openspec/changes/release-model-inference-memory/.openspec.yaml

@@ -0,0 +1,2 @@
+schema: spec-driven
+created: 2026-09-01

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 产物无需处理。

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 协议、模型文件、输出文件与外部依赖保持不变。

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，并返回既有的可读资源错误

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.

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(

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",

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(

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)
 
 

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,
     )

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)

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:]

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)

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):

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,

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(

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)

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"),

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)

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()

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):

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(

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()