diff --git a/README.md b/README.md index 1b7b170..b005763 100644 --- a/README.md +++ b/README.md @@ -268,11 +268,50 @@ The `mflux-generate-controlnet` command supports most of the same arguments as ` See the [Controlnet](#%EF%B8%8F-controlnet) section for more details on how to use this feature effectively. -#### 📜 Batch Image Generation Arguments +#### Dynamic Prompts with `--prompt-file` -- **`--prompts-file`** (required, `str`): Local path for a file that holds a batch of prompts. +MFlux supports dynamic prompt updates through the `--prompt-file` option. Instead of providing a fixed prompt with `--prompt`, you can specify a plain text file containing your prompt. The file is re-read before each generation, allowing you to modify prompts between iterations without restarting. -- **`--global-seed`** (optional, `int`): Entropy Seed used for all prompts in the batch. +##### Key Benefits + +- **Real-time Prompt Editing**: Edit prompts between generations while the program continues running +- **Iterative Refinement**: Fine-tune prompts based on previous results in a batch +- **Editor Integration**: Work with complex prompts in your preferred text editor +- **Consistent Parameters**: Experiment with prompt variations while keeping all other settings constant + +##### Example Usage + +```bash +# Generate 10 images using a prompt from a file +mflux-generate --prompt-file my_prompt.txt --auto-seeds 10 +``` + +##### Workflow Example + +1. Create a prompt file: + ``` + echo "a surreal landscape with floating islands" > my_prompt.txt + ``` + +2. Start a batch generation: + ``` + mflux-generate --prompt-file my_prompt.txt --auto-seeds 10 + ``` + +3. After reviewing initial results, edit the prompt file while generation continues: + ``` + echo "a surreal landscape with floating islands and waterfalls" > my_prompt.txt + ``` + +4. Subsequent generations will automatically use your updated prompt + +5. Continue refining as needed between generations + +##### Notes + +- `--prompt` and `--prompt-file` are mutually exclusive options +- Empty prompt files or non-existent files will raise appropriate errors +- Each generated image's metadata will contain the actual prompt used for that specific generation #### 📜 Training Arguments diff --git a/src/mflux/error/exceptions.py b/src/mflux/error/exceptions.py index e2693ac..f3bd743 100644 --- a/src/mflux/error/exceptions.py +++ b/src/mflux/error/exceptions.py @@ -14,6 +14,10 @@ class MFluxUserException(MFluxException): """an exception raised by user behavior or intention.""" +class PromptFileReadError(MFluxUserException): + """Exception raised for errors in reading the prompt file.""" + + class StopImageGenerationException(MFluxUserException): """user has requested to stop a image generation in progress.""" diff --git a/src/mflux/generate.py b/src/mflux/generate.py index dbd7765..3128bac 100644 --- a/src/mflux/generate.py +++ b/src/mflux/generate.py @@ -2,7 +2,9 @@ from mflux import Config, Flux1, ModelConfig, StopImageGenerationException from mflux.callbacks.callback_registry import CallbackRegistry from mflux.callbacks.instances.memory_saver import MemorySaver from mflux.callbacks.instances.stepwise_handler import StepwiseHandler +from mflux.error.exceptions import PromptFileReadError from mflux.ui.cli.parsers import CommandLineParser +from mflux.ui.prompt_utils import get_effective_prompt def main(): @@ -44,7 +46,7 @@ def main(): # 3. Generate an image for each seed value image = flux.generate_image( seed=seed, - prompt=args.prompt, + prompt=get_effective_prompt(args), config=Config( num_inference_steps=args.steps, height=args.height, @@ -56,8 +58,8 @@ def main(): ) # 4. Save the image image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata) - except StopImageGenerationException as stop_exc: - print(stop_exc) + except (StopImageGenerationException, PromptFileReadError) as exc: + print(exc) finally: if memory_saver: print(memory_saver.memory_stats()) diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py index 41dfd79..38ffd73 100644 --- a/src/mflux/generate_controlnet.py +++ b/src/mflux/generate_controlnet.py @@ -3,7 +3,9 @@ from mflux.callbacks.callback_registry import CallbackRegistry from mflux.callbacks.instances.canny_saver import CannyImageSaver from mflux.callbacks.instances.memory_saver import MemorySaver from mflux.callbacks.instances.stepwise_handler import StepwiseHandler +from mflux.error.exceptions import PromptFileReadError from mflux.ui.cli.parsers import CommandLineParser +from mflux.ui.prompt_utils import get_effective_prompt def main(): @@ -47,7 +49,7 @@ def main(): # 3. Generate an image for each seed value image = flux.generate_image( seed=seed, - prompt=args.prompt, + prompt=get_effective_prompt(args), controlnet_image_path=args.controlnet_image_path, config=Config( num_inference_steps=args.steps, @@ -60,8 +62,8 @@ def main(): # 4. Save the image image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata) - except StopImageGenerationException as stop_exc: - print(stop_exc) + except (StopImageGenerationException, PromptFileReadError) as exc: + print(exc) finally: if memory_saver: print(memory_saver.memory_stats()) diff --git a/src/mflux/generate_depth.py b/src/mflux/generate_depth.py index 4860bc0..eb34ec6 100644 --- a/src/mflux/generate_depth.py +++ b/src/mflux/generate_depth.py @@ -3,8 +3,10 @@ from mflux.callbacks.callback_registry import CallbackRegistry from mflux.callbacks.instances.depth_saver import DepthImageSaver from mflux.callbacks.instances.memory_saver import MemorySaver from mflux.callbacks.instances.stepwise_handler import StepwiseHandler +from mflux.error.exceptions import PromptFileReadError from mflux.flux_tools.depth.flux_depth import Flux1Depth from mflux.ui.cli.parsers import CommandLineParser +from mflux.ui.prompt_utils import get_effective_prompt def main(): @@ -51,7 +53,7 @@ def main(): # 3. Generate an image for each seed value image = flux.generate_image( seed=seed, - prompt=args.prompt, + prompt=get_effective_prompt(args), config=Config( num_inference_steps=args.steps, height=args.height, @@ -64,8 +66,8 @@ def main(): # 4. Save the image image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata) - except StopImageGenerationException as stop_exc: - print(stop_exc) + except (StopImageGenerationException, PromptFileReadError) as exc: + print(exc) finally: if memory_saver: print(memory_saver.memory_stats()) diff --git a/src/mflux/generate_fill.py b/src/mflux/generate_fill.py index ae9e131..968d407 100644 --- a/src/mflux/generate_fill.py +++ b/src/mflux/generate_fill.py @@ -2,8 +2,10 @@ from mflux import Config, StopImageGenerationException from mflux.callbacks.callback_registry import CallbackRegistry from mflux.callbacks.instances.memory_saver import MemorySaver from mflux.callbacks.instances.stepwise_handler import StepwiseHandler +from mflux.error.exceptions import PromptFileReadError from mflux.flux_tools.fill.flux_fill import Flux1Fill from mflux.ui.cli.parsers import CommandLineParser +from mflux.ui.prompt_utils import get_effective_prompt def main(): @@ -48,7 +50,7 @@ def main(): # 3. Generate an image for each seed value image = flux.generate_image( seed=seed, - prompt=args.prompt, + prompt=get_effective_prompt(args), config=Config( num_inference_steps=args.steps, height=args.height, @@ -61,8 +63,8 @@ def main(): # 4. Save the image image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata) - except StopImageGenerationException as stop_exc: - print(stop_exc) + except (StopImageGenerationException, PromptFileReadError) as exc: + print(exc) finally: if memory_saver: print(memory_saver.memory_stats()) diff --git a/src/mflux/generate_in_context.py b/src/mflux/generate_in_context.py index 8ad2cb9..fefa443 100644 --- a/src/mflux/generate_in_context.py +++ b/src/mflux/generate_in_context.py @@ -7,7 +7,9 @@ from mflux.callbacks.instances.stepwise_handler import StepwiseHandler from mflux.community.in_context_lora.flux_in_context_lora import Flux1InContextLoRA from mflux.community.in_context_lora.in_context_loras import LORA_REPO_ID, get_lora_filename from mflux.config.model_config import ModelConfig +from mflux.error.exceptions import PromptFileReadError from mflux.ui.cli.parsers import CommandLineParser +from mflux.ui.prompt_utils import get_effective_prompt def main(): @@ -56,7 +58,7 @@ def main(): # 3. Generate an image for each seed value image = flux.generate_image( seed=seed, - prompt=args.prompt, + prompt=get_effective_prompt(args), config=Config( num_inference_steps=args.steps, height=args.height, @@ -70,8 +72,8 @@ def main(): image.get_right_half().save(path=output_path, export_json_metadata=args.metadata) if args.save_full_image: image.save(path=output_path.with_stem(output_path.stem + "_full")) - except StopImageGenerationException as stop_exc: - print(stop_exc) + except (StopImageGenerationException, PromptFileReadError) as exc: + print(exc) finally: if memory_saver: print(memory_saver.memory_stats()) diff --git a/src/mflux/generate_redux.py b/src/mflux/generate_redux.py index b52ed3f..8b77b05 100644 --- a/src/mflux/generate_redux.py +++ b/src/mflux/generate_redux.py @@ -4,8 +4,10 @@ from mflux import Config, ModelConfig, StopImageGenerationException from mflux.callbacks.callback_registry import CallbackRegistry from mflux.callbacks.instances.memory_saver import MemorySaver from mflux.callbacks.instances.stepwise_handler import StepwiseHandler +from mflux.error.exceptions import PromptFileReadError from mflux.flux_tools.redux.flux_redux import Flux1Redux from mflux.ui.cli.parsers import CommandLineParser +from mflux.ui.prompt_utils import get_effective_prompt def main(): @@ -53,7 +55,7 @@ def main(): # 3. Generate an image for each seed value image = flux.generate_image( seed=seed, - prompt=args.prompt, + prompt=get_effective_prompt(args), config=Config( num_inference_steps=args.steps, height=args.height, @@ -66,8 +68,8 @@ def main(): # 4. Save the image image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata) - except StopImageGenerationException as stop_exc: - print(stop_exc) + except (StopImageGenerationException, PromptFileReadError) as exc: + print(exc) finally: if memory_saver: print(memory_saver.memory_stats()) diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py index 343c00c..366ca15 100644 --- a/src/mflux/ui/cli/parsers.py +++ b/src/mflux/ui/cli/parsers.py @@ -71,7 +71,9 @@ class CommandLineParser(argparse.ArgumentParser): 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, supports_metadata_config=False) -> None: - self.add_argument("--prompt", type=str, required=(not supports_metadata_config), default=None, help="The textual description of the image to generate.") + prompt_group = self.add_mutually_exclusive_group(required=(not supports_metadata_config)) + prompt_group.add_argument("--prompt", type=str, help="The textual description of the image to generate.") + prompt_group.add_argument("--prompt-file", type=Path, help="Path to a file containing the prompt text. The file will be re-read before each generation, allowing you to edit the prompt between iterations when using multiple seeds without restarting the program.") self.add_argument("--seed", type=int, default=None, nargs='+', help="Specify 1+ Entropy Seeds (Default is 1 time-based random-seed)") self.add_argument("--auto-seeds", type=int, default=-1, help="Auto generate N Entropy Seeds (random ints between 0 and 1 billion") self._add_image_generator_common_arguments() @@ -84,7 +86,7 @@ class CommandLineParser(argparse.ArgumentParser): self.add_argument("--image-strength", type=float, required=False, default=ui_defaults.IMAGE_STRENGTH, help=f"Controls how strongly the init image influences the output image. A value of 0.0 means no influence. (Default is {ui_defaults.IMAGE_STRENGTH})") def add_batch_image_generator_arguments(self) -> None: - self.add_argument("--prompts-file", type=Path, required=True, default=argparse.SUPPRESS, help="Local path for a file that holds a batch of prompts.") + self.add_argument("--batch-prompts-file", type=Path, required=True, default=argparse.SUPPRESS, help="Local path for a file that holds a batch of prompts.") self.add_argument("--global-seed", type=int, default=argparse.SUPPRESS, help="Entropy Seed (used for all prompts in the batch)") self._add_image_generator_common_arguments() @@ -220,9 +222,9 @@ class CommandLineParser(argparse.ArgumentParser): output_path = Path(namespace.output) namespace.output = str(output_path.with_stem(output_path.stem + "_seed_{seed}")) - if self.supports_image_generation and namespace.prompt is None: - # not supplied by CLI and not supplied by metadata config file - self.error("--prompt argument required or 'prompt' required in metadata config file") + if self.supports_image_generation and namespace.prompt is None and namespace.prompt_file is None: + # when metadata config is supported but neither prompt nor prompt-file is provided + self.error("Either --prompt or --prompt-file argument is required, or 'prompt' required in metadata config file") if self.supports_image_generation and namespace.steps is None: namespace.steps = ui_defaults.MODEL_INFERENCE_STEPS.get(namespace.model, 14) diff --git a/src/mflux/ui/prompt_utils.py b/src/mflux/ui/prompt_utils.py new file mode 100644 index 0000000..38080ae --- /dev/null +++ b/src/mflux/ui/prompt_utils.py @@ -0,0 +1,30 @@ +import typing as t +from pathlib import Path + +from mflux.error.exceptions import PromptFileReadError + + +def read_prompt_file(prompt_file_path: Path) -> str: + # Check if file exists + if not prompt_file_path.exists(): + raise PromptFileReadError(f"Prompt file does not exist: {prompt_file_path}") + + try: + with open(prompt_file_path, "rt") as f: + content: str = f.read().strip() + + # Validate content + if not content: + raise PromptFileReadError(f"Prompt file is empty: {prompt_file_path}") + + print(f"Using prompt from file: {prompt_file_path}") + return content + except (IOError, OSError) as e: + raise PromptFileReadError(f"Error reading prompt file '{prompt_file_path}': {e}") + + +def get_effective_prompt(args: t.Any) -> str: + if args.prompt_file is not None: + prompt = read_prompt_file(args.prompt_file) + return prompt + return args.prompt diff --git a/tests/arg_parser/test_cli_argparser.py b/tests/arg_parser/test_cli_argparser.py index 8ed12d9..cde36e1 100644 --- a/tests/arg_parser/test_cli_argparser.py +++ b/tests/arg_parser/test_cli_argparser.py @@ -222,6 +222,32 @@ def test_prompt_arg(mflux_generate_parser, mflux_generate_minimal_argv, base_met assert args.prompt == cli_prompt +def test_prompt_file_arg(mflux_generate_parser, mflux_generate_minimal_argv, temp_dir): + # Create a prompt file + prompt_content = "prompt from a file being re-read for each generation" + prompt_file = temp_dir / "prompt.txt" + with prompt_file.open("wt") as pf: + pf.write(prompt_content) + + # Test that prompt-file is correctly read + with patch('sys.argv', ["mflux-generate", "--prompt-file", prompt_file.as_posix()]): # fmt: off + args = mflux_generate_parser.parse_args() + assert args.prompt_file == prompt_file + assert args.prompt is None # prompt should be None since we're using prompt-file + + +def test_prompt_and_prompt_file_mutually_exclusive(mflux_generate_parser, temp_dir): + # Create a prompt file + prompt_file = temp_dir / "prompt.txt" + with prompt_file.open("wt") as pf: + pf.write("some prompt content") + + # Test that using both --prompt and --prompt-file raises an error + with pytest.raises(SystemExit): + with patch('sys.argv', ["mflux-generate", "--prompt", "direct prompt", "--prompt-file", prompt_file.as_posix()]): # fmt: off + mflux_generate_parser.parse_args() + + def test_guidance_arg(mflux_generate_parser, mflux_generate_minimal_argv, base_metadata_dict, temp_dir): # fmt: off metadata_file = temp_dir / "guidance.json" with metadata_file.open("wt") as m: