Support prompt text via stdin. e.g. LLMs can generate prompts that directly feed into mflux. (#204)

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

.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

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

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.

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

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

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