mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-21 03:01:27 -03:00
feat(backend): add Other model type (VAE/upscaler/text encoder) scanner, service and routes
This commit is contained in:
@@ -0,0 +1,277 @@
|
||||
"""Tests for other-model path handling in py/config.py."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from py import config as config_module
|
||||
from py.services.settings_manager import get_settings_manager
|
||||
|
||||
|
||||
def _normalize(path: str) -> str:
|
||||
return os.path.normpath(path).replace(os.sep, "/")
|
||||
|
||||
|
||||
def _make_config(**overrides) -> config_module.Config:
|
||||
"""Create a bare Config instance for _prepare_other_paths tests."""
|
||||
config = config_module.Config.__new__(config_module.Config)
|
||||
config._path_mappings = {}
|
||||
config._preview_root_paths = set()
|
||||
config._cached_fingerprint = None
|
||||
config.base_models_roots = []
|
||||
config.embeddings_roots = []
|
||||
for key, value in overrides.items():
|
||||
setattr(config, key, value)
|
||||
return config
|
||||
|
||||
|
||||
class TestPrepareOtherPaths:
|
||||
"""Unit tests for Config._prepare_other_paths."""
|
||||
|
||||
def test_maps_each_folder_key_to_sub_type(self, tmp_path):
|
||||
roots = {
|
||||
"vae": tmp_path / "vae",
|
||||
"upscale_models": tmp_path / "upscale_models",
|
||||
"text_encoders": tmp_path / "text_encoders",
|
||||
"clip": tmp_path / "clip",
|
||||
"clip_vision": tmp_path / "clip_vision",
|
||||
"controlnet": tmp_path / "controlnet",
|
||||
}
|
||||
for root in roots.values():
|
||||
root.mkdir()
|
||||
|
||||
config = _make_config()
|
||||
unique, sub_type_map, per_key = config._prepare_other_paths(
|
||||
{key: [str(root)] for key, root in roots.items()}
|
||||
)
|
||||
|
||||
assert len(unique) == 6
|
||||
assert sub_type_map[_normalize(str(roots["vae"]))] == "vae"
|
||||
assert sub_type_map[_normalize(str(roots["upscale_models"]))] == "upscaler"
|
||||
assert sub_type_map[_normalize(str(roots["text_encoders"]))] == "text_encoder"
|
||||
# Legacy ComfyUI 'clip' key maps to text_encoder as well
|
||||
assert sub_type_map[_normalize(str(roots["clip"]))] == "text_encoder"
|
||||
assert sub_type_map[_normalize(str(roots["clip_vision"]))] == "clip_vision"
|
||||
assert sub_type_map[_normalize(str(roots["controlnet"]))] == "controlnet"
|
||||
assert per_key["vae"] == [_normalize(str(roots["vae"]))]
|
||||
assert per_key["controlnet"] == [_normalize(str(roots["controlnet"]))]
|
||||
|
||||
def test_missing_or_unknown_keys_are_skipped(self, tmp_path):
|
||||
vae_root = tmp_path / "vae"
|
||||
vae_root.mkdir()
|
||||
|
||||
config = _make_config()
|
||||
unique, sub_type_map, per_key = config._prepare_other_paths(
|
||||
{
|
||||
"vae": [str(vae_root)],
|
||||
"does_not_exist_key": [str(tmp_path / "nope_dir")],
|
||||
"upscale_models": [],
|
||||
}
|
||||
)
|
||||
|
||||
assert unique == [_normalize(str(vae_root))]
|
||||
assert set(per_key.keys()) == {"vae"}
|
||||
|
||||
def test_nonexistent_paths_are_filtered(self, tmp_path):
|
||||
config = _make_config()
|
||||
unique, _, _ = config._prepare_other_paths(
|
||||
{"vae": [str(tmp_path / "missing_vae")]}
|
||||
)
|
||||
assert unique == []
|
||||
|
||||
def test_cross_category_overlap_warns_and_keeps_first(
|
||||
self, tmp_path, caplog
|
||||
):
|
||||
"""The same physical folder under two categories warns; first wins."""
|
||||
shared = tmp_path / "shared"
|
||||
shared.mkdir()
|
||||
|
||||
config = _make_config()
|
||||
with caplog.at_level(logging.WARNING, logger=config_module.logger.name):
|
||||
unique, sub_type_map, per_key = config._prepare_other_paths(
|
||||
{
|
||||
"vae": [str(shared)],
|
||||
"upscale_models": [str(shared)],
|
||||
}
|
||||
)
|
||||
|
||||
assert unique == [_normalize(str(shared))]
|
||||
assert sub_type_map[_normalize(str(shared))] == "vae"
|
||||
assert "upscale_models" not in per_key
|
||||
|
||||
warnings = [
|
||||
record.message
|
||||
for record in caplog.records
|
||||
if record.levelname == "WARNING"
|
||||
and "multiple other-model categories" in record.message
|
||||
]
|
||||
assert len(warnings) == 1
|
||||
|
||||
def test_cross_scanner_overlap_warns_but_keeps_path(self, tmp_path, caplog):
|
||||
"""An other root overlapping a checkpoint root warns but stays managed."""
|
||||
shared = tmp_path / "shared_models"
|
||||
shared.mkdir()
|
||||
|
||||
config = _make_config(base_models_roots=[_normalize(str(shared))])
|
||||
with caplog.at_level(logging.WARNING, logger=config_module.logger.name):
|
||||
unique, sub_type_map, _ = config._prepare_other_paths(
|
||||
{"vae": [str(shared)]}
|
||||
)
|
||||
|
||||
# Kept on purpose: dropping would silently unmanage the files
|
||||
assert unique == [_normalize(str(shared))]
|
||||
assert sub_type_map[_normalize(str(shared))] == "vae"
|
||||
|
||||
warnings = [
|
||||
record.message
|
||||
for record in caplog.records
|
||||
if record.levelname == "WARNING"
|
||||
and "overlaps an existing checkpoints/embeddings root" in record.message
|
||||
]
|
||||
assert len(warnings) == 1
|
||||
|
||||
def test_no_warning_for_disjoint_roots(self, tmp_path, caplog):
|
||||
checkpoints_root = tmp_path / "checkpoints"
|
||||
checkpoints_root.mkdir()
|
||||
vae_root = tmp_path / "vae"
|
||||
vae_root.mkdir()
|
||||
|
||||
config = _make_config(base_models_roots=[_normalize(str(checkpoints_root))])
|
||||
with caplog.at_level(logging.WARNING, logger=config_module.logger.name):
|
||||
unique, _, _ = config._prepare_other_paths({"vae": [str(vae_root)]})
|
||||
|
||||
assert unique == [_normalize(str(vae_root))]
|
||||
warnings = [
|
||||
record.message
|
||||
for record in caplog.records
|
||||
if record.levelname == "WARNING" and "overlap" in record.message.lower()
|
||||
]
|
||||
assert warnings == []
|
||||
|
||||
|
||||
class TestInitOtherPaths:
|
||||
"""Config._init_other_paths with mocked folder_paths (plugin + standalone modes).
|
||||
|
||||
Config only depends on ``folder_paths.get_folder_paths(key)``: ComfyUI in
|
||||
plugin mode, or MockFolderPaths serving ``settings.json.folder_paths`` in
|
||||
standalone mode. A dict-backed stub therefore covers both.
|
||||
"""
|
||||
|
||||
def _stub_folder_paths(self, monkeypatch, mapping):
|
||||
def get_folder_paths(key):
|
||||
value = mapping.get(key, [])
|
||||
return [value] if isinstance(value, str) else list(value)
|
||||
|
||||
monkeypatch.setattr(
|
||||
config_module.folder_paths, "get_folder_paths", get_folder_paths
|
||||
)
|
||||
|
||||
def test_default_enabled_keys_exclude_controlnet(self, monkeypatch, tmp_path):
|
||||
dirs = {}
|
||||
for key in (
|
||||
"vae",
|
||||
"upscale_models",
|
||||
"text_encoders",
|
||||
"clip",
|
||||
"clip_vision",
|
||||
"controlnet",
|
||||
):
|
||||
path = tmp_path / key
|
||||
path.mkdir()
|
||||
dirs[key] = str(path)
|
||||
|
||||
self._stub_folder_paths(monkeypatch, dirs)
|
||||
|
||||
config = _make_config()
|
||||
roots = config._init_other_paths()
|
||||
|
||||
assert _normalize(dirs["controlnet"]) not in roots
|
||||
assert _normalize(dirs["controlnet"]) not in config.other_root_subtypes
|
||||
for key in ("vae", "upscale_models", "text_encoders", "clip", "clip_vision"):
|
||||
assert _normalize(dirs[key]) in roots
|
||||
|
||||
def test_controlnet_opt_in_via_setting(self, monkeypatch, tmp_path):
|
||||
controlnet_dir = tmp_path / "controlnet"
|
||||
controlnet_dir.mkdir()
|
||||
|
||||
self._stub_folder_paths(monkeypatch, {"controlnet": str(controlnet_dir)})
|
||||
get_settings_manager().set("enabled_other_folders", ["controlnet"])
|
||||
|
||||
config = _make_config()
|
||||
roots = config._init_other_paths()
|
||||
|
||||
assert _normalize(str(controlnet_dir)) in roots
|
||||
assert (
|
||||
config.other_root_subtypes[_normalize(str(controlnet_dir))]
|
||||
== "controlnet"
|
||||
)
|
||||
|
||||
def test_unknown_opt_in_keys_are_ignored(self, monkeypatch, tmp_path):
|
||||
vae_dir = tmp_path / "vae"
|
||||
vae_dir.mkdir()
|
||||
|
||||
self._stub_folder_paths(monkeypatch, {"vae": str(vae_dir)})
|
||||
get_settings_manager().set("enabled_other_folders", ["not_a_real_key", 42])
|
||||
|
||||
config = _make_config()
|
||||
roots = config._init_other_paths()
|
||||
|
||||
assert roots == [_normalize(str(vae_dir))]
|
||||
|
||||
def test_apply_library_paths_picks_up_other_keys(self, monkeypatch, tmp_path):
|
||||
vae_dir = tmp_path / "vae"
|
||||
vae_dir.mkdir()
|
||||
|
||||
config = _make_config()
|
||||
monkeypatch.setattr(config, "_initialize_symlink_mappings", lambda: None)
|
||||
|
||||
config._apply_library_paths(
|
||||
{
|
||||
"loras": [],
|
||||
"checkpoints": [],
|
||||
"unet": [],
|
||||
"embeddings": [],
|
||||
"vae": [str(vae_dir)],
|
||||
}
|
||||
)
|
||||
|
||||
assert config.other_roots == [_normalize(str(vae_dir))]
|
||||
assert config.other_root_subtypes == {
|
||||
_normalize(str(vae_dir)): "vae"
|
||||
}
|
||||
assert config.other_folder_roots == {"vae": [_normalize(str(vae_dir))]}
|
||||
|
||||
|
||||
class TestOtherRootsWiring:
|
||||
"""other_roots participates in symlink and preview root bookkeeping."""
|
||||
|
||||
def test_symlink_roots_include_other_roots(self):
|
||||
config = _make_config()
|
||||
config.loras_roots = ["/loras"]
|
||||
config.embeddings_roots = ["/embeddings"]
|
||||
config.other_roots = ["/vae"]
|
||||
config.extra_loras_roots = []
|
||||
config.extra_checkpoints_roots = []
|
||||
config.extra_unet_roots = []
|
||||
config.extra_embeddings_roots = []
|
||||
|
||||
assert "/vae" in config._symlink_roots()
|
||||
|
||||
def test_preview_roots_include_other_roots(self, tmp_path):
|
||||
vae_dir = tmp_path / "vae"
|
||||
vae_dir.mkdir()
|
||||
|
||||
config = _make_config()
|
||||
config.loras_roots = []
|
||||
config.embeddings_roots = []
|
||||
config.other_roots = [_normalize(str(vae_dir))]
|
||||
config.extra_loras_roots = []
|
||||
config.extra_checkpoints_roots = []
|
||||
config.extra_unet_roots = []
|
||||
config.extra_embeddings_roots = []
|
||||
config.recipes_path = ""
|
||||
|
||||
config._rebuild_preview_roots()
|
||||
|
||||
assert config.is_preview_path_allowed(str(vae_dir / "model.preview.png"))
|
||||
@@ -132,6 +132,7 @@ async def test_lora_manager_lifecycle(monkeypatch: pytest.MonkeyPatch, tmp_path:
|
||||
"lora": _DummyScanner("lora"),
|
||||
"checkpoint": _DummyScanner("checkpoint"),
|
||||
"embedding": _DummyScanner("embedding"),
|
||||
"other": _DummyScanner("other"),
|
||||
"recipe": _DummyScanner("recipe"),
|
||||
}
|
||||
|
||||
@@ -147,6 +148,7 @@ async def test_lora_manager_lifecycle(monkeypatch: pytest.MonkeyPatch, tmp_path:
|
||||
monkeypatch.setattr(lora_manager.ServiceRegistry, "get_lora_scanner", lambda: _stub("lora_scanner", scanners["lora"]))
|
||||
monkeypatch.setattr(lora_manager.ServiceRegistry, "get_checkpoint_scanner", lambda: _stub("checkpoint_scanner", scanners["checkpoint"]))
|
||||
monkeypatch.setattr(lora_manager.ServiceRegistry, "get_embedding_scanner", lambda: _stub("embedding_scanner", scanners["embedding"]))
|
||||
monkeypatch.setattr(lora_manager.ServiceRegistry, "get_other_scanner", lambda: _stub("other_scanner", scanners["other"]))
|
||||
monkeypatch.setattr(lora_manager.ServiceRegistry, "get_recipe_scanner", lambda: _stub("recipe_scanner", scanners["recipe"]))
|
||||
|
||||
migration_calls: list[bool] = []
|
||||
@@ -205,7 +207,7 @@ async def test_lora_manager_lifecycle(monkeypatch: pytest.MonkeyPatch, tmp_path:
|
||||
await asyncio.gather(*pending)
|
||||
|
||||
task_names = {task.get_name() for task in scheduled_tasks}
|
||||
assert {"lora_cache_init", "checkpoint_cache_init", "embedding_cache_init", "recipe_cache_init", "post_init_tasks", "cleanup_bak_files"}.issubset(task_names)
|
||||
assert {"lora_cache_init", "checkpoint_cache_init", "embedding_cache_init", "other_cache_init", "recipe_cache_init", "post_init_tasks", "cleanup_bak_files"}.issubset(task_names)
|
||||
|
||||
# Startup sweep: an expired pending-delete purge task is spawned during
|
||||
# service initialization (covers both plugin and standalone modes).
|
||||
@@ -219,4 +221,4 @@ async def test_lora_manager_lifecycle(monkeypatch: pytest.MonkeyPatch, tmp_path:
|
||||
for root in (loras_root, checkpoints_root, embeddings_root):
|
||||
assert not any(path.suffix == ".bak" for path in root.rglob("*")), f"Backup files remain in {root}"
|
||||
|
||||
assert {"civitai_client", "download_manager", "websocket_manager", "lora_scanner", "checkpoint_scanner", "embedding_scanner", "recipe_scanner"}.issubset(registry_calls)
|
||||
assert {"civitai_client", "download_manager", "websocket_manager", "lora_scanner", "checkpoint_scanner", "embedding_scanner", "other_scanner", "recipe_scanner"}.issubset(registry_calls)
|
||||
|
||||
@@ -0,0 +1,118 @@
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
|
||||
from py.routes.other_routes import OtherRoutes
|
||||
from py.services.other_model_service import OtherModelService
|
||||
|
||||
|
||||
class DummyRequest:
|
||||
def __init__(self, *, match_info=None):
|
||||
self.match_info = match_info or {}
|
||||
|
||||
|
||||
class StubOtherModelService:
|
||||
def __init__(self):
|
||||
self.info = {}
|
||||
|
||||
async def get_model_info_by_name(self, name):
|
||||
value = self.info.get(name)
|
||||
if isinstance(value, Exception):
|
||||
raise value
|
||||
return value
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def routes():
|
||||
handler = OtherRoutes()
|
||||
handler.service = StubOtherModelService() # pyright: ignore[reportAttributeAccessIssue]
|
||||
return handler
|
||||
|
||||
|
||||
def test_common_and_specific_routes_registered():
|
||||
"""Registration smoke test: /api/lm/other/* surface plus the /other page."""
|
||||
app = web.Application()
|
||||
OtherRoutes().setup_routes(app)
|
||||
|
||||
registered = {(route.method, route.resource.canonical) for route in app.router.routes()}
|
||||
|
||||
assert ("GET", "/other") in registered
|
||||
assert ("GET", "/api/lm/other/list") in registered
|
||||
assert ("GET", "/api/lm/other/model-types") in registered
|
||||
assert ("GET", "/api/lm/other/roots") in registered
|
||||
assert ("POST", "/api/lm/other/fetch-civitai") in registered
|
||||
assert ("POST", "/api/lm/other/delete") in registered
|
||||
assert ("POST", "/api/lm/other/move_model") in registered
|
||||
assert ("GET", "/api/lm/other/info/{name}") in registered
|
||||
|
||||
|
||||
def test_template_name_is_other_page():
|
||||
assert OtherRoutes().template_name == "other.html"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_type",
|
||||
["VAE", "Upscaler", "TextEncoder", "CLIP", "CLIPVision", "Controlnet", "Other"],
|
||||
)
|
||||
def test_validate_civitai_model_type_accepts_other_types(model_type):
|
||||
assert OtherRoutes()._validate_civitai_model_type(model_type) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_type", ["Lora", "Checkpoint", "TextualInversion"])
|
||||
def test_validate_civitai_model_type_rejects_foreign_types(model_type):
|
||||
assert OtherRoutes()._validate_civitai_model_type(model_type) is False
|
||||
|
||||
|
||||
def test_get_expected_model_types_mentions_supported_types():
|
||||
expected = OtherRoutes()._get_expected_model_types()
|
||||
for name in ("VAE", "Upscaler", "TextEncoder", "CLIPVision", "Controlnet"):
|
||||
assert name in expected
|
||||
|
||||
|
||||
async def test_get_other_model_info_success(routes):
|
||||
routes.service.info["demo"] = {"name": "demo"}
|
||||
response = await routes.get_other_model_info(DummyRequest(match_info={"name": "demo"}))
|
||||
payload = json.loads(response.text)
|
||||
assert payload == {"name": "demo"}
|
||||
|
||||
|
||||
async def test_get_other_model_info_missing(routes):
|
||||
response = await routes.get_other_model_info(DummyRequest(match_info={"name": "missing"}))
|
||||
payload = json.loads(response.text)
|
||||
assert response.status == 404
|
||||
assert payload == {"error": "Model not found"}
|
||||
|
||||
|
||||
async def test_get_other_model_info_error(routes):
|
||||
routes.service.info["demo"] = RuntimeError("boom")
|
||||
response = await routes.get_other_model_info(DummyRequest(match_info={"name": "demo"}))
|
||||
payload = json.loads(response.text)
|
||||
assert response.status == 500
|
||||
assert payload == {"error": "boom"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_services_builds_other_model_service(monkeypatch):
|
||||
from py.services.service_registry import ServiceRegistry
|
||||
|
||||
sentinel_scanner = object()
|
||||
sentinel_update_service = object()
|
||||
|
||||
async def fake_scanner():
|
||||
return sentinel_scanner
|
||||
|
||||
async def fake_update_service():
|
||||
return sentinel_update_service
|
||||
|
||||
monkeypatch.setattr(ServiceRegistry, "get_other_scanner", staticmethod(fake_scanner))
|
||||
monkeypatch.setattr(
|
||||
ServiceRegistry, "get_model_update_service", staticmethod(fake_update_service)
|
||||
)
|
||||
|
||||
handler = OtherRoutes()
|
||||
await handler.initialize_services()
|
||||
|
||||
assert isinstance(handler.service, OtherModelService)
|
||||
assert handler.service.model_type == "other"
|
||||
assert handler.service.scanner is sentinel_scanner
|
||||
@@ -0,0 +1,337 @@
|
||||
"""Tests for OtherScanner: root aggregation, sub_type derivation, lazy hash."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from py import config as config_module
|
||||
from py.services import model_scanner
|
||||
from py.services.model_scanner import ModelScanner
|
||||
from py.services.other_scanner import OtherScanner
|
||||
from py.utils.models import OtherModelMetadata
|
||||
|
||||
|
||||
def _normalize(path) -> str:
|
||||
return str(path).replace(os.sep, "/")
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_model_scanner_singletons():
|
||||
ModelScanner._instances.clear()
|
||||
ModelScanner._locks.clear()
|
||||
yield
|
||||
ModelScanner._instances.clear()
|
||||
ModelScanner._locks.clear()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def other_config(monkeypatch, tmp_path):
|
||||
"""Point the global config at a synthetic set of other-model roots."""
|
||||
vae_root = tmp_path / "vae"
|
||||
upscaler_root = tmp_path / "upscale_models"
|
||||
te_root = tmp_path / "text_encoders"
|
||||
clip_root = tmp_path / "clip"
|
||||
clip_vision_root = tmp_path / "clip_vision"
|
||||
for root in (vae_root, upscaler_root, te_root, clip_root, clip_vision_root):
|
||||
root.mkdir()
|
||||
|
||||
roots = [
|
||||
_normalize(vae_root),
|
||||
_normalize(upscaler_root),
|
||||
_normalize(te_root),
|
||||
_normalize(clip_root),
|
||||
_normalize(clip_vision_root),
|
||||
]
|
||||
subtypes = {
|
||||
_normalize(vae_root): "vae",
|
||||
_normalize(upscaler_root): "upscaler",
|
||||
_normalize(te_root): "text_encoder",
|
||||
_normalize(clip_root): "text_encoder",
|
||||
_normalize(clip_vision_root): "clip_vision",
|
||||
}
|
||||
monkeypatch.setattr(config_module.config, "other_roots", roots)
|
||||
monkeypatch.setattr(config_module.config, "other_root_subtypes", subtypes)
|
||||
return {
|
||||
"roots": roots,
|
||||
"subtypes": subtypes,
|
||||
"vae": _normalize(vae_root),
|
||||
"upscaler": _normalize(upscaler_root),
|
||||
"text_encoders": _normalize(te_root),
|
||||
"clip": _normalize(clip_root),
|
||||
"clip_vision": _normalize(clip_vision_root),
|
||||
}
|
||||
|
||||
|
||||
def _make_scanner() -> OtherScanner:
|
||||
"""Create a scanner without __init__ to avoid async initialization."""
|
||||
scanner = object.__new__(OtherScanner)
|
||||
scanner.model_type = "other"
|
||||
scanner.model_class = OtherModelMetadata
|
||||
scanner.file_extensions = {".safetensors", ".pt", ".bin"}
|
||||
scanner._hash_index = MagicMock()
|
||||
return scanner
|
||||
|
||||
|
||||
class TestOtherScannerRoots:
|
||||
"""Root aggregation and sub_type resolution."""
|
||||
|
||||
def test_get_model_roots_aggregates_and_dedupes(self, other_config, monkeypatch):
|
||||
scanner = _make_scanner()
|
||||
monkeypatch.setattr(
|
||||
config_module.config,
|
||||
"other_roots",
|
||||
other_config["roots"] + [other_config["vae"]],
|
||||
)
|
||||
roots = scanner.get_model_roots()
|
||||
assert roots == other_config["roots"]
|
||||
|
||||
def test_get_model_roots_empty_when_unconfigured(self, monkeypatch):
|
||||
scanner = _make_scanner()
|
||||
monkeypatch.setattr(config_module.config, "other_roots", None)
|
||||
assert scanner.get_model_roots() == []
|
||||
|
||||
def test_resolve_sub_type_for_each_default_category(self, other_config):
|
||||
scanner = _make_scanner()
|
||||
cases = [
|
||||
(other_config["vae"], "vae"),
|
||||
(other_config["upscaler"], "upscaler"),
|
||||
(other_config["text_encoders"], "text_encoder"),
|
||||
# Legacy ComfyUI 'clip' key maps to text_encoder as well
|
||||
(other_config["clip"], "text_encoder"),
|
||||
(other_config["clip_vision"], "clip_vision"),
|
||||
]
|
||||
for root, expected in cases:
|
||||
file_path = f"{root}/model.safetensors"
|
||||
assert scanner.resolve_sub_type_for_path(file_path) == expected
|
||||
|
||||
def test_resolve_sub_type_longest_prefix_wins(self, monkeypatch, tmp_path):
|
||||
"""A nested root (controlnet inside vae) resolves to the inner category."""
|
||||
outer = tmp_path / "vae"
|
||||
inner = outer / "controlnet"
|
||||
inner.mkdir(parents=True)
|
||||
monkeypatch.setattr(
|
||||
config_module.config,
|
||||
"other_root_subtypes",
|
||||
{_normalize(outer): "vae", _normalize(inner): "controlnet"},
|
||||
)
|
||||
scanner = _make_scanner()
|
||||
assert (
|
||||
scanner.resolve_sub_type_for_path(f"{_normalize(inner)}/cn.safetensors")
|
||||
== "controlnet"
|
||||
)
|
||||
assert (
|
||||
scanner.resolve_sub_type_for_path(f"{_normalize(outer)}/vae.safetensors")
|
||||
== "vae"
|
||||
)
|
||||
|
||||
def test_resolve_sub_type_none_for_unknown_or_empty(self, other_config):
|
||||
scanner = _make_scanner()
|
||||
assert scanner.resolve_sub_type_for_path(None) is None
|
||||
assert scanner.resolve_sub_type_for_path("") is None
|
||||
assert scanner.resolve_sub_type_for_path("/unrelated/model.safetensors") is None
|
||||
|
||||
def test_adjust_metadata_sets_sub_type(self, other_config):
|
||||
scanner = _make_scanner()
|
||||
metadata = OtherModelMetadata(
|
||||
file_name="te",
|
||||
model_name="te",
|
||||
file_path=f"{other_config['text_encoders']}/te.safetensors",
|
||||
size=1,
|
||||
modified=0.0,
|
||||
sha256="",
|
||||
base_model="Unknown",
|
||||
preview_url="",
|
||||
)
|
||||
result = scanner.adjust_metadata(
|
||||
metadata,
|
||||
metadata.file_path,
|
||||
other_config["text_encoders"],
|
||||
)
|
||||
assert result.sub_type == "text_encoder"
|
||||
|
||||
def test_adjust_cached_entry_rederives_sub_type(self, other_config):
|
||||
"""Persisted sub_type is never trusted: it is re-derived from location."""
|
||||
scanner = _make_scanner()
|
||||
entry = {
|
||||
"file_path": f"{other_config['clip']}/legacy.safetensors",
|
||||
"sub_type": "vae", # stale value from an old snapshot
|
||||
}
|
||||
result = scanner.adjust_cached_entry(entry)
|
||||
assert result["sub_type"] == "text_encoder"
|
||||
|
||||
def test_adjust_cached_entry_keeps_value_when_root_unknown(self, other_config):
|
||||
scanner = _make_scanner()
|
||||
entry = {
|
||||
"file_path": "/gone/model.safetensors",
|
||||
"sub_type": "upscaler",
|
||||
}
|
||||
result = scanner.adjust_cached_entry(entry)
|
||||
assert result["sub_type"] == "upscaler"
|
||||
|
||||
|
||||
class TestOtherScannerLazyHash:
|
||||
"""Lazy hashing: pending by default, singleflight on-demand calculation."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_metadata_has_pending_hash(self, other_config):
|
||||
vae_file = Path(other_config["vae"]) / "vae_model.safetensors"
|
||||
vae_file.write_text("fake vae content", encoding="utf-8")
|
||||
|
||||
scanner = OtherScanner()
|
||||
metadata = await scanner._create_default_metadata(_normalize(vae_file))
|
||||
|
||||
assert metadata is not None
|
||||
assert metadata.sha256 == ""
|
||||
assert metadata.hash_status == "pending"
|
||||
assert metadata.from_civitai is False
|
||||
assert metadata.sub_type == "vae"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_metadata_sub_type_from_location(self, other_config):
|
||||
te_file = Path(other_config["text_encoders"]) / "t5.safetensors"
|
||||
te_file.write_text("fake text encoder", encoding="utf-8")
|
||||
|
||||
scanner = OtherScanner()
|
||||
metadata = await scanner._create_default_metadata(_normalize(te_file))
|
||||
|
||||
assert metadata is not None
|
||||
assert metadata.sub_type == "text_encoder"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_calculate_hash_for_model_completes_pending(self, other_config):
|
||||
model_file = Path(other_config["upscaler"]) / "upscaler.safetensors"
|
||||
model_file.write_text("fake upscaler content", encoding="utf-8")
|
||||
normalized_file = _normalize(model_file)
|
||||
|
||||
scanner = OtherScanner()
|
||||
metadata = await scanner._create_default_metadata(normalized_file)
|
||||
assert metadata is not None and metadata.hash_status == "pending"
|
||||
|
||||
hash_result = await scanner.calculate_hash_for_model(normalized_file)
|
||||
|
||||
assert hash_result is not None
|
||||
assert len(hash_result) == 64
|
||||
|
||||
metadata_file = model_file.with_suffix(".metadata.json")
|
||||
saved_data = json.loads(metadata_file.read_text(encoding="utf-8"))
|
||||
assert saved_data["sha256"] == hash_result
|
||||
assert saved_data["hash_status"] == "completed"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_calculate_hash_singleflight_same_file(self, other_config):
|
||||
"""Concurrent calls for the same file share one SHA256 task."""
|
||||
model_file = Path(other_config["vae"]) / "shared.safetensors"
|
||||
model_file.write_text("fake content", encoding="utf-8")
|
||||
normalized_file = _normalize(model_file)
|
||||
real_file = os.path.realpath(normalized_file)
|
||||
|
||||
scanner = OtherScanner()
|
||||
metadata = await scanner._create_default_metadata(normalized_file)
|
||||
assert metadata is not None
|
||||
|
||||
calls = []
|
||||
|
||||
async def fake_calculate_sha256(file_path: str) -> str:
|
||||
calls.append(file_path)
|
||||
await asyncio.sleep(0.01)
|
||||
return "a" * 64
|
||||
|
||||
with patch(
|
||||
"py.utils.file_utils.calculate_sha256", side_effect=fake_calculate_sha256
|
||||
):
|
||||
results = await asyncio.gather(
|
||||
*[scanner.calculate_hash_for_model(normalized_file) for _ in range(8)]
|
||||
)
|
||||
|
||||
assert calls == [real_file]
|
||||
assert results == ["a" * 64] * 8
|
||||
assert scanner._hash_calculation_tasks == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_calculate_hash_skips_completed(self, other_config):
|
||||
model_file = Path(other_config["clip_vision"]) / "cv.safetensors"
|
||||
model_file.write_text("fake content", encoding="utf-8")
|
||||
normalized_file = _normalize(model_file)
|
||||
|
||||
scanner = OtherScanner()
|
||||
metadata = await scanner._create_default_metadata(normalized_file)
|
||||
assert metadata is not None
|
||||
# Simulate an already-completed hash
|
||||
metadata.sha256 = "existing_hash"
|
||||
metadata.hash_status = "completed"
|
||||
from py.utils.metadata_manager import MetadataManager
|
||||
|
||||
await MetadataManager.save_metadata(normalized_file, metadata)
|
||||
|
||||
with patch("py.utils.file_utils.calculate_sha256") as mock_calc:
|
||||
hash_result = await scanner.calculate_hash_for_model(normalized_file)
|
||||
|
||||
assert hash_result == "existing_hash"
|
||||
mock_calc.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_calculate_all_pending_hashes(self, other_config):
|
||||
for index in range(3):
|
||||
model_file = Path(other_config["vae"]) / f"model_{index}.safetensors"
|
||||
model_file.write_text(f"content {index}", encoding="utf-8")
|
||||
|
||||
scanner = OtherScanner()
|
||||
for index in range(3):
|
||||
model_file = Path(other_config["vae"]) / f"model_{index}.safetensors"
|
||||
await scanner._create_default_metadata(_normalize(model_file))
|
||||
|
||||
progress_calls = []
|
||||
|
||||
async def progress_callback(current, total, file_path):
|
||||
progress_calls.append((current, total, file_path))
|
||||
|
||||
result = await scanner.calculate_all_pending_hashes(progress_callback)
|
||||
|
||||
assert result["total"] == 3
|
||||
assert result["completed"] == 3
|
||||
assert result["failed"] == 0
|
||||
assert len(progress_calls) == 3
|
||||
|
||||
|
||||
class TestOtherModelMetadataFromCivitai:
|
||||
"""CivitAI type mapping in OtherModelMetadata.from_civitai_info."""
|
||||
|
||||
def _build(self, civitai_type: str) -> OtherModelMetadata:
|
||||
return OtherModelMetadata.from_civitai_info(
|
||||
{
|
||||
"type": civitai_type,
|
||||
"baseModel": "SDXL",
|
||||
"model": {"name": "Model", "tags": ["tag"], "description": "desc"},
|
||||
},
|
||||
{"name": "model.safetensors", "sizeKB": 1, "hashes": {"SHA256": "AB"}},
|
||||
"/tmp/model.safetensors",
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"civitai_type,expected",
|
||||
[
|
||||
("VAE", "vae"),
|
||||
("Upscaler", "upscaler"),
|
||||
("TextEncoder", "text_encoder"),
|
||||
("CLIP", "text_encoder"),
|
||||
("CLIPVision", "clip_vision"),
|
||||
("Controlnet", "controlnet"),
|
||||
("Other", "vae"), # unknown types fall back to the placeholder
|
||||
],
|
||||
)
|
||||
def test_civitai_type_mapping(self, civitai_type, expected):
|
||||
metadata = self._build(civitai_type)
|
||||
assert metadata.sub_type == expected
|
||||
assert metadata.sha256 == "ab"
|
||||
assert metadata.tags == ["tag"]
|
||||
|
||||
|
||||
def test_page_type_maps_to_other():
|
||||
"""The WS progress page type for the other scanner is 'other'."""
|
||||
assert model_scanner.PAGE_TYPE_MAP["other"] == "other"
|
||||
scanner = _make_scanner()
|
||||
assert scanner.page_type == "other"
|
||||
@@ -7,6 +7,7 @@ from unittest.mock import MagicMock, AsyncMock
|
||||
from py.services.lora_service import LoraService
|
||||
from py.services.checkpoint_service import CheckpointService
|
||||
from py.services.embedding_service import EmbeddingService
|
||||
from py.services.other_model_service import OtherModelService
|
||||
|
||||
|
||||
class TestLoraServiceFormatResponse:
|
||||
@@ -206,6 +207,89 @@ class TestEmbeddingServiceFormatResponse:
|
||||
assert "model_type" not in result # Removed in refactoring
|
||||
|
||||
|
||||
class TestOtherModelServiceFormatResponse:
|
||||
"""Test OtherModelService.format_response includes sub_type."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_scanner(self):
|
||||
scanner = MagicMock()
|
||||
scanner._hash_index = MagicMock()
|
||||
return scanner
|
||||
|
||||
@pytest.fixture
|
||||
def other_service(self, mock_scanner):
|
||||
return OtherModelService(mock_scanner)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_format_response_includes_sub_type(self, other_service):
|
||||
"""format_response should include sub_type field."""
|
||||
other_data = {
|
||||
"model_name": "Test VAE",
|
||||
"file_name": "test_vae",
|
||||
"preview_url": "test.webp",
|
||||
"preview_nsfw_level": 0,
|
||||
"base_model": "SDXL",
|
||||
"folder": "",
|
||||
"sha256": "abc123",
|
||||
"file_path": "/models/vae/test_vae.safetensors",
|
||||
"size": 1000,
|
||||
"modified": 1234567890.0,
|
||||
"tags": [],
|
||||
"from_civitai": True,
|
||||
"notes": "",
|
||||
"favorite": False,
|
||||
"sub_type": "vae",
|
||||
"civitai": {},
|
||||
}
|
||||
|
||||
result = await other_service.format_response(other_data)
|
||||
|
||||
assert "sub_type" in result
|
||||
assert result["sub_type"] == "vae"
|
||||
assert "model_type" not in result # Removed in refactoring
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_format_response_defaults_to_vae(self, other_service):
|
||||
"""format_response should default to 'vae' if no sub_type field."""
|
||||
other_data = {
|
||||
"model_name": "Test Upscaler",
|
||||
"file_name": "test_upscaler",
|
||||
"preview_url": "test.webp",
|
||||
"preview_nsfw_level": 0,
|
||||
"base_model": "SD1.5",
|
||||
"folder": "",
|
||||
"sha256": "abc123",
|
||||
"file_path": "/models/upscale_models/test.pth",
|
||||
"size": 1000,
|
||||
"modified": 1234567890.0,
|
||||
"tags": [],
|
||||
"from_civitai": True,
|
||||
"civitai": {},
|
||||
}
|
||||
|
||||
result = await other_service.format_response(other_data)
|
||||
|
||||
assert result["sub_type"] == "vae"
|
||||
assert "model_type" not in result # Removed in refactoring
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_format_response_returns_none_on_missing_file_path(self, other_service):
|
||||
"""format_response returns None when file_path is missing (corrupted row)."""
|
||||
other_data = {
|
||||
"model_name": "Test",
|
||||
"file_name": "test",
|
||||
"file_path": None, # corrupted: missing file_path
|
||||
"folder": "",
|
||||
"sha256": "abc",
|
||||
"tags": [],
|
||||
"from_civitai": True,
|
||||
"civitai": {},
|
||||
"sub_type": "text_encoder",
|
||||
}
|
||||
result = await other_service.format_response(other_data)
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestFormatResponseCorruptedEntries:
|
||||
"""Test format_response handles corrupted cache entries gracefully (issue #730).
|
||||
|
||||
|
||||
Reference in New Issue
Block a user