From 0125103634ff25e30df04f7b2d3141672252d35d Mon Sep 17 00:00:00 2001 From: Anthony Wu <462072+anthonywu@users.noreply.github.com> Date: Fri, 27 Jun 2025 06:48:51 -0700 Subject: [PATCH] =?UTF-8?q?=E2=9E=A1=EF=B8=8F=E2=9E=A1=EF=B8=8F=20Tab=20Ta?= =?UTF-8?q?b=20for=20zsh:=20CLI=20completion=20generator=20(#208)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Anthony Wu Co-authored-by: filipstrand --- README.md | 26 +++ pyproject.toml | 1 + src/mflux/completions/README.md | 144 +++++++++++++ src/mflux/completions/__init__.py | 1 + src/mflux/completions/generator.py | 328 +++++++++++++++++++++++++++++ src/mflux/completions/install.py | 203 ++++++++++++++++++ 6 files changed, 703 insertions(+) create mode 100644 src/mflux/completions/README.md create mode 100644 src/mflux/completions/__init__.py create mode 100644 src/mflux/completions/generator.py create mode 100644 src/mflux/completions/install.py diff --git a/README.md b/README.md index cfa43ee..a5d216f 100644 --- a/README.md +++ b/README.md @@ -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 -- +``` + +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: diff --git a/pyproject.toml b/pyproject.toml index 115fb67..954a3a7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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"] diff --git a/src/mflux/completions/README.md b/src/mflux/completions/README.md new file mode 100644 index 0000000..39580a5 --- /dev/null +++ b/src/mflux/completions/README.md @@ -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 -- + ``` + +## 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. diff --git a/src/mflux/completions/__init__.py b/src/mflux/completions/__init__.py new file mode 100644 index 0000000..7dfa6a7 --- /dev/null +++ b/src/mflux/completions/__init__.py @@ -0,0 +1 @@ +"""ZSH completion generator for mflux CLI commands.""" diff --git a/src/mflux/completions/generator.py b/src/mflux/completions/generator.py new file mode 100644 index 0000000..02d4c98 --- /dev/null +++ b/src/mflux/completions/generator.py @@ -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()) diff --git a/src/mflux/completions/install.py b/src/mflux/completions/install.py new file mode 100644 index 0000000..1b0393a --- /dev/null +++ b/src/mflux/completions/install.py @@ -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 --") + 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()