From d3e71328b24ba514750f8e2d720513d1d7e7fd66 Mon Sep 17 00:00:00 2001 From: Anthony Wu <462072+anthonywu@users.noreply.github.com> Date: Wed, 4 Jun 2025 07:33:31 -0700 Subject: [PATCH] LORA_LIBRARY_PATH for Unix-style resource discovery (#198) Co-authored-by: Anthony Wu Co-authored-by: filipstrand --- README.md | 35 +++ pyproject.toml | 1 + src/mflux/lora_library.py | 133 ++++++++++++ src/mflux/ui/cli/parsers.py | 12 ++ src/mflux/weights/lora_library.py | 87 ++++++++ .../test_lora_library_integration.py | 149 +++++++++++++ tests/weights/__init__.py | 0 tests/weights/test_lora_library.py | 199 ++++++++++++++++++ 8 files changed, 616 insertions(+) create mode 100644 src/mflux/lora_library.py create mode 100644 src/mflux/weights/lora_library.py create mode 100644 tests/arg_parser/test_lora_library_integration.py create mode 100644 tests/weights/__init__.py create mode 100644 tests/weights/test_lora_library.py diff --git a/README.md b/README.md index 00613a6..d43cb26 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,7 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black - [🎨 Image-to-Image](#-image-to-image) - [🔌 LoRA](#-lora) * [Multi-LoRA](#multi-lora) + * [LoRA Library Path](#lora-library-path) * [Supported LoRA formats (updated)](#supported-lora-formats-updated) - [🎭 In-Context LoRA](#-in-context-lora) * [Available Styles](#available-styles) @@ -766,6 +767,40 @@ mflux-generate \ Just to see the difference, this image displays the four cases: One of having both adapters fully active, partially active and no LoRA at all. The example above also show the usage of `--lora-scales` flag. +#### LoRA Library Path + +MFLUX supports a convenient LoRA library feature that allows you to reference LoRA files by their basename instead of full paths. This is particularly useful when you have a collection of LoRA files organized in one or more directories. + +To use this feature, set the `LORA_LIBRARY_PATH` environment variable to point to your LoRA directories. You can specify multiple directories separated by colons (`:`): + +```sh +export LORA_LIBRARY_PATH="/path/to/loras:/another/path/to/more/loras" +``` + +Once set, MFLUX will automatically discover all `.safetensors` files in these directories (including subdirectories) and allow you to reference them by their basename: + +```sh +# Instead of using full paths: +mflux-generate \ + --prompt "a portrait" \ + --lora-paths "/path/to/loras/style1.safetensors" "/another/path/to/more/loras/style2.safetensors" + +# You can simply use basenames: +mflux-generate \ + --prompt "a portrait" \ + --lora-paths "style1" "style2" +``` + +
+Notes on organizing your LoRA files + +- The basename is the filename without the `.safetensors` extension +- If multiple files have the same basename, the first directory in `LORA_LIBRARY_PATH` takes precedence + - to workaround this, rename or symlink to another name your `.safetensors` files to avoid conflicts +- Full paths still work as before, making this feature fully backwards compatible +- The library paths are scanned recursively, so LoRAs in subdirectories are also discovered. However, we do not recommend setting the library paths to a directory with a large number of files, as it can slow down the scanning process on every run. +
+ #### Supported LoRA formats (updated) Since different fine-tuning services can use different implementations of FLUX, the corresponding diff --git a/pyproject.toml b/pyproject.toml index b82c226..b2fd4d0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,6 +65,7 @@ mflux-save = "mflux.save:main" mflux-save-depth = "mflux.save_depth:main" mflux-train = "mflux.train:main" mflux-upscale = "mflux.upscale:main" +mflux-lora-library = "mflux.lora_library:main" [tool.setuptools.packages.find] where = ["src"] diff --git a/src/mflux/lora_library.py b/src/mflux/lora_library.py new file mode 100644 index 0000000..11158c3 --- /dev/null +++ b/src/mflux/lora_library.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +"""CLI tool for managing the MFLUX LoRA library.""" + +import argparse +import os +import sys +from collections import defaultdict +from pathlib import Path + +from mflux.weights.lora_library import _discover_lora_files + + +def list_loras(paths: list[str] | None = None) -> int: + """List all discovered LoRA files from specified paths or LORA_LIBRARY_PATH. + + Args: + paths: Optional list of paths to use instead of LORA_LIBRARY_PATH + + Returns: + Exit code (0 for success, 1 for error) + """ + if paths: + # Use provided paths + library_paths = [Path(p.strip()) for p in paths] + else: + # Use environment variable + library_path_env = os.environ.get("LORA_LIBRARY_PATH") + + if not library_path_env: + print("LORA_LIBRARY_PATH environment variable is not set.", file=sys.stderr) + print("Set it to one or more colon-separated directories containing .safetensors files.", file=sys.stderr) + print("Alternatively, use --paths to specify directories directly.", file=sys.stderr) + return 1 + + # Parse library paths from environment + library_paths = [Path(p.strip()) for p in library_path_env.split(":") if p.strip()] + + # Check which paths exist + valid_paths = [] + for path in library_paths: + if path.exists() and path.is_dir(): + valid_paths.append(path) + else: + print(f"Warning: Path does not exist or is not a directory: {path}", file=sys.stderr) + + if not valid_paths: + print("No valid directories found in LORA_LIBRARY_PATH.", file=sys.stderr) + return 1 + + # Discover all LoRA files + lora_registry = _discover_lora_files(valid_paths) + + if not lora_registry: + print("No .safetensors files found in the specified directories.") + return 0 + + # Sort by basename for consistent output + sorted_items = sorted(lora_registry.items()) + + # Print all discovered LoRAs + print("Discovered LoRA files:") + print("-" * 80) + for basename, full_path in sorted_items: + print(f"{basename} -> {full_path}") + + # Calculate statistics per top-level directory + stats = defaultdict(int) + for full_path in lora_registry.values(): + # Find which library path this file belongs to + for lib_path in valid_paths: + try: + # Check if full_path is relative to lib_path + full_path.relative_to(lib_path.resolve()) + stats[str(lib_path)] += 1 + break + except ValueError: + # Not relative to this lib_path, continue + continue + + # Print summary + print("-" * 80) + print(f"\nTotal LoRA files found: {len(lora_registry)}") + + if len(valid_paths) > 1: + print("\nBreakdown by library path:") + for lib_path in valid_paths: + count = stats.get(str(lib_path), 0) + print(f" {lib_path}: {count} files") + + return 0 + + +def main(): + """Main entry point for the CLI.""" + parser = argparse.ArgumentParser( + description="MFLUX LoRA Library management tool", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Environment Variables: + LORA_LIBRARY_PATH Colon-separated list of directories containing .safetensors files + Example: /path/to/loras:/another/path/to/loras + +Examples: + # List all discovered LoRA files using LORA_LIBRARY_PATH + mflux-lora-library list + + # With environment variable set + LORA_LIBRARY_PATH=/home/user/loras:/opt/shared/loras mflux-lora-library list + + # Override with specific paths + mflux-lora-library list --paths /path/to/loras /another/path/to/loras +""", + ) + + subparsers = parser.add_subparsers(dest="command", help="Available commands") + + # Add 'list' command + list_parser = subparsers.add_parser("list", help="List all discovered LoRA files") + list_parser.add_argument( + "--paths", nargs="+", help="Override LORA_LIBRARY_PATH with these directories (space-separated)" + ) + + args = parser.parse_args() + + if args.command == "list": + return list_loras(paths=args.paths) + else: + parser.print_help() + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py index e12f88d..a6d512c 100644 --- a/src/mflux/ui/cli/parsers.py +++ b/src/mflux/ui/cli/parsers.py @@ -10,6 +10,7 @@ from mflux.ui import ( box_values, defaults as ui_defaults, ) +from mflux.weights.lora_library import get_lora_path class ModelSpecAction(argparse.Action): @@ -237,4 +238,15 @@ class CommandLineParser(argparse.ArgumentParser): namespace.image_outpaint_padding = box_values.parse_box_value(namespace.image_outpaint_padding) print(f"{namespace.image_outpaint_padding=}") + # Resolve lora paths from library if needed + if self.supports_lora and hasattr(namespace, "lora_paths") and namespace.lora_paths: + resolved_paths = [] + for lora_path in namespace.lora_paths: + try: + resolved_path = get_lora_path(lora_path) + resolved_paths.append(resolved_path) + except FileNotFoundError as e: # noqa: PERF203 + self.error(str(e)) + namespace.lora_paths = resolved_paths + return namespace diff --git a/src/mflux/weights/lora_library.py b/src/mflux/weights/lora_library.py new file mode 100644 index 0000000..cde3ac8 --- /dev/null +++ b/src/mflux/weights/lora_library.py @@ -0,0 +1,87 @@ +import os +from pathlib import Path + + +def _discover_lora_files(library_paths: list[Path]) -> dict[str, Path]: + """ + Discover all .safetensors files in the library paths and their subdirectories. + Earlier paths in the list have higher precedence for duplicate basenames. + + Args: + library_paths: List of paths to LORA library directories (in precedence order) + + Returns: + Dictionary mapping basename (without extension) to full path + """ + lora_files = {} + + # Process paths in reverse order so earlier paths overwrite later ones + for library_path in reversed(library_paths): + if not library_path.exists() or not library_path.is_dir(): + continue + + # Find all .safetensors files recursively + for safetensor_path in library_path.rglob("*.safetensors"): + # Use the basename without extension as the key + basename = safetensor_path.stem + + # Skip files with digit-only names (0-9) in transformer directories + if basename.isdigit() and safetensor_path.parent.name == "transformer": + continue + + # Earlier paths in the list take precedence (overwrite) + lora_files[basename] = safetensor_path.resolve() + + return lora_files + + +# Global registry that will be populated on module import +_LORA_REGISTRY: dict[str, Path] = {} + + +def _initialize_registry() -> None: + """Initialize the global LORA registry from LORA_LIBRARY_PATH environment variable.""" + global _LORA_REGISTRY + + library_path_env = os.environ.get("LORA_LIBRARY_PATH") + if library_path_env: + # Split by colon to support multiple paths + library_paths = [Path(p.strip()) for p in library_path_env.split(":") if p.strip()] + _LORA_REGISTRY = _discover_lora_files(library_paths) + + +def get_lora_path(path_or_name: str) -> str: + """ + Get the full path for a LORA file, resolving from library if needed. + + Args: + path_or_name: Either a full path or a basename that exists in the library + + Returns: + The resolved path as a string + + Raises: + FileNotFoundError: If the file cannot be found either as a path or in the registry + """ + # If it's already a path that exists, return it as-is + path = Path(path_or_name) + if path.exists(): + return str(path) + + # Otherwise, check if it's in the registry + if path_or_name in _LORA_REGISTRY: + return str(_LORA_REGISTRY[path_or_name]) + + # If not found, raise FileNotFoundError + raise FileNotFoundError( + f"LoRA file not found: '{path_or_name}'. File does not exist and is not in the LoRA library." + ) + + +def get_registry() -> dict[str, Path]: + """Get a copy of the current LORA registry.""" + return _LORA_REGISTRY.copy() + + +# Initialize the registry when the module is imported +_initialize_registry() diff --git a/tests/arg_parser/test_lora_library_integration.py b/tests/arg_parser/test_lora_library_integration.py new file mode 100644 index 0000000..2c5a155 --- /dev/null +++ b/tests/arg_parser/test_lora_library_integration.py @@ -0,0 +1,149 @@ +import os +import sys +import tempfile +from pathlib import Path +from unittest import mock + +import pytest + +from mflux.ui.cli.parsers import CommandLineParser +from mflux.weights import lora_library + + +@pytest.fixture +def temp_lora_library(): + """Create a temporary lora library for testing.""" + with tempfile.TemporaryDirectory() as temp_dir: + lib_path = Path(temp_dir) / "lora_library" + lib_path.mkdir() + + # Create some test .safetensors files + (lib_path / "style1.safetensors").touch() + (lib_path / "style2.safetensors").touch() + (lib_path / "subdir").mkdir() + (lib_path / "subdir" / "style3.safetensors").touch() + + yield lib_path + + +def test_parser_resolves_lora_paths_from_library(temp_lora_library): + """Test that parser resolves lora names to full paths from library.""" + # Set up the environment variable + with mock.patch.dict(os.environ, {"LORA_LIBRARY_PATH": str(temp_lora_library)}): + # Re-initialize the registry + lora_library._initialize_registry() + + parser = CommandLineParser() + parser.add_model_arguments() + parser.add_lora_arguments() + parser.add_image_generator_arguments() + + # Test with lora names that should be resolved + # Simulate command line arguments + test_args = [ + "mflux-generate", # argv[0] is program name + "--model", + "dev", + "--prompt", + "test prompt", + "--lora-paths", + "style1", + "style3", + "--lora-scales", + "0.5", + "0.8", + ] + + with mock.patch.object(sys, "argv", test_args): + args = parser.parse_args() + + assert len(args.lora_paths) == 2 + assert args.lora_paths[0] == str((temp_lora_library / "style1.safetensors").resolve()) + assert args.lora_paths[1] == str((temp_lora_library / "subdir" / "style3.safetensors").resolve()) + assert args.lora_scales == [0.5, 0.8] + + +def test_parser_preserves_full_paths(temp_lora_library): + """Test that parser preserves full paths that already exist.""" + # Create a lora file outside the library + with tempfile.NamedTemporaryFile(suffix=".safetensors", delete=False) as tmp_file: + external_lora = tmp_file.name + + try: + parser = CommandLineParser() + parser.add_model_arguments() + parser.add_lora_arguments() + parser.add_image_generator_arguments() + + test_args = ["mflux-generate-fill", "--model", "dev", "--prompt", "test prompt", "--lora-paths", external_lora] + with mock.patch.object(sys, "argv", test_args): + args = parser.parse_args() + + assert len(args.lora_paths) == 1 + assert args.lora_paths[0] == external_lora + finally: + os.unlink(external_lora) + + +def test_parser_mixed_lora_paths(temp_lora_library): + """Test parser with mix of library names and full paths.""" + # Create an external lora file + with tempfile.NamedTemporaryFile(suffix=".safetensors", delete=False) as tmp_file: + external_lora = tmp_file.name + + try: + with mock.patch.dict(os.environ, {"LORA_LIBRARY_PATH": str(temp_lora_library)}): + lora_library._initialize_registry() + + parser = CommandLineParser() + parser.add_model_arguments() + parser.add_lora_arguments() + parser.add_image_generator_arguments() + + test_args = [ + "mflux-generate-redux", + "--model", + "dev", + "--prompt", + "test prompt", + "--lora-paths", + "style1", + external_lora, + "style2", + ] + with mock.patch.object(sys, "argv", test_args): + args = parser.parse_args() + + assert len(args.lora_paths) == 3 + assert args.lora_paths[0] == str((temp_lora_library / "style1.safetensors").resolve()) + assert args.lora_paths[1] == external_lora + assert args.lora_paths[2] == str((temp_lora_library / "style2.safetensors").resolve()) + finally: + os.unlink(external_lora) + + +def test_parser_unknown_lora_names_error(): + """Test that unknown lora names raise an error.""" + # No library path set + with mock.patch.dict(os.environ, {}, clear=True): + lora_library._initialize_registry() + + parser = CommandLineParser() + parser.add_model_arguments() + parser.add_lora_arguments() + parser.add_image_generator_arguments() + + # Should raise an error for unknown LoRA names + test_args = [ + "mflux-generate-depth", + "--model", + "dev", + "--prompt", + "test prompt", + "--lora-paths", + "unknown_style", + ] + + with mock.patch.object(sys, "argv", test_args): + with pytest.raises(SystemExit): + parser.parse_args() diff --git a/tests/weights/__init__.py b/tests/weights/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/weights/test_lora_library.py b/tests/weights/test_lora_library.py new file mode 100644 index 0000000..716d4d0 --- /dev/null +++ b/tests/weights/test_lora_library.py @@ -0,0 +1,199 @@ +import os +import tempfile +from pathlib import Path +from unittest import mock + +import pytest + +from mflux.weights import lora_library + + +@pytest.fixture(autouse=True) +def reset_lora_registry(): + """Reset the global registry before and after each test.""" + original = lora_library._LORA_REGISTRY.copy() + lora_library._LORA_REGISTRY.clear() + yield + lora_library._LORA_REGISTRY.clear() + lora_library._LORA_REGISTRY.update(original) + + +@pytest.fixture +def temp_lora_dirs(): + """Create temporary directories with fake .safetensors files for testing.""" + with tempfile.TemporaryDirectory() as temp_root: + # Create two library directories + lib1 = Path(temp_root) / "library1" + lib2 = Path(temp_root) / "library2" + lib1.mkdir() + lib2.mkdir() + + # Create some .safetensors files + (lib1 / "style_a.safetensors").touch() + (lib1 / "style_b.safetensors").touch() + (lib1 / "subdir").mkdir() + (lib1 / "subdir" / "style_c.safetensors").touch() + + # Create transformer directory with digit-named files (should be ignored) + (lib1 / "transformer").mkdir() + (lib1 / "transformer" / "0.safetensors").touch() + (lib1 / "transformer" / "1.safetensors").touch() + (lib1 / "transformer" / "style_valid.safetensors").touch() # Should not be ignored + + (lib2 / "style_b.safetensors").touch() # Duplicate name + (lib2 / "style_d.safetensors").touch() + + yield lib1, lib2 + + +def test_discover_lora_files_single_path(temp_lora_dirs): + """Test discovering files in a single library path.""" + lib1, lib2 = temp_lora_dirs + + result = lora_library._discover_lora_files([lib1]) + + assert len(result) == 4 # Should not include digit-named files in transformer/ + assert "style_a" in result + assert "style_b" in result + assert "style_c" in result + assert "style_valid" in result + assert "0" not in result # Digit-named files in transformer/ should be excluded + assert "1" not in result + assert result["style_a"] == (lib1 / "style_a.safetensors").resolve() + assert result["style_c"] == (lib1 / "subdir" / "style_c.safetensors").resolve() + assert result["style_valid"] == (lib1 / "transformer" / "style_valid.safetensors").resolve() + + +def test_discover_lora_files_multiple_paths_precedence(temp_lora_dirs): + """Test that earlier paths take precedence for duplicate names.""" + lib1, lib2 = temp_lora_dirs + + # lib1 should take precedence over lib2 + result = lora_library._discover_lora_files([lib1, lib2]) + + assert len(result) == 5 # style_a, style_b, style_c, style_valid, style_d + assert "style_a" in result + assert "style_b" in result + assert "style_c" in result + assert "style_valid" in result + assert "style_d" in result + assert "0" not in result # Digit-named files should still be excluded + assert "1" not in result + + # style_b should come from lib1 (first in list) + assert result["style_b"] == (lib1 / "style_b.safetensors").resolve() + assert result["style_d"] == (lib2 / "style_d.safetensors").resolve() + + +def test_discover_lora_files_nonexistent_path(): + """Test handling of nonexistent paths.""" + result = lora_library._discover_lora_files([Path("/nonexistent/path")]) + assert result == {} + + +def test_get_lora_path_existing_file(temp_lora_dirs): + """Test that existing file paths are returned as-is.""" + lib1, _ = temp_lora_dirs + existing_file = lib1 / "style_a.safetensors" + + result = lora_library.get_lora_path(str(existing_file)) + assert result == str(existing_file) + + +def test_get_lora_path_from_registry(temp_lora_dirs): + """Test resolving paths from the registry.""" + lib1, lib2 = temp_lora_dirs + + # Mock the registry + with mock.patch.object( + lora_library, + "_LORA_REGISTRY", + { + "style_a": lib1 / "style_a.safetensors", + "style_b": lib1 / "style_b.safetensors", + }, + ): + # Should resolve from registry + assert lora_library.get_lora_path("style_a") == str(lib1 / "style_a.safetensors") + assert lora_library.get_lora_path("style_b") == str(lib1 / "style_b.safetensors") + + # Should raise FileNotFoundError if not in registry + with pytest.raises(FileNotFoundError, match="LoRA file not found: 'unknown_style'"): + lora_library.get_lora_path("unknown_style") + + +def test_initialize_registry_from_env(temp_lora_dirs): + """Test that registry is initialized from LORA_LIBRARY_PATH env var.""" + lib1, lib2 = temp_lora_dirs + + # Set up environment variable with colon-delimited paths + env_path = f"{lib1}:{lib2}" + + with mock.patch.dict(os.environ, {"LORA_LIBRARY_PATH": env_path}): + # Re-initialize the registry + lora_library._initialize_registry() + registry = lora_library.get_registry() + + assert len(registry) == 5 # Includes style_valid from transformer/ + assert "style_a" in registry + assert "style_b" in registry + assert "style_c" in registry + assert "style_valid" in registry + assert "style_d" in registry + assert "0" not in registry # Digit files should be excluded + assert "1" not in registry + + # Check precedence - style_b should be from lib1 + assert registry["style_b"] == (lib1 / "style_b.safetensors").resolve() + + +def test_initialize_registry_empty_env(): + """Test that registry is empty when env var is not set.""" + # Save the original registry + original_registry = lora_library._LORA_REGISTRY.copy() + try: + with mock.patch.dict(os.environ, {}, clear=True): + lora_library._initialize_registry() + registry = lora_library.get_registry() + assert registry == {} + finally: + # Restore the original registry + lora_library._LORA_REGISTRY = original_registry + + +def test_get_registry_returns_copy(): + """Test that get_registry returns a copy, not the original.""" + with mock.patch.object(lora_library, "_LORA_REGISTRY", {"test": Path("/test")}): + registry1 = lora_library.get_registry() + registry2 = lora_library.get_registry() + + # Should be equal but different objects + assert registry1 == registry2 + assert registry1 is not registry2 + + # Modifying the copy shouldn't affect the original + registry1["new"] = Path("/new") + assert "new" not in lora_library._LORA_REGISTRY + + +def test_transformer_digit_filtering(temp_lora_dirs): + """Test that digit-named files in transformer directories are filtered out.""" + lib1, _ = temp_lora_dirs + + # Create additional test cases + (lib1 / "9.safetensors").touch() # Should be included (not in transformer/) + (lib1 / "other_dir").mkdir() + (lib1 / "other_dir" / "5.safetensors").touch() # Should be included (not in transformer/) + + result = lora_library._discover_lora_files([lib1]) + + # Digit files NOT in transformer/ should be included + assert "9" in result + assert "5" in result + + # Digit files in transformer/ should be excluded + assert "0" not in result + assert "1" not in result + + # Non-digit files in transformer/ should be included + assert "style_valid" in result