icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Implement unified memory management for ONNX model inference

Browse Source 
- Introduced a new memory management policy for ONNX Runtime to release temporary GPU memory after inference completion.
- Established a shared inference entry point for all ONNX models, differentiating between long-lived sessions and temporary memory arenas.
- Updated CUDA sessions to utilize a shrinkable memory arena, reclaiming GPU memory post-inference.
- Ensured consistent memory release semantics across BEN2, DA3, MoGe-2, and Upscale models while reusing existing sessions.
- Added lifecycle tests to validate temporary memory release and session reuse.
- Modified model adapters to integrate the new memory management features, ensuring backward compatibility with existing inference behavior.
- Enhanced unit tests to cover new memory release scenarios and verify correct behavior across different execution providers.
This commit is contained in:
马南
2026-09-01 13:22:35 +08:00
parent 4d4b892d64
commit 2ff7b04f56
24 changed files with 767 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/anyimage/server/jobs/debug_predictions.py
View File

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

1
src/anyimage/server/jobs/upscale.py
View File

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

8
src/anyimage/server/models/ben2.py
View File

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

21
src/anyimage/server/models/da3.py
View File

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

16
src/anyimage/server/models/moge2.py
View File

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

11
src/anyimage/server/models/onnx_ben2.py
View File

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

8
src/anyimage/server/models/onnx_da3.py
View File

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

10
src/anyimage/server/models/onnx_moge2.py
View File

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

44
src/anyimage/server/models/onnx_runtime.py
View File

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

23
src/anyimage/server/models/onnx_upscale.py
View File

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

8
src/anyimage/server/models/upscale.py
View File

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