AnyImage/src/anyimage/operators/upscale.py

import bpy

from ..preferences import configured_max_ai_input_size
from ..properties import configured_upscale_model, upscale_model_label
from ..server.media.resolution import (
    DEFAULT_MAX_AI_INPUT_SIZE,
    MAX_MAX_AI_INPUT_SIZE,
    MIN_MAX_AI_INPUT_SIZE,
)
from ..runtime import JobOperatorBase
from ..common.ai import (
    require_input_path,
    invoke_ai_setup_if_needed,
    production_device,
    require_environment,
    require_model,
)
from ..common.image import (
    cleanup_image_input,
    default_extract_frames,
    is_image_empty,
    load_image_edit_result,
    prepare_image_input,
    source_start_frame,
    replace_empty_image,
    require_image_empty,
    video_total_frames,
)

_HIDDEN = {"HIDDEN", "SKIP_SAVE"}

UPSCALE_FACTOR = 2


def upscale_output_size(image):
    width, height = (int(value) for value in getattr(image, "size", (0, 0)))
    if width <= 0 or height <= 0:
        return None
    return width * UPSCALE_FACTOR, height * UPSCALE_FACTOR


def upscale_input_size_allowed(image, max_input_size):
    width, height = (int(value) for value in getattr(image, "size", (0, 0)))
    if width <= 0 or height <= 0:
        return True
    return max(width, height) < int(max_input_size)


class UpscaleImage(JobOperatorBase, bpy.types.Operator):
    bl_idname = "anyimage.upscale_image"
    bl_label = "Upscale"
    bl_description = "Upscale this Image Empty by 2x 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_input_size: bpy.props.IntProperty(
        name="Max AI Input Size",
        description="Maximum input long edge used for AI upscale analysis",
        default=DEFAULT_MAX_AI_INPUT_SIZE,
        min=MIN_MAX_AI_INPUT_SIZE,
        max=MAX_MAX_AI_INPUT_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

        source_object = getattr(context, "object", None)
        return (
            is_image_empty(source_object)
            and upscale_input_size_allowed(
                source_object.data,
                configured_max_ai_input_size(context),
            )
            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_input_size = configured_max_ai_input_size(context)
        source_object = context.object
        if not is_image_empty(source_object):
            self.report({"ERROR"}, "Select an Image Empty to upscale")
            return {"CANCELLED"}
        source_image = source_object.data
        if not upscale_input_size_allowed(source_image, self.max_input_size):
            self.report(
                {"ERROR"},
                "Image long edge must be smaller than Maximum AI Input Size "
                f"({self.max_input_size} px)",
            )
            return {"CANCELLED"}
        self.source_image_name = source_image.name
        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_image_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(require_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_input_size": int(self.max_input_size),
        }

    def response(self, _context, result):
        source_object = require_image_empty(self.source_object_name)
        source_image = bpy.data.images.get(self.source_image_name)
        if source_image is None:
            raise RuntimeError("The source image is no longer available")
        result_image = load_image_edit_result(
            source_object,
            result.directory,
            frame_prefix="upscale.",
        )
        replace_empty_image(source_object, result_image)
        model_label = upscale_model_label(self.model)
        return f"Image is upscaled with {model_label}"

    def cleanup(self):
        cleanup_image_input(self.input_path, self.delete_input)