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.

src/anyimage/server/jobs/background.py

@@ -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": "."}

src/anyimage/server/jobs/debug_upscale.py

@@ -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(

src/anyimage/server/jobs/foreground.py

@@ -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"
     )

src/anyimage/server/jobs/upscale.py

@@ -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": "."}