Add support for --prompt-file for dynamic prompt updates in large batches (#185)

Co-authored-by: Anthony Wu <pls-file-gh-issue@users.noreply.github.com>
Co-authored-by: filipstrand <strand.filip@gmail.com>

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
 

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."""
 

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())

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())

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())

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())

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())

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())

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)

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

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: