mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-27 05:54:08 -03:00
fix(linking): support external-source linking for Other-model roots
set_hf_url rejected files under config.other_roots (VAEs, text encoders, upscalers, ...) with 'File is not within any configured model directory' because neither _find_matching_root nor _infer_model_type knew about the Other category. Include other_roots in both, so linking routes the cache update to the Other scanner instead of falling back to the LoRA one, and downloads into Other roots keep the lazy-hash metadata path. Also make the root prefix match boundary-aware so /models/vae no longer swallows a sibling like /models/vae-old.
This commit is contained in:
@@ -39,7 +39,12 @@ from ...services.settings_manager import get_settings_manager
|
||||
from ...services.service_registry import ServiceRegistry
|
||||
from ...services.websocket_manager import ws_manager
|
||||
from ...utils.metadata_manager import MetadataManager
|
||||
from ...utils.models import LoraMetadata, CheckpointMetadata, EmbeddingMetadata
|
||||
from ...utils.models import (
|
||||
LoraMetadata,
|
||||
CheckpointMetadata,
|
||||
EmbeddingMetadata,
|
||||
OtherModelMetadata,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -78,6 +83,11 @@ def _infer_model_type(model_root: str) -> tuple[Any, str]:
|
||||
if os.path.normpath(p).replace(os.sep, "/") == norm:
|
||||
return EmbeddingMetadata, "get_embedding_scanner"
|
||||
|
||||
# Other-model roots (VAE, text encoders, upscalers, ...)
|
||||
for p in config.other_roots or []:
|
||||
if os.path.normpath(p).replace(os.sep, "/") == norm:
|
||||
return OtherModelMetadata, "get_other_scanner"
|
||||
|
||||
# Fallback — should not happen in normal use
|
||||
logger.warning(
|
||||
"Could not determine model type for root '%s'; defaulting to LoRA",
|
||||
@@ -213,12 +223,14 @@ def _find_matching_root(dest_dir: str) -> str | None:
|
||||
config.extra_unet_roots or [],
|
||||
config.embeddings_roots or [],
|
||||
config.extra_embeddings_roots or [],
|
||||
config.other_roots or [],
|
||||
):
|
||||
all_roots.extend([os.path.normpath(p).replace(os.sep, "/") for p in root_list])
|
||||
# Find the longest matching prefix
|
||||
# Find the longest matching prefix. The boundary check prevents a root like
|
||||
# `/models/vae` from swallowing a sibling directory like `/models/vae-old`.
|
||||
match: str | None = None
|
||||
for root in all_roots:
|
||||
if norm.startswith(root):
|
||||
if norm == root or norm.startswith(root + "/"):
|
||||
if match is None or len(root) > len(match):
|
||||
match = root
|
||||
return match
|
||||
|
||||
@@ -329,6 +329,106 @@ async def test_get_model_sources_lists_capabilities():
|
||||
assert all(s["example_url"] for s in sources)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Root matching (linking must recognise every configured model category)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def isolated_config_roots(monkeypatch, tmp_path):
|
||||
"""Clear every configured root so root matching is fully test-controlled."""
|
||||
from py.config import config
|
||||
|
||||
for attr in (
|
||||
"loras_roots",
|
||||
"extra_loras_roots",
|
||||
"checkpoints_roots",
|
||||
"extra_checkpoints_roots",
|
||||
"unet_roots",
|
||||
"extra_unet_roots",
|
||||
"embeddings_roots",
|
||||
"extra_embeddings_roots",
|
||||
"other_roots",
|
||||
):
|
||||
monkeypatch.setattr(config, attr, [])
|
||||
return config
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_hf_url_links_model_in_other_root(
|
||||
isolated_config_roots, tmp_path, monkeypatch
|
||||
):
|
||||
""""Other" category files live under config.other_roots; linking one to an
|
||||
external source must not fail with 'not within any configured model
|
||||
directory' and must update the Other scanner's cache, not the LoRA one."""
|
||||
vae_root = tmp_path / "vae"
|
||||
vae_root.mkdir()
|
||||
isolated_config_roots.other_roots = [str(vae_root)]
|
||||
model_path = vae_root / "some_vae.safetensors"
|
||||
await _write_plain_model(model_path)
|
||||
|
||||
other_scanner = SimpleNamespace(update_single_model_cache=AsyncMock())
|
||||
lora_scanner = SimpleNamespace(update_single_model_cache=AsyncMock())
|
||||
monkeypatch.setattr(
|
||||
ServiceRegistry, "get_other_scanner", AsyncMock(return_value=other_scanner)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ServiceRegistry, "get_lora_scanner", AsyncMock(return_value=lora_scanner)
|
||||
)
|
||||
|
||||
response = await ModelSourceHandler().set_hf_url(
|
||||
FakeRequest(
|
||||
json_data={
|
||||
"file_path": str(model_path),
|
||||
"source_url": "https://huggingface.co/user/repo",
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
assert response.status == 200
|
||||
payload = _json_payload(response)
|
||||
assert payload["success"] is True
|
||||
assert payload["source_platform"] == "huggingface"
|
||||
|
||||
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
|
||||
assert saved["source_platform"] == "huggingface"
|
||||
assert saved["hf_url"] == "https://huggingface.co/user/repo"
|
||||
|
||||
other_scanner.update_single_model_cache.assert_awaited_once()
|
||||
lora_scanner.update_single_model_cache.assert_not_awaited()
|
||||
|
||||
|
||||
def test_find_matching_root_includes_other_roots(isolated_config_roots, tmp_path):
|
||||
vae_root = tmp_path / "vae"
|
||||
isolated_config_roots.other_roots = [str(vae_root)]
|
||||
|
||||
assert model_source_handlers._find_matching_root(str(vae_root / "sub")) == str(
|
||||
vae_root
|
||||
)
|
||||
|
||||
|
||||
def test_find_matching_root_respects_path_boundaries(isolated_config_roots, tmp_path):
|
||||
"""A root must not swallow a sibling that merely shares its name prefix."""
|
||||
vae_root = tmp_path / "vae"
|
||||
isolated_config_roots.other_roots = [str(vae_root)]
|
||||
|
||||
assert (
|
||||
model_source_handlers._find_matching_root(str(tmp_path / "vae-old")) is None
|
||||
)
|
||||
|
||||
|
||||
def test_infer_model_type_recognises_other_roots(isolated_config_roots, tmp_path):
|
||||
from py.utils.models import OtherModelMetadata
|
||||
|
||||
vae_root = tmp_path / "vae"
|
||||
isolated_config_roots.other_roots = [str(vae_root)]
|
||||
|
||||
assert model_source_handlers._infer_model_type(str(vae_root)) == (
|
||||
OtherModelMetadata,
|
||||
"get_other_scanner",
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# File listing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user