Image Outpaint Feature: prepare the CLI and ImageUtil helpers

src/mflux/post_processing/image_util.py

@@ -7,9 +7,11 @@ import mlx.core as mx
 import numpy as np
 import piexif
 import PIL.Image
+import PIL.ImageDraw
 
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.post_processing.generated_image import GeneratedImage
+from mflux.ui.box_values import AbsoluteBoxValues, BoxValues
 
 log = logging.getLogger(__name__)
 
@@ -101,6 +103,80 @@ class ImageUtil:
     def load_image(path: str | pathlib.Path) -> PIL.Image.Image:
         return PIL.Image.open(path)
 
+    @staticmethod
+    def expand_image(
+        image: PIL.Image.Image,
+        box_values: AbsoluteBoxValues = None,
+        top: int | str = 0,
+        right: int | str = 0,
+        bottom: int | str = 0,
+        left: int | str = 0,
+        fill_color: tuple = (255, 255, 255),
+    ) -> PIL.Image.Image:
+        """
+        Expand the image by padding it with the top/right/bottom/left box values specified
+        in either pixels or percentages relative to original image dimensions.
+        """
+        if box_values is None:
+            box_values = BoxValues(top=top, right=right, bottom=bottom, left=left).normalize_to_dimensions(
+                image.width, image.height
+            )  # Create new image with expanded dimensions, paste the original image into it
+
+        new_width = image.width + box_values.left + box_values.right
+        new_height = image.height + box_values.top + box_values.bottom
+        expanded_image = PIL.Image.new(image.mode, (new_width, new_height), fill_color)
+        expanded_image.paste(image, (box_values.left, box_values.top))
+        return expanded_image
+
+    @staticmethod
+    def create_outpaint_mask_image(orig_width: int, orig_height: int, **create_bordered_image_kwargs):
+        """
+        Create an outpaint mask image that is black in the middle representing the original image dimensions
+        and a white border on the outside paddings representing the areas to be painted over.
+        """
+        return ImageUtil.create_bordered_image(
+            orig_width,
+            orig_height,
+            border_color=(255, 255, 255),
+            content_color=(0, 0, 0),
+            **create_bordered_image_kwargs,
+        )
+
+    @staticmethod
+    def create_bordered_image(
+        orig_width: int,
+        orig_height: int,
+        border_color: tuple,
+        content_color: tuple,
+        box_values: AbsoluteBoxValues = None,
+        top: int | str = 0,
+        right: int | str = 0,
+        bottom: int | str = 0,
+        left: int | str = 0,
+    ) -> PIL.Image.Image:
+        """
+        Create an image with border color and a content/fill-colored center based on CSS box model values.
+        """
+        if box_values is None:
+            box_values = BoxValues(top=top, right=right, bottom=bottom, left=left).normalize_to_dimensions(
+                orig_width, orig_height
+            )
+
+        # Create a new white image
+        new_width = orig_width + box_values.right + box_values.left
+        new_height = orig_height + box_values.top + box_values.bottom
+
+        result = PIL.Image.new("RGB", (new_width, new_height), border_color)
+        draw = PIL.ImageDraw.Draw(result)
+
+        # Draw black rectangle in the center
+        draw.rectangle(
+            [(box_values.left, box_values.top), (box_values.left + orig_width, box_values.top + orig_height)],
+            fill=content_color,
+        )
+
+        return result
+
     @staticmethod
     def scale_to_dimensions(
         image: PIL.Image.Image,

src/mflux/ui/box_values.py

@@ -0,0 +1,80 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class AbsoluteBoxValues:
+    top: int
+    right: int
+    bottom: int
+    left: int
+
+
+@dataclass
+class BoxValues:
+    top: int | str
+    right: int | str
+    bottom: int | str
+    left: int | str
+
+    def normalize_to_dimensions(self, width, height) -> AbsoluteBoxValues:
+        parts = []
+        dimension_base = [height, width, height, width]
+        for index, part in enumerate([self.top, self.right, self.bottom, self.left]):
+            if isinstance(part, str) and part.endswith("%"):
+                parts.append(int(int(part.strip("%")) / 100 * dimension_base[index]))
+            else:
+                # simple integer value
+                parts.append(int(part))
+        return AbsoluteBoxValues(*parts)
+
+
+class BoxValueError(ValueError):
+    pass
+
+
+def parse_box_value(value, delimiter=","):
+    """
+    Parse CSS-style padding values into a dictionary.
+
+    Accepts formats:
+    - Single value (all sides)
+    - Two values (vertical | horizontal)
+    - Three values (top | horizontal | bottom)
+    - Four values (top | right | bottom | left)
+
+    Args:
+        value: A string containing CSS-style box values
+
+    Returns:
+        BoxValues dataclass with 'top', 'right', 'bottom', and 'left' attributes
+    """
+    parts = []
+    for part_value in value.strip().split(delimiter):
+        try:
+            part = int(part_value.strip())
+            parts.append(part)
+        except ValueError:  # noqa: PERF203
+            # not an int - is it a %?
+            if (part_value := part_value.strip()).endswith("%"):
+                parts.append(part_value)
+            else:
+                raise BoxValueError(f"Invalid padding value: {part_value}")
+
+    if len(parts) == 1:
+        # If only one value is provided, apply to all sides
+        return BoxValues(top=parts[0], right=parts[0], bottom=parts[0], left=parts[0])
+    elif len(parts) == 2:
+        # If two values: first is top/bottom, second is left/right
+        return BoxValues(top=parts[0], right=parts[1], bottom=parts[0], left=parts[1])
+    elif len(parts) == 3:
+        # If three values: top, left/right, bottom
+        return BoxValues(top=parts[0], right=parts[1], bottom=parts[2], left=parts[1])
+    elif len(parts) == 4:
+        # If four values: top, right, bottom, left
+        return BoxValues(top=parts[0], right=parts[1], bottom=parts[2], left=parts[3])
+    else:
+        raise BoxValueError(
+            "Invalid outpaint padding box value format: {value} "
+            "Expected: 1 (all-sides), 2 (top/bottom, left/right) "
+            "or 4 (top, right, bottom, left)values of int or percentages. e.g. 10px, 20%"
+        )

src/mflux/ui/cli/parsers.py

@@ -6,7 +6,10 @@ import typing as t
 from pathlib import Path
 
 from mflux.community.in_context_lora.in_context_loras import LORA_NAME_MAP, LORA_REPO_ID
-from mflux.ui import defaults as ui_defaults
+from mflux.ui import (
+    box_values,
+    defaults as ui_defaults,
+)
 
 
 class ModelSpecAction(argparse.Action):
@@ -33,6 +36,7 @@ class CommandLineParser(argparse.ArgumentParser):
         self.supports_image_generation = False
         self.supports_controlnet = False
         self.supports_image_to_image = False
+        self.supports_image_outpaint = False
         self.supports_lora = False
 
     def add_general_arguments(self) -> None:
@@ -88,6 +92,10 @@ class CommandLineParser(argparse.ArgumentParser):
         self.add_argument("--output", type=str, default="image.png", help="The filename for the output image. Default is \"image.png\".")
         self.add_argument('--stepwise-image-output-dir', type=str, default=None, help='[EXPERIMENTAL] Output dir to write step-wise images and their final composite image to. This feature may change in future versions.')
 
+    def add_image_outpaint_arguments(self, required=False) -> None:
+        self.supports_image_outpaint = True
+        self.add_argument("--image-outpaint-padding", type=str, default=None, required=required, help="For outpainting mode: CSS-style box padding values to extend the canvas of image specified by--image-path. E.g. '20', '50%%'")
+
     def add_controlnet_arguments(self) -> None:
         self.supports_controlnet = True
         self.add_argument("--controlnet-image-path", type=str, required=False, help="Local path of the image to use as input for controlnet.")
@@ -171,8 +179,13 @@ class CommandLineParser(argparse.ArgumentParser):
                 if namespace.controlnet_save_canny == self.get_default("controlnet_save_canny") and (cnet_canny_from_metadata := prior_gen_metadata.get("controlnet_save_canny", None)):
                     namespace.controlnet_save_canny = cnet_canny_from_metadata
 
+
+            if self.supports_image_outpaint:
+                if namespace.image_outpaint_padding is None:
+                    namespace.image_outpaint_padding = prior_gen_metadata.get("image_outpaint_padding", None)
+
         # Only require model if we're not in training mode
-        if namespace.model is None and not has_training_args:
+        if hasattr(namespace, "model") and namespace.model is None and not has_training_args:
             self.error("--model / -m must be provided, or 'model' must be specified in the config file.")
 
         if self.supports_image_generation and namespace.seed is None and namespace.auto_seeds > 0:
@@ -196,7 +209,12 @@ class CommandLineParser(argparse.ArgumentParser):
         if self.supports_image_generation and namespace.steps is None:
             namespace.steps = ui_defaults.MODEL_INFERENCE_STEPS.get(namespace.model, None)
 
-        if namespace.low_ram and len(namespace.seed) > 1:
+        if self.supports_image_outpaint and namespace.image_outpaint_padding is not None:
+            # parse and normalize any acceptable 1,2,3,4-tuple box value to 4-tuple
+            namespace.image_outpaint_padding = box_values.parse_box_value(namespace.image_outpaint_padding)
+            print(f"{namespace.image_outpaint_padding=}")
+
+        if hasattr(namespace, "low_ram") and namespace.low_ram is not None and len(namespace.seed) > 1:
             self.error("--low-ram cannot be used with multiple seeds")
 
         return namespace

tests/arg_parser/test_cli_argparser.py

@@ -5,6 +5,7 @@ from unittest.mock import patch
 
 import pytest
 
+from mflux.ui.box_values import BoxValues
 from mflux.ui.cli.parsers import CommandLineParser
 
 
@@ -15,6 +16,7 @@ def _create_mflux_generate_parser(with_controlnet=False) -> CommandLineParser:
     parser.add_image_generator_arguments(supports_metadata_config=True)
     parser.add_lora_arguments()
     parser.add_image_to_image_arguments(required=False)
+    parser.add_image_outpaint_arguments()
     if with_controlnet:
         parser.add_controlnet_arguments()
     parser.add_output_arguments()
@@ -328,6 +330,40 @@ def test_image_to_image_args(mflux_generate_parser, mflux_generate_minimal_argv,
         assert args.image_strength == 0.4  # default
 
 
+def test_image_outpaint_args(mflux_generate_parser, mflux_generate_minimal_argv, base_metadata_dict, temp_dir):  # fmt: off
+    metadata_file = temp_dir / "image_outpaint.json"
+    test_padding = "10,20,30,40"
+    with metadata_file.open("wt") as m:
+        base_metadata_dict["image_outpaint_padding"] = test_padding
+        json.dump(base_metadata_dict, m, indent=4)
+
+    # test user default value
+    with patch("sys.argv", mflux_generate_minimal_argv + ["-m", "dev"]):
+        args = mflux_generate_parser.parse_args()
+        assert args.image_outpaint_padding is None
+
+    # test metadata config accepted
+    with patch('sys.argv', mflux_generate_minimal_argv + ['--config-from-metadata', metadata_file.as_posix()]):  # fmt: off
+        args = mflux_generate_parser.parse_args()
+        assert args.image_outpaint_padding == BoxValues(10, 20, 30, 40)
+
+    # test outpaint padding override in 4-value format
+    with patch('sys.argv', mflux_generate_minimal_argv + ['--image-outpaint-padding', '5,15,25,35', '--config-from-metadata', metadata_file.as_posix()]):  # fmt: off
+        args = mflux_generate_parser.parse_args()
+        assert args.image_outpaint_padding == BoxValues(5, 15, 25, 35)
+
+    # test outpaint padding override in percentages in two-value format
+    with patch('sys.argv', mflux_generate_minimal_argv + ['--image-outpaint-padding', '10%,20%', '--config-from-metadata', metadata_file.as_posix()]):  # fmt: off
+        args = mflux_generate_parser.parse_args()
+        assert args.image_outpaint_padding == BoxValues("10%", "20%", "10%", "20%")
+
+    # test outpaint padding override in percentages in three-value format, mixed int/percentages
+    # also allowing whitespace between the box values
+    with patch('sys.argv', mflux_generate_minimal_argv + ['--image-outpaint-padding', '10%, 50,   20%', '--config-from-metadata', metadata_file.as_posix()]):  # fmt: off
+        args = mflux_generate_parser.parse_args()
+        assert args.image_outpaint_padding == BoxValues("10%", 50, "20%", 50)
+
+
 def test_controlnet_args(mflux_generate_controlnet_parser, mflux_generate_controlnet_minimal_argv, base_metadata_dict, temp_dir):  # fmt: off
     test_path = "/some/cnet/1.safetensors"
     metadata_file = temp_dir / "cnet_args.json"

tests/image_generation/test_image_util.py

@@ -0,0 +1,134 @@
+import PIL.Image
+import pytest
+
+from mflux.post_processing.image_util import ImageUtil
+
+
+@pytest.fixture
+def test_image():
+    # Create a simple test image
+    return PIL.Image.new("RGB", (100, 80), color="blue")
+
+
+def test_expand_image_with_pixels(test_image):
+    # Test expanding with pixel values
+    expanded = ImageUtil.expand_image(
+        test_image,
+        top=20,
+        right=30,
+        bottom=40,
+        left=10,
+        fill_color=(255, 0, 0),  # red
+    )
+
+    # Check dimensions
+    assert expanded.width == 100 + 30 + 10  # original + right + left
+    assert expanded.height == 80 + 20 + 40  # original + top + bottom
+
+    # Check color at corners (should be the fill color)
+    assert expanded.getpixel((0, 0)) == (255, 0, 0)  # top-left
+    assert expanded.getpixel((139, 0)) == (255, 0, 0)  # top-right
+    assert expanded.getpixel((0, 139)) == (255, 0, 0)  # bottom-left
+    assert expanded.getpixel((139, 139)) == (255, 0, 0)  # bottom-right
+
+    # Check original image is preserved in the middle
+    assert expanded.getpixel((15, 25)) == (0, 0, 255)  # blue
+
+
+def test_expand_image_with_percentages(test_image):
+    # Test expanding with percentage values
+    expanded = ImageUtil.expand_image(
+        test_image,
+        top="25%",
+        right="30%",
+        bottom="50%",
+        left="10%",
+        fill_color=(0, 255, 0),  # green
+    )
+
+    # Calculate expected dimensions
+    expected_top = int(0.25 * 80)  # 25% of height
+    expected_right = int(0.3 * 100)  # 30% of width
+    expected_bottom = int(0.5 * 80)  # 50% of height
+    expected_left = int(0.1 * 100)  # 10% of width
+
+    expected_width = 100 + expected_right + expected_left
+    expected_height = 80 + expected_top + expected_bottom
+
+    # Check dimensions
+    assert expanded.width == expected_width
+    assert expanded.height == expected_height
+
+    # Check fill color
+    assert expanded.getpixel((0, 0)) == (0, 255, 0)
+
+
+def test_expand_image_with_invalid_values(test_image):
+    # Test with invalid percentage string (a string that does not end with %"
+    with pytest.raises(ValueError):
+        ImageUtil.expand_image(test_image, top="25#", right="30", bottom="50%", left="10")
+
+    with pytest.raises(ValueError):
+        ImageUtil.expand_image(test_image, top="25", right="30$", bottom="50%", left="10")
+
+    with pytest.raises(ValueError):
+        ImageUtil.expand_image(test_image, top="25", right="30", bottom="50#", left="10")
+
+    with pytest.raises(ValueError):
+        ImageUtil.expand_image(test_image, top="25", right="30", bottom="50#", left="10*")
+
+
+def test_create_outpaint_mask_image():
+    # Test with various padding values
+    orig_width = 100
+    orig_height = 80
+    top_padding = 20
+    right_padding = 30
+    bottom_padding = 40
+    left_padding = 10
+
+    mask = ImageUtil.create_outpaint_mask_image(
+        orig_width=orig_width,
+        orig_height=orig_height,
+        top=top_padding,
+        right=right_padding,
+        bottom=bottom_padding,
+        left=left_padding,
+    )
+
+    # Check dimensions of the mask
+    expected_width = orig_width + right_padding + left_padding
+    expected_height = orig_height + bottom_padding + top_padding
+    assert mask.width == expected_width
+    assert mask.height == expected_height
+
+    # Check padding areas are white (255, 255, 255)
+    # Top-left corner
+    assert mask.getpixel((5, 5)) == (255, 255, 255)
+
+    # Top-right corner
+    assert mask.getpixel((expected_width - 5, 5)) == (255, 255, 255)
+
+    # Bottom-left corner
+    assert mask.getpixel((5, expected_height - 5)) == (255, 255, 255)
+
+    # Bottom-right corner
+    assert mask.getpixel((expected_width - 5, expected_height - 5)) == (255, 255, 255)
+
+    # Check center is black (0, 0, 0)
+    center_x = left_padding + (orig_width // 2)
+    center_y = top_padding + (orig_height // 2)
+    assert mask.getpixel((center_x, center_y)) == (0, 0, 0)
+
+    # Check boundaries
+    # Top edge of center box
+    assert mask.getpixel((left_padding + 10, top_padding)) == (0, 0, 0)
+
+    # Right edge of center box
+    assert mask.getpixel((left_padding + orig_width - 1, top_padding + 10)) == (0, 0, 0)
+
+    # Bottom edge of center box
+    assert mask.getpixel((left_padding + 10, top_padding + orig_height - 1)) == (0, 0, 0)
+
+    # Left edge of center box
+    assert mask.getpixel((left_padding, top_padding + 10)) == (0, 0, 0)

tools/create_outpaint_image_canvas_and_mask.py

@@ -0,0 +1,69 @@
+import pathlib
+import sys
+
+from mflux.post_processing.image_util import ImageUtil
+from mflux.ui.box_values import AbsoluteBoxValues, BoxValues
+from mflux.ui.cli.parsers import CommandLineParser
+
+
+def main():
+    parser = CommandLineParser(description="Create expanded canvas and mask for outpainting")
+    parser.add_argument("image_path", type=pathlib.Path, help="Path to the input image file")
+    parser.add_image_outpaint_arguments(required=True)
+    args = parser.parse_args()
+
+    if not args.image_path.exists():
+        print(f"Input image not found: {args.image_path}")
+        return 1
+
+    try:
+        # Load the image
+        image = ImageUtil.load_image(args.image_path)
+        # Get original dimensions
+        orig_width, orig_height = image.width, image.height
+        print(f"Loaded original image: {args.image_path} ({orig_width}x{orig_height})")
+
+        # Calculate padding values
+        padding: BoxValues = args.image_outpaint_padding
+        print(f"{padding=}")
+
+        # Create expanded canvas
+        expanded_image = ImageUtil.expand_image(
+            image=image,
+            top=padding.top,
+            right=padding.right,
+            bottom=padding.bottom,
+            left=padding.left,
+        )
+
+        abs_padding: AbsoluteBoxValues = args.image_outpaint_padding.normalize_to_dimensions(orig_width, orig_height)
+        # Create mask image
+        mask_image = ImageUtil.create_outpaint_mask_image(
+            orig_width=orig_width,
+            orig_height=orig_height,
+            top=abs_padding.top,
+            right=abs_padding.right,
+            bottom=abs_padding.bottom,
+            left=abs_padding.left,
+        )
+
+        # Construct output paths
+        expanded_output_path = args.image_path.with_stem(f"{args.image_path.stem}_expanded")
+        mask_output_path = args.image_path.with_stem(f"{args.image_path.stem}_mask")
+
+        # Save the images
+        ImageUtil.save_image(expanded_image, expanded_output_path)
+        ImageUtil.save_image(mask_image, mask_output_path)
+
+        print(f"Saved expanded canvas: {expanded_output_path}")
+        print(f"Saved mask image: {mask_output_path}")
+
+        return 0
+
+    except Exception as e:  # noqa: BLE001
+        print(f"Error processing image: {e}")
+        raise e
+
+
+if __name__ == "__main__":
+    sys.exit(main())