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

Implement BiRefNet and HAT Sharper models with background removal capabilities

Browse Source 
- Added validation documentation for AI model tiers.
- Introduced `background.py` for background removal logic, including inference functions for alpha and foreground extraction.
- Created `onnx_birefnet.py` to handle BiRefNet model sessions and alpha inference.
- Developed tests for background models, ensuring proper dispatch and session reuse.
- Added tests for BiRefNet ONNX model to validate alpha output and error handling for invalid values.
- Implemented model export scripts for BiRefNet and HAT Sharper, ensuring proper weight compression and licensing.
- Created `deform_sampling.py` for deformable convolution operations used in BiRefNet.
- Added utility functions for compressing ONNX weights without altering graph arithmetic.
- Introduced a reference plane node for depth field adjustments in geometry processing.
This commit is contained in:
马南
2026-09-10 20:50:44 +08:00
parent 8ae3c1d5b7
commit e8129fe495
70 changed files with 1890 additions and 1208 deletions
Download Patch File Download Diff File Expand all files Collapse all files

BIN
src/anyimage/assets/O_AnyImage.blend LFS
View File

Binary file not shown.

2
src/anyimage/blender_manifest.toml
View File

@@ -9,7 +9,7 @@ type = "add-on"
blender_version_min = "4.5.0"
platforms = ["windows-x64"]
wheels = ["./wheels/blendjob-0.1.16-py3-none-any.whl", "./wheels/scipy-1.15.3-cp311-cp311-win_amd64.whl", "./wheels/scipy-1.15.3-cp313-cp313-win_amd64.whl"]
wheels = ["./wheels/blendjob-0.1.17-py3-none-any.whl", "./wheels/scipy-1.15.3-cp311-cp311-win_amd64.whl", "./wheels/scipy-1.15.3-cp313-cp313-win_amd64.whl"]
license = ["SPDX:GPL-3.0-or-later"]
[permissions]

8
src/anyimage/operators/ai_setup.py
View File

@@ -2,8 +2,8 @@ import bpy
from ..preferences import ADDON_PACKAGE
from ..runtime import JobOperatorBase, runtime
from ..server.model_catalog import DEFAULT_MODEL_KEYS
from ..properties import (
REQUIRED_MODEL_KEYS,
ai_status,
model_download_size,
model_record,
@@ -42,7 +42,7 @@ class SetupAIEnvironment(bpy.types.Operator):
if not status["environment_ready"]:
required_size = sum(
model_download_size(model_key)
for model_key in REQUIRED_MODEL_KEYS
for model_key in DEFAULT_MODEL_KEYS
)
layout.label(
text="AnyImage AI features run in an isolated Python environment."
@@ -52,7 +52,7 @@ class SetupAIEnvironment(bpy.types.Operator):
)
layout.label(
text=(
"It also downloads MoGe-2 ViT-S Normal, BEN2 Base and "
"It also downloads MoGe-2 ViT-S Normal, BiRefNet Lite and "
"Real-ESRGAN General WDN x4v3 "
f"({format_download_size(required_size)})."
)
@@ -154,7 +154,7 @@ class DownloadRequiredModels(JobOperatorBase, bpy.types.Operator):
require_environment()
if not bpy.app.online_access:
raise RuntimeError("Online Access is disabled in Blender Preferences")
return {"models": list(REQUIRED_MODEL_KEYS)}
return {"models": list(DEFAULT_MODEL_KEYS)}
class ClearModels(bpy.types.Operator):

7
src/anyimage/operators/convert_to_panorama/operators.py
View File

@@ -6,7 +6,8 @@ from ...common.ai import invoke_ai_setup_if_needed, production_device, report_ai
from ...common.image import is_image_empty, video_total_frames
from ...common.color_image import prepare_material_color_input, material_analysis_input, cleanup_material_color_input
from ...preferences import addon_preferences, configured_max_ai_input_size
from ...properties import DEFAULT_MOGE2_MODEL_KEY
from ...preferences import configured_geometry_model
from ...server.model_catalog import DEFAULT_GEOMETRY_MODEL_KEY
from ...runtime import JobOperatorBase
from ...server.media.resolution import DEFAULT_MAX_AI_INPUT_SIZE, MAX_MAX_AI_INPUT_SIZE, MIN_MAX_AI_INPUT_SIZE
from .object import create_panorama_from_result
@@ -26,7 +27,7 @@ class GeneratePanorama(JobOperatorBase, bpy.types.Operator):
source_identity: bpy.props.StringProperty(options=_HIDDEN)
image_identity: bpy.props.StringProperty(options=_HIDDEN)
mesh_detail: bpy.props.IntProperty(default=7, min=0, max=8, options=_HIDDEN)
model: bpy.props.StringProperty(default=DEFAULT_MOGE2_MODEL_KEY, options=_HIDDEN)
model: bpy.props.StringProperty(default=DEFAULT_GEOMETRY_MODEL_KEY, options=_HIDDEN)
resolution_level: bpy.props.IntProperty(default=5, min=0, max=9, options=_HIDDEN)
max_input_size: bpy.props.IntProperty(default=DEFAULT_MAX_AI_INPUT_SIZE,
min=MIN_MAX_AI_INPUT_SIZE, max=MAX_MAX_AI_INPUT_SIZE, options=_HIDDEN)
@@ -82,7 +83,7 @@ class ConvertToPanorama(bpy.types.Operator):
input_path=str(path), color_path=str(color_path), source_object_name=source.name,
source_identity=str(source.as_pointer()),
image_identity=str(image.as_pointer()), mesh_detail=self.mesh_detail,
model=getattr(preferences, "geometry_model", DEFAULT_MOGE2_MODEL_KEY),
model=configured_geometry_model(context),
resolution_level=int(getattr(preferences, "geometry_resolution_level", 5)),
max_input_size=configured_max_ai_input_size(context),
)

8
src/anyimage/operators/convert_to_plane/operators.py
View File

@@ -1,7 +1,7 @@
import bpy
from ...preferences import addon_preferences, configured_max_ai_input_size
from ...properties import DEFAULT_MOGE2_MODEL_KEY
from ...preferences import addon_preferences, configured_max_ai_input_size, configured_geometry_model
from ...server.model_catalog import DEFAULT_GEOMETRY_MODEL_KEY
from ...runtime import JobOperatorBase
from ...server.media.resolution import (
DEFAULT_MAX_AI_INPUT_SIZE,
@@ -37,7 +37,7 @@ _HIDDEN = {"HIDDEN", "SKIP_SAVE"}
def depth_plane_settings(context):
preferences = addon_preferences(context)
return (
getattr(preferences, "geometry_model", DEFAULT_MOGE2_MODEL_KEY),
configured_geometry_model(context),
int(getattr(preferences, "geometry_resolution_level", 5)),
)
@@ -59,7 +59,7 @@ class GenerateDepthPlane(JobOperatorBase, bpy.types.Operator):
options=_HIDDEN,
)
mesh_detail: bpy.props.IntProperty(default=6, min=0, max=10, options=_HIDDEN)
model: bpy.props.StringProperty(default=DEFAULT_MOGE2_MODEL_KEY, options=_HIDDEN)
model: bpy.props.StringProperty(default=DEFAULT_GEOMETRY_MODEL_KEY, options=_HIDDEN)
resolution_level: bpy.props.IntProperty(default=5, min=0, max=9, options=_HIDDEN)
max_input_size: bpy.props.IntProperty(
default=DEFAULT_MAX_AI_INPUT_SIZE,

7
src/anyimage/operators/cutout_tool/operators.py
View File

@@ -6,15 +6,16 @@ import bpy
from ...preferences import (
addon_preferences,
configured_geometry_model,
configured_max_ai_input_size,
configured_max_cutout_mesh_resolution,
)
from ...properties import (
DEFAULT_MOGE2_MODEL_KEY,
cutout_alpha_threshold_property,
cutout_fine_outline_property,
)
from ...runtime import JobOperatorBase
from ...server.model_catalog import DEFAULT_GEOMETRY_MODEL_KEY
from ...server.media.resolution import (
DEFAULT_MAX_AI_INPUT_SIZE,
MAX_MAX_AI_INPUT_SIZE,
@@ -169,7 +170,7 @@ def _generated_color_result(input_path, source_bounds):
def cutout_geometry_settings(context):
preferences = addon_preferences(context)
model = getattr(preferences, "geometry_model", DEFAULT_MOGE2_MODEL_KEY)
model = configured_geometry_model(context)
return (
model,
int(getattr(preferences, "geometry_resolution_level", 5)),
@@ -196,7 +197,7 @@ class CutoutSelectionToShape(JobOperatorBase, bpy.types.Operator):
options=_HIDDEN,
)
model: bpy.props.StringProperty(
default=DEFAULT_MOGE2_MODEL_KEY,
default=DEFAULT_GEOMETRY_MODEL_KEY,
options=_HIDDEN,
)
resolution_level: bpy.props.IntProperty(

9
src/anyimage/operators/remove_background.py
View File

@@ -1,6 +1,6 @@
import bpy
from ..properties import BEN2_MODEL_KEY
from ..preferences import configured_background_model
from ..runtime import JobOperatorBase
from ..common.ai import (
require_input_path,
@@ -22,7 +22,7 @@ _HIDDEN = {"HIDDEN", "SKIP_SAVE"}
class RemoveImageBackground(bpy.types.Operator):
bl_idname = "anyimage.remove_image_background"
bl_label = "Remove Background"
bl_description = "Remove the image background with BEN2"
bl_description = "Remove the image background with the selected AI model"
bl_options = {"REGISTER"}
@classmethod
@@ -84,11 +84,12 @@ class RunBackgroundRemoval(JobOperatorBase, bpy.types.Operator):
def request(self, context):
require_environment()
require_model(BEN2_MODEL_KEY)
model = configured_background_model(context)
require_model(model)
settings = context.scene.anyimage_settings
parameters = {
"input": str(require_input_path(self.input_path)),
"model": BEN2_MODEL_KEY,
"model": model,
"device": production_device().lower(),
"video_frames": settings.video_frames,
"start_frame": self._image_target.start_frame,

3
src/anyimage/operators/upscale.py
View File

@@ -1,7 +1,8 @@
import bpy
from ..preferences import configured_max_ai_input_size
from ..properties import configured_upscale_model, upscale_model_label
from ..preferences import configured_upscale_model
from ..properties import upscale_model_label
from ..server.media.resolution import (
DEFAULT_MAX_AI_INPUT_SIZE,
MAX_MAX_AI_INPUT_SIZE,

46
src/anyimage/preferences.py
View File

@@ -8,12 +8,14 @@ from .server.media.resolution import (
MIN_MAX_AI_INPUT_SIZE,
)
from .server.model_catalog import (
BACKGROUND_MODELS,
MOGE_MODELS,
DEFAULT_BACKGROUND_MODEL_KEY,
DEFAULT_GEOMETRY_MODEL_KEY,
DEFAULT_UPSCALE_MODEL_KEY,
UPSCALE_MODELS,
)
DEFAULT_MOGE2_MODEL_KEY = "MOGE2_VITS_NORMAL"
DEFAULT_BACKGROUND_MODEL_KEY = "BEN2_BASE"
DEFAULT_MAX_CUTOUT_MESH_RESOLUTION = 1024
MIN_CUTOUT_MESH_RESOLUTION = 4
MAX_CUTOUT_MESH_RESOLUTION = 1024
@@ -24,16 +26,13 @@ MAX_FRAME_RESOLUTION = 8192
ADDON_PACKAGE = __package__
DEFAULT_STORAGE_ROOT = Path.home() / ".anyimage"
UPSCALE_MODEL_ITEMS = tuple(
(model.key, model.label, model.description, index)
for index, model in enumerate(UPSCALE_MODELS.values())
(model.key, model.label, model.description, number)
for model, number in zip(UPSCALE_MODELS.values(), (0, 2, 4))
)
BACKGROUND_MODEL_ITEMS = (
(
DEFAULT_BACKGROUND_MODEL_KEY,
"BEN2 Base",
"Model used for AI Remove Background",
),
BACKGROUND_MODEL_ITEMS = tuple(
(model.key, model.label, model.description, number)
for model, number in zip(BACKGROUND_MODELS.values(), (1, 0, 2))
)
BALLOON_PROFILE_ALGORITHM_ITEMS = (
@@ -61,6 +60,27 @@ def configured_storage_root():
return DEFAULT_STORAGE_ROOT
def _configured_model(context, property_name, models, default):
preferences = addon_preferences(context)
model = getattr(preferences, property_name, default)
if model not in models:
setattr(preferences, property_name, default)
return default
return model
def configured_geometry_model(context=None):
return _configured_model(context, "geometry_model", MOGE_MODELS, DEFAULT_GEOMETRY_MODEL_KEY)
def configured_background_model(context=None):
return _configured_model(context, "background_model", BACKGROUND_MODELS, DEFAULT_BACKGROUND_MODEL_KEY)
def configured_upscale_model(context=None):
return _configured_model(context, "upscale_model", UPSCALE_MODELS, DEFAULT_UPSCALE_MODEL_KEY)
def configured_material_view_adaptation(context=None):
return bool(
getattr(addon_preferences(context), "adapt_material_to_view_transform", True)
@@ -239,7 +259,7 @@ class AnyImagePreferences(bpy.types.AddonPreferences):
ai,
"geometry_model",
models,
DEFAULT_MOGE2_MODEL_KEY,
DEFAULT_GEOMETRY_MODEL_KEY,
"AI Generate Depth Map",
)
self._draw_model_property(
@@ -286,8 +306,10 @@ class AnyImagePreferences(bpy.types.AddonPreferences):
ready_label,
):
model_key = getattr(self, property_name, default_model_key)
if model_key not in models:
choices = {"geometry_model": MOGE_MODELS, "background_model": BACKGROUND_MODELS, "upscale_model": UPSCALE_MODELS}
if model_key not in choices[property_name]:
model_key = default_model_key
setattr(self, property_name, model_key)
split = layout.split(factor=0.58, align=True)
split.column(align=True).prop(self, property_name, text="")
self._draw_model_status(

28
src/anyimage/properties.py
View File

@@ -3,18 +3,9 @@ import sys
import bpy
from .preferences import addon_preferences
from .runtime import runtime
from .server import model_catalog as shared_model_catalog
BEN2_MODEL_KEY = "BEN2_BASE"
REALESRGAN_GENERAL_WDN_MODEL_KEY = shared_model_catalog.DEFAULT_UPSCALE_MODEL_KEY
DEFAULT_MOGE2_MODEL_KEY = "MOGE2_VITS_NORMAL"
REQUIRED_MODEL_KEYS = (
DEFAULT_MOGE2_MODEL_KEY,
BEN2_MODEL_KEY,
REALESRGAN_GENERAL_WDN_MODEL_KEY,
)
DEVICE_LABELS = {
@@ -87,7 +78,7 @@ def ai_status():
environment_ready = runtime.environment_ready()
ready_models = {model["key"] for model in model_catalog() if model.get("ready")}
missing_models = tuple(
model_key for model_key in REQUIRED_MODEL_KEYS if model_key not in ready_models
model_key for model_key in shared_model_catalog.DEFAULT_MODEL_KEYS if model_key not in ready_models
)
return {
"environment_ready": environment_ready,
@@ -127,14 +118,15 @@ def model_record(model_key):
def _model_items(models):
numbers = {"MOGE2_VITS_NORMAL": 0, "MOGE2_VITB_NORMAL": 1, "MOGE3_VITL": 3}
return [
(
model["key"],
model["label"],
f"{model['description']}; license: {model['license']}",
index,
numbers[model["key"]],
)
for index, model in enumerate(models)
for model in models
]
@@ -143,18 +135,6 @@ def geometry_model_enum_items(_owner, _context):
return _model_items(models)
def configured_upscale_model(context=None):
preferences = addon_preferences(context)
model = getattr(
preferences,
"upscale_model",
shared_model_catalog.DEFAULT_UPSCALE_MODEL_KEY,
)
if model in shared_model_catalog.UPSCALE_MODELS:
return model
return shared_model_catalog.DEFAULT_UPSCALE_MODEL_KEY
def upscale_model_label(model_key):
model = shared_model_catalog.UPSCALE_MODELS.get(model_key)
return model.label if model is not None else model_key

13
src/anyimage/runtime.py
View File

@@ -3,6 +3,7 @@ from pathlib import Path
from blendjob import JobRuntime
from .preferences import configured_storage_root
from .server.model_catalog import DEFAULT_MODEL_KEYS, get_downloadable_model
ENVIRONMENT = {
@@ -24,19 +25,11 @@ ENVIRONMENT = {
def post_install(job_runtime):
errors = []
models = (
("MOGE2_VITS_NORMAL", "MoGe-2 ViT-S Normal"),
("BEN2_BASE", "BEN2 Base"),
(
"REALESRGAN_GENERAL_WDN_X4V3",
"Real-ESRGAN General WDN x4v3",
),
)
for model, label in models:
for model in DEFAULT_MODEL_KEYS:
try:
job_runtime.request("download-model", {"model": model})
except RuntimeError as error:
errors.append(f"{label}: {error}")
errors.append(f"{get_downloadable_model(model).label}: {error}")
if errors:
raise RuntimeError("; ".join(errors))

12
src/anyimage/server/jobs/remove_background.py
View File

@@ -2,12 +2,15 @@ import time
from ..media import input_frames
from ..media.input import open_image
from ..models.ben2 import infer_alpha, infer_foreground, infer_foreground_frames
from ..model_catalog import BACKGROUND_MODELS
from ..models.background import infer_alpha, infer_foreground, infer_foreground_frames, inference_device
def run(context, parameters):
"""Run BEN2 and write transparent foreground images into the Job directory."""
"""Write foreground images or continuous alpha into the Job directory."""
total_started = time.perf_counter()
if parameters["model"] not in BACKGROUND_MODELS:
raise ValueError(f"Unknown background removal model: {parameters['model']}")
context.progress(0.02, "Starting Environment")
context.check_cancelled()
@@ -16,12 +19,13 @@ def run(context, parameters):
device = parameters["device"]
model_manager = context.resource("model_manager")
model_directory = model_manager.directory(parameters["model"])
device = inference_device(model_directory, device)
output_kind = parameters.get("output_kind", "foreground")
if output_kind not in {"foreground", "alpha"}:
raise ValueError(f"Unknown background removal output: {output_kind}")
with input_frames(context, parameters) as images:
context.progress(0.12, f"Loading BEN2 for {device.upper()}")
context.progress(0.12, f"Loading {parameters['model']} for {device.upper()}")
if output_kind == "alpha":
import numpy as np
@@ -53,7 +57,7 @@ def run(context, parameters):
compress_level=1,
)
if not output_path.is_file():
raise RuntimeError("BEN2 did not write a foreground image")
raise RuntimeError("Background removal did not write a foreground image")
else:
load_ms, inference_ms = infer_foreground_frames(
images,

129
src/anyimage/server/model_catalog.py
View File

@@ -16,14 +16,31 @@ class ModelSpec:
r2_files: tuple[tuple[str, int, str], ...] = ()
BIREFNET_LITE = ModelSpec(
key="BIREFNET_LITE",
label="Fast | BiRefNet Lite",
directory_name="birefnet-lite",
repository="ZhengPeng7/BiRefNet_lite",
license_name="MIT",
description="Compact 1024px background removal with continuous alpha",
download_patterns=("model.onnx", "LICENSE.txt"),
ready_patterns=("model.onnx", "LICENSE.txt"),
r2_directory="birefnet-lite",
r2_files=(
("model.onnx", 112968598, "6e2a9844c57b080e95a402722ddcae5660e6f254a04d3e003347f06f587b71fe"),
("LICENSE.txt", 1066, "92a7089e0915fc32bc40067560b398f1e6a7a5958abd7d04eda393629a5acefb"),
),
)
BEN2_MODEL = ModelSpec(
key="BEN2_BASE",
label="BEN2 Base",
label="General | BEN2 Base",
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",
description="General-purpose FP16 ONNX background removal model",
download_patterns=("onnx/model_fp16.onnx",),
ready_patterns=("onnx/model_fp16.onnx",),
r2_directory="BEN2-ONNX",
@@ -37,7 +54,7 @@ BEN2_MODEL = ModelSpec(
REALESRGAN_GENERAL_WDN_X4V3 = ModelSpec(
key="REALESRGAN_GENERAL_WDN_X4V3",
label="Real-ESRGAN General WDN x4v3",
label="Fast | Real-ESRGAN General WDN x4v3",
directory_name="realesr-general-wdn-x4v3",
repository="xinntao/Real-ESRGAN",
license_name="BSD 3-Clause",
@@ -52,26 +69,10 @@ REALESRGAN_GENERAL_WDN_X4V3 = ModelSpec(
),),
)
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 2x output from a real-world 4x ONNX model",
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",
label="General | Real-ESRGAN x4plus",
directory_name="realesrgan-x4plus",
repository="jonathanst29/tinier-upscale-models",
license_name="BSD 3-Clause",
@@ -88,25 +89,21 @@ REALESRGAN_X4PLUS = ModelSpec(
)
PISA_SR = ModelSpec(
key="PISA_SR",
label="PiSA-SR",
directory_name="pisa-sr",
repository="csslc/PiSA-SR",
license_name="Apache 2.0 / CreativeML Open RAIL++-M",
description="Generative 2x upscale with fixed detail strength 0.7",
download_patterns=("encoder.onnx", "decoder.onnx", "unet-pixel.onnx", "unet-full.onnx", "noise.npy", "LICENSE-APACHE.txt", "LICENSE-MODEL.txt", "NOTICE.txt"),
ready_patterns=("encoder.onnx", "decoder.onnx", "unet-pixel.onnx", "unet-full.onnx", "noise.npy", "LICENSE-APACHE.txt", "LICENSE-MODEL.txt", "NOTICE.txt"),
r2_directory="pisa-sr",
HAT_GAN_X4_SHARPER = ModelSpec(
key="HAT_GAN_X4_SHARPER",
label="Detailed | HAT Sharper",
directory_name="hat-gan-x4-sharper",
repository="XPixelGroup/HAT",
license_name="Apache 2.0",
description="Detailed texture enhancement with 2x output; higher processing cost",
download_patterns=("model.onnx", "LICENSE.txt"),
ready_patterns=("model.onnx", "LICENSE.txt"),
r2_directory="hat-gan-x4-sharper",
r2_files=(
("encoder.onnx", 68426712, "4b93e4a3d81c3349a78ab96619501e39ee8f815d14a20b73fc4f969e30a1ceeb"),
("decoder.onnx", 99084589, "c9021384c402e267654907527fa5f0ee86cab1814bfd96233216a15ae8e83c4c"),
("unet-pixel.onnx", 1732713007, "4f381574ad584050d0534c8cca794fb5828cb8e99e148f8efc1aaac20833c0b1"),
("unet-full.onnx", 1732713007, "3b0fd44ff1093608a65865c3f26d0060775337e1fc117d8ac29aacbe748d3f9c"),
("noise.npy", 32896, "a94ed01e2e963022cfa8ec01e04d2e96e599408ef0c385878347eaacd25a663b"),
("LICENSE-APACHE.txt", 11358, "cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30"),
("LICENSE-MODEL.txt", 14206, "e514164be7f2ca2c0fb39e7283b01ab4a06d802837eaf171aa4645cfd7ee1157"),
("NOTICE.txt", 1081, "ccb55fcb98cce42401368b88e3ccd555f4ff49ac1347c6ce063061dd48726050"),
("model.onnx", 45874438, "032d8b03c704220a08d91465a822aca252c1d6e63db20c3cf23c5fb67f3d2366"),
("LICENSE.txt", 11342, "a91d57ebad8955a1757be7891ca2da8e07493e662f0e34e1e1926571e05f37fc"),
),
)
@@ -115,9 +112,8 @@ UPSCALE_MODELS = {
model.key: model
for model in (
REALESRGAN_GENERAL_WDN_X4V3,
SAFMN_MODEL,
REALESRGAN_X4PLUS,
PISA_SR,
HAT_GAN_X4_SHARPER,
)
}
@@ -128,7 +124,7 @@ DEFAULT_UPSCALE_MODEL_KEY = REALESRGAN_GENERAL_WDN_X4V3.key
MOGE2_MODELS = {
"MOGE2_VITS_NORMAL": ModelSpec(
key="MOGE2_VITS_NORMAL",
label="MoGe-2 ViT-S Normal",
label="Fast | MoGe-2 ViT-S Normal",
directory_name="moge-2-vits-normal-onnx",
repository="Ruicheng/moge-2-vits-normal-onnx",
license_name="MIT",
@@ -145,7 +141,7 @@ MOGE2_MODELS = {
),
"MOGE2_VITB_NORMAL": ModelSpec(
key="MOGE2_VITB_NORMAL",
label="MoGe-2 ViT-B Normal",
label="Balanced | MoGe-2 ViT-B Normal",
directory_name="moge-2-vitb-normal-onnx",
repository="Ruicheng/moge-2-vitb-normal-onnx",
license_name="MIT",
@@ -160,27 +156,10 @@ MOGE2_MODELS = {
"bbf14e07a30f11e69d36ab861590123f5598ababcbc8946a063eb4a966f35a21",
),),
),
"MOGE2_VITL_NORMAL": ModelSpec(
key="MOGE2_VITL_NORMAL",
label="MoGe-2 ViT-L Normal",
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",
download_patterns=("model.onnx",),
ready_patterns=("model.onnx",),
r2_directory="moge-2-vitl-normal-onnx",
r2_files=((
"model.onnx",
1324265014,
"afbc4ccc3450298f3afb35b90f015f4c4f552dea21dc6470d5f7b78b77e2d751",
),),
),
}
MOGE3_VITL = ModelSpec(
key="MOGE3_VITL",
label="MoGe-3 ViT-L",
label="Detailed | MoGe-3 ViT-L",
directory_name="moge-3-vitl-onnx",
repository="Ruicheng/moge-3-vitl",
license_name="MIT",
@@ -195,9 +174,33 @@ MOGE3_VITL = ModelSpec(
)
MOGE_MODELS = {**MOGE2_MODELS, MOGE3_VITL.key: MOGE3_VITL}
BIREFNET_HR_MATTING = ModelSpec(
key="BIREFNET_HR_MATTING",
label="Detailed | BiRefNet HR-matting (CPU)",
directory_name="birefnet-hr-matting",
repository="ZhengPeng7/BiRefNet_HR-matting",
license_name="MIT",
description="2048px matting with continuous alpha; CPU only, high memory and processing cost",
download_patterns=("model.onnx", "LICENSE.txt"),
ready_patterns=("model.onnx", "LICENSE.txt"),
r2_directory="birefnet-hr-matting",
r2_files=(
("model.onnx", 529467341, "b3abc34843a4d96f669ca9fb1a2a16b3417ab8b91c987573fdbda36d20514133"),
("LICENSE.txt", 1066, "92a7089e0915fc32bc40067560b398f1e6a7a5958abd7d04eda393629a5acefb"),
),
)
BACKGROUND_MODELS = {model.key: model for model in (BIREFNET_LITE, BEN2_MODEL, BIREFNET_HR_MATTING)}
DEFAULT_GEOMETRY_MODEL_KEY = next(iter(MOGE_MODELS))
DEFAULT_BACKGROUND_MODEL_KEY = next(iter(BACKGROUND_MODELS))
DEFAULT_MODEL_KEYS = (
DEFAULT_GEOMETRY_MODEL_KEY,
DEFAULT_BACKGROUND_MODEL_KEY,
DEFAULT_UPSCALE_MODEL_KEY,
)
DOWNLOADABLE_MODELS = {
BEN2_MODEL.key: BEN2_MODEL,
**BACKGROUND_MODELS,
**UPSCALE_MODELS,
**MOGE_MODELS,
}
@@ -220,8 +223,8 @@ def model_record(model, ready=False):
"family": (
"moge"
if model.key in MOGE_MODELS
else "ben2"
if model.key == BEN2_MODEL.key
else "background"
if model.key in BACKGROUND_MODELS
else "upscale"
),
"ready": bool(ready),

57
src/anyimage/server/model_manager.py
View File

@@ -12,28 +12,30 @@ class ModelSessionCache:
"""Cache ONNX sessions shared by Jobs in one Server process."""
def __init__(self):
self.ben2_model = None
self.ben2_key = None
self.background_model = None
self.background_key = None
self.moge_model = None
self.moge_key = None
self.upscale_model = None
self.upscale_key = None
def get_ben2(self, model_dir, device):
requested_device = getattr(device, "type", device)
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
if self.ben2_model is not None and self.ben2_key == key:
return self.ben2_model, 0.0
self.release_ben2()
return self._load_ben2(model_dir, requested_device, key)
def get_background(self, model_dir, device):
from .models.background import inference_device
def _load_ben2(self, model_dir, requested_device, key):
from .models.onnx_ben2 import create_session
requested_device = inference_device(model_dir, device)
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
if self.background_model is not None and self.background_key == key:
return self.background_model, 0.0
self.release_background()
return self._load_background(model_dir, requested_device, key)
def _load_background(self, model_dir, requested_device, key):
from .models.background import model_adapter
started = time.perf_counter()
model = create_session(model_dir, requested_device)
self.ben2_model = model
self.ben2_key = key
model = model_adapter(model_dir).create_session(model_dir, requested_device)
self.background_model = model
self.background_key = key
elapsed_ms = (time.perf_counter() - started) * 1000.0
return model, elapsed_ms
@@ -56,9 +58,6 @@ class ModelSessionCache:
def get_upscale(self, model_dir, requested_device):
from .models.onnx_upscale import create_session
if Path(model_dir).name == "pisa-sr":
from .models.onnx_pisa_sr import PisaSession as create_session
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
if self.upscale_model is not None and self.upscale_key == key:
return self.upscale_model
@@ -68,9 +67,9 @@ class ModelSessionCache:
self.upscale_key = key
return model
def release_ben2(self):
self.ben2_model = None
self.ben2_key = None
def release_background(self):
self.background_model = None
self.background_key = None
self._release_runtime_cache()
def release_geometry(self):
@@ -79,21 +78,15 @@ class ModelSessionCache:
self._release_runtime_cache()
def release_upscale(self):
close = getattr(self.upscale_model, "close", None)
if close is not None:
close()
self.upscale_model = None
self.upscale_key = None
self._release_runtime_cache()
def close(self):
close = getattr(self.upscale_model, "close", None)
if close is not None:
close()
self.ben2_model = None
self.background_model = None
self.moge_model = None
self.upscale_model = None
self.ben2_key = None
self.background_key = None
self.moge_key = None
self.upscale_key = None
self._release_runtime_cache()
@@ -101,7 +94,7 @@ class ModelSessionCache:
def snapshot(self):
geometry_key = self.moge_key
return {
"ben2": Path(self.ben2_key[0]).name if self.ben2_key else "",
"background": Path(self.background_key[0]).name if self.background_key else "",
"geometry": Path(geometry_key[0]).name if geometry_key else "",
"upscale": (Path(self.upscale_key[0]).name if self.upscale_key else ""),
}
@@ -125,7 +118,7 @@ class ModelManager:
def ready(self, model_key):
model = model_catalog.get_downloadable_model(model_key)
directory = self.directory(model_key)
if model_key in {"PISA_SR", "MOGE3_VITL"}:
if model_key in {"MOGE3_VITL", "HAT_GAN_X4_SHARPER", *model_catalog.BACKGROUND_MODELS}:
for filename, size, checksum in model.r2_files:
path = directory / filename
try:
@@ -148,8 +141,8 @@ class ModelManager:
return True
return all(any(directory.glob(pattern)) for pattern in model.ready_patterns)
def get_ben2(self, model_dir, device):
return self.sessions.get_ben2(model_dir, device)
def get_background(self, model_dir, device):
return self.sessions.get_background(model_dir, device)
def get_moge(self, model_dir, device):

4
src/anyimage/server/models/__init__.py
View File

@@ -1,7 +1,7 @@
from . import ben2, moge, upscale
from . import background, moge, upscale
__all__ = (
"ben2",
"background",
"moge",
"upscale",
)

30
src/anyimage/server/models/ben2.py → src/anyimage/server/models/background.py
View File

@@ -1,10 +1,26 @@
import time
from importlib import import_module
from pathlib import Path
from ..media import check_cancelled
from ..media.input import open_image, validate_input_path
def inference_device(model_dir, requested_device):
if Path(model_dir).name == "birefnet-hr-matting":
return "cpu"
return getattr(requested_device, "type", requested_device)
def model_adapter(model_dir):
name = Path(model_dir).name
if name == "BEN2-ONNX":
return import_module(".onnx_ben2", __package__)
if name in {"birefnet-lite", "birefnet-hr-matting"}:
return import_module(".onnx_birefnet", __package__)
raise ValueError(f"Unknown background model directory: {name}")
def infer_alpha(
image,
model_dir,
@@ -16,16 +32,16 @@ def infer_alpha(
from PIL import Image
check_cancelled(cancel_check)
from .onnx_ben2 import create_session, infer_alpha as predict_alpha
adapter = model_adapter(model_dir)
owns_model = model_cache is None
requested_device = getattr(device, "type", device)
requested_device = inference_device(model_dir, device)
if owns_model:
load_started = time.perf_counter()
model = create_session(model_dir, requested_device)
model = adapter.create_session(model_dir, requested_device)
load_ms = (time.perf_counter() - load_started) * 1000.0
else:
model, load_ms = model_cache.get_ben2(model_dir, device)
model, load_ms = model_cache.get_background(model_dir, requested_device)
check_cancelled(cancel_check)
source_rgba = image.convert("RGBA")
@@ -36,7 +52,7 @@ def infer_alpha(
source_rgba.getchannel("A"),
)
inference_started = time.perf_counter()
alpha = predict_alpha(
alpha = adapter.infer_alpha(
model,
image,
release_memory=release_memory,
@@ -90,7 +106,7 @@ def infer_foreground_frames(
output_prefix="foreground",
preserve_input_alpha=False,
):
"""Run BEN2 on every frame and write numbered transparent PNG files."""
"""Remove each frame's background and write numbered transparent PNG files."""
total_load_ms = 0.0
total_inference_ms = 0.0
for index, image_path in enumerate(images):
@@ -115,7 +131,7 @@ def infer_foreground_frames(
compress_level=1,
)
if not output_path.is_file():
raise RuntimeError(f"BEN2 did not write: {output_path}")
raise RuntimeError(f"Background removal did not write: {output_path}")
total_load_ms += load_ms
total_inference_ms += inference_ms
return total_load_ms, total_inference_ms

41
src/anyimage/server/models/onnx_birefnet.py Normal file
View File

@@ -0,0 +1,41 @@
"""BiRefNet probability alpha with official RGB normalization."""
from pathlib import Path
from .onnx_runtime import create_session as create_runtime_session
from .onnx_runtime import run_session
def create_session(model_directory, requested_device="auto"):
return create_runtime_session(Path(model_directory) / "model.onnx", requested_device)
def preprocess(image, size):
import numpy as np
from PIL import Image
resized = image.convert("RGB").resize((size, size), Image.Resampling.BILINEAR)
array = np.asarray(resized, dtype=np.float32) / 255.0
array = (array - np.array((0.485, 0.456, 0.406), dtype=np.float32)) / np.array(
(0.229, 0.224, 0.225), dtype=np.float32,
)
return np.ascontiguousarray(array.transpose(2, 0, 1)[None])
def infer_alpha(session, image, *, release_memory=True):
import numpy as np
from PIL import Image
input_spec = session.get_inputs()[0]
size = input_spec.shape[-1]
if size not in (1024, 2048) or input_spec.shape != [1, 3, size, size]:
raise RuntimeError("BiRefNet returned an invalid input contract")
alpha = np.asarray(run_session(
session, [session.get_outputs()[0].name],
{input_spec.name: preprocess(image, size)}, release_memory=release_memory,
)[0], dtype=np.float32)
if (alpha.shape != (1, 1, size, size) or not np.isfinite(alpha).all()
or alpha.min() < 0 or alpha.max() > 1):
raise RuntimeError("BiRefNet returned an invalid alpha mask")
mask = Image.fromarray(alpha[0, 0]).resize(image.size, Image.Resampling.BILINEAR)
return np.clip(np.asarray(mask, dtype=np.float32), 0.0, 1.0)

153
src/anyimage/server/models/onnx_pisa_sr.py
View File

@@ -1,153 +0,0 @@
"""Fixed-strength PiSA-SR inference with one resident ONNX component."""
import gc
from pathlib import Path
from tempfile import TemporaryDirectory
import numpy as np
from PIL import Image
from .onnx_runtime import create_session, run_session
from .onnx_upscale import RESOURCE_ERROR_TOKENS, UpscaleResourceError
TILE_SIZE = 512
OVERLAP = 128
LATENT_SCALE = np.float16(0.18215)
SEMANTIC_STRENGTH = np.float16(0.7)
def tile_starts(length):
padded = max(TILE_SIZE, (length + 7) // 8 * 8)
starts = list(range(0, padded - TILE_SIZE + 1, TILE_SIZE - OVERLAP))
if starts[-1] != padded - TILE_SIZE:
starts.append(padded - TILE_SIZE)
return starts, padded
def correct_color(pixels, source):
target = pixels.astype(np.float32) / 255
reference = np.asarray(source, dtype=np.float32) / 255
axes = (0, 1)
deviation = np.sqrt(target.var(axis=axes, ddof=1) + 1e-5)
reference_deviation = np.sqrt(reference.var(axis=axes, ddof=1) + 1e-5)
corrected = (target - target.mean(axis=axes)) / deviation
corrected = corrected * reference_deviation + reference.mean(axis=axes)
return Image.fromarray((corrected.clip(0, 1) * 255).astype(np.uint8))
class PisaSession:
def __init__(self, directory, device):
if device == "coreml":
raise RuntimeError("PiSA-SR CoreML is not validated; select CPU for this model")
self.directory = Path(directory)
self.device = device
self.session = None
self.component = None
self.noise = np.load(self.directory / "noise.npy", allow_pickle=False)
if self.noise.shape != (1, 4, 64, 64) or not np.isfinite(self.noise).all():
raise ValueError("Invalid PiSA-SR noise asset")
def close(self):
self.session = None
self.component = None
gc.collect()
def activate(self, component):
if self.component == component:
return
self.close()
self.session = create_session(self.directory / f"{component}.onnx", self.device)
self.session.disable_fallback()
if "CoreMLExecutionProvider" in self.session.get_providers():
self.close()
raise RuntimeError("PiSA-SR CoreML is not validated; select CPU for this model")
self.component = component
def run(self, value, final=False):
try:
result = run_session(
self.session, ["output"],
{"input": np.ascontiguousarray(value, dtype=np.float16)},
release_memory=final,
)[0]
except Exception as error:
if "CUDAExecutionProvider" in self.session.get_providers():
raise RuntimeError(
"PiSA-SR failed on CUDA; select DirectML or CPU. " + str(error)
) from error
raise
if not np.isfinite(result).all():
raise RuntimeError("PiSA-SR returned non-finite pixels on the selected device")
return result
def infer(model, image, cancel_check=None, *, release_memory=True):
def check():
if cancel_check is not None and cancel_check():
raise RuntimeError("Task cancelled")
source = image.convert("RGB").resize(
(image.width * 2, image.height * 2), Image.Resampling.BICUBIC,
)
rows, height = tile_starts(source.height)
columns, width = tile_starts(source.width)
coordinates = [(y, x) for y in rows for x in columns]
array = np.asarray(source)
array = np.pad(array, ((0, height-source.height), (0, width-source.width), (0, 0)),
mode="reflect" if min(source.size) > 1 else "edge")
try:
with TemporaryDirectory(prefix="anyimage-pisa-") as temporary:
directory = Path(temporary)
check()
model.activate("encoder")
for index, (y, x) in enumerate(coordinates):
check()
tile = array[y:y+TILE_SIZE, x:x+TILE_SIZE].astype(np.float32) / 255
tensor = (tile.transpose(2, 0, 1)[None] * 2 - 1).astype(np.float16)
mean, logvar = np.split(model.run(tensor, index == len(coordinates)-1), 2, axis=1)
noise = np.take(np.take(model.noise, (np.arange(64)+y//8) % 64, axis=2),
(np.arange(64)+x//8) % 64, axis=3)
std = np.exp(logvar.clip(-30, 20) * np.float16(0.5)).astype(np.float16)
latent = ((mean + std * noise) * LATENT_SCALE).astype(np.float16)
np.save(directory / f"latent-{index}.npy", latent)
for component in ("unet-pixel", "unet-full"):
check()
model.activate(component)
for index in range(len(coordinates)):
check()
latent = np.load(directory / f"latent-{index}.npy")
prediction = model.run(latent, index == len(coordinates)-1)
pixel_path = directory / f"pixel-{index}.npy"
if component == "unet-pixel":
np.save(pixel_path, prediction)
else:
pixel = np.load(pixel_path)
corrected = latent - (pixel + SEMANTIC_STRENGTH * (prediction-pixel))
np.save(directory / f"latent-{index}.npy", corrected)
pixel_path.unlink()
check()
model.activate("decoder")
result = np.zeros((height, width, 3), dtype=np.float32)
weights = np.zeros((height, width, 1), dtype=np.float32)
ramp = np.minimum(np.arange(TILE_SIZE)+1, np.arange(TILE_SIZE, 0, -1))
ramp = np.minimum(ramp / OVERLAP, 1).astype(np.float32)
blend = (ramp[:, None] * ramp[None, :])[..., None]
for index, (y, x) in enumerate(coordinates):
check()
latent = np.load(directory / f"latent-{index}.npy")
decoded = model.run(latent, release_memory and index == len(coordinates)-1)
pixels = decoded[0].transpose(1, 2, 0).astype(np.float32)
result[y:y+TILE_SIZE, x:x+TILE_SIZE] += pixels * blend
weights[y:y+TILE_SIZE, x:x+TILE_SIZE] += blend
check()
result = (result / weights)[:source.height, :source.width]
pixels = ((result.clip(-1, 1)*0.5+0.5)*255).clip(0, 255).astype(np.uint8)
return correct_color(pixels, source)
except (MemoryError, UnicodeDecodeError) as error:
model.close()
raise UpscaleResourceError("PiSA-SR exhausted available inference memory") from error
except Exception as error:
model.close()
if any(token in str(error).lower() for token in RESOURCE_ERROR_TOKENS):
raise UpscaleResourceError("PiSA-SR exhausted available inference memory") from error
raise

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

@@ -79,6 +79,8 @@ def _run(session, input_tensor, *, release_memory=True):
raise RuntimeError(
f"Upscale model returned an unexpected output shape: {prediction.shape}"
)
if not np.isfinite(prediction).all():
raise RuntimeError("Upscale model returned non-finite pixels")
return prediction
@@ -89,32 +91,34 @@ def _prediction_pixels(prediction):
return np.rint(pixels * 255.0).astype(np.uint8)
def _pad_source(source, extra_height, extra_width):
def _pad_source(source, extra_height, extra_width, tile_border):
import numpy as np
mode = "reflect" if source.shape[0] > 1 and source.shape[1] > 1 else "edge"
return np.pad(
source,
(
(TILE_BORDER, TILE_BORDER + extra_height),
(TILE_BORDER, TILE_BORDER + extra_width),
(tile_border, tile_border + extra_height),
(tile_border, tile_border + extra_width),
(0, 0),
),
mode=mode,
)
def infer(session, image, cancel_check=None, *, release_memory=True):
def infer(session, image, cancel_check=None, *, release_memory=True, tile_border=TILE_BORDER):
import numpy as np
from PIL import Image
try:
source = np.asarray(image.convert("RGB"), dtype=np.uint8)
height, width = source.shape[:2]
core_size = TILE_SIZE - 2 * TILE_BORDER
if not 0 <= tile_border < TILE_SIZE // 2:
raise ValueError("Invalid upscale tile border")
core_size = TILE_SIZE - 2 * tile_border
extra_height = (-height) % core_size
extra_width = (-width) % core_size
padded = _pad_source(source, extra_height, extra_width)
padded = _pad_source(source, extra_height, extra_width, tile_border)
output = np.empty((height * SCALE, width * SCALE, 3), dtype=np.uint8)
for top in range(0, height, core_size):
core_height = min(core_size, height - top)
@@ -138,7 +142,7 @@ def infer(session, image, cancel_check=None, *, release_memory=True):
release_memory=final_tile,
)
)
border = TILE_BORDER * SCALE
border = tile_border * SCALE
output[
top * SCALE : (top + core_height) * SCALE,
left * SCALE : (left + core_width) * SCALE,

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

@@ -1,5 +1,6 @@
from ..media.input import open_image
from ..media.resolution import DEFAULT_MAX_AI_INPUT_SIZE, limit_image
from ..model_catalog import UPSCALE_MODELS
def _source_alpha(image):
@@ -19,8 +20,8 @@ def infer_one(
from .onnx_upscale import UpscaleResourceError, infer
if parameters["model"] == "PISA_SR":
from .onnx_pisa_sr import infer
if parameters["model"] not in UPSCALE_MODELS:
raise ValueError(f"Unknown upscale model: {parameters['model']}")
model_directory = model_manager.directory(parameters["model"])
model = model_manager.get_upscale(
@@ -39,6 +40,7 @@ def infer_one(
source,
cancel_check=cancel_check,
release_memory=release_memory,
tile_border=48 if parameters["model"] == "HAT_GAN_X4_SHARPER" else 16,
)
prediction = prediction.resize(
(source.width * 2, source.height * 2),