mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-18 19:41:26 -03:00
Will Miao
fc3f3f3bdb
The Checkpoint/Unet Loader (LoraManager) nodes now support ComfyUI's built-in control_after_generate mechanism on the ckpt_name/unet_name combos, letting users pick a random model on every queue with the selected model written back into the widget (visible, and lockable via the 'fixed' mode). A base_model input narrows the random pool: a front-end extension fetches the name/base_model mapping from the new /api/lm/checkpoints/loader-pool endpoint and filters the combo options, wired through the node callback, the refreshComboInNodes extension hook, and a graph.onConfigure hook installed from onAdded (onNodeCreated fires before the node is attached to a graph, so the graph reference is unavailable there).
90 lines
2.6 KiB
Python
90 lines
2.6 KiB
Python
"""Tests for the loader-pool endpoint backing the Random Checkpoint/Unet
|
|
Loader nodes' front-end base_model filtering.
|
|
"""
|
|
|
|
import json
|
|
|
|
import pytest
|
|
|
|
from py.routes.checkpoint_routes import CheckpointRoutes
|
|
from py.services.service_registry import ServiceRegistry
|
|
|
|
|
|
class _FakeCache:
|
|
def __init__(self, raw_data):
|
|
self.raw_data = raw_data
|
|
|
|
|
|
class _FakeScanner:
|
|
def __init__(self, raw_data, model_roots):
|
|
self._raw_data = raw_data
|
|
self._model_roots = model_roots
|
|
|
|
async def get_cached_data(self, force_refresh=False):
|
|
return _FakeCache(self._raw_data)
|
|
|
|
def get_model_roots(self):
|
|
return self._model_roots
|
|
|
|
|
|
class DummyRequest:
|
|
def __init__(self, query=None):
|
|
self.query = query or {}
|
|
|
|
|
|
@pytest.fixture
|
|
def routes(tmp_path, monkeypatch):
|
|
existing = tmp_path / "flux.safetensors"
|
|
existing.write_bytes(b"x")
|
|
missing = tmp_path / "missing.safetensors" # referenced but never created
|
|
|
|
raw_data = [
|
|
{"sub_type": "checkpoint", "file_path": str(existing), "base_model": "Flux.1 D"},
|
|
{"sub_type": "checkpoint", "file_path": str(missing), "base_model": "SDXL 1.0"},
|
|
{
|
|
"sub_type": "diffusion_model",
|
|
"file_path": str(existing),
|
|
"base_model": "Flux.1 D",
|
|
},
|
|
]
|
|
|
|
async def _fake_scanner():
|
|
return _FakeScanner(raw_data, [str(tmp_path)])
|
|
|
|
monkeypatch.setattr(ServiceRegistry, "get_checkpoint_scanner", _fake_scanner)
|
|
return CheckpointRoutes()
|
|
|
|
|
|
async def test_loader_pool_checkpoint_subtype(routes):
|
|
response = await routes.get_loader_pool(DummyRequest(query={"sub_type": "checkpoint"}))
|
|
assert response.status == 200
|
|
payload = json.loads(response.text)
|
|
assert payload == {
|
|
"items": [{"name": "flux.safetensors", "base_model": "Flux.1 D"}]
|
|
}
|
|
|
|
|
|
async def test_loader_pool_diffusion_model_subtype(routes):
|
|
response = await routes.get_loader_pool(
|
|
DummyRequest(query={"sub_type": "diffusion_model"})
|
|
)
|
|
assert response.status == 200
|
|
payload = json.loads(response.text)
|
|
assert payload == {
|
|
"items": [{"name": "flux.safetensors", "base_model": "Flux.1 D"}]
|
|
}
|
|
|
|
|
|
async def test_loader_pool_default_subtype_is_checkpoint(routes):
|
|
response = await routes.get_loader_pool(DummyRequest())
|
|
assert response.status == 200
|
|
payload = json.loads(response.text)
|
|
assert payload == {
|
|
"items": [{"name": "flux.safetensors", "base_model": "Flux.1 D"}]
|
|
}
|
|
|
|
|
|
async def test_loader_pool_invalid_subtype(routes):
|
|
response = await routes.get_loader_pool(DummyRequest(query={"sub_type": "lora"}))
|
|
assert response.status == 400
|