From c31326ff336d08b29af699d14bc552cf5044930b Mon Sep 17 00:00:00 2001 From: Anthony Wu <462072+anthonywu@users.noreply.github.com> Date: Tue, 10 Jun 2025 22:31:46 -0700 Subject: [PATCH] Support prompt text via stdin. e.g. LLMs can generate prompts that directly feed into mflux. (#204) Co-authored-by: Anthony Wu --- .pre-commit-config.yaml | 2 +- README.md | 10 +- ignore_pr.md | 62 ++++++++ src/mflux/ui/prompt_utils.py | 24 ++- tests/arg_parser/test_stdin_prompt.py | 77 +++++++++ .../test_stdin_prompt_integration.py | 148 ++++++++++++++++++ 6 files changed, 318 insertions(+), 5 deletions(-) create mode 100644 ignore_pr.md create mode 100644 tests/arg_parser/test_stdin_prompt.py create mode 100644 tests/image_generation/test_stdin_prompt_integration.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 440eef8..0d5f5fd 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,7 +1,7 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit # Ruff version. - rev: v0.11.2 + rev: v0.11.13 hooks: # Run the linter. - id: ruff diff --git a/README.md b/README.md index 2121644..03c2545 100644 --- a/README.md +++ b/README.md @@ -150,6 +150,14 @@ This example uses the more powerful `dev` model with 25 time steps: mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2 -q 8 ``` +You can also pipe prompts from other commands using stdin: + +```sh +echo "A majestic mountain landscape" | mflux-generate --prompt - --model schnell --steps 2 +``` + +This is useful for integrating MFLUX into shell scripts or dynamically generating prompts using LLM inference tools such as [`llm`](https://llm.datasette.io/en/stable/), [`mlx-lm`](https://github.com/ml-explore/mlx-lm), [`ollama`](https://ollama.ai/), etc. + ⚠️ *If the specific model is not already downloaded on your machine, it will start the download process and fetch the model weights (~34GB in size for the Schnell or Dev model respectively). See the [quantization](#%EF%B8%8F-quantization) section for running compressed versions of the model.* ⚠️ *By default, model files are downloaded to the `.cache` folder within your home directory. For example, in my setup, the path looks like this:* @@ -164,7 +172,7 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2 #### 📜 Full list of Command-Line Arguments -- **`--prompt`** (required, `str`): Text description of the image to generate. +- **`--prompt`** (required, `str`): Text description of the image to generate. Use `-` to read the prompt from stdin (e.g., `echo "A beautiful sunset" | mflux-generate --prompt -`). - **`--model`** or **`-m`** (required, `str`): Model to use for generation. Can be one of the official models (`"schnell"` or `"dev"`) or a HuggingFace repository ID for a compatible third-party model (e.g., `"Freepik/flux.1-lite-8B-alpha"`). diff --git a/ignore_pr.md b/ignore_pr.md new file mode 100644 index 0000000..f02e1c0 --- /dev/null +++ b/ignore_pr.md @@ -0,0 +1,62 @@ +# Add stdin support for prompt input + +## Summary + +This PR adds support for reading prompts from stdin when `--prompt -` is specified, following Unix/Linux conventions. This enables piping prompts from other commands, making mflux more composable in shell scripts and workflows. + +## Changes + +### Feature Implementation +- **Added stdin support in `prompt_utils.py`**: When `--prompt -` is provided, the prompt is read from stdin +- **Follows Unix conventions**: Uses `-` as the special marker for stdin input, consistent with tools like `cat`, `grep`, etc. +- **Proper error handling**: Raises `PromptFileReadError` with clear message when stdin is empty +- **Preserves existing behavior**: Regular prompts and `--prompt-file` continue to work as before + +### Code Improvements +- **Converted print statements to logging**: Changed `print()` calls to `logger.info()` for better control over output +- **Improved type annotations**: Changed `get_effective_prompt(args: t.Any)` to use proper `argparse.Namespace` type +- **Added comprehensive tests**: Both unit tests and integration tests verify the functionality + +## Usage Examples + +```bash +# Pipe prompt from echo +echo "A beautiful sunset over mountains" | mflux-generate --prompt - --model dev --steps 20 + +# Pipe prompt from a file +cat prompt.txt | mflux-generate --prompt - --model dev --output result.png + +# Use in a pipeline +generate-prompt-script.py | mflux-generate --prompt - --model dev --metadata +``` + +## Testing + +### Unit Tests (`tests/arg_parser/test_stdin_prompt.py`) +- ✅ Basic stdin parsing (`--prompt -` returns "-") +- ✅ Regular prompt still works unchanged +- ✅ Whitespace trimming from stdin +- ✅ `--prompt-file` takes precedence over stdin + +### Integration Tests (`tests/image_generation/test_stdin_prompt_integration.py`) +- ✅ Single-line stdin prompt generates image with correct metadata +- ✅ Multi-line stdin prompt is preserved correctly +- ✅ Empty stdin shows appropriate error message +- ✅ Shell piping with echo command works correctly + +All tests verify that the stdin prompt is correctly captured in the output metadata when `--metadata` flag is used. + +## Technical Details + +- Stdin reading happens in `get_effective_prompt()` when `args.prompt == "-"` +- Uses `sys.stdin.read().strip()` to read and clean the input +- Logging messages indicate when prompt is read from stdin vs file +- Error handling for `IOError`, `OSError`, and `KeyboardInterrupt` + +## Breaking Changes + +None. This is a purely additive feature that doesn't affect existing functionality. + +## Future Considerations + +This implementation could be extended to other input arguments if needed, following the same `-` convention for stdin input. \ No newline at end of file diff --git a/src/mflux/ui/prompt_utils.py b/src/mflux/ui/prompt_utils.py index 38080ae..97a9fb7 100644 --- a/src/mflux/ui/prompt_utils.py +++ b/src/mflux/ui/prompt_utils.py @@ -1,8 +1,12 @@ -import typing as t +import logging +import sys +from argparse import Namespace from pathlib import Path from mflux.error.exceptions import PromptFileReadError +logger = logging.getLogger(__name__) + def read_prompt_file(prompt_file_path: Path) -> str: # Check if file exists @@ -17,14 +21,28 @@ def read_prompt_file(prompt_file_path: Path) -> str: if not content: raise PromptFileReadError(f"Prompt file is empty: {prompt_file_path}") - print(f"Using prompt from file: {prompt_file_path}") + logger.info(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: +def get_effective_prompt(args: Namespace) -> str: + # Handle stdin input when prompt is "-" + if args.prompt == "-": + try: + content = sys.stdin.read().strip() + if not content: + raise PromptFileReadError("No prompt provided via stdin") + logger.info("Using prompt from stdin") + return content + except (IOError, OSError, KeyboardInterrupt) as e: + raise PromptFileReadError(f"Error reading from stdin: {e}") + + # Handle prompt file if args.prompt_file is not None: prompt = read_prompt_file(args.prompt_file) return prompt + + # Return regular prompt return args.prompt diff --git a/tests/arg_parser/test_stdin_prompt.py b/tests/arg_parser/test_stdin_prompt.py new file mode 100644 index 0000000..53cf589 --- /dev/null +++ b/tests/arg_parser/test_stdin_prompt.py @@ -0,0 +1,77 @@ +from io import StringIO +from pathlib import Path +from unittest.mock import patch + +import pytest + +from mflux.ui.cli.parsers import CommandLineParser +from mflux.ui.prompt_utils import get_effective_prompt + + +@pytest.fixture +def mflux_generate_parser() -> CommandLineParser: + parser = CommandLineParser(description="Generate an image based on a prompt.") + parser.add_general_arguments() + parser.add_model_arguments(require_model_arg=False) + 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() + parser.add_output_arguments() + return parser + + +@pytest.fixture +def temp_output_dir(tmp_path_factory) -> Path: + return tmp_path_factory.mktemp("mflux_stdin_test") + + +def test_prompt_from_stdin(mflux_generate_parser): + """Test that --prompt - reads from stdin correctly.""" + stdin_content = "A beautiful sunset over the ocean" + + # Simulate stdin input + with patch("sys.stdin", StringIO(stdin_content)): + with patch("sys.argv", ["mflux-generate", "--prompt", "-", "--model", "dev"]): + args = mflux_generate_parser.parse_args() + + # The parser returns the raw args, get_effective_prompt handles stdin + assert args.prompt == "-" + + +def test_prompt_stdin_vs_regular(mflux_generate_parser): + """Test that regular prompt still works when not using stdin.""" + regular_prompt = "A regular prompt not from stdin" + + with patch("sys.argv", ["mflux-generate", "--prompt", regular_prompt, "--model", "dev"]): + args = mflux_generate_parser.parse_args() + assert args.prompt == regular_prompt + assert get_effective_prompt(args) == regular_prompt + + +def test_prompt_stdin_with_whitespace(mflux_generate_parser): + """Test that stdin prompt with surrounding whitespace is properly stripped.""" + stdin_content = "\n\n A prompt with whitespace \n\n" + expected_prompt = "A prompt with whitespace" + + with patch("sys.stdin", StringIO(stdin_content)): + with patch("sys.argv", ["mflux-generate", "--prompt", "-", "--model", "dev"]): + args = mflux_generate_parser.parse_args() + effective_prompt = get_effective_prompt(args) + assert effective_prompt == expected_prompt + + +def test_prompt_file_takes_precedence_over_stdin(mflux_generate_parser, temp_output_dir): + """Test that --prompt-file still works and takes precedence over stdin detection. + because --prompt is not used in this scenario.""" + # Create a prompt file + prompt_file = temp_output_dir / "prompt.txt" + file_prompt = "Prompt from file" + prompt_file.write_text(file_prompt) + stdin_content = "This should not be used because --prompt is not provided in the command." + + with patch("sys.stdin", StringIO(stdin_content)): + with patch("sys.argv", ["mflux-generate", "--prompt-file", str(prompt_file), "--model", "dev"]): + args = mflux_generate_parser.parse_args() + effective_prompt = get_effective_prompt(args) + assert effective_prompt == file_prompt diff --git a/tests/image_generation/test_stdin_prompt_integration.py b/tests/image_generation/test_stdin_prompt_integration.py new file mode 100644 index 0000000..e7d4322 --- /dev/null +++ b/tests/image_generation/test_stdin_prompt_integration.py @@ -0,0 +1,148 @@ +import json +import subprocess +import sys +from pathlib import Path + +import pytest + + +@pytest.fixture +def temp_output_dir(tmp_path) -> Path: + return tmp_path + + +def test_stdin_prompt_with_actual_generation(temp_output_dir): + """Test that stdin prompt is correctly used in actual image generation and saved in metadata.""" + stdin_prompt = "A beautiful mountain landscape with snow" + output_image = temp_output_dir / "test_stdin.png" + metadata_file = temp_output_dir / "test_stdin.json" + + # Run the actual mflux.generate module with stdin + cmd = [ + sys.executable, + "-m", + "mflux.generate", + "--prompt", + "-", + "--model", + "dev", + "--steps", + "1", + "--seed", + "42", + "--output", + str(output_image), + "--metadata", + ] + + # Execute the command with stdin input + process = subprocess.run( + cmd, + input=stdin_prompt, + text=True, + capture_output=True, + ) + + # Check that the command succeeded + assert process.returncode == 0, f"Command failed with stderr: {process.stderr}" + + # Check that output files were created + assert output_image.exists(), "Output image was not created" + assert metadata_file.exists(), "Metadata file was not created" + + # Load and verify metadata + with open(metadata_file, "r") as f: + metadata = json.load(f) + + # Verify the prompt in metadata matches our stdin input + assert metadata["prompt"] == stdin_prompt + + +def test_stdin_prompt_multiline_with_actual_generation(temp_output_dir): + """Test that multiline stdin prompt is correctly preserved in metadata.""" + stdin_prompt = """A fantasy scene with: +- Dragons flying in the sky +- A castle on a mountain +- Magical aurora lights""" + + output_image = temp_output_dir / "test_multiline.png" + metadata_file = temp_output_dir / "test_multiline.json" + + cmd = [ + sys.executable, + "-m", + "mflux.generate", + "--prompt", + "-", + "--model", + "dev", + "--steps", + "1", + "--seed", + "123", + "--output", + str(output_image), + "--metadata", + ] + + process = subprocess.run(cmd, input=stdin_prompt, text=True, capture_output=True) + + assert process.returncode == 0, f"Command failed with stderr: {process.stderr}" + assert metadata_file.exists(), "Metadata file was not created" + + with open(metadata_file, "r") as f: + metadata = json.load(f) + + # Verify the multiline prompt is preserved (strip() removes leading/trailing whitespace) + assert metadata["prompt"] == stdin_prompt.strip() + + +def test_empty_stdin_fails_generation(temp_output_dir): + """Test that empty stdin causes generation to fail with appropriate error.""" + output_image = temp_output_dir / "test_empty.png" + + cmd = [ + sys.executable, + "-m", + "mflux.generate", + "--prompt", + "-", + "--model", + "dev", + "--steps", + "1", + "--output", + str(output_image), + ] + + process = subprocess.run( + cmd, + input="", # Empty stdin + text=True, + capture_output=True, + ) + + # The application prints the error but exits with 0 (by design) + assert process.returncode == 0 + # Error message should be in stdout (as it's printed, not written to stderr) + assert "No prompt provided via stdin" in process.stdout + + +def test_pipe_from_echo_command(temp_output_dir): + """Test using echo command to pipe prompt, simulating real usage.""" + prompt = "A serene lake at sunset" + output_image = temp_output_dir / "test_echo.png" + metadata_file = temp_output_dir / "test_echo.json" + + # Simulate: echo "prompt" | mflux-generate --prompt - ... + echo_cmd = f'echo "{prompt}" | {sys.executable} -m mflux.generate --prompt - --model dev --steps 1 --output {output_image} --metadata' + + process = subprocess.run(echo_cmd, shell=True, capture_output=True, text=True) + + assert process.returncode == 0, f"Command failed with stderr: {process.stderr}" + assert metadata_file.exists(), "Metadata file was not created" + + with open(metadata_file, "r") as f: + metadata = json.load(f) + + assert metadata["prompt"] == prompt