monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

➡️➡️ Tab Tab for zsh: CLI completion generator (#208)

Browse Source 
Co-authored-by: Anthony Wu <pls-file-gh-issue@users.noreply.github.com>
Co-authored-by: filipstrand <strand.filip@gmail.com>
This commit is contained in:
Anthony Wu 2025-06-27 06:48:51 -07:00 committed by GitHub
parent 6c3ba1e065
commit 0125103634
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 703 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
README.md
View File

@ -12,6 +12,7 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black
- [Philosophy](#philosophy)
- [💿 Installation](#-installation)
- [🚀 Shell Completions (Quick Start)](#-shell-completions-quick-start)
- [🖼️ Generating an image](#%EF%B8%8F-generating-an-image)
* [📜 Full list of Command-Line Arguments](#-full-list-of-command-line-arguments)
- [⏱️ Image generation speed (updated)](#%EF%B8%8F-image-generation-speed-updated)
@ -139,6 +140,31 @@ pip install -U mflux
*If you have trouble installing MFLUX, please see the [installation related issues section](https://github.com/filipstrand/mflux/issues?q=is%3Aissue+install+).*
### ⌨️ Shell Completions (Quick Start)
MFLUX supports ZSH (default on macOS) shell completions for all CLI commands.
To enable completions:
```sh
# Install completions
mflux-completions
# Reload your shell
exec zsh
# Now try tab completion!
mflux-generate --<TAB>
```
If you encounter issues (ZSH completion setup can vary by system), check your installation:
```sh
mflux-completions --check
```
For more details and troubleshooting, see the [completions documentation](src/mflux/completions/README.md).
### 🖼️ Generating an image
Run the command `mflux-generate` by specifying a prompt and the model and some optional arguments. For example, here we use a quantized version of the `schnell` model for 2 steps:

1
pyproject.toml
View File

@ -71,6 +71,7 @@ mflux-save-depth = "mflux.save_depth:main"
mflux-train = "mflux.train:main"
mflux-upscale = "mflux.upscale:main"
mflux-lora-library = "mflux.lora_library:main"
mflux-completions = "mflux.completions.install:main"
[tool.setuptools.packages.find]
where = ["src"]

144
src/mflux/completions/README.md Normal file
View File

@ -0,0 +1,144 @@
# mflux Shell Completions
This module provides ZSH shell completion support for all mflux CLI commands.
## Quick Start
1. **Install completions:**
```bash
mflux-completions
```
2. **Reload your shell:**
```bash
exec zsh
```
3. **Test it:**
```bash
mflux-generate --<TAB>
```
## Features
- **Auto-completion for all 15 mflux commands**: Generate, train, save, upscale, and more
- **Smart completions**:
- Model names (dev, schnell, dev-fill, or HuggingFace repos)
- Quantization levels (3, 4, 6, 8)
- LoRA styles (storyboard, portrait, etc.)
- File paths with native completion
- Common percentage values for battery limit
- **Dynamic generation**: Completions stay in sync with code changes
- **Mutually exclusive options**: Properly handles -m/--model style options
## Installation
### Automatic Installation
The simplest way is to let mflux-completions auto-detect the best location:
```bash
mflux-completions
```
### Custom Directory
Install to a specific directory:
```bash
mflux-completions --dir ~/.config/zsh/completions
```
### Update Existing Installation
Update an existing completion file:
```bash
mflux-completions --update
```
## Troubleshooting
### Check Installation
Verify that completions are properly installed:
```bash
mflux-completions --check
```
This will:
- Check if the completion file exists in your $fpath
- Verify that compinit is available
- Provide troubleshooting steps if needed
### Common Issues
1. **Completions not working after installation:**
- Start a new shell: `exec zsh`
2. **Completions not updated after mflux upgrade:**
- Update completions: `mflux-completions --update`
- Then force reload: `compinit -u`
3. **Permission denied:**
- Use `--dir` to specify a user-writable directory
4. **Completion shows \} or other artifacts:**
- Update to the latest version: `mflux-completions --update`
For Zsh-specific configuration issues (fpath, compinit, etc.), please refer to the [official Zsh documentation](https://zsh.sourceforge.io/Doc/Release/Completion-System.html).
### Manual Installation
If automatic installation fails:
```bash
# Generate completion script
mflux-completions --print > _mflux
# Move to a directory in your $fpath
mkdir -p ~/.zsh/completions
mv _mflux ~/.zsh/completions/
# Restart your shell
exec zsh
```
## How It Works
The completion system:
1. **Introspects argparse**: Dynamically extracts all command-line options from the parsers
2. **Generates ZSH functions**: Creates completion functions for each mflux command
3. **Provides context-aware completions**: Different completion strategies based on argument type
4. **Handles special cases**: Custom completions for models, LoRA styles, etc.
## Development
### Adding New Commands
When adding a new mflux CLI command:
1. Create the command with argparse as usual
2. Add it to the `commands` list in `generator.py`
3. Add the parser creation logic to `create_parser_for_command()`
4. Completions will be automatically generated!
### Testing Completions
Test the generator without installing:
```bash
python -m mflux.completions.generator
```
### File Structure
- `generator.py`: Core logic for generating ZSH completion scripts
- `install.py`: Installation utility with user-friendly CLI
- `README.md`: This documentation
## Platform Support
Currently supports ZSH (the default macOS shell since 2019). Supporting other shells is possible but is considered out of scope for the official repository.

1
src/mflux/completions/__init__.py Normal file
View File

@ -0,0 +1 @@
"""ZSH completion generator for mflux CLI commands."""

328
src/mflux/completions/generator.py Normal file
View File

@ -0,0 +1,328 @@
"""Generate ZSH completion scripts for mflux commands."""
import argparse
from pathlib import Path
from mflux.community.in_context.utils.in_context_loras import LORA_NAME_MAP
from mflux.ui import defaults as ui_defaults
from mflux.ui.cli.parsers import CommandLineParser
class CompletionGenerator:
"""Generate ZSH completion scripts by introspecting argparse parsers."""
def __init__(self):
self.commands = [
"mflux-generate",
"mflux-generate-controlnet",
"mflux-generate-in-context",
"mflux-generate-in-context-edit",
"mflux-generate-in-context-catvton",
"mflux-generate-fill",
"mflux-generate-depth",
"mflux-generate-redux",
"mflux-concept",
"mflux-concept-from-image",
"mflux-save",
"mflux-save-depth",
"mflux-train",
"mflux-upscale",
"mflux-lora-library",
]
def create_parser_for_command(self, command: str) -> CommandLineParser:
"""Create the appropriate parser for a given command."""
parser = CommandLineParser(prog=command, add_help=False)
if command == "mflux-generate":
parser.add_general_arguments()
parser.add_model_arguments()
parser.add_lora_arguments()
parser.add_image_generator_arguments(supports_metadata_config=True)
parser.add_image_to_image_arguments()
parser.add_image_outpaint_arguments()
parser.add_output_arguments()
elif command == "mflux-generate-controlnet":
parser.add_general_arguments()
parser.add_model_arguments()
parser.add_lora_arguments()
parser.add_controlnet_arguments()
parser.add_image_generator_arguments()
parser.add_output_arguments()
elif command == "mflux-generate-in-context":
parser.add_general_arguments()
parser.add_model_arguments()
parser.add_image_generator_arguments()
parser.add_image_to_image_arguments(required=True)
parser.add_output_arguments()
# Add save-full-image manually since it's not in parsers.py
parser.add_argument("--save-full-image", action="store_true")
elif command == "mflux-generate-in-context-edit":
parser.add_general_arguments()
parser.add_model_arguments(require_model_arg=False)
parser.add_lora_arguments()
parser.add_image_generator_arguments(supports_metadata_config=False, require_prompt=False)
parser.add_in_context_edit_arguments()
parser.add_in_context_arguments()
parser.add_output_arguments()
elif command == "mflux-generate-in-context-catvton":
parser.add_general_arguments()
parser.add_model_arguments(require_model_arg=False)
parser.add_lora_arguments()
parser.add_image_generator_arguments(supports_metadata_config=False, require_prompt=False)
parser.add_catvton_arguments()
parser.add_in_context_arguments()
parser.add_output_arguments()
elif command == "mflux-generate-fill":
parser.add_general_arguments()
parser.add_model_arguments()
parser.add_lora_arguments()
parser.add_image_generator_arguments()
parser.add_fill_arguments()
parser.add_output_arguments()
elif command == "mflux-generate-depth":
parser.add_general_arguments()
parser.add_model_arguments()
parser.add_lora_arguments()
parser.add_image_generator_arguments()
parser.add_depth_arguments()
parser.add_output_arguments()
elif command == "mflux-generate-redux":
parser.add_general_arguments()
parser.add_model_arguments()
parser.add_lora_arguments()
parser.add_image_generator_arguments()
parser.add_redux_arguments()
parser.add_output_arguments()
elif command == "mflux-concept":
parser.add_general_arguments()
parser.add_model_arguments()
parser.add_concept_attention_arguments()
parser.add_image_generator_arguments()
parser.add_output_arguments()
elif command == "mflux-concept-from-image":
parser.add_general_arguments()
parser.add_model_arguments()
parser.add_concept_from_image_arguments()
parser.add_image_generator_arguments()
parser.add_output_arguments()
elif command == "mflux-save":
parser.add_model_arguments(path_type="save")
elif command == "mflux-save-depth":
parser.add_save_depth_arguments()
parser.add_output_arguments()
elif command == "mflux-train":
parser.add_model_arguments(require_model_arg=False)
parser.add_training_arguments()
elif command == "mflux-upscale":
parser.add_general_arguments()
parser.add_model_arguments()
parser.add_image_generator_arguments(supports_metadata_config=True)
parser.add_image_to_image_arguments(required=True)
parser.add_output_arguments()
elif command == "mflux-lora-library":
# Special case: has subcommands
subparsers = parser.add_subparsers(dest="subcommand")
list_parser = subparsers.add_parser("list")
list_parser.add_argument("--paths", action="store_true")
return parser
def escape_description(self, desc: str) -> str:
"""Escape special characters in descriptions for ZSH."""
if not desc:
return ""
# Escape brackets, quotes, and other special characters
desc = desc.replace("\\", "\\\\") # Escape backslashes first
desc = desc.replace("[", "\\[").replace("]", "\\]")
desc = desc.replace("'", "") # Remove single quotes entirely
desc = desc.replace('"', '\\"')
desc = desc.replace("`", "\\`")
desc = desc.replace("$", "\\$")
return desc
def format_argument_spec(self, action: argparse.Action) -> list[str]:
"""Format an argparse action into ZSH completion syntax."""
specs = []
# Handle options with both short and long forms
if action.option_strings:
opts = action.option_strings
desc = self.escape_description(action.help or "")
# For store_true/store_false actions, no value spec needed
is_flag = isinstance(action, (argparse._StoreTrueAction, argparse._StoreFalseAction))
if len(opts) == 2:
# Both short and long form - create exclusion group
opt_spec = f"'({opts[0]} {opts[1]})'{{{opts[0]},{opts[1]}}}"
else:
# Only one form
opt_spec = f"'{opts[0]}'"
opt_spec += f"'[{desc}]"
# Add value spec for non-flag arguments
if not is_flag and (action.type or action.choices or action.nargs):
opt_spec += self.get_value_spec(action)
opt_spec += "'"
specs.append(opt_spec)
return specs
def get_value_spec(self, action: argparse.Action) -> str:
"""Get the value specification for an argument."""
# Check for special cases first
if action.dest == "model":
return ":model:_mflux_models"
elif action.dest == "quantize":
return ":quantization:_mflux_quantize"
elif action.dest == "lora_style":
return ":style:_mflux_lora_styles"
if action.choices:
# Fixed choices
choices = " ".join(str(c) for c in action.choices)
return f":value:({choices})"
elif action.type == Path or (action.dest and ("path" in action.dest or "file" in action.dest)):
# File completion
return ":file:_files"
elif action.type is int:
# Integer value
return ":number:"
elif action.type is float:
# Float value
return ":float:"
elif callable(action.type):
# Callable type (like lambda) - usually returns int or float
# Check the dest name for hints
if "percentage" in (action.dest or "") or "percent" in (action.dest or ""):
return ":percentage:(1 5 10 15 20 25 30 40 50 60 70 80 90 95 99)"
return ":value:"
elif action.nargs in ["*", "+"]:
# Multiple values
if action.type == Path or (action.dest and ("path" in action.dest)):
return ":files:_files"
return ":values:"
elif action.type is str or action.type is None:
# String value
return f":{action.dest or 'value'}:"
return ""
def generate_command_function(self, command: str, parser: CommandLineParser) -> str:
"""Generate the ZSH completion function for a specific command."""
func_name = command.replace("-", "_")
lines = [f"_{func_name}() {{"]
lines.append(" local -a args")
lines.append(" args=(")
# Extract all actions from the parser
for action in parser._actions:
if isinstance(action, argparse._HelpAction):
continue
arg_specs = self.format_argument_spec(action)
lines.extend(f" {spec}" for spec in arg_specs)
lines.append(" )")
lines.append(" _arguments -s -S $args")
lines.append("}")
lines.append("")
return "\n".join(lines)
def generate_header(self) -> str:
"""Generate the completion script header."""
commands = " ".join(self.commands)
return f"""#compdef {commands}
# ZSH completion for mflux commands
# Generated by mflux-completions
"""
def generate_helper_functions(self) -> str:
"""Generate helper functions for common completions."""
helpers = []
# Model completion helper
model_choices = " ".join(f"'{m}[{m} model]'" for m in ui_defaults.MODEL_CHOICES)
helpers.append(f"""_mflux_models() {{
_values 'model' \\
{model_choices} \\
'*:Hugging Face repo:(org/model)'
}}
""")
# LoRA style completion helper
if LORA_NAME_MAP:
lora_choices = " ".join(f"'{k}[{v}]'" for k, v in LORA_NAME_MAP.items())
helpers.append(f"""_mflux_lora_styles() {{
_values 'lora style' \\
{lora_choices}
}}
""")
# Quantization choices
quant_choices = " ".join(f"'{q}[{q}-bit quantization]'" for q in ui_defaults.QUANTIZE_CHOICES)
helpers.append(f"""_mflux_quantize() {{
_values 'quantization' \\
{quant_choices}
}}
""")
return "\n".join(helpers)
def generate_main_function(self) -> str:
"""Generate the main completion dispatcher."""
lines = ["# Main completion dispatcher", "_mflux() {", " local cmd=$words[1]", " case $cmd in"]
for command in self.commands:
func_name = command.replace("-", "_")
lines.append(f" {command})")
lines.append(f" _{func_name}")
lines.append(" ;;")
lines.extend([" *)", " _default", " ;;", " esac", "}", "", "_mflux"])
return "\n".join(lines)
def generate(self) -> str:
"""Generate the complete ZSH completion script."""
script = [self.generate_header()]
script.append(self.generate_helper_functions())
# Generate completion functions for each command
for command in self.commands:
parser = self.create_parser_for_command(command)
script.append(self.generate_command_function(command, parser))
script.append(self.generate_main_function())
return "\n".join(script)
if __name__ == "__main__":
# Test the generator by printing out the completion function to stdout
generator = CompletionGenerator()
print(generator.generate())

203
src/mflux/completions/install.py Normal file
View File

@ -0,0 +1,203 @@
#!/usr/bin/env python3
"""Install ZSH completions for mflux commands."""
import argparse
import os
import subprocess
import sys
from pathlib import Path
from mflux.completions.generator import CompletionGenerator
def get_zsh_fpath():
"""Get the ZSH fpath directories."""
try:
result = subprocess.run(
["zsh", "-c", "echo $fpath"],
capture_output=True,
text=True,
check=True,
)
return result.stdout.strip().split()
except subprocess.CalledProcessError:
return []
def find_completion_dir():
"""Find appropriate directory for completion files."""
# Common completion directories in order of preference
candidates = [
Path.home() / ".zsh" / "completions",
Path.home() / ".config" / "zsh" / "completions",
Path("/usr/local/share/zsh/site-functions"),
Path("/opt/homebrew/share/zsh/site-functions"),
]
# Check fpath directories
fpath = get_zsh_fpath()
for path_str in fpath:
path = Path(path_str)
if path.exists() and os.access(path, os.W_OK):
candidates.insert(0, path)
# Return first writable directory
for candidate in candidates:
if candidate.exists() and os.access(candidate, os.W_OK):
return candidate
# Default to user directory
user_dir = Path.home() / ".zsh" / "completions"
user_dir.mkdir(parents=True, exist_ok=True)
return user_dir
def check_installation():
"""Check if completions are properly installed and accessible."""
print("Checking mflux completions installation...\n")
# Check if completion file exists
found_locations = []
fpath = get_zsh_fpath()
for path_str in fpath:
path = Path(path_str)
completion_file = path / "_mflux"
if completion_file.exists():
found_locations.append(completion_file)
if not found_locations:
print("❌ No mflux completion file found in $fpath")
print("\nSearched in:")
for path_str in fpath[:5]: # Show first 5 locations
print(f" - {path_str}")
print("\nRun 'mflux-completions' to install completions")
return False
print("✓ Found completion file(s):")
for location in found_locations:
print(f" - {location}")
# Check if compinit is loaded in interactive shell
try:
result = subprocess.run(
["zsh", "-i", "-c", "type compinit"],
capture_output=True,
text=True,
check=False,
)
if result.returncode == 0:
print("\n✓ compinit is available")
else:
print("\n⚠️ compinit may not be initialized")
print(" If completions don't work, add to ~/.zshrc:")
print(" autoload -U compinit && compinit")
except (subprocess.SubprocessError, OSError):
print("\n⚠️ Could not verify compinit")
# Test a completion
print("\n✓ Installation appears correct!")
print("\nTo test completions:")
print(" 1. Start a new shell or run: exec zsh")
print(" 2. Type: mflux-generate --<TAB>")
print("\nIf completions don't work, try:")
print(" rm -f ~/.zcompdump && compinit")
return True
def main():
"""Main entry point for completion installation."""
parser = argparse.ArgumentParser(
description="Install or generate ZSH completions for mflux commands",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
# Install to default location
mflux-completions
# Generate to stdout
mflux-completions --print
# Install to specific directory
mflux-completions --dir ~/.config/zsh/completions
# Update existing completion
mflux-completions --update
# Check if completions are properly installed
mflux-completions --check
""",
)
parser.add_argument(
"--print",
action="store_true",
help="Print completion script to stdout instead of installing",
)
parser.add_argument(
"--dir",
type=Path,
help="Directory to install completion file (default: auto-detect)",
)
parser.add_argument(
"--update",
action="store_true",
help="Update existing completion file",
)
parser.add_argument(
"--check",
action="store_true",
help="Check if completions are properly installed",
)
args = parser.parse_args()
# Handle --check flag
if args.check:
sys.exit(0 if check_installation() else 1)
# Generate completion script
generator = CompletionGenerator()
completion_script = generator.generate()
if args.print:
# Just print to stdout
print(completion_script)
return
# Determine installation directory
if args.dir:
install_dir = args.dir
install_dir.mkdir(parents=True, exist_ok=True)
else:
install_dir = find_completion_dir()
# Install completion file
completion_file = install_dir / "_mflux"
if completion_file.exists() and not args.update:
print(f"Error: Completion file already exists at {completion_file}")
print("Use --update to overwrite")
sys.exit(1)
try:
completion_file.write_text(completion_script)
print(f"✓ Installed ZSH completions to: {completion_file}")
print()
print("To activate completions:")
print("1. Ensure this directory is in your $fpath:")
print(f" echo 'fpath=({install_dir} $fpath)' >> ~/.zshrc")
print("2. Reload your shell:")
print(" exec zsh")
print()
print("If completions don't work immediately, you may need to:")
print(" rm -f ~/.zcompdump && compinit")
except OSError as e:
print(f"Error installing completion file: {e}")
sys.exit(1)
if __name__ == "__main__":
main()