- 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
53 lines
2.0 KiB
Python
53 lines
2.0 KiB
Python
"""Logic for the Checkpoint Folder Iterator node."""
|
|
|
|
import os
|
|
from typing import List, Sequence, Tuple, TypeVar
|
|
|
|
from ...base.model_folders import filter_by_pattern, scan_model_entries
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
def list_checkpoint_entries(
|
|
folder_path: str, include_pattern: str = "", exclude_pattern: str = ""
|
|
) -> List[Tuple[str, str]]:
|
|
"""(name, absolute_path) pairs for every checkpoint, filtered, in natural order.
|
|
|
|
The absolute path is what gets loaded, so folders outside the configured
|
|
checkpoints roots work too; the name is relative to the root (or to the
|
|
folder itself when it lives outside every root) and is used for labels.
|
|
"""
|
|
entries = scan_model_entries(folder_path, "checkpoints")
|
|
names = [name for name, _ in entries]
|
|
kept = set(filter_by_pattern(names, include_pattern, exclude_pattern))
|
|
return [(name, path) for name, path in entries if name in kept]
|
|
|
|
|
|
def list_checkpoints(
|
|
folder_path: str, include_pattern: str = "", exclude_pattern: str = ""
|
|
) -> List[str]:
|
|
"""Checkpoint names (relative to the checkpoints roots) in natural order."""
|
|
entries = list_checkpoint_entries(folder_path, include_pattern, exclude_pattern)
|
|
return [name for name, _ in entries]
|
|
|
|
|
|
def pick_checkpoint(items: Sequence[T], index: int) -> T:
|
|
"""Return items[index], or raise when the sweep has run past the end.
|
|
|
|
Wrapping would silently restart the sweep (and, downstream, archive the
|
|
grid that just finished), so an over-queued run fails loudly instead.
|
|
"""
|
|
if not items:
|
|
raise ValueError("No checkpoints found")
|
|
if not 0 <= index < len(items):
|
|
raise ValueError(
|
|
f"Index {index} is past the end: only {len(items)} checkpoints matched. "
|
|
"The sweep is complete; reset index to 0 to start a new one."
|
|
)
|
|
return items[index]
|
|
|
|
|
|
def label_for(ckpt_name: str) -> str:
|
|
"""Label drawn under a tile for this checkpoint. Default: the file stem."""
|
|
return os.path.splitext(os.path.basename(ckpt_name))[0]
|