mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-21 21:22:11 -03:00
feat: Add extra folder paths support for LoRA Manager
Introduce extra_folder_paths feature to allow users to add additional model roots that are managed by LoRA Manager but not shared with ComfyUI. Changes: - Add extra_folder_paths support in SettingsManager (stored per library) - Add extra path attributes in Config class (extra_loras_roots, etc.) - Merge folder_paths with extra_folder_paths when applying library settings - Update LoraScanner, CheckpointScanner, EmbeddingScanner to include extra paths in their model roots - Add comprehensive tests for the new functionality This enables users to manage models from additional directories without modifying ComfyUI's model folder configuration.
This commit is contained in:
60
py/config.py
60
py/config.py
@@ -91,6 +91,11 @@ class Config:
|
||||
self.embeddings_roots = None
|
||||
self.base_models_roots = self._init_checkpoint_paths()
|
||||
self.embeddings_roots = self._init_embedding_paths()
|
||||
# Extra paths (only for LoRA Manager, not shared with ComfyUI)
|
||||
self.extra_loras_roots: List[str] = []
|
||||
self.extra_checkpoints_roots: List[str] = []
|
||||
self.extra_unet_roots: List[str] = []
|
||||
self.extra_embeddings_roots: List[str] = []
|
||||
# Scan symbolic links during initialization
|
||||
self._initialize_symlink_mappings()
|
||||
|
||||
@@ -250,6 +255,11 @@ class Config:
|
||||
roots.extend(self.loras_roots or [])
|
||||
roots.extend(self.base_models_roots or [])
|
||||
roots.extend(self.embeddings_roots or [])
|
||||
# Include extra paths for scanning symlinks
|
||||
roots.extend(self.extra_loras_roots or [])
|
||||
roots.extend(self.extra_checkpoints_roots or [])
|
||||
roots.extend(self.extra_unet_roots or [])
|
||||
roots.extend(self.extra_embeddings_roots or [])
|
||||
return roots
|
||||
|
||||
def _build_symlink_fingerprint(self) -> Dict[str, object]:
|
||||
@@ -570,6 +580,15 @@ class Config:
|
||||
preview_roots.update(self._expand_preview_root(root))
|
||||
for root in self.embeddings_roots or []:
|
||||
preview_roots.update(self._expand_preview_root(root))
|
||||
# Include extra paths for preview access
|
||||
for root in self.extra_loras_roots or []:
|
||||
preview_roots.update(self._expand_preview_root(root))
|
||||
for root in self.extra_checkpoints_roots or []:
|
||||
preview_roots.update(self._expand_preview_root(root))
|
||||
for root in self.extra_unet_roots or []:
|
||||
preview_roots.update(self._expand_preview_root(root))
|
||||
for root in self.extra_embeddings_roots or []:
|
||||
preview_roots.update(self._expand_preview_root(root))
|
||||
|
||||
for target, link in self._path_mappings.items():
|
||||
preview_roots.update(self._expand_preview_root(target))
|
||||
@@ -577,11 +596,11 @@ class Config:
|
||||
|
||||
self._preview_root_paths = {path for path in preview_roots if path.is_absolute()}
|
||||
logger.debug(
|
||||
"Preview roots rebuilt: %d paths from %d lora roots, %d checkpoint roots, %d embedding roots, %d symlink mappings",
|
||||
"Preview roots rebuilt: %d paths from %d lora roots (%d extra), %d checkpoint roots (%d extra), %d embedding roots (%d extra), %d symlink mappings",
|
||||
len(self._preview_root_paths),
|
||||
len(self.loras_roots or []),
|
||||
len(self.base_models_roots or []),
|
||||
len(self.embeddings_roots or []),
|
||||
len(self.loras_roots or []), len(self.extra_loras_roots or []),
|
||||
len(self.base_models_roots or []), len(self.extra_checkpoints_roots or []),
|
||||
len(self.embeddings_roots or []), len(self.extra_embeddings_roots or []),
|
||||
len(self._path_mappings),
|
||||
)
|
||||
|
||||
@@ -692,7 +711,11 @@ class Config:
|
||||
|
||||
return unique_paths
|
||||
|
||||
def _apply_library_paths(self, folder_paths: Mapping[str, Iterable[str]]) -> None:
|
||||
def _apply_library_paths(
|
||||
self,
|
||||
folder_paths: Mapping[str, Iterable[str]],
|
||||
extra_folder_paths: Optional[Mapping[str, Iterable[str]]] = None,
|
||||
) -> None:
|
||||
self._path_mappings.clear()
|
||||
self._preview_root_paths = set()
|
||||
|
||||
@@ -705,6 +728,20 @@ class Config:
|
||||
self.base_models_roots = self._prepare_checkpoint_paths(checkpoint_paths, unet_paths)
|
||||
self.embeddings_roots = self._prepare_embedding_paths(embedding_paths)
|
||||
|
||||
# Process extra paths (only for LoRA Manager, not shared with ComfyUI)
|
||||
extra_paths = extra_folder_paths or {}
|
||||
extra_lora_paths = extra_paths.get('loras', []) or []
|
||||
extra_checkpoint_paths = extra_paths.get('checkpoints', []) or []
|
||||
extra_unet_paths = extra_paths.get('unet', []) or []
|
||||
extra_embedding_paths = extra_paths.get('embeddings', []) or []
|
||||
|
||||
self.extra_loras_roots = self._prepare_lora_paths(extra_lora_paths)
|
||||
self.extra_checkpoints_roots = self._prepare_checkpoint_paths(extra_checkpoint_paths, extra_unet_paths)
|
||||
self.extra_embeddings_roots = self._prepare_embedding_paths(extra_embedding_paths)
|
||||
# extra_unet_roots is set by _prepare_checkpoint_paths (access unet_roots before it's reset)
|
||||
unet_roots_value: List[str] = getattr(self, 'unet_roots', None) or []
|
||||
self.extra_unet_roots = unet_roots_value
|
||||
|
||||
self._initialize_symlink_mappings()
|
||||
|
||||
def _init_lora_paths(self) -> List[str]:
|
||||
@@ -864,16 +901,19 @@ class Config:
|
||||
def apply_library_settings(self, library_config: Mapping[str, object]) -> None:
|
||||
"""Update runtime paths to match the provided library configuration."""
|
||||
folder_paths = library_config.get('folder_paths') if isinstance(library_config, Mapping) else {}
|
||||
extra_folder_paths = library_config.get('extra_folder_paths') if isinstance(library_config, Mapping) else None
|
||||
if not isinstance(folder_paths, Mapping):
|
||||
folder_paths = {}
|
||||
if not isinstance(extra_folder_paths, Mapping):
|
||||
extra_folder_paths = None
|
||||
|
||||
self._apply_library_paths(folder_paths)
|
||||
self._apply_library_paths(folder_paths, extra_folder_paths)
|
||||
|
||||
logger.info(
|
||||
"Applied library settings with %d lora roots, %d checkpoint roots, and %d embedding roots",
|
||||
len(self.loras_roots or []),
|
||||
len(self.base_models_roots or []),
|
||||
len(self.embeddings_roots or []),
|
||||
"Applied library settings with %d lora roots (%d extra), %d checkpoint roots (%d extra), and %d embedding roots (%d extra)",
|
||||
len(self.loras_roots or []), len(self.extra_loras_roots or []),
|
||||
len(self.base_models_roots or []), len(self.extra_checkpoints_roots or []),
|
||||
len(self.embeddings_roots or []), len(self.extra_embeddings_roots or []),
|
||||
)
|
||||
|
||||
def get_library_registry_snapshot(self) -> Dict[str, object]:
|
||||
|
||||
@@ -51,5 +51,16 @@ class CheckpointScanner(ModelScanner):
|
||||
return entry
|
||||
|
||||
def get_model_roots(self) -> List[str]:
|
||||
"""Get checkpoint root directories"""
|
||||
return config.base_models_roots
|
||||
"""Get checkpoint root directories (including extra paths)"""
|
||||
roots: List[str] = []
|
||||
roots.extend(config.base_models_roots or [])
|
||||
roots.extend(config.extra_checkpoints_roots or [])
|
||||
roots.extend(config.extra_unet_roots or [])
|
||||
# Remove duplicates while preserving order
|
||||
seen: set = set()
|
||||
unique_roots: List[str] = []
|
||||
for root in roots:
|
||||
if root not in seen:
|
||||
seen.add(root)
|
||||
unique_roots.append(root)
|
||||
return unique_roots
|
||||
|
||||
@@ -22,5 +22,15 @@ class EmbeddingScanner(ModelScanner):
|
||||
)
|
||||
|
||||
def get_model_roots(self) -> List[str]:
|
||||
"""Get embedding root directories"""
|
||||
return config.embeddings_roots
|
||||
"""Get embedding root directories (including extra paths)"""
|
||||
roots: List[str] = []
|
||||
roots.extend(config.embeddings_roots or [])
|
||||
roots.extend(config.extra_embeddings_roots or [])
|
||||
# Remove duplicates while preserving order
|
||||
seen: set = set()
|
||||
unique_roots: List[str] = []
|
||||
for root in roots:
|
||||
if root and root not in seen:
|
||||
seen.add(root)
|
||||
unique_roots.append(root)
|
||||
return unique_roots
|
||||
|
||||
@@ -25,8 +25,18 @@ class LoraScanner(ModelScanner):
|
||||
)
|
||||
|
||||
def get_model_roots(self) -> List[str]:
|
||||
"""Get lora root directories"""
|
||||
return config.loras_roots
|
||||
"""Get lora root directories (including extra paths)"""
|
||||
roots: List[str] = []
|
||||
roots.extend(config.loras_roots or [])
|
||||
roots.extend(config.extra_loras_roots or [])
|
||||
# Remove duplicates while preserving order
|
||||
seen: set = set()
|
||||
unique_roots: List[str] = []
|
||||
for root in roots:
|
||||
if root and root not in seen:
|
||||
seen.add(root)
|
||||
unique_roots.append(root)
|
||||
return unique_roots
|
||||
|
||||
async def diagnose_hash_index(self):
|
||||
"""Diagnostic method to verify hash index functionality"""
|
||||
|
||||
@@ -54,6 +54,7 @@ DEFAULT_SETTINGS: Dict[str, Any] = {
|
||||
"base_model_path_mappings": {},
|
||||
"download_path_templates": {},
|
||||
"folder_paths": {},
|
||||
"extra_folder_paths": {},
|
||||
"example_images_path": "",
|
||||
"optimize_example_images": True,
|
||||
"auto_download_example_images": False,
|
||||
@@ -402,6 +403,7 @@ class SettingsManager:
|
||||
active_library = libraries.get(active_name, {})
|
||||
folder_paths = copy.deepcopy(active_library.get("folder_paths", {}))
|
||||
self.settings["folder_paths"] = folder_paths
|
||||
self.settings["extra_folder_paths"] = copy.deepcopy(active_library.get("extra_folder_paths", {}))
|
||||
self.settings["default_lora_root"] = active_library.get("default_lora_root", "")
|
||||
self.settings["default_checkpoint_root"] = active_library.get("default_checkpoint_root", "")
|
||||
self.settings["default_unet_root"] = active_library.get("default_unet_root", "")
|
||||
@@ -417,6 +419,7 @@ class SettingsManager:
|
||||
self,
|
||||
*,
|
||||
folder_paths: Optional[Mapping[str, Iterable[str]]] = None,
|
||||
extra_folder_paths: Optional[Mapping[str, Iterable[str]]] = None,
|
||||
default_lora_root: Optional[str] = None,
|
||||
default_checkpoint_root: Optional[str] = None,
|
||||
default_unet_root: Optional[str] = None,
|
||||
@@ -432,6 +435,11 @@ class SettingsManager:
|
||||
else:
|
||||
payload.setdefault("folder_paths", {})
|
||||
|
||||
if extra_folder_paths is not None:
|
||||
payload["extra_folder_paths"] = self._normalize_folder_paths(extra_folder_paths)
|
||||
else:
|
||||
payload.setdefault("extra_folder_paths", {})
|
||||
|
||||
if default_lora_root is not None:
|
||||
payload["default_lora_root"] = default_lora_root
|
||||
else:
|
||||
@@ -546,6 +554,7 @@ class SettingsManager:
|
||||
self,
|
||||
*,
|
||||
folder_paths: Optional[Mapping[str, Iterable[str]]] = None,
|
||||
extra_folder_paths: Optional[Mapping[str, Iterable[str]]] = None,
|
||||
default_lora_root: Optional[str] = None,
|
||||
default_checkpoint_root: Optional[str] = None,
|
||||
default_unet_root: Optional[str] = None,
|
||||
@@ -565,6 +574,12 @@ class SettingsManager:
|
||||
library["folder_paths"] = normalized_paths
|
||||
changed = True
|
||||
|
||||
if extra_folder_paths is not None:
|
||||
normalized_extra_paths = self._normalize_folder_paths(extra_folder_paths)
|
||||
if library.get("extra_folder_paths") != normalized_extra_paths:
|
||||
library["extra_folder_paths"] = normalized_extra_paths
|
||||
changed = True
|
||||
|
||||
if default_lora_root is not None and library.get("default_lora_root") != default_lora_root:
|
||||
library["default_lora_root"] = default_lora_root
|
||||
changed = True
|
||||
@@ -816,12 +831,14 @@ class SettingsManager:
|
||||
defaults['download_path_templates'] = {}
|
||||
defaults['priority_tags'] = DEFAULT_PRIORITY_TAG_CONFIG.copy()
|
||||
defaults.setdefault('folder_paths', {})
|
||||
defaults.setdefault('extra_folder_paths', {})
|
||||
defaults['auto_organize_exclusions'] = []
|
||||
defaults['metadata_refresh_skip_paths'] = []
|
||||
|
||||
library_name = defaults.get("active_library") or "default"
|
||||
default_library = self._build_library_payload(
|
||||
folder_paths=defaults.get("folder_paths", {}),
|
||||
extra_folder_paths=defaults.get("extra_folder_paths", {}),
|
||||
default_lora_root=defaults.get("default_lora_root"),
|
||||
default_checkpoint_root=defaults.get("default_checkpoint_root"),
|
||||
default_embedding_root=defaults.get("default_embedding_root"),
|
||||
@@ -927,6 +944,35 @@ class SettingsManager:
|
||||
self._save_settings()
|
||||
return skip_paths
|
||||
|
||||
def get_extra_folder_paths(self) -> Dict[str, List[str]]:
|
||||
"""Get extra folder paths for the active library.
|
||||
|
||||
These paths are only used by LoRA Manager and not shared with ComfyUI.
|
||||
Returns a dictionary with keys like 'loras', 'checkpoints', 'embeddings', 'unet'.
|
||||
"""
|
||||
extra_paths = self.settings.get("extra_folder_paths", {})
|
||||
if not isinstance(extra_paths, dict):
|
||||
return {}
|
||||
return self._normalize_folder_paths(extra_paths)
|
||||
|
||||
def update_extra_folder_paths(
|
||||
self,
|
||||
extra_folder_paths: Mapping[str, Iterable[str]],
|
||||
) -> None:
|
||||
"""Update extra folder paths for the active library.
|
||||
|
||||
These paths are only used by LoRA Manager and not shared with ComfyUI.
|
||||
Validates that extra paths don't overlap with other libraries' paths.
|
||||
"""
|
||||
active_name = self.get_active_library_name()
|
||||
self._validate_folder_paths(active_name, extra_folder_paths)
|
||||
|
||||
normalized_paths = self._normalize_folder_paths(extra_folder_paths)
|
||||
self.settings["extra_folder_paths"] = normalized_paths
|
||||
self._update_active_library_entry(extra_folder_paths=normalized_paths)
|
||||
self._save_settings()
|
||||
logger.info("Updated extra folder paths for library '%s'", active_name)
|
||||
|
||||
def get_startup_messages(self) -> List[Dict[str, Any]]:
|
||||
return [message.copy() for message in self._startup_messages]
|
||||
|
||||
@@ -973,6 +1019,8 @@ class SettingsManager:
|
||||
self._prepare_portable_switch(value)
|
||||
if key == 'folder_paths' and isinstance(value, Mapping):
|
||||
self._update_active_library_entry(folder_paths=value) # type: ignore[arg-type]
|
||||
elif key == 'extra_folder_paths' and isinstance(value, Mapping):
|
||||
self._update_active_library_entry(extra_folder_paths=value) # type: ignore[arg-type]
|
||||
elif key == 'default_lora_root':
|
||||
self._update_active_library_entry(default_lora_root=str(value))
|
||||
elif key == 'default_checkpoint_root':
|
||||
@@ -1284,6 +1332,7 @@ class SettingsManager:
|
||||
library_name: str,
|
||||
*,
|
||||
folder_paths: Optional[Mapping[str, Iterable[str]]] = None,
|
||||
extra_folder_paths: Optional[Mapping[str, Iterable[str]]] = None,
|
||||
default_lora_root: Optional[str] = None,
|
||||
default_checkpoint_root: Optional[str] = None,
|
||||
default_unet_root: Optional[str] = None,
|
||||
@@ -1300,11 +1349,15 @@ class SettingsManager:
|
||||
if folder_paths is not None:
|
||||
self._validate_folder_paths(name, folder_paths)
|
||||
|
||||
if extra_folder_paths is not None:
|
||||
self._validate_folder_paths(name, extra_folder_paths)
|
||||
|
||||
libraries = self.settings.setdefault("libraries", {})
|
||||
existing = libraries.get(name, {})
|
||||
|
||||
payload = self._build_library_payload(
|
||||
folder_paths=folder_paths if folder_paths is not None else existing.get("folder_paths"),
|
||||
extra_folder_paths=extra_folder_paths if extra_folder_paths is not None else existing.get("extra_folder_paths"),
|
||||
default_lora_root=default_lora_root if default_lora_root is not None else existing.get("default_lora_root"),
|
||||
default_checkpoint_root=(
|
||||
default_checkpoint_root
|
||||
@@ -1343,6 +1396,7 @@ class SettingsManager:
|
||||
library_name: str,
|
||||
*,
|
||||
folder_paths: Mapping[str, Iterable[str]],
|
||||
extra_folder_paths: Optional[Mapping[str, Iterable[str]]] = None,
|
||||
default_lora_root: str = "",
|
||||
default_checkpoint_root: str = "",
|
||||
default_unet_root: str = "",
|
||||
@@ -1359,6 +1413,7 @@ class SettingsManager:
|
||||
return self.upsert_library(
|
||||
library_name,
|
||||
folder_paths=folder_paths,
|
||||
extra_folder_paths=extra_folder_paths,
|
||||
default_lora_root=default_lora_root,
|
||||
default_checkpoint_root=default_checkpoint_root,
|
||||
default_unet_root=default_unet_root,
|
||||
@@ -1417,6 +1472,7 @@ class SettingsManager:
|
||||
self,
|
||||
folder_paths: Mapping[str, Iterable[str]],
|
||||
*,
|
||||
extra_folder_paths: Optional[Mapping[str, Iterable[str]]] = None,
|
||||
default_lora_root: Optional[str] = None,
|
||||
default_checkpoint_root: Optional[str] = None,
|
||||
default_unet_root: Optional[str] = None,
|
||||
@@ -1428,6 +1484,7 @@ class SettingsManager:
|
||||
self.upsert_library(
|
||||
active_name,
|
||||
folder_paths=folder_paths,
|
||||
extra_folder_paths=extra_folder_paths,
|
||||
default_lora_root=default_lora_root,
|
||||
default_checkpoint_root=default_checkpoint_root,
|
||||
default_unet_root=default_unet_root,
|
||||
|
||||
@@ -215,3 +215,110 @@ def test_save_paths_removes_template_default_library(monkeypatch, tmp_path):
|
||||
)
|
||||
assert payload["metadata"] == {"display_name": "ComfyUI", "source": "comfyui"}
|
||||
assert payload["activate"] is True
|
||||
|
||||
|
||||
def test_apply_library_settings_merges_extra_paths(monkeypatch, tmp_path):
|
||||
"""Test that apply_library_settings correctly merges folder_paths with extra_folder_paths."""
|
||||
loras_dir = tmp_path / "loras"
|
||||
extra_loras_dir = tmp_path / "extra_loras"
|
||||
checkpoints_dir = tmp_path / "checkpoints"
|
||||
extra_checkpoints_dir = tmp_path / "extra_checkpoints"
|
||||
embeddings_dir = tmp_path / "embeddings"
|
||||
extra_embeddings_dir = tmp_path / "extra_embeddings"
|
||||
|
||||
for directory in (loras_dir, extra_loras_dir, checkpoints_dir, extra_checkpoints_dir, embeddings_dir, extra_embeddings_dir):
|
||||
directory.mkdir()
|
||||
|
||||
config_instance = config_module.Config()
|
||||
|
||||
folder_paths = {
|
||||
"loras": [str(loras_dir)],
|
||||
"checkpoints": [str(checkpoints_dir)],
|
||||
"unet": [],
|
||||
"embeddings": [str(embeddings_dir)],
|
||||
}
|
||||
extra_folder_paths = {
|
||||
"loras": [str(extra_loras_dir)],
|
||||
"checkpoints": [str(extra_checkpoints_dir)],
|
||||
"unet": [],
|
||||
"embeddings": [str(extra_embeddings_dir)],
|
||||
}
|
||||
|
||||
library_config = {
|
||||
"folder_paths": folder_paths,
|
||||
"extra_folder_paths": extra_folder_paths,
|
||||
}
|
||||
|
||||
config_instance.apply_library_settings(library_config)
|
||||
|
||||
assert str(loras_dir) in config_instance.loras_roots
|
||||
assert str(extra_loras_dir) in config_instance.extra_loras_roots
|
||||
assert str(checkpoints_dir) in config_instance.base_models_roots
|
||||
assert str(extra_checkpoints_dir) in config_instance.extra_checkpoints_roots
|
||||
assert str(embeddings_dir) in config_instance.embeddings_roots
|
||||
assert str(extra_embeddings_dir) in config_instance.extra_embeddings_roots
|
||||
|
||||
|
||||
def test_apply_library_settings_without_extra_paths(monkeypatch, tmp_path):
|
||||
"""Test that apply_library_settings works when extra_folder_paths is not provided."""
|
||||
loras_dir = tmp_path / "loras"
|
||||
checkpoints_dir = tmp_path / "checkpoints"
|
||||
embeddings_dir = tmp_path / "embeddings"
|
||||
|
||||
for directory in (loras_dir, checkpoints_dir, embeddings_dir):
|
||||
directory.mkdir()
|
||||
|
||||
config_instance = config_module.Config()
|
||||
|
||||
folder_paths = {
|
||||
"loras": [str(loras_dir)],
|
||||
"checkpoints": [str(checkpoints_dir)],
|
||||
"unet": [],
|
||||
"embeddings": [str(embeddings_dir)],
|
||||
}
|
||||
|
||||
library_config = {
|
||||
"folder_paths": folder_paths,
|
||||
}
|
||||
|
||||
config_instance.apply_library_settings(library_config)
|
||||
|
||||
assert str(loras_dir) in config_instance.loras_roots
|
||||
assert config_instance.extra_loras_roots == []
|
||||
assert str(checkpoints_dir) in config_instance.base_models_roots
|
||||
assert config_instance.extra_checkpoints_roots == []
|
||||
assert str(embeddings_dir) in config_instance.embeddings_roots
|
||||
assert config_instance.extra_embeddings_roots == []
|
||||
|
||||
|
||||
def test_extra_paths_deduplication(monkeypatch, tmp_path):
|
||||
"""Test that extra paths are stored separately from main paths in Config."""
|
||||
loras_dir = tmp_path / "loras"
|
||||
extra_loras_dir = tmp_path / "extra_loras"
|
||||
loras_dir.mkdir()
|
||||
extra_loras_dir.mkdir()
|
||||
|
||||
config_instance = config_module.Config()
|
||||
|
||||
folder_paths = {
|
||||
"loras": [str(loras_dir)],
|
||||
"checkpoints": [],
|
||||
"unet": [],
|
||||
"embeddings": [],
|
||||
}
|
||||
extra_folder_paths = {
|
||||
"loras": [str(extra_loras_dir)],
|
||||
"checkpoints": [],
|
||||
"unet": [],
|
||||
"embeddings": [],
|
||||
}
|
||||
|
||||
library_config = {
|
||||
"folder_paths": folder_paths,
|
||||
"extra_folder_paths": extra_folder_paths,
|
||||
}
|
||||
|
||||
config_instance.apply_library_settings(library_config)
|
||||
|
||||
assert config_instance.loras_roots == [str(loras_dir)]
|
||||
assert config_instance.extra_loras_roots == [str(extra_loras_dir)]
|
||||
|
||||
@@ -132,4 +132,59 @@ async def test_persisted_cache_restores_model_type(tmp_path: Path, monkeypatch):
|
||||
assert types_by_path[normalized_unet_file] == "diffusion_model"
|
||||
|
||||
assert ws_stub.payloads
|
||||
assert ws_stub.payloads[-1]["stage"] == "loading_cache"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checkpoint_scanner_get_model_roots_includes_extra_paths(monkeypatch, tmp_path):
|
||||
"""Test that get_model_roots includes both main and extra paths."""
|
||||
checkpoints_root = tmp_path / "checkpoints"
|
||||
extra_checkpoints_root = tmp_path / "extra_checkpoints"
|
||||
unet_root = tmp_path / "unet"
|
||||
extra_unet_root = tmp_path / "extra_unet"
|
||||
|
||||
for directory in (checkpoints_root, extra_checkpoints_root, unet_root, extra_unet_root):
|
||||
directory.mkdir()
|
||||
|
||||
normalized_checkpoints = _normalize(checkpoints_root)
|
||||
normalized_extra_checkpoints = _normalize(extra_checkpoints_root)
|
||||
normalized_unet = _normalize(unet_root)
|
||||
normalized_extra_unet = _normalize(extra_unet_root)
|
||||
|
||||
monkeypatch.setattr(
|
||||
model_scanner.config,
|
||||
"base_models_roots",
|
||||
[normalized_checkpoints, normalized_unet],
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
model_scanner.config,
|
||||
"checkpoints_roots",
|
||||
[normalized_checkpoints],
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
model_scanner.config,
|
||||
"unet_roots",
|
||||
[normalized_unet],
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
model_scanner.config,
|
||||
"extra_checkpoints_roots",
|
||||
[normalized_extra_checkpoints],
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
model_scanner.config,
|
||||
"extra_unet_roots",
|
||||
[normalized_extra_unet],
|
||||
raising=False,
|
||||
)
|
||||
|
||||
scanner = CheckpointScanner()
|
||||
roots = scanner.get_model_roots()
|
||||
|
||||
assert normalized_checkpoints in roots
|
||||
assert normalized_unet in roots
|
||||
assert normalized_extra_checkpoints in roots
|
||||
assert normalized_extra_unet in roots
|
||||
|
||||
@@ -470,6 +470,100 @@ def test_upsert_library_creates_entry_and_activates(manager, tmp_path):
|
||||
assert str(lora_dir).replace(os.sep, "/") in normalized_stored_paths
|
||||
|
||||
|
||||
def test_extra_folder_paths_stored_separately(manager, tmp_path):
|
||||
lora_dir = tmp_path / "loras"
|
||||
extra_dir = tmp_path / "extra_loras"
|
||||
lora_dir.mkdir()
|
||||
extra_dir.mkdir()
|
||||
|
||||
manager.upsert_library(
|
||||
"test_library",
|
||||
folder_paths={"loras": [str(lora_dir)]},
|
||||
extra_folder_paths={"loras": [str(extra_dir)]},
|
||||
activate=True,
|
||||
)
|
||||
|
||||
libraries = manager.get_libraries()
|
||||
lib = libraries["test_library"]
|
||||
|
||||
# Verify folder_paths contains main path
|
||||
assert str(lora_dir) in lib["folder_paths"]["loras"]
|
||||
# Verify extra_folder_paths contains extra path
|
||||
assert str(extra_dir) in lib["extra_folder_paths"]["loras"]
|
||||
# Verify they are separate
|
||||
assert str(extra_dir) not in lib["folder_paths"]["loras"]
|
||||
|
||||
|
||||
def test_get_extra_folder_paths(manager, tmp_path):
|
||||
extra_dir = tmp_path / "extra_loras"
|
||||
extra_dir.mkdir()
|
||||
|
||||
manager.update_extra_folder_paths({"loras": [str(extra_dir)]})
|
||||
|
||||
extra_paths = manager.get_extra_folder_paths()
|
||||
assert str(extra_dir) in extra_paths.get("loras", [])
|
||||
|
||||
|
||||
def test_library_switch_preserves_extra_paths(manager, tmp_path):
|
||||
"""Test that switching libraries preserves each library's extra paths."""
|
||||
lora_dir1 = tmp_path / "lib1_loras"
|
||||
extra_dir1 = tmp_path / "lib1_extra"
|
||||
lora_dir2 = tmp_path / "lib2_loras"
|
||||
extra_dir2 = tmp_path / "lib2_extra"
|
||||
|
||||
for directory in (lora_dir1, extra_dir1, lora_dir2, extra_dir2):
|
||||
directory.mkdir()
|
||||
|
||||
manager.create_library(
|
||||
"library1",
|
||||
folder_paths={"loras": [str(lora_dir1)]},
|
||||
extra_folder_paths={"loras": [str(extra_dir1)]},
|
||||
activate=True,
|
||||
)
|
||||
|
||||
manager.create_library(
|
||||
"library2",
|
||||
folder_paths={"loras": [str(lora_dir2)]},
|
||||
extra_folder_paths={"loras": [str(extra_dir2)]},
|
||||
)
|
||||
|
||||
assert manager.get_active_library_name() == "library1"
|
||||
lib1 = manager.get_active_library()
|
||||
assert str(lora_dir1) in lib1["folder_paths"]["loras"]
|
||||
assert str(extra_dir1) in lib1["extra_folder_paths"]["loras"]
|
||||
|
||||
manager.activate_library("library2")
|
||||
|
||||
assert manager.get_active_library_name() == "library2"
|
||||
lib2 = manager.get_active_library()
|
||||
assert str(lora_dir2) in lib2["folder_paths"]["loras"]
|
||||
assert str(extra_dir2) in lib2["extra_folder_paths"]["loras"]
|
||||
|
||||
|
||||
def test_extra_paths_validation_no_overlap_with_other_libraries(manager, tmp_path):
|
||||
"""Test that extra paths cannot overlap with other libraries' paths."""
|
||||
lora_dir1 = tmp_path / "lib1_loras"
|
||||
lora_dir1.mkdir()
|
||||
|
||||
manager.create_library(
|
||||
"library1",
|
||||
folder_paths={"loras": [str(lora_dir1)]},
|
||||
activate=True,
|
||||
)
|
||||
|
||||
extra_dir = tmp_path / "extra_loras"
|
||||
extra_dir.mkdir()
|
||||
|
||||
manager.create_library(
|
||||
"library2",
|
||||
folder_paths={"loras": [str(extra_dir)]},
|
||||
activate=True,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="already assigned to library"):
|
||||
manager.update_extra_folder_paths({"loras": [str(lora_dir1)]})
|
||||
|
||||
|
||||
def test_delete_library_switches_active(manager, tmp_path):
|
||||
other_dir = tmp_path / "other"
|
||||
other_dir.mkdir()
|
||||
|
||||
Reference in New Issue
Block a user