diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index 563d9fa..81bf4a3 100644 --- a/src/mflux/post_processing/image_util.py +++ b/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, diff --git a/src/mflux/ui/box_values.py b/src/mflux/ui/box_values.py new file mode 100644 index 0000000..c1808ad --- /dev/null +++ b/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%" + ) diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py index ad898df..1e601f8 100644 --- a/src/mflux/ui/cli/parsers.py +++ b/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 diff --git a/tests/arg_parser/test_cli_argparser.py b/tests/arg_parser/test_cli_argparser.py index fcc91f1..54b1490 100644 --- a/tests/arg_parser/test_cli_argparser.py +++ b/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" diff --git a/tests/image_generation/test_image_util.py b/tests/image_generation/test_image_util.py new file mode 100644 index 0000000..5a16a83 --- /dev/null +++ b/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) diff --git a/tools/create_outpaint_image_canvas_and_mask.py b/tools/create_outpaint_image_canvas_and_mask.py new file mode 100755 index 0000000..4b59f02 --- /dev/null +++ b/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())