nan
2ff7b04f56
- 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.
56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
from ..media.resolution import DEFAULT_MAX_AI_INPUT_SIZE, limit_image
|
|
|
|
|
|
def _source_alpha(image):
|
|
if "A" not in image.getbands() and "transparency" not in image.info:
|
|
return None
|
|
return image.convert("RGBA").getchannel("A")
|
|
|
|
|
|
def infer_one(
|
|
model_manager,
|
|
parameters,
|
|
image_path,
|
|
cancel_check=None,
|
|
preserve_alpha=False,
|
|
release_memory=True,
|
|
):
|
|
from PIL import Image
|
|
|
|
from .onnx_upscale import UpscaleResourceError, infer
|
|
|
|
model_directory = model_manager.directory(parameters["model"])
|
|
model = model_manager.get_upscale(
|
|
model_directory,
|
|
parameters["device"],
|
|
)
|
|
try:
|
|
with Image.open(image_path) as opened:
|
|
source = limit_image(
|
|
opened,
|
|
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,
|
|
release_memory=release_memory,
|
|
)
|
|
if source_alpha is None:
|
|
return prediction
|
|
alpha = source_alpha.resize(
|
|
prediction.size,
|
|
Image.Resampling.LANCZOS,
|
|
)
|
|
prediction = prediction.convert("RGBA")
|
|
prediction.putalpha(alpha)
|
|
return prediction
|
|
except (MemoryError, UnicodeDecodeError, UpscaleResourceError) as error:
|
|
model_manager.release_upscale()
|
|
device = str(parameters.get("device", "selected device")).upper()
|
|
raise RuntimeError(
|
|
f"Upscale could not finish on {device}: {error}. "
|
|
"Reduce Maximum AI Input Size or select CPU, then try again."
|
|
) from error
|