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

11
src/anyimage/operators/__init__.py
View File

@@ -4,12 +4,7 @@ from .balloon import (
ExtractBalloon,
AnalyzeBalloon,
)
from .debug import (
DebugBen2,
DebugDa3,
DebugMoge2,
DebugRealesrganGeneralX4v3,
)
from .debug import DebugBen2, DebugDa3, DebugMoge2, DebugUpscale
from .image_edit import CLASSES as IMAGE_EDIT_CLASSES
from .depth_plane import (
ConvertToDepthPlane,
@@ -27,6 +22,7 @@ from .models import (
)
from .mesh_plane import ConvertToMeshPlane
from .surface import ConvertToSurface, GenerateDa3Surface, GenerateMoge2Surface
from .upscale import UpscaleImage
CLASSES = (
@@ -37,7 +33,7 @@ CLASSES = (
DebugDa3,
DebugMoge2,
DebugBen2,
DebugRealesrganGeneralX4v3,
DebugUpscale,
GenerateDa3Depth,
GenerateMoge2Depth,
GenerateDa3Surface,
@@ -51,6 +47,7 @@ CLASSES = (
AnalyzeBalloon,
ActivateBalloonTool,
RemoveImageBackground,
UpscaleImage,
ConvertToMeshPlane,
ConvertToDepthPlane,
ConvertToSurface,

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

@@ -79,7 +79,14 @@ def apply_background_response(result, source_object_name):
source_object = bpy.data.objects.get(source_object_name)
if not is_image_empty(source_object):
raise RuntimeError("The source image empty is no longer available")
create_color_image(source_object, result.directory)
source_image = source_object.data
source_name = source_image.name
original_fake_user = source_image.use_fake_user
result_image = create_color_image(source_object, result.directory)
source_image.use_fake_user = original_fake_user
if not original_fake_user and int(source_image.users) == 0:
bpy.data.images.remove(source_image)
result_image.name = source_name
push_undo("Remove Background")

46
src/anyimage/operators/balloon/operator.py
View File

@@ -151,24 +151,6 @@ def data_name_matches(name, base_name):
return suffix != name and suffix.isdigit()
def material_uses_balloon_texture(material, color_name):
node_tree = getattr(material, "node_tree", None)
if node_tree is None:
return False
has_texture_layer = any(
getattr(node, "type", None) == "GROUP"
and getattr(getattr(node, "node_tree", None), "name", None)
== TEXTURE_MATERIAL_NODE_GROUP_NAME
for node in node_tree.nodes
)
has_color_image = any(
getattr(node, "type", None) == "TEX_IMAGE"
and getattr(node, "image", None) is not None
and data_name_matches(node.image.name, color_name)
for node in node_tree.nodes
)
return has_texture_layer and has_color_image
def marked_balloon_material(source_image, normal_required=False):
return next(
@@ -194,30 +176,8 @@ def mark_balloon_material(material, source_image, normal_image=None):
material[BALLOON_NORMAL_IMAGE_PROPERTY] = normal_image
def reusable_balloon_material(
material_name,
color_name,
source_image=None,
normal_required=False,
):
marked = (
marked_balloon_material(source_image, normal_required)
if source_image is not None
else None
)
if marked is not None:
return marked
if normal_required:
return None
return next(
(
material
for material in bpy.data.materials
if data_name_matches(material.name, material_name)
and material_uses_balloon_texture(material, color_name)
),
None,
)
def reusable_balloon_material(source_image, normal_required=False):
return marked_balloon_material(source_image, normal_required)
def load_balloon_normal_image(path, source_object):
@@ -331,8 +291,6 @@ def extract_balloon_object(
uses_normal = normal_path is not None
material = (
reusable_balloon_material(
material_name,
color_name,
image,
normal_required=uses_normal,
)

5
src/anyimage/operators/common.py
View File

@@ -4,6 +4,7 @@ from pathlib import Path
import bpy
from .image import extract_frame_range
from ..preferences import ADDON_PACKAGE, addon_preferences
from ..runtime import runtime
from ..properties import ai_ready, model_ready, model_record
@@ -13,7 +14,6 @@ BEN2_VIDEO_SUFFIXES = {
".mp4", ".mov", ".avi", ".mkv", ".webm", ".m4v", ".mts", ".mpg", ".mpeg"
}
HIDDEN = {"HIDDEN", "SKIP_SAVE"}
ADDON_PACKAGE = __package__.partition(".")[0]
def input_path(value):
@@ -40,8 +40,7 @@ def source_start_frame(source_object_name):
def production_device():
addon = bpy.context.preferences.addons.get(ADDON_PACKAGE)
return getattr(getattr(addon, "preferences", None), "device", "AUTO")
return getattr(addon_preferences(), "device", "AUTO")
def require_environment():

27
src/anyimage/operators/debug.py
View File

@@ -1,10 +1,7 @@
import bpy
from ..runtime import JobOperatorBase
from ..properties import (
BEN2_MODEL_KEY,
REALESRGAN_MODEL_KEY,
)
from ..properties import BEN2_MODEL_KEY, REALESRGAN_GENERAL_WDN_MODEL_KEY
from .common import (
background_input_path,
da3_parameters,
@@ -94,19 +91,24 @@ class DebugBen2(JobOperatorBase, bpy.types.Operator):
return f"Removed background to {result.directory}"
class DebugRealesrganGeneralX4v3(JobOperatorBase, bpy.types.Operator):
bl_idname = "anyimage.debug_realesrgan_general_x4v3"
bl_label = "Debug Real-ESRGAN General x4v3"
job_type = "debug-realesrgan-general-x4v3"
starting_message = "Starting Real-ESRGAN General x4v3 inference"
class DebugUpscale(JobOperatorBase, bpy.types.Operator):
bl_idname = "anyimage.debug_upscale"
bl_label = "Debug Upscale"
job_type = "debug-upscale"
starting_message = "Starting Upscale inference"
model: bpy.props.StringProperty(
default=REALESRGAN_GENERAL_WDN_MODEL_KEY,
options={"HIDDEN", "SKIP_SAVE"},
)
def request(self, context):
require_environment()
settings = context.scene.anyimage_settings
require_model(REALESRGAN_MODEL_KEY)
require_model(self.model)
return {
"input": str(input_path(settings.input_path)),
"model": REALESRGAN_MODEL_KEY,
"model": self.model,
"device": settings.device.lower(),
"video_frames": settings.video_frames,
"start_frame": 1,
@@ -114,5 +116,4 @@ class DebugRealesrganGeneralX4v3(JobOperatorBase, bpy.types.Operator):
def response(self, _context, result):
open_job_directory(result.directory)
return f"Saved Real-ESRGAN API prediction to {result.directory}"
return f"Saved upscale API prediction to {result.directory}"

4
src/anyimage/operators/depth_plane/build.py
View File

@@ -267,10 +267,10 @@ def build_image_depth_plane(
frames_image = create_sequence_image(
depth_data_name(source_object),
depth_frames_directory,
frame_prefix="color.",
frame_prefix="image.",
)
frame_count = len(
sorted_frame_paths(depth_frames_directory, "color.")
sorted_frame_paths(depth_frames_directory, "image.")
)
(
source_frame_start,

11
src/anyimage/operators/depth_plane/data.py
View File

@@ -9,7 +9,7 @@ INTRINSICS_KEY_CANDIDATES = ("intrinsics", "ixts")
def create_depth_sequence_image(
name,
depth_frames,
scene,
_scene,
frame_directory,
frame_prefix="depth",
):
@@ -27,11 +27,10 @@ def create_depth_sequence_image(
float_buffer=True,
)
image.colorspace_settings.name = "Non-Color"
image_settings = scene.render.image_settings
original_file_format = image_settings.file_format
original_file_format = image.file_format
frame_paths = []
try:
image_settings.file_format = "OPEN_EXR"
image.file_format = "OPEN_EXR"
for frame_index, depth_values in enumerate(depth_frames):
# NPZ arrays use a top-left origin, while Blender image pixels start below.
depth_pixels = np.flipud(depth_values)
@@ -44,12 +43,12 @@ def create_depth_sequence_image(
frame_path = (
frame_directory / f"{frame_prefix}.{frame_index + 1:04d}.exr"
)
image.save_render(str(frame_path), scene=scene)
image.save(filepath=str(frame_path), save_copy=True)
if not frame_path.is_file():
raise RuntimeError(f"Blender did not write: {frame_path}")
frame_paths.append(frame_path)
finally:
image_settings.file_format = original_file_format
image.file_format = original_file_format
bpy.data.images.remove(image, do_unlink=True)
if not frame_paths:
raise RuntimeError("Unable to write the depth frame sequence")

26
src/anyimage/operators/image.py
View File

@@ -155,8 +155,8 @@ def copy_depth_image(source_object):
return depth_image
def color_data_name(source_object):
return f"{image_base_name(source_object.data)}.color"
def color_data_name(source_object, suffix=".color"):
return f"{image_base_name(source_object.data)}{suffix}"
def material_data_name(source_object):
@@ -204,18 +204,17 @@ def draw_process_resolution(layout, owner, image):
)
def save_packed_image(image, scene):
def save_packed_image(image, _scene):
output_path = Path(tempfile.mkdtemp(prefix="anyimage-input-")) / "input.png"
image_settings = scene.render.image_settings
original_file_format = image_settings.file_format
original_file_format = image.file_format
try:
image_settings.file_format = "PNG"
image.save_render(str(output_path), scene=scene)
image.file_format = "PNG"
image.save(filepath=str(output_path), save_copy=True)
except (OSError, RuntimeError) as error:
output_path.unlink(missing_ok=True)
raise RuntimeError("Unable to export the packed image") from error
finally:
image_settings.file_format = original_file_format
image.file_format = original_file_format
if not output_path.is_file():
raise RuntimeError("Blender did not export the packed image")
return output_path
@@ -377,14 +376,17 @@ def create_color_image(
source_object,
output_directory,
assign_to_source=True,
*,
frame_prefix="image.",
):
output_directory = Path(output_directory)
color_paths = sorted_frame_paths(output_directory, "color.")
color_paths = sorted_frame_paths(output_directory, frame_prefix)
if not color_paths:
raise RuntimeError(
f"Color result has no images: {output_directory}"
)
source_image = source_object.data
source_color_space = source_image.colorspace_settings.name
animated = is_animated_image(source_image)
try:
result_image = bpy.data.images.load(
@@ -392,16 +394,17 @@ def create_color_image(
check_existing=False,
)
result_image.alpha_mode = "CHANNEL_PACKED"
result_image.colorspace_settings.name = source_color_space
if animated:
force_sequence_source(result_image)
except (OSError, RuntimeError) as error:
except (OSError, RuntimeError, TypeError, ValueError) as error:
if "result_image" in locals():
bpy.data.images.remove(result_image, do_unlink=True)
raise RuntimeError("Unable to load the color result") from error
original_fake_user = source_image.use_fake_user
try:
result_image.name = color_data_name(source_object)
result_name = color_data_name(source_object)
if animated:
pack_sequence_if_possible(result_image, len(color_paths))
else:
@@ -427,6 +430,7 @@ def create_color_image(
source_image.use_fake_user = original_fake_user
bpy.data.images.remove(result_image, do_unlink=True)
raise RuntimeError("Unable to prepare the color result") from error
result_image.name = result_name
return result_image

55
src/anyimage/operators/image_edit/common.py
View File

@@ -13,6 +13,15 @@ from ..common import production_device, require_environment, require_model
from ..image import image_base_name, is_image_empty, plane_bounds
class ImageEditWarning(ValueError):
pass
def report_image_edit_exception(operator, error):
level = "WARNING" if isinstance(error, ImageEditWarning) else "ERROR"
operator.report({level}, str(error))
def image_edit_poll(context):
return context.area is not None and context.area.type == "VIEW_3D"
@@ -94,7 +103,7 @@ def crop_pixels_to_alpha(pixels, bounds):
)
visible_y, visible_x = np.nonzero(rgba[:, :, 3] > 0.0)
if not len(visible_x):
raise ValueError("The edit contains no visible image pixels")
raise ImageEditWarning("The edit contains no visible image pixels")
content_left = int(visible_x.min())
content_top = int(visible_y.min())
content_right = int(visible_x.max()) + 1
@@ -108,10 +117,10 @@ def crop_pixels_to_alpha(pixels, bounds):
)
def align_empty_display_to_bounds(
def empty_display_alignment(
source_object, source_size, output_size, placement_bounds
):
"""Place a new image canvas without changing the Empty transform."""
"""Calculate image canvas placement without changing the Empty transform."""
source_width, source_height = (float(value) for value in source_size)
output_width, output_height = (float(value) for value in output_size)
left, top, right, bottom = (float(value) for value in placement_bounds)
@@ -134,21 +143,36 @@ def align_empty_display_to_bounds(
display_height = display_size
center_x = (desired_x_min + desired_x_max) * 0.5
center_y = (desired_y_min + desired_y_max) * 0.5
source_object.empty_display_size = display_size
source_object.empty_image_offset = (
return display_size, (
center_x / display_width - 0.5,
center_y / display_height - 0.5,
)
def write_image_pixels(image, pixels):
image.alpha_mode = "STRAIGHT"
image.pixels.foreach_set(pixels)
image.pack()
image.update()
def replace_empty_image_pixels(
source_object, pixels, output_size, placement_bounds, suffix=".edit"
source_object, pixels, output_size, placement_bounds, suffix=""
):
source_image = source_object.data
source_size = tuple(source_image.size)
output_width, output_height = (int(value) for value in output_size)
if output_width <= 0 or output_height <= 0:
raise ValueError("The edited image dimensions must be positive")
if len(pixels) != output_width * output_height * 4:
raise ValueError("The edited pixels do not match the image dimensions")
display_size, image_offset = empty_display_alignment(
source_object,
source_size,
(output_width, output_height),
placement_bounds,
)
image = bpy.data.images.new(
f"{image_base_name(source_image)}{suffix}",
width=output_width,
@@ -157,20 +181,17 @@ def replace_empty_image_pixels(
)
try:
image.colorspace_settings.name = source_image.colorspace_settings.name
image.alpha_mode = "STRAIGHT"
image.pixels.foreach_set(pixels)
image.pack()
image.update()
align_empty_display_to_bounds(
source_object,
source_size,
(output_width, output_height),
placement_bounds,
)
source_object.data = image
write_image_pixels(image, pixels)
except Exception:
bpy.data.images.remove(image, do_unlink=True)
raise
source_name = source_image.name
source_object.data = image
if not source_image.use_fake_user and int(source_image.users) == 0:
bpy.data.images.remove(source_image)
image.name = source_name
source_object.empty_display_size = display_size
source_object.empty_image_offset = image_offset
return image

18
src/anyimage/operators/image_edit/mask.py
View File

@@ -30,10 +30,12 @@ from .common import (
draw_dashed_line,
draw_image_mask_operator_settings,
draw_polygon_fill,
ImageEditWarning,
image_edit_poll,
image_mask_refine_request,
polygon_bounds,
replace_empty_image_pixels,
report_image_edit_exception,
resolve_image_edit_click,
)
@@ -141,7 +143,6 @@ def edit_image_object(source_object, selection_mask, invert=False):
pixels,
(right - left, bottom - top),
crop_bounds,
suffix=".mask",
)
@@ -389,7 +390,8 @@ class EditImageMask(bpy.types.Operator):
else:
path = self._screen_polygon()
if len(path) < 3:
raise ValueError("Draw a larger image region")
self._finish(context)
return {"CANCELLED"}
polygon = screen_path_to_image_pixels(
path,
context.region,
@@ -444,7 +446,7 @@ class EditImageMask(bpy.types.Operator):
try:
return self._complete(context)
except (RuntimeError, TypeError, ValueError) as error:
self.report({"ERROR"}, str(error))
report_image_edit_exception(self, error)
self._finish(context)
return {"CANCELLED"}
return {"PASS_THROUGH"}
@@ -482,7 +484,7 @@ class EditImageMask(bpy.types.Operator):
try:
return self._complete(context)
except (RuntimeError, TypeError, ValueError) as error:
self.report({"ERROR"}, str(error))
report_image_edit_exception(self, error)
self._finish(context)
return {"CANCELLED"}
@@ -518,7 +520,7 @@ class EditImageMask(bpy.types.Operator):
self.invert,
)
except (json.JSONDecodeError, RuntimeError, TypeError, ValueError) as error:
self.report({"ERROR"}, str(error))
report_image_edit_exception(self, error)
return {"CANCELLED"}
return {"FINISHED"}
@@ -559,7 +561,11 @@ class RefineImageMask(JobOperatorBase, bpy.types.Operator):
)
finally:
bpy.data.images.remove(refined_image, do_unlink=True)
edit_image_object(source_object, mask, self.invert)
try:
edit_image_object(source_object, mask, self.invert)
except ImageEditWarning as error:
self.report({"WARNING"}, str(error))
return str(error)
push_undo("Edit Image Mask")
return "Image selection refined"

17
src/anyimage/operators/image_edit/perspective.py
View File

@@ -24,6 +24,7 @@ from .common import (
image_edit_poll,
polygon_bounds,
replace_empty_image_pixels,
report_image_edit_exception,
resolve_image_edit_click,
)
@@ -484,7 +485,6 @@ def create_perspective_result(
pixels,
output_size,
placement_bounds,
suffix=".perspective",
)
@@ -592,9 +592,14 @@ class EditImagePerspective(bpy.types.Operator):
context.area.tag_redraw()
def _remove_preview(self):
if self._preview_image is not None:
bpy.data.images.remove(self._preview_image, do_unlink=True)
self._preview_image = None
preview_image = self._preview_image
self._preview_image = None
if preview_image is None:
return
try:
bpy.data.images.remove(preview_image, do_unlink=True)
except ReferenceError:
pass
def _finish(self, context):
if self._handle is not None:
@@ -771,7 +776,7 @@ class EditImagePerspective(bpy.types.Operator):
self._phase = "ASPECT"
self._cursor = None
except (RuntimeError, TypeError, ValueError) as error:
self.report({"ERROR"}, str(error))
report_image_edit_exception(self, error)
self._finish(context)
return {"CANCELLED"}
self._set_aspect_status(context)
@@ -825,7 +830,7 @@ class EditImagePerspective(bpy.types.Operator):
TypeError,
ValueError,
) as error:
self.report({"ERROR"}, str(error))
report_image_edit_exception(self, error)
return {"CANCELLED"}
return {"FINISHED"}

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

@@ -51,7 +51,8 @@ class SetupAIEnvironment(bpy.types.Operator):
)
layout.label(
text=(
"It also downloads MoGe-2 ViT-S Normal and BEN2 Base "
"It also downloads MoGe-2 ViT-S Normal, BEN2 Base and "
"Real-ESRGAN General WDN x4v3 "
f"({format_download_size(required_size)})."
)
)

158
src/anyimage/operators/upscale.py Normal file
View File

@@ -0,0 +1,158 @@
import bpy
from ..preferences import (
DEFAULT_UPSCALE_MAX_OUTPUT_SIZE,
MAX_UPSCALE_OUTPUT_SIZE,
)
from ..properties import (
configured_upscale_max_output_size,
configured_upscale_model,
upscale_model_label,
)
from ..runtime import JobOperatorBase
from .common import (
HIDDEN,
cleanup_temporary_input,
input_path,
invoke_ai_setup_if_needed,
production_device,
require_environment,
require_model,
source_start_frame,
)
from .image import (
create_color_image,
default_extract_frames,
image_base_name,
is_image_empty,
prepare_image_input,
video_total_frames,
)
class UpscaleImage(JobOperatorBase, bpy.types.Operator):
bl_idname = "anyimage.upscale_image"
bl_label = "Upscale"
bl_description = "Upscale this Image Empty by 4x with an AI model"
bl_options = {"UNDO"}
job_type = "upscale-image"
starting_message = "Starting Image Empty upscale"
model: bpy.props.StringProperty(
default="REALESRGAN_GENERAL_WDN_X4V3",
options=HIDDEN,
)
max_output_size: bpy.props.IntProperty(
name="Max Output Size",
description=(
"Maximum output long edge; larger inputs are resized before 4x "
"upscaling to keep memory use bounded"
),
default=DEFAULT_UPSCALE_MAX_OUTPUT_SIZE,
min=1024,
max=MAX_UPSCALE_OUTPUT_SIZE,
step=256,
subtype="PIXEL",
options=HIDDEN,
)
input_path: bpy.props.StringProperty(options=HIDDEN)
source_object_name: bpy.props.StringProperty(options=HIDDEN)
source_image_name: bpy.props.StringProperty(options=HIDDEN)
delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN)
@classmethod
def poll(cls, context):
from ..runtime import runtime
return is_image_empty(context.object) and not runtime.server_busy()
def invoke(self, context, _event):
setup_result = invoke_ai_setup_if_needed()
if setup_result is not None:
return setup_result
return self.execute(context)
def execute(self, context):
self.model = configured_upscale_model(context)
self.max_output_size = configured_upscale_max_output_size(context)
source_object = context.object
if not is_image_empty(source_object):
self.report({"ERROR"}, "Select an Image Empty to upscale")
return {"CANCELLED"}
current_image = source_object.data
try:
original_image = getattr(
current_image,
"anyimage_upscale_source",
None,
)
except ReferenceError:
original_image = None
source_image = original_image or current_image
self.source_image_name = source_image.name
self._rebuild_from_original = source_image is not current_image
if self._rebuild_from_original:
self.report(
{"WARNING"},
"This image was already upscaled; rebuilding from the original",
)
if video_total_frames(source_image) > 1:
context.scene.anyimage_settings.video_frames = (
default_extract_frames(source_object)
)
try:
source_path, delete_input = prepare_image_input(
source_image,
context.scene,
)
except RuntimeError as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
self.input_path = str(source_path)
self.source_object_name = source_object.name
self.delete_input = delete_input
result = super().execute(context)
if "CANCELLED" in result:
cleanup_temporary_input(self.input_path, self.delete_input)
return result
def request(self, context):
require_environment()
require_model(self.model)
settings = context.scene.anyimage_settings
return {
"input": str(input_path(self.input_path)),
"model": self.model,
"device": production_device().lower(),
"video_frames": settings.video_frames,
"start_frame": source_start_frame(self.source_object_name),
"max_output_size": int(self.max_output_size),
}
def response(self, _context, result):
source_object = bpy.data.objects.get(self.source_object_name)
if not is_image_empty(source_object):
raise RuntimeError("The source image empty is no longer available")
source_image = bpy.data.images.get(self.source_image_name)
if source_image is None:
raise RuntimeError("The original source image is no longer available")
replaced_image = source_object.data
result_image = create_color_image(
source_object,
result.directory,
frame_prefix="upscale.",
)
if replaced_image is not source_image:
replaced_image.use_fake_user = False
if int(replaced_image.users) == 0:
bpy.data.images.remove(replaced_image, do_unlink=True)
result_image.name = f"{image_base_name(source_image)}.upscale"
result_image.anyimage_upscale_source = source_image
model_label = upscale_model_label(self.model)
if getattr(self, "_rebuild_from_original", False):
return f"Image is rebuilt from the original with {model_label}"
return f"Image is upscaled with {model_label}"
def cleanup(self):
cleanup_temporary_input(self.input_path, self.delete_input)