diff --git a/src/mflux/generate.py b/src/mflux/generate.py index 2184ef7..b89b17b 100644 --- a/src/mflux/generate.py +++ b/src/mflux/generate.py @@ -1,37 +1,19 @@ -import argparse import time from pathlib import Path from mflux import Config, Flux1, ModelConfig, StopImageGenerationException +from mflux.ui.cli.parsers import CommandLineParser def main(): # fmt: off - parser = argparse.ArgumentParser(description="Generate an image based on a prompt.") - parser.add_argument("--prompt", type=str, required=True, help="The textual description of the image to generate.") - parser.add_argument("--output", type=str, default="image.png", help="The filename for the output image. Default is \"image.png\".") - parser.add_argument("--model", "-m", type=str, required=True, choices=["dev", "schnell"], help="The model to use (\"schnell\" or \"dev\").") - parser.add_argument("--seed", type=int, default=None, help="Entropy Seed (Default is time-based random-seed)") - parser.add_argument("--height", type=int, default=1024, help="Image height (Default is 1024)") - parser.add_argument("--width", type=int, default=1024, help="Image width (Default is 1024)") - parser.add_argument("--steps", type=int, default=None, help="Inference Steps") - parser.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.') - parser.add_argument("--guidance", type=float, default=3.5, help="Guidance Scale (Default is 3.5)") - parser.add_argument("--quantize", "-q", type=int, choices=[4, 8], default=None, help="Quantize the model (4 or 8, Default is None)") - parser.add_argument("--path", type=str, default=None, help="Local path for loading a model from disk") - parser.add_argument("--lora-paths", type=str, nargs="*", default=None, help="Local safetensors for applying LORA from disk") - parser.add_argument("--lora-scales", type=float, nargs="*", default=None, help="Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.") - parser.add_argument("--metadata", action="store_true", help="Export image metadata as a JSON file.") - # fmt: on - + parser = CommandLineParser(description="Generate an image based on a prompt.") + parser.add_model_arguments() + parser.add_lora_arguments() + parser.add_image_generator_arguments() + parser.add_output_arguments() args = parser.parse_args() - if args.path and args.model is None: - parser.error("--model must be specified when using --path") - - if args.steps is None: - args.steps = 4 if args.model == "schnell" else 14 - # Load the model flux = Flux1( model_config=ModelConfig.from_alias(args.model), diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py index 28c1a09..6495fcf 100644 --- a/src/mflux/generate_controlnet.py +++ b/src/mflux/generate_controlnet.py @@ -1,40 +1,19 @@ -import argparse import time from pathlib import Path from mflux import ConfigControlnet, Flux1Controlnet, ModelConfig, StopImageGenerationException +from mflux.ui.cli.parsers import CommandLineParser def main(): - # fmt: off - parser = argparse.ArgumentParser(description="Generate an image based on a prompt.") - parser.add_argument("--prompt", type=str, required=True, help="The textual description of the image to generate.") - parser.add_argument("--controlnet-image-path", type=str, required=True, help="Local path of the image to use as input for controlnet.") - parser.add_argument("--controlnet-strength", type=float, default=0.4, help="Controls how strongly the control image influences the output image. A value of 0.0 means no influence. (Default is 0.4)") - parser.add_argument("--controlnet-save-canny", action="store_true", help="If set, save the Canny edge detection reference input image.") - parser.add_argument("--output", type=str, default="image.png", help="The filename for the output image. Default is \"image.png\".") - parser.add_argument("--model", "-m", type=str, required=True, choices=["dev", "schnell"], help="The model to use (\"schnell\" or \"dev\").") - parser.add_argument("--seed", type=int, default=None, help="Entropy Seed (Default is time-based random-seed)") - parser.add_argument("--height", type=int, default=1024, help="Image height (Default is 1024)") - parser.add_argument("--width", type=int, default=1024, help="Image width (Default is 1024)") - parser.add_argument("--steps", type=int, default=None, help="Inference Steps") - parser.add_argument('--stepwise-image-output-dir', type=str, default=None, help='Output dir to write step-wise images and their final composite image to.') - parser.add_argument("--guidance", type=float, default=3.5, help="Guidance Scale (Default is 3.5)") - parser.add_argument("--quantize", "-q", type=int, choices=[4, 8], default=None, help="Quantize the model (4 or 8, Default is None)") - parser.add_argument("--path", type=str, default=None, help="Local path for loading a model from disk") - parser.add_argument("--lora-paths", type=str, nargs="*", default=None, help="Local safetensors for applying LORA from disk") - parser.add_argument("--lora-scales", type=float, nargs="*", default=None, help="Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.") - parser.add_argument("--metadata", action="store_true", help="Export image metadata as a JSON file.") - # fmt: on - + parser = CommandLineParser(description="Generate an image based on a prompt and a controlnet reference image.") # fmt: off + parser.add_model_arguments() + parser.add_lora_arguments() + parser.add_image_generator_arguments() + parser.add_controlnet_arguments() + parser.add_output_arguments() args = parser.parse_args() - if args.path and args.model is None: - parser.error("--model must be specified when using --path") - - if args.steps is None: - args.steps = 4 if args.model == "schnell" else 14 - # Load the model flux = Flux1Controlnet( model_config=ModelConfig.from_alias(args.model), diff --git a/src/mflux/save.py b/src/mflux/save.py index 641d06f..f26db41 100644 --- a/src/mflux/save.py +++ b/src/mflux/save.py @@ -1,18 +1,11 @@ -import argparse - from mflux import Flux1, ModelConfig +from mflux.ui.cli.parsers import CommandLineParser def main(): - # fmt: off - parser = argparse.ArgumentParser(description="Save a quantized version of Flux.1 to disk.") - parser.add_argument("--path", type=str, required=True, help="Local path for loading a model from disk") - parser.add_argument("--model", "-m", type=str, required=True, choices=["dev", "schnell"], help="The model to use (\"schnell\" or \"dev\").") - parser.add_argument("--quantize", "-q", type=int, choices=[4, 8], default=8, help="Quantize the model (4 or 8, Default is 8)") - parser.add_argument("--lora-paths", type=str, nargs="*", default=None, help="Local safetensors for applying LORA from disk") - parser.add_argument("--lora-scales", type=float, nargs="*", default=None, help="Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.") - # fmt: on - + parser = CommandLineParser(description="Save a quantized version of Flux.1 to disk.") # fmt: off + parser.add_model_arguments() + parser.add_lora_arguments() args = parser.parse_args() print(f"Saving model {args.model} with quantization level {args.quantize}\n") diff --git a/src/mflux/ui/__init__.py b/src/mflux/ui/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/ui/cli/__init__.py b/src/mflux/ui/cli/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py new file mode 100644 index 0000000..eae8c49 --- /dev/null +++ b/src/mflux/ui/cli/parsers.py @@ -0,0 +1,51 @@ +import argparse +from pathlib import Path + +from mflux.ui import defaults as ui_defaults + + +# fmt: off +class CommandLineParser(argparse.ArgumentParser): + + def add_model_arguments(self): + self.add_argument("--model", "-m", type=str, required=True, choices=ui_defaults.MODEL_CHOICES, help=f"The model to use ({' or '.join(ui_defaults.MODEL_CHOICES)}).") + self.add_argument("--path", type=str, default=None, help="Local path for loading a model from disk") + 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_lora_arguments(self): + self.add_argument("--lora-paths", type=str, nargs="*", default=None, help="Local safetensors for applying LORA from disk") + self.add_argument("--lora-scales", type=float, nargs="*", default=None, help="Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.") + + def _add_image_generator_common_arguments(self): + self.add_argument("--height", type=int, default=ui_defaults.HEIGHT, help=f"Image height (Default is {ui_defaults.HEIGHT})") + self.add_argument("--width", type=int, default=ui_defaults.WIDTH, help=f"Image width (Default is {ui_defaults.HEIGHT})") + self.add_argument("--steps", type=int, default=None, help="Inference Steps") + self.add_argument("--guidance", type=float, default=ui_defaults.GUIDANCE_SCALE, help=f"Guidance Scale (Default is {ui_defaults.GUIDANCE_SCALE})") + + def add_image_generator_arguments(self): + self.add_argument("--prompt", type=str, required=True, help="The textual description of the image to generate.") + self.add_argument("--seed", type=int, default=None, help="Entropy Seed (Default is time-based random-seed)") + self._add_image_generator_common_arguments() + + def add_batch_image_generator_arguments(self): + self.add_argument("--prompts-file", type=Path, required=True, help="Local path for a file that holds a batch of prompts.") + self.add_argument("--global-seed", type=int, default=None, help="Entropy Seed (used for all prompts in the batch)") + self._add_image_generator_common_arguments() + + def add_output_arguments(self): + self.add_argument("--metadata", action="store_true", help="Export image metadata as a JSON file.") + 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_controlnet_arguments(self): + self.add_argument("--controlnet-image-path", type=str, required=True, help="Local path of the image to use as input for controlnet.") + self.add_argument("--controlnet-strength", type=float, default=ui_defaults.CONTROLNET_STRENGTH, help=f"Controls how strongly the control image influences the output image. A value of 0.0 means no influence. (Default is {ui_defaults.CONTROLNET_STRENGTH})") + self.add_argument("--controlnet-save-canny", action="store_true", help="If set, save the Canny edge detection reference input image.") + + def parse_args(self): + namespace = super().parse_args() + if hasattr(namespace, "path") and namespace.path is not None and namespace.model is None: + namespace.error("--model must be specified when using --path") + if hasattr(namespace, "steps") and namespace.steps is None: + namespace.steps = ui_defaults.MODEL_INFERENCE_STEPS.get(namespace.model, None) + return namespace diff --git a/src/mflux/ui/defaults.py b/src/mflux/ui/defaults.py new file mode 100644 index 0000000..fec8d7d --- /dev/null +++ b/src/mflux/ui/defaults.py @@ -0,0 +1,9 @@ +CONTROLNET_STRENGTH = 0.4 +GUIDANCE_SCALE = 3.5 +HEIGHT, WIDTH = 1024, 1024 +MODEL_CHOICES = ["dev", "schnell"] +MODEL_INFERENCE_STEPS = { + "dev": 14, + "schnell": 4, +} +QUANTIZE_CHOICES = [4, 8]