mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-18 19:41:26 -03:00
feat(loaders): add control_after_generate random model selection to checkpoint/unet loaders
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).
This commit is contained in:
@@ -0,0 +1,89 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user