Restore edited output to original image size

gradio_app/edit.py

@@ -11,7 +11,7 @@ from gradio_app.config import (
 from gradio_app.logging_utils import setup_logging
 from gradio_app.pipeline import load_lora
 from pixelsmile.linear_conditioning import compute_text_embeddings
-from pixelsmile.utils.image import resize
+from pixelsmile.utils.image import build_edit_image_bundle, restore_edited_image
 
 
 logger = setup_logging()
@@ -39,8 +39,7 @@ def prepare_input_image(image: Image.Image) -> Image.Image:
         raise gr.Error("Please upload an input image.")
     if not isinstance(image, Image.Image):
         image = Image.fromarray(image)
-    image = image.convert("RGB")
-    return resize(image, (DEFAULT_WIDTH, DEFAULT_HEIGHT), DEFAULT_RESIZE_MODE)
+    return image.convert("RGB")
 
 
 def run_edit(
@@ -70,7 +69,12 @@ def run_edit(
     pipe = load_lora(weight_version)
     if progress is not None:
         progress(0.35, desc="Preparing input image...")
-    input_image = prepare_input_image(image)
+    source_image = prepare_input_image(image)
+    input_image, restore_meta = build_edit_image_bundle(
+        source_image,
+        (DEFAULT_WIDTH, DEFAULT_HEIGHT),
+        DEFAULT_RESIZE_MODE,
+    )
     logger.info("Input image prepared at size: %s", input_image.size)
     edit_condition = build_edit_condition(subject, expression, float(scale))
     logger.info("Edit condition: %s", edit_condition)
@@ -107,4 +111,6 @@ def run_edit(
     logger.info("Pipeline inference finished")
     if progress is not None:
         progress(0.95, desc="Finalizing output...")
-    return output.images[0]
+    restored = restore_edited_image(output.images[0], restore_meta)
+    logger.info("Restored edited image to original size: %s", restored.size)
+    return restored

pixelsmile/utils/image.py

@@ -145,6 +145,119 @@ def resize(
     else:
         raise ValueError(f"Resize mode error: {resize_mode}")
 
+
+def build_edit_image_bundle(
+    image: Image.Image,
+    target_size: tuple,
+    resize_mode: str,
+    box: list = None,
+    vertical_bias: float = 0.5,
+):
+    original = image.convert("RGB")
+    original_size = original.size
+    target_w, target_h = target_size
+
+    if resize_mode == "direct":
+        edited_input = original.resize((target_w, target_h), resample=Image.LANCZOS)
+        meta = {
+            "mode": "direct",
+            "original_size": original_size,
+            "target_size": target_size,
+        }
+        return edited_input, meta
+
+    if resize_mode == "padding":
+        w, h = original.size
+        scale = min(target_w / w, target_h / h)
+        new_w = int(w * scale)
+        new_h = int(h * scale)
+        resized = original.resize((new_w, new_h), resample=Image.LANCZOS)
+
+        canvas = Image.new("RGB", (target_w, target_h))
+        pad_left = (target_w - new_w) // 2
+        pad_top = (target_h - new_h) // 2
+        canvas.paste(resized, (pad_left, pad_top))
+
+        meta = {
+            "mode": "padding",
+            "original_size": original_size,
+            "target_size": target_size,
+            "content_box": (pad_left, pad_top, pad_left + new_w, pad_top + new_h),
+        }
+        return canvas, meta
+
+    if resize_mode == "crop":
+        w, h = original.size
+        scale = max(target_w / w, target_h / h)
+        new_w = int(w * scale)
+        new_h = int(h * scale)
+        resized = original.resize((new_w, new_h), resample=Image.LANCZOS)
+
+        if box:
+            nx1, ny1, nx2, ny2 = [c * scale for c in box]
+            box_cx = (nx1 + nx2) / 2
+            box_cy = (ny1 + ny2) / 2
+            box_w = nx2 - nx1
+            box_h = ny2 - ny1
+
+            left = box_cx - target_w / 2
+            if box_h <= target_h:
+                min_top_to_contain = ny2 - target_h
+                max_top_to_contain = ny1
+                top = max_top_to_contain - (
+                    max_top_to_contain - min_top_to_contain
+                ) * vertical_bias
+            else:
+                top = box_cy - target_h / 2
+        else:
+            left = (new_w - target_w) / 2
+            top = (new_h - target_h) * vertical_bias
+
+        left = max(0, min(left, new_w - target_w))
+        top = max(0, min(top, new_h - target_h))
+        left = int(left)
+        top = int(top)
+        right = left + target_w
+        bottom = top + target_h
+
+        edited_input = resized.crop((left, top, right, bottom))
+        meta = {
+            "mode": "crop",
+            "original_size": original_size,
+            "target_size": target_size,
+            "resized_size": (new_w, new_h),
+            "crop_box": (left, top, right, bottom),
+            "resized_background": resized,
+        }
+        return edited_input, meta
+
+    raise ValueError(f"Resize mode error: {resize_mode}")
+
+
+def restore_edited_image(
+    edited_image: Image.Image,
+    meta: dict,
+):
+    mode = meta["mode"]
+    original_size = tuple(meta["original_size"])
+
+    if mode == "direct":
+        return edited_image.resize(original_size, resample=Image.LANCZOS)
+
+    if mode == "padding":
+        left, top, right, bottom = meta["content_box"]
+        content = edited_image.crop((left, top, right, bottom))
+        return content.resize(original_size, resample=Image.LANCZOS)
+
+    if mode == "crop":
+        canvas = meta["resized_background"].copy()
+        left, top, right, bottom = meta["crop_box"]
+        pasted = edited_image.resize((right - left, bottom - top), resample=Image.LANCZOS)
+        canvas.paste(pasted, (left, top))
+        return canvas.resize(original_size, resample=Image.LANCZOS)
+
+    raise ValueError(f"Unsupported restore mode: {mode}")
+
 def scale_fun(x: float) -> float:
     y = max(0.0, min(1.0, x))
     if x <= 0.6: