From 7cc40999956aaca2fde21c2a308cfe429cf5d4e1 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Fri, 25 Apr 2025 10:34:31 +0200 Subject: [PATCH] Fix command line args --- src/mflux/save_depth.py | 2 +- src/mflux/ui/cli/parsers.py | 5 +- tests/arg_parser/test_cli_argparser.py | 88 +++++++++++++++++++++++++- 3 files changed, 92 insertions(+), 3 deletions(-) diff --git a/src/mflux/save_depth.py b/src/mflux/save_depth.py index a657133..1e6e7fd 100644 --- a/src/mflux/save_depth.py +++ b/src/mflux/save_depth.py @@ -6,7 +6,7 @@ from mflux.ui.cli.parsers import CommandLineParser def main(): # 0. Parse command line arguments parser = CommandLineParser(description="Save a depth map of an input image") - parser.add_depth_arguments() + parser.add_save_depth_arguments() args = parser.parse_args() # 1. Create and save the depth map diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py index e873732..f1ae894 100644 --- a/src/mflux/ui/cli/parsers.py +++ b/src/mflux/ui/cli/parsers.py @@ -96,7 +96,10 @@ class CommandLineParser(argparse.ArgumentParser): self.add_argument("--image-path", type=Path, required=False, help="Local path to the source image") self.add_argument("--depth-image-path", type=Path, required=False, help="Local path to the depth image") self.add_argument("--save-depth-map", action="store_true", required=False, help="If set, save the depth map created from the source image.") - self.add_argument("--quantize", "-q", type=int, choices=ui_defaults.QUANTIZE_CHOICES, default=None, help=f"Quantize the model ({' or '.join(map(str, ui_defaults.QUANTIZE_CHOICES))}, Default is None)") + + def add_save_depth_arguments(self) -> None: + self.add_argument("--image-path", type=Path, required=True, help="Local path to the source image") + self.add_argument("--quantize", "-q", type=int, choices=ui_defaults.QUANTIZE_CHOICES, default=None, required=False, help=f"Quantize the model ({' or '.join(map(str, ui_defaults.QUANTIZE_CHOICES))}, Default is None)") def add_redux_arguments(self) -> None: self.add_argument("--redux-image-paths", type=Path, nargs="*", required=True, help="Local path to the source image") diff --git a/tests/arg_parser/test_cli_argparser.py b/tests/arg_parser/test_cli_argparser.py index 8287034..f16f328 100644 --- a/tests/arg_parser/test_cli_argparser.py +++ b/tests/arg_parser/test_cli_argparser.py @@ -100,7 +100,44 @@ def mflux_fill_parser() -> CommandLineParser: @pytest.fixture def mflux_fill_minimal_argv() -> list[str]: - return ["mflux-fill", "--prompt", "meaning of life", "--image-path", "image.png", "--masked-image-path", "mask.png"] + return [ + "mflux-generate-fill", + "--prompt", + "meaning of life", + "--image-path", + "image.png", + "--masked-image-path", + "mask.png", + ] + + +@pytest.fixture +def mflux_save_depth_parser() -> CommandLineParser: + parser = CommandLineParser(description="Save depth map from an image.") + parser.add_general_arguments() + parser.add_save_depth_arguments() + parser.add_output_arguments() + return parser + + +@pytest.fixture +def mflux_save_depth_minimal_argv() -> list[str]: + return ["mflux-save-depth", "--image-path", "image.png"] + + +@pytest.fixture +def mflux_redux_parser() -> CommandLineParser: + parser = CommandLineParser(description="Generate redux images.") + parser.add_general_arguments() + parser.add_model_arguments(require_model_arg=False) + parser.add_redux_arguments() + parser.add_output_arguments() + return parser + + +@pytest.fixture +def mflux_redux_minimal_argv() -> list[str]: + return ["mflux-generate-redux", "--redux-image-paths", "image1.png", "image2.png"] def test_model_path_requires_model_arg(mflux_generate_parser): @@ -529,3 +566,52 @@ def test_fill_default_guidance(): # Now check that guidance is correctly set to 30 assert args.guidance == 30 + + +def test_save_depth_args(mflux_save_depth_parser, mflux_save_depth_minimal_argv): + # Test required arguments + with patch("sys.argv", mflux_save_depth_minimal_argv): + args = mflux_save_depth_parser.parse_args() + assert args.image_path == Path("image.png") + assert hasattr(args, "quantize") + + # Test with quantized argument + with patch("sys.argv", mflux_save_depth_minimal_argv + ["--quantize", "4"]): + args = mflux_save_depth_parser.parse_args() + assert args.image_path == Path("image.png") + assert args.quantize == 4 + + # Test with output argument + with patch("sys.argv", mflux_save_depth_minimal_argv + ["--output", "depth_map.png"]): + args = mflux_save_depth_parser.parse_args() + assert args.image_path == Path("image.png") + assert args.output == "depth_map.png" + + +def test_redux_args(mflux_redux_parser, mflux_redux_minimal_argv): + # Test required arguments + with patch("sys.argv", mflux_redux_minimal_argv): + args = mflux_redux_parser.parse_args() + assert len(args.redux_image_paths) == 2 + assert args.redux_image_paths[0] == Path("image1.png") + assert args.redux_image_paths[1] == Path("image2.png") + + # Test with more image paths + with patch("sys.argv", ["mflux-redux", "--redux-image-paths", "image1.png", "image2.png", "image3.png"]): + args = mflux_redux_parser.parse_args() + assert len(args.redux_image_paths) == 3 + assert args.redux_image_paths[0] == Path("image1.png") + assert args.redux_image_paths[1] == Path("image2.png") + assert args.redux_image_paths[2] == Path("image3.png") + + # Test with model argument + with patch("sys.argv", mflux_redux_minimal_argv + ["--model", "dev"]): + args = mflux_redux_parser.parse_args() + assert len(args.redux_image_paths) == 2 + assert args.model == "dev" + + # Test with output argument + with patch("sys.argv", mflux_redux_minimal_argv + ["--output", "redux_result.png"]): + args = mflux_redux_parser.parse_args() + assert len(args.redux_image_paths) == 2 + assert args.output == "redux_result.png"