Vito Sansevero
c5e00c130b
- Checkpoint Folder Iterator loads one checkpoint per queued run from a folder (relative to models/checkpoints or absolute), so a whole folder can be swept with only one model in memory at a time - Model Compare Sheet saves each run's image as a tile under output/model_compare/<run_name>/ and rebuilds a labeled contact sheet from disk every run, so the grid survives between independent prompts - Starting a sweep at index 0 archives stale tiles to a timestamped folder - Running past the last checkpoint raises a clear error instead of wrapping, so an over-queued batch never restarts and archives a finished sweep - Checkpoints load from the scanned absolute path, so folders outside the configured roots work as the tooltip promises - Extract the LoRA Folder Batch folder-scanning, natural-sort and regex filter helpers into kikotools/base/model_folders.py so both tools share one implementation
72 lines
2.9 KiB
Python
72 lines
2.9 KiB
Python
"""Tests for the Checkpoint Folder Iterator node."""
|
|
|
|
import sys
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from kikotools.tools.checkpoint_folder_iterator import CheckpointFolderIteratorNode
|
|
|
|
NODE = "kikotools.tools.checkpoint_folder_iterator.node"
|
|
ENTRIES = [
|
|
("X/a_v1.safetensors", "/models/checkpoints/X/a_v1.safetensors"),
|
|
("X/b_v2.safetensors", "/models/checkpoints/X/b_v2.safetensors"),
|
|
]
|
|
|
|
|
|
def _fake_comfy_modules():
|
|
fp = MagicMock()
|
|
fp.get_folder_paths = MagicMock(return_value=["/models/embeddings"])
|
|
comfy = MagicMock()
|
|
comfy.sd.load_checkpoint_guess_config = MagicMock(
|
|
return_value=("MODEL", "CLIP", "VAE", "extra")
|
|
)
|
|
return (
|
|
patch.dict(
|
|
sys.modules, {"folder_paths": fp, "comfy": comfy, "comfy.sd": comfy.sd}
|
|
),
|
|
fp,
|
|
comfy,
|
|
)
|
|
|
|
|
|
class TestCheckpointFolderIteratorNode:
|
|
def test_input_index_increments_after_generate(self):
|
|
index = CheckpointFolderIteratorNode.INPUT_TYPES()["required"]["index"]
|
|
assert index[0] == "INT"
|
|
assert index[1]["control_after_generate"] is True
|
|
|
|
def test_loads_the_indexed_checkpoint_from_its_scanned_path(self):
|
|
ctx, fp, comfy = _fake_comfy_modules()
|
|
with ctx, patch(f"{NODE}.list_checkpoint_entries", return_value=ENTRIES):
|
|
out = CheckpointFolderIteratorNode().load("X", 1, "", "")
|
|
comfy.sd.load_checkpoint_guess_config.assert_called_once()
|
|
assert comfy.sd.load_checkpoint_guess_config.call_args.args == (
|
|
"/models/checkpoints/X/b_v2.safetensors",
|
|
)
|
|
fp.get_full_path_or_raise.assert_not_called()
|
|
assert out == ("MODEL", "CLIP", "VAE", "X/b_v2.safetensors", "b_v2", 1, 2)
|
|
|
|
def test_absolute_folder_outside_models_dir_loads_by_real_path(self):
|
|
ctx, _, comfy = _fake_comfy_modules()
|
|
external = [("ext_v1.safetensors", "/mnt/drive/ckpts/ext_v1.safetensors")]
|
|
with ctx, patch(f"{NODE}.list_checkpoint_entries", return_value=external):
|
|
out = CheckpointFolderIteratorNode().load("/mnt/drive/ckpts", 0, "", "")
|
|
assert comfy.sd.load_checkpoint_guess_config.call_args.args == (
|
|
"/mnt/drive/ckpts/ext_v1.safetensors",
|
|
)
|
|
assert out[3:] == ("ext_v1.safetensors", "ext_v1", 0, 1)
|
|
|
|
def test_index_past_the_end_raises_instead_of_restarting(self):
|
|
ctx, _, comfy = _fake_comfy_modules()
|
|
with ctx, patch(f"{NODE}.list_checkpoint_entries", return_value=ENTRIES):
|
|
with pytest.raises(ValueError, match="past the end"):
|
|
CheckpointFolderIteratorNode().load("X", 2, "", "")
|
|
comfy.sd.load_checkpoint_guess_config.assert_not_called()
|
|
|
|
def test_empty_folder_raises_clear_error(self):
|
|
ctx, _, _ = _fake_comfy_modules()
|
|
with ctx, patch(f"{NODE}.list_checkpoint_entries", return_value=[]):
|
|
with pytest.raises(ValueError, match="No checkpoints"):
|
|
CheckpointFolderIteratorNode().load("Empty", 0, "", "")
|