feat: Implement upscale functionality with ONNX models

- Added upscale operator to handle image upscaling by 4x using AI models.
- Introduced new jobs for upscale processing in the server, including debug and standard upscale jobs.
- Created a new ONNX model handler for upscale operations, supporting tiled inference and resource management.
- Updated tests to cover new upscale functionality, including input validation and output verification.
- Added documentation for the upscale feature, detailing usage and model options.
This commit is contained in:
2026-08-23 20:34:35 +08:00
parent 1613bd3194
commit 0cf04a13c6
66 changed files with 2156 additions and 1088 deletions

View File

@@ -10,6 +10,7 @@ class ModelSpec:
license_name: str
description: str
supports_gs: bool
huggingface_repository: str = ""
download_patterns: tuple[str, ...] = ()
ready_patterns: tuple[str, ...] = ("config.json", "*.safetensors")
enabled: bool = True
@@ -24,6 +25,7 @@ MODELS = {
directory_name="DA3-SMALL",
repository="depth-anything/DA3-SMALL",
license_name="Apache 2.0",
huggingface_repository="depth-anything/DA3-SMALL",
description="Fastest any-view model using ONNX Runtime",
supports_gs=False,
download_patterns=("model.onnx",),
@@ -39,6 +41,7 @@ MODELS = {
directory_name="DA3-BASE",
repository="depth-anything/DA3-BASE",
license_name="Apache 2.0",
huggingface_repository="depth-anything/DA3-BASE",
description="Balanced any-view model using ONNX Runtime",
supports_gs=False,
download_patterns=("model.onnx",),
@@ -76,6 +79,7 @@ BEN2_MODEL = ModelSpec(
directory_name="BEN2-ONNX",
repository="onnx-community/BEN2-ONNX",
license_name="Apache 2.0",
huggingface_repository="onnx-community/BEN2-ONNX",
description="Fast high-quality FP16 ONNX background removal model",
supports_gs=False,
download_patterns=("onnx/model_fp16.onnx",),
@@ -89,30 +93,75 @@ BEN2_MODEL = ModelSpec(
)
REALESRGAN_MODEL = ModelSpec(
key="REALESRGAN_GENERAL_X4V3",
label="Real-ESRGAN General x4v3",
directory_name="realesr-general-x4v3",
REALESRGAN_GENERAL_WDN_X4V3 = ModelSpec(
key="REALESRGAN_GENERAL_WDN_X4V3",
label="Real-ESRGAN General WDN x4v3",
directory_name="realesr-general-wdn-x4v3",
repository="xinntao/Real-ESRGAN",
license_name="BSD 3-Clause",
description="Fast and conservative 4x image upscaler using ONNX Runtime",
supports_gs=False,
download_patterns=("realesr-general-wdn-x4v3.onnx",),
ready_patterns=("realesr-general-wdn-x4v3.onnx",),
r2_directory="realesr-general-wdn-x4v3",
r2_files=((
"realesr-general-wdn-x4v3.onnx",
4867419,
"b021a837f7a0d3f287f9a39e367f581f1e6c3a759a08dc0c1c79c22442729774",
),),
)
SAFMN_MODEL = ModelSpec(
key="SAFMN_REAL_X4",
label="SAFMN Real x4",
directory_name="safmn-real-x4",
repository="sunny2109/SAFMN",
license_name="Apache 2.0",
description="Balanced real-world 4x image upscaler using ONNX Runtime",
supports_gs=False,
download_patterns=("safmn-real-x4.onnx",),
ready_patterns=("safmn-real-x4.onnx",),
r2_directory="safmn-real-x4",
r2_files=((
"safmn-real-x4.onnx",
22657477,
"8f13867b66c0300941b902bf6a5c6e0b545794faabe8b0478bf55668bc517621",
),),
)
REALESRGAN_X4PLUS = ModelSpec(
key="REALESRGAN_X4PLUS",
label="Real-ESRGAN x4plus",
directory_name="realesrgan-x4plus",
repository="jonathanst29/tinier-upscale-models",
license_name="BSD 3-Clause",
description="Fast general-purpose 4x image upscaler using ONNX Runtime",
huggingface_repository="jonathanst29/tinier-upscale-models",
description="Detailed general-purpose 4x image upscaler using ONNX Runtime",
supports_gs=False,
download_patterns=("realesr-general-x4v3.onnx",),
ready_patterns=("realesr-general-x4v3.onnx",),
r2_directory="realesr-general-x4v3",
download_patterns=("realesrgan-x4plus.onnx",),
ready_patterns=("realesrgan-x4plus.onnx",),
r2_directory="realesrgan-x4plus",
r2_files=((
"realesr-general-x4v3.onnx",
4866396,
"924ebad6532777303582d4ce7811849b88869231a3ae7093e0f21200249df8d5",
"realesrgan-x4plus.onnx",
67051618,
"4ed6a45a8185bde6c8d7c7790a6e80e3c5b9f608946f634021068dd0fd0473c8",
),),
)
UPSCALE_MODELS = {
REALESRGAN_MODEL.key: REALESRGAN_MODEL,
model.key: model
for model in (
REALESRGAN_GENERAL_WDN_X4V3,
SAFMN_MODEL,
REALESRGAN_X4PLUS,
)
}
DEFAULT_UPSCALE_MODEL_KEY = REALESRGAN_GENERAL_WDN_X4V3.key
MOGE2_MODELS = {
"MOGE2_VITS_NORMAL": ModelSpec(
key="MOGE2_VITS_NORMAL",
@@ -120,6 +169,7 @@ MOGE2_MODELS = {
directory_name="moge-2-vits-normal-onnx",
repository="Ruicheng/moge-2-vits-normal-onnx",
license_name="MIT",
huggingface_repository="Ruicheng/moge-2-vits-normal-onnx",
description="Fastest MoGe-2 model using ONNX Runtime",
supports_gs=False,
download_patterns=("model.onnx",),
@@ -137,6 +187,7 @@ MOGE2_MODELS = {
directory_name="moge-2-vitb-normal-onnx",
repository="Ruicheng/moge-2-vitb-normal-onnx",
license_name="MIT",
huggingface_repository="Ruicheng/moge-2-vitb-normal-onnx",
description="Balanced MoGe-2 model using ONNX Runtime",
supports_gs=False,
download_patterns=("model.onnx",),
@@ -154,6 +205,7 @@ MOGE2_MODELS = {
directory_name="moge-2-vitl-normal-onnx",
repository="Ruicheng/moge-2-vitl-normal-onnx",
license_name="MIT",
huggingface_repository="Ruicheng/moge-2-vitl-normal-onnx",
description="Highest-quality MoGe-2 model using ONNX Runtime",
supports_gs=False,
download_patterns=("model.onnx",),