LORA_LIBRARY_PATH for Unix-style resource discovery (#198)
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
GitHub
parent
0003bbe2cc
commit
d3e71328b2
35
README.md
35
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"
|
||||
```
|
||||
|
||||
<details>
|
||||
<summary>Notes on organizing your LoRA files</summary>
|
||||
|
||||
- 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.
|
||||
</details>
|
||||
|
||||
#### Supported LoRA formats (updated)
|
||||
|
||||
Since different fine-tuning services can use different implementations of FLUX, the corresponding
|
||||
|
||||
@ -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"]
|
||||
|
||||
133
src/mflux/lora_library.py
Normal file
133
src/mflux/lora_library.py
Normal file
@ -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())
|
||||
@ -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
|
||||
|
||||
87
src/mflux/weights/lora_library.py
Normal file
87
src/mflux/weights/lora_library.py
Normal file
@ -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()
|
||||
149
tests/arg_parser/test_lora_library_integration.py
Normal file
149
tests/arg_parser/test_lora_library_integration.py
Normal file
@ -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()
|
||||
0
tests/weights/__init__.py
Normal file
0
tests/weights/__init__.py
Normal file
199
tests/weights/test_lora_library.py
Normal file
199
tests/weights/test_lora_library.py
Normal file
@ -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
|
||||
Loading…
Reference in New Issue
Block a user