ethanfel-ComfyUI-Lora-Sweeper/lora_sweeper.py

import os
import re
import hashlib

import folder_paths
import comfy.utils
import comfy.sd

# Matches the trailing training-step token used by ai-toolkit / kohya style
# checkpoints, e.g. "my_lora_000002000.safetensors" -> 2000
STEP_RE = re.compile(r"_(\d{4,})\.safetensors$", re.IGNORECASE)
SAFE_EXT = ".safetensors"


def _resolve_folder(folder_path):
    """Resolve an absolute path, a ~ path, or a name relative to the loras dir."""
    folder_path = (folder_path or "").strip().strip('"').strip("'")
    if not folder_path:
        raise ValueError("LoraEpochSweeper: 'folder_path' is empty.")

    expanded = os.path.expanduser(folder_path)
    if os.path.isdir(expanded):
        return expanded

    # Allow typing just a sub-folder name that lives inside a loras root.
    try:
        roots = folder_paths.get_folder_paths("loras")
    except Exception:
        roots = []
    for root in roots:
        candidate = os.path.join(root, folder_path)
        if os.path.isdir(candidate):
            return candidate

    raise FileNotFoundError(f"LoraEpochSweeper: folder not found: {folder_path}")


def _basename(fn):
    """LoRA name without the step suffix/extension (the grouping key)."""
    m = STEP_RE.search(fn)
    if m:
        return fn[:m.start()]
    if fn.lower().endswith(SAFE_EXT):
        return fn[:-len(SAFE_EXT)]
    return fn


def _scan(folder):
    """Return [(step_or_None, filename)] for every .safetensors in the folder."""
    entries = []
    for fn in os.listdir(folder):
        if not fn.lower().endswith(SAFE_EXT):
            continue
        m = STEP_RE.search(fn)
        step = int(m.group(1)) if m else None
        entries.append((step, fn))
    return entries


def _group(entries):
    """Group entries by LoRA basename -> [(step_or_None, filename)]."""
    groups = {}
    for step, fn in entries:
        groups.setdefault(_basename(fn), []).append((step, fn))
    return groups


def _choose_group(groups, lora_basename):
    """Pick which LoRA in the folder to sweep."""
    names = sorted(groups.keys())
    lb = (lora_basename or "").strip()
    if lb:
        if lb in groups:
            return lb
        matches = [n for n in names if lb in n]
        if len(matches) == 1:
            return matches[0]
        if not matches:
            raise ValueError(
                f"LoraEpochSweeper: no LoRA name matching '{lb}'. Available: {names}")
        raise ValueError(
            f"LoraEpochSweeper: '{lb}' is ambiguous, matches {matches}. Be more specific.")

    # Auto: prefer the LoRA with the most numbered epochs.
    def score(n):
        ents = groups[n]
        return (sum(1 for s, _ in ents if s is not None), len(ents))

    return max(names, key=score)


def _by_interval(in_range, step_interval):
    """Keep ~one checkpoint per `step_interval` training steps.

    Targets the multiples of the interval inside the range and snaps each to the
    nearest available checkpoint, so clean schedules (e.g. saves every 250, want
    every 500) yield the round steps, while odd intervals still return something.
    """
    interval = int(step_interval or 0)
    if interval <= 0 or not in_range:
        return list(in_range)

    steps = [s for s, _ in in_range]
    lo, hi = steps[0], steps[-1]
    first_mult = ((lo + interval - 1) // interval) * interval
    targets = list(range(first_mult, hi + 1, interval))
    if not targets:
        return [in_range[0]]  # interval wider than the whole range

    chosen = set()
    for t in targets:
        best = min(range(len(in_range)), key=lambda i: (abs(steps[i] - t), steps[i]))
        chosen.add(best)
    return [in_range[i] for i in sorted(chosen)]


def _select(entries, start_step, end_step, step_interval, include_base, max_count):
    """Filter one LoRA's epochs by step range, then thin by training-step interval."""
    numbered = sorted((e for e in entries if e[0] is not None), key=lambda e: e[0])
    base = sorted((e for e in entries if e[0] is None), key=lambda e: e[1])

    hi = end_step if end_step and end_step > 0 else None
    in_range = [e for e in numbered if e[0] >= start_step and (hi is None or e[0] <= hi)]

    selected = _by_interval(in_range, step_interval)
    if include_base:
        selected = selected + base  # the suffix-less "final" file goes last

    if max_count and max_count > 0:
        selected = selected[:max_count]
    return selected


class LoraEpochSweeper:
    """
    Load a range of one LoRA's training epochs from a folder and emit them as
    lists, so ComfyUI samples one image per epoch from a single queue.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL", {"tooltip": "Base diffusion model to apply each epoch to."}),
                "folder_path": ("STRING", {
                    "default": "",
                    "tooltip": "Absolute path to a folder of LoRA epochs "
                               "(or a sub-folder name inside the loras dir).",
                }),
                "lora_basename": ("STRING", {
                    "default": "",
                    "tooltip": "Which LoRA in the folder to sweep (substring of its name). "
                               "Leave empty to auto-pick the LoRA with the most epochs.",
                }),
                "start_step": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "step": 1,
                                       "tooltip": "Lowest training step to include (inclusive)."}),
                "end_step": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "step": 1,
                                     "tooltip": "Highest training step to include. 0 = no upper limit."}),
                "step_interval": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "step": 1,
                                          "tooltip": "Thin the range to ~one checkpoint per this many training "
                                                     "steps (e.g. 500). 0 = keep every epoch."}),
                "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
                "strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
            },
            "optional": {
                "clip": ("CLIP", {"tooltip": "Optional CLIP/text-encoder to also patch."}),
                "include_base": ("BOOLEAN", {"default": True,
                                             "tooltip": "Also include the final, suffix-less .safetensors."}),
                "max_count": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1,
                                      "tooltip": "Cap the number of epochs swept. 0 = unlimited."}),
            },
        }

    RETURN_TYPES = ("MODEL", "CLIP", "STRING", "STRING")
    RETURN_NAMES = ("model", "clip", "step_label", "lora_name")
    OUTPUT_IS_LIST = (True, True, True, True)
    OUTPUT_TOOLTIPS = (
        "List of models, one per selected epoch.",
        "List of CLIPs, one per selected epoch (passthrough if no clip given).",
        "Training step of each epoch as a string ('final' for the suffix-less file).",
        "Filename (without extension) of each epoch.",
    )
    FUNCTION = "sweep"
    CATEGORY = "loaders/LoraSweeper"
    DESCRIPTION = ("Loads a range of one LoRA's epochs from a folder and outputs them as lists, "
                   "so a single queue samples one image per epoch.")

    def sweep(self, model, folder_path, lora_basename, start_step, end_step, step_interval,
              strength_model, strength_clip, clip=None, include_base=True, max_count=0):
        folder = _resolve_folder(folder_path)
        entries = _scan(folder)
        if not entries:
            raise ValueError(f"LoraEpochSweeper: no .safetensors files in {folder}")

        groups = _group(entries)
        chosen = _choose_group(groups, lora_basename)
        if len(groups) > 1:
            others = sorted(n for n in groups if n != chosen)
            print(f"[LoraEpochSweeper] {len(groups)} LoRAs in folder; sweeping '{chosen}'. "
                  f"Others available: {others}")

        selected = _select(groups[chosen], start_step, end_step, step_interval, include_base, max_count)
        if not selected:
            raise ValueError(
                f"LoraEpochSweeper: no epochs of '{chosen}' match start/end/step_interval.")

        models, clips, labels, names = [], [], [], []
        print(f"[LoraEpochSweeper] sweeping {len(selected)} epoch(s) of '{chosen}' from {folder}")
        for step, fn in selected:
            path = os.path.join(folder, fn)
            label = str(step) if step is not None else "final"
            print(f"  - {label}: {fn}")

            if strength_model == 0 and strength_clip == 0:
                m, c = model, clip
            else:
                lora, meta = comfy.utils.load_torch_file(path, safe_load=True, return_metadata=True)
                try:
                    m, c = comfy.sd.load_lora_for_models(
                        model, clip, lora, strength_model, strength_clip, lora_metadata=meta)
                except TypeError:
                    # Older ComfyUI without the lora_metadata kwarg.
                    m, c = comfy.sd.load_lora_for_models(
                        model, clip, lora, strength_model, strength_clip)
                del lora

            models.append(m)
            clips.append(c)
            labels.append(label)
            names.append(os.path.splitext(fn)[0])

        return (models, clips, labels, names)

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        folder_path = kwargs.get("folder_path", "")
        try:
            folder = _resolve_folder(folder_path)
            entries = sorted(_scan(folder), key=lambda e: e[1])
            sig = [(fn, os.path.getmtime(os.path.join(folder, fn))) for _, fn in entries]
        except Exception as e:
            sig = str(e)

        keys = ["folder_path", "lora_basename", "start_step", "end_step", "step_interval",
                "strength_model", "strength_clip", "include_base", "max_count"]
        payload = repr([sig] + [kwargs.get(k) for k in keys])
        return hashlib.sha256(payload.encode("utf-8")).hexdigest()


NODE_CLASS_MAPPINGS = {"LoraEpochSweeper": LoraEpochSweeper}
NODE_DISPLAY_NAME_MAPPINGS = {"LoraEpochSweeper": "Lora Epoch Sweeper"}