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:
158
src/anyimage/operators/upscale.py
Normal file
158
src/anyimage/operators/upscale.py
Normal 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)
|
||||
Reference in New Issue
Block a user