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

feat: Implement upscale functionality with ONNX models

Browse Source 
- 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
Download Patch File Download Diff File Expand all files Collapse all files

94
src/anyimage/server/__init__.py
View File

@@ -1,90 +1,12 @@
"""AnyImage Job Server without Blender dependencies."""
from blendjob import JobServer
from .jobs import generate_depth as _generate_depth
from .jobs import generate_surface as _generate_surface
from .jobs.background import run as _remove_background
from .jobs.balloon import run as _analyze_balloon
from .jobs.debug_ben2 import run as _debug_ben2
from .jobs.debug_da3 import run as _debug_da3
from .jobs.debug_moge2 import run as _debug_moge2
from .jobs.debug_realesrgan import run as _debug_realesrgan
from .jobs.frames import run as _extract_frames
from .jobs.image_mask import run as _refine_image_mask
from .model_manager import ModelManager, ModelSessionCache
"""Lightweight package boundary shared by Blender and the Job Server."""
server = JobServer("AnyImage Job Server")
def __getattr__(name):
if name != "server":
raise AttributeError(name)
from .app import server
return server
@server.resource("model_manager")
def create_model_manager(job_server):
return ModelManager(job_server.storage_root / "models")
@server.job("download-model")
def download_model(context, parameters):
manager = context.resource("model_manager")
return manager.download(context, parameters["model"])
@server.job("download-required-models")
def download_required_models(context, parameters):
manager = context.resource("model_manager")
return manager.download_missing(context, parameters["models"])
@server.job("analyze-balloon")
def analyze_balloon(context, parameters):
return _analyze_balloon(context, parameters)
@server.job("refine-image-mask")
def refine_image_mask(context, parameters):
return _refine_image_mask(context, parameters)
@server.job("remove-background")
def remove_background(context, parameters):
return _remove_background(context, parameters)
@server.job("extract-frames")
def extract_frames(context, parameters):
return _extract_frames(context, parameters)
@server.job("generate-depth")
def generate_depth(context, parameters):
return _generate_depth(context, parameters)
@server.job("generate-surface")
def generate_surface(context, parameters):
return _generate_surface(context, parameters)
@server.job("debug-da3")
def debug_da3(context, parameters):
return _debug_da3(context, parameters)
@server.job("debug-moge2")
def debug_moge2(context, parameters):
return _debug_moge2(context, parameters)
@server.job("debug-ben2")
def debug_ben2(context, parameters):
return _debug_ben2(context, parameters)
@server.job("debug-realesrgan-general-x4v3")
def debug_realesrgan(context, parameters):
return _debug_realesrgan(context, parameters)
__all__ = ("ModelManager", "ModelSessionCache", "server")
__all__ = ("server",)

92
src/anyimage/server/app.py Normal file
View File

@@ -0,0 +1,92 @@
"""AnyImage Job Server without Blender dependencies."""
from blendjob import JobServer
from .jobs import generate_depth as _generate_depth
from .jobs import generate_surface as _generate_surface
from .jobs.background import run as _remove_background
from .jobs.balloon import run as _analyze_balloon
from .jobs.debug_ben2 import run as _debug_ben2
from .jobs.debug_da3 import run as _debug_da3
from .jobs.debug_moge2 import run as _debug_moge2
from .jobs.debug_upscale import run as _debug_upscale
from .jobs.frames import run as _extract_frames
from .jobs.image_mask import run as _refine_image_mask
from .jobs.upscale import run as _upscale_image
from .model_manager import ModelManager
server = JobServer("AnyImage Job Server")
@server.resource("model_manager")
def create_model_manager(job_server):
return ModelManager(job_server.storage_root / "models")
@server.job("download-model")
def download_model(context, parameters):
manager = context.resource("model_manager")
return manager.download(context, parameters["model"])
@server.job("download-required-models")
def download_required_models(context, parameters):
manager = context.resource("model_manager")
return manager.download_missing(context, parameters["models"])
@server.job("analyze-balloon")
def analyze_balloon(context, parameters):
return _analyze_balloon(context, parameters)
@server.job("refine-image-mask")
def refine_image_mask(context, parameters):
return _refine_image_mask(context, parameters)
@server.job("remove-background")
def remove_background(context, parameters):
return _remove_background(context, parameters)
@server.job("upscale-image")
def upscale_image(context, parameters):
return _upscale_image(context, parameters)
@server.job("extract-frames")
def extract_frames(context, parameters):
return _extract_frames(context, parameters)
@server.job("generate-depth")
def generate_depth(context, parameters):
return _generate_depth(context, parameters)
@server.job("generate-surface")
def generate_surface(context, parameters):
return _generate_surface(context, parameters)
@server.job("debug-da3")
def debug_da3(context, parameters):
return _debug_da3(context, parameters)
@server.job("debug-moge2")
def debug_moge2(context, parameters):
return _debug_moge2(context, parameters)
@server.job("debug-ben2")
def debug_ben2(context, parameters):
return _debug_ben2(context, parameters)
@server.job("debug-upscale")
def debug_upscale(context, parameters):
return _debug_upscale(context, parameters)

8
src/anyimage/server/jobs/background.py
View File

@@ -53,7 +53,7 @@ def run(context, parameters):
)
context.check_cancelled()
context.progress(0.92, "Writing transparent PNG")
output_path = output_directory / "color.0001.png"
output_path = output_directory / "image.0001.png"
foreground.convert("RGBA").save(
output_path,
format="PNG",
@@ -72,7 +72,7 @@ def run(context, parameters):
0.7,
model_cache=model_manager,
cancel_check=context.check_cancelled,
output_prefix="color",
output_prefix="image",
preserve_input_alpha=True,
)
finally:
@@ -86,6 +86,6 @@ def run(context, parameters):
f"Background removed{frame_text} in {inference_ms:.0f} ms "
f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)"
)
print(message)
context.log(message)
context.progress(1.0, message)
return {"color": "."}
return {"image": "."}

3
src/anyimage/server/jobs/debug_realesrgan.py → src/anyimage/server/jobs/debug_upscale.py
View File

@@ -1,5 +1,5 @@
from ..media import copy_frames
from ..models.realesrgan import infer_one
from ..models.upscale import infer_one
from .common import input_frames
@@ -13,6 +13,7 @@ def run(context, parameters):
context.resource("model_manager"),
parameters,
image,
cancel_check=context.check_cancelled,
)
suffix = "" if len(images) == 1 else f".{index + 1:04d}"
prediction.save(

8
src/anyimage/server/jobs/foreground.py
View File

@@ -25,7 +25,7 @@ def write_optional(
cancel_check=context.check_cancelled,
preserve_input_alpha=True,
)
path = context.directory / "color.0001.png"
path = context.directory / "image.0001.png"
foreground.convert("RGBA").save(
path,
format="PNG",
@@ -43,11 +43,11 @@ def write_optional(
progress_span,
model_cache=model_manager,
cancel_check=context.check_cancelled,
output_prefix="color",
output_prefix="image",
preserve_input_alpha=True,
)
paths = sorted(context.directory.glob("color.*.png"))
print(
paths = sorted(context.directory.glob("image.*.png"))
context.log(
f"BEN2 inference {inference_ms:.0f} ms; "
f"model load {load_ms:.0f} ms"
)

34
src/anyimage/server/jobs/upscale.py Normal file
View File

@@ -0,0 +1,34 @@
from ..models.upscale import infer_one
from .common import frame_progress_message, input_frames
def run(context, parameters):
context.log(
f"Upscale model: {parameters['model']}; "
f"maximum output size: {parameters.get('max_output_size', 4096)} px"
)
context.progress(0.02, "Preparing upscale input")
context.directory.mkdir(parents=True, exist_ok=True)
with input_frames(context, parameters) as images:
total = len(images)
for index, image in enumerate(images):
context.check_cancelled()
context.progress(
0.08 + index / total * 0.84,
frame_progress_message("Upscaling image", index, total),
)
prediction = infer_one(
context.resource("model_manager"),
parameters,
image,
cancel_check=context.check_cancelled,
preserve_alpha=True,
)
prediction.save(
context.directory / f"upscale.{index + 1:04d}.png",
format="PNG",
compress_level=1,
)
context.check_cancelled()
context.progress(1.0, "Image upscale complete")
return {"upscale": "."}

4
src/anyimage/server/media/input.py
View File

@@ -144,11 +144,11 @@ def prepare_background_input(
def copy_color_images(images, output_directory):
"""Copy input images as the canonical numbered color layer."""
"""Copy input images as the canonical numbered image layer."""
for index, image_path in enumerate(images):
suffix = Path(image_path).suffix.lower() or ".jpg"
destination = (
Path(output_directory) / f"color.{index + 1:04d}{suffix}"
Path(output_directory) / f"image.{index + 1:04d}{suffix}"
)
if Path(image_path).resolve() == destination.resolve():
continue

76
src/anyimage/server/model_catalog.py
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",),

14
src/anyimage/server/model_download.py
View File

@@ -171,15 +171,15 @@ def download_model(args):
def _download_declared_or_fail(args, model_dir, label, start, span):
"""Download declared files from Hugging Face, then the R2 mirror."""
sources = (
(HF_ENDPOINT, download_from_hf, f"Downloading {label}"),
(
"""Download declared files from Hugging Face when available, then R2."""
sources = []
if args.repository:
sources.append((HF_ENDPOINT, download_from_hf, f"Downloading {label}"))
sources.append((
R2_ENDPOINT,
download_from_r2,
f"Downloading {label} (mirror)",
),
)
))
failures = []
for source, downloader, message in sources:
args.progress(start, message)
@@ -431,6 +431,8 @@ def download_from_hf(
)
def format_rate(rate):
size = float(rate)
if size <= 0:

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

@@ -17,15 +17,15 @@ class ModelSessionCache:
self.da3_key = None
self.moge_model = None
self.moge_key = None
self.realesrgan_model = None
self.realesrgan_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()
self.close()
return self._load_ben2(model_dir, requested_device, key)
def _load_ben2(self, model_dir, requested_device, key):
@@ -41,11 +41,9 @@ class ModelSessionCache:
def get_da3(self, model_dir, device):
requested_device = getattr(device, "type", device)
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
if self.moge_model is not None:
self.release_moge()
if self.da3_model is not None and self.da3_key == key:
return self.da3_model
self.release_da3()
self.close()
return self._load_da3(model_dir, requested_device, key)
def _load_da3(self, model_dir, requested_device, key):
@@ -60,26 +58,24 @@ class ModelSessionCache:
from .models.onnx_moge2 import create_session
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
if self.da3_model is not None:
self.release_da3()
if self.moge_model is not None and self.moge_key == key:
return self.moge_model
self.release_moge()
self.close()
model = create_session(model_dir, requested_device)
self.moge_model = model
self.moge_key = key
return model
def get_realesrgan(self, model_dir, requested_device):
from .models.onnx_realesrgan import create_session
def get_upscale(self, model_dir, requested_device):
from .models.onnx_upscale import create_session
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
if self.realesrgan_model is not None and self.realesrgan_key == key:
return self.realesrgan_model
self.release_realesrgan()
if self.upscale_model is not None and self.upscale_key == key:
return self.upscale_model
self.close()
model = create_session(model_dir, requested_device)
self.realesrgan_model = model
self.realesrgan_key = key
self.upscale_model = model
self.upscale_key = key
return model
def release_ben2(self):
@@ -97,20 +93,20 @@ class ModelSessionCache:
self.moge_key = None
self._release_runtime_cache()
def release_realesrgan(self):
self.realesrgan_model = None
self.realesrgan_key = None
def release_upscale(self):
self.upscale_model = None
self.upscale_key = None
self._release_runtime_cache()
def close(self):
self.ben2_model = None
self.da3_model = None
self.moge_model = None
self.realesrgan_model = None
self.upscale_model = None
self.ben2_key = None
self.da3_key = None
self.moge_key = None
self.realesrgan_key = None
self.upscale_key = None
self._release_runtime_cache()
def snapshot(self):
@@ -118,9 +114,9 @@ class ModelSessionCache:
"ben2": Path(self.ben2_key[0]).name if self.ben2_key else "",
"da3": Path(self.da3_key[0]).name if self.da3_key else "",
"moge2": Path(self.moge_key[0]).name if self.moge_key else "",
"realesrgan": (
Path(self.realesrgan_key[0]).name
if self.realesrgan_key
"upscale": (
Path(self.upscale_key[0]).name
if self.upscale_key
else ""
),
}
@@ -138,22 +134,7 @@ class ModelManager:
def directory(self, model_key):
model = model_catalog.get_downloadable_model(model_key)
directory = self.models_directory / model.directory_name
old_moge_directory = self.models_directory / "moge2_base_normal"
if (
model_key == "MOGE2_VITB_NORMAL"
and not directory.is_dir()
and old_moge_directory.is_dir()
):
old_moge_directory.rename(directory)
legacy_directory = self.models_directory / model_key.lower()
if (
not directory.is_dir()
and legacy_directory.is_dir()
and legacy_directory != directory
):
legacy_directory.rename(directory)
return directory
return self.models_directory / model.directory_name
def ready(self, model_key):
model = model_catalog.get_downloadable_model(model_key)
@@ -169,8 +150,11 @@ class ModelManager:
def get_moge2(self, model_dir, device):
return self.sessions.get_moge2(model_dir, device)
def get_realesrgan(self, model_dir, device):
return self.sessions.get_realesrgan(model_dir, device)
def get_upscale(self, model_dir, device):
return self.sessions.get_upscale(model_dir, device)
def release_upscale(self):
self.sessions.release_upscale()
def snapshot(self):
return {"loaded": self.sessions.snapshot()}
@@ -201,7 +185,7 @@ class ModelManager:
def _download(self, context, model_key, progress):
model = model_catalog.get_downloadable_model(model_key)
args = SimpleNamespace(
repository=model.repository,
repository=model.huggingface_repository,
model_dir=str(self.directory(model_key)),
progress=progress,
cancel_check=context.check_cancelled,

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

@@ -1,4 +1,4 @@
from . import ben2, da3, moge2, realesrgan
from . import ben2, da3, moge2, upscale
def geometry_adapter(model_key):
@@ -9,6 +9,6 @@ __all__ = (
"ben2",
"da3",
"moge2",
"realesrgan",
"upscale",
"geometry_adapter",
)

51
src/anyimage/server/models/onnx_realesrgan.py
View File

@@ -1,51 +0,0 @@
"""Real-ESRGAN General x4v3 ONNX inference."""
from pathlib import Path
try:
from .onnx_runtime import create_session as create_runtime_session
except ImportError:
from onnx_runtime import create_session as create_runtime_session
ONNX_FILENAME = "realesr-general-x4v3.onnx"
SCALE = 4
def create_session(model_directory, requested_device="auto"):
path = Path(model_directory) / ONNX_FILENAME
return create_runtime_session(path, requested_device)
def preprocess(image):
import numpy as np
image_array = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0
return np.ascontiguousarray(image_array.transpose(2, 0, 1)[None])
def infer(session, image):
import numpy as np
from PIL import Image
input_tensor = preprocess(image)
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
prediction = np.asarray(
session.run([output_name], {input_name: input_tensor})[0],
dtype=np.float32,
)
expected_shape = (
1,
3,
input_tensor.shape[2] * SCALE,
input_tensor.shape[3] * SCALE,
)
if prediction.shape != expected_shape:
raise RuntimeError(
"Real-ESRGAN General x4v3 returned an unexpected output shape: "
f"{prediction.shape}"
)
pixels = np.clip(prediction[0].transpose(1, 2, 0), 0.0, 1.0)
pixels = np.rint(pixels * 255.0).astype(np.uint8)
return Image.fromarray(pixels, mode="RGB")

142
src/anyimage/server/models/onnx_upscale.py Normal file
View File

@@ -0,0 +1,142 @@
"""Tiled ONNX inference shared by the Upscale models."""
from pathlib import Path
try:
from .onnx_runtime import create_session as create_runtime_session
except ImportError:
from onnx_runtime import create_session as create_runtime_session
SCALE = 4
TILE_SIZE = 256
TILE_BORDER = 16
RESOURCE_ERROR_TOKENS = (
"out of memory",
"bad allocation",
"bad_alloc",
"failed to allocate",
"not enough memory",
"insufficient memory",
)
class UpscaleResourceError(RuntimeError):
"""The execution provider could not allocate resources for one tile."""
def create_session(model_directory, requested_device="auto"):
paths = tuple(Path(model_directory).glob("*.onnx"))
if len(paths) != 1:
raise RuntimeError(
f"Expected exactly one ONNX model in {model_directory}, found {len(paths)}"
)
path = paths[0]
return create_runtime_session(path, requested_device)
def preprocess(image):
import numpy as np
image_array = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0
return np.ascontiguousarray(image_array.transpose(2, 0, 1)[None])
def _run(session, input_tensor):
import numpy as np
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
try:
raw_prediction = session.run(
[output_name],
{input_name: input_tensor},
)[0]
except (MemoryError, UnicodeDecodeError) as error:
raise UpscaleResourceError(
"the ONNX execution provider failed, usually because GPU or "
"system memory was exhausted"
) from error
except RuntimeError as error:
if any(token in str(error).lower() for token in RESOURCE_ERROR_TOKENS):
raise UpscaleResourceError(
"the ONNX execution provider failed because GPU or system "
"memory was exhausted"
) from error
raise
prediction = np.asarray(raw_prediction, dtype=np.float32)
expected_shape = (
1,
3,
input_tensor.shape[2] * SCALE,
input_tensor.shape[3] * SCALE,
)
if prediction.shape != expected_shape:
raise RuntimeError(
f"Upscale model returned an unexpected output shape: {prediction.shape}"
)
return prediction
def _prediction_pixels(prediction):
import numpy as np
pixels = np.clip(prediction[0].transpose(1, 2, 0), 0.0, 1.0)
return np.rint(pixels * 255.0).astype(np.uint8)
def _pad_source(source, extra_height, extra_width):
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),
(0, 0),
),
mode=mode,
)
def infer(session, image, cancel_check=None):
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
extra_height = (-height) % core_size
extra_width = (-width) % core_size
padded = _pad_source(source, extra_height, extra_width)
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)
for left in range(0, width, core_size):
core_width = min(core_size, width - left)
if cancel_check is not None and cancel_check():
raise RuntimeError("Task cancelled")
tile = Image.fromarray(
padded[top : top + TILE_SIZE, left : left + TILE_SIZE],
mode="RGB",
)
pixels = _prediction_pixels(_run(session, preprocess(tile)))
border = TILE_BORDER * SCALE
output[
top * SCALE : (top + core_height) * SCALE,
left * SCALE : (left + core_width) * SCALE,
] = pixels[
border : border + core_height * SCALE,
border : border + core_width * SCALE,
]
return Image.fromarray(output, mode="RGB")
except UpscaleResourceError:
raise
except MemoryError as error:
raise UpscaleResourceError(
"the complete upscale result exceeded available system memory"
) from error

12
src/anyimage/server/models/realesrgan.py
View File

@@ -1,12 +0,0 @@
def infer_one(model_manager, parameters, image_path):
from PIL import Image
from .onnx_realesrgan import infer
model_directory = model_manager.directory(parameters["model"])
model = model_manager.get_realesrgan(
model_directory,
parameters["device"],
)
with Image.open(image_path) as opened:
return infer(model, opened)

76
src/anyimage/server/models/upscale.py Normal file
View File

@@ -0,0 +1,76 @@
import math
SCALE = 4
MAX_OUTPUT_SIZE = 8192
MAX_OUTPUT_PIXELS = 32 * 1024 * 1024
def _limited_input(image, max_output_size):
from PIL import Image
max_output_size = int(max_output_size)
if max_output_size < 1 or max_output_size > MAX_OUTPUT_SIZE:
raise ValueError(
f"Max Output Size must be between 1 and {MAX_OUTPUT_SIZE} pixels"
)
width, height = image.size
scale = min(
1.0,
max_output_size / (max(width, height) * SCALE),
math.sqrt(MAX_OUTPUT_PIXELS / (width * height * SCALE * SCALE)),
)
if scale >= 1.0:
return image
size = (
max(1, int(width * scale)),
max(1, int(height * scale)),
)
return image.resize(size, Image.Resampling.LANCZOS)
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,
):
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 = opened
if "max_output_size" in parameters:
source = _limited_input(source, parameters["max_output_size"])
source_alpha = _source_alpha(source) if preserve_alpha else None
prediction = infer(model, source, cancel_check=cancel_check)
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 Max Output Size or select CPU, then try again."
) from error