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:
@@ -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",),
|
||||
|
||||
Reference in New Issue
Block a user