86 lines
2.8 KiB
Python
86 lines
2.8 KiB
Python
import os
|
||
import math
|
||
import yaml
|
||
from PIL import Image
|
||
import folder_paths
|
||
|
||
def collect_dirs(key_name: str):
|
||
"""
|
||
key_name: "checkpoints" or "loras"
|
||
Returns a list of absolute dirs:
|
||
1) ComfyUI/models/<key_name>
|
||
2) For each section in extra_model_paths.yaml:
|
||
section.base_path + '/' + section[key_name]
|
||
"""
|
||
ROOT = os.getcwd()
|
||
dirs = [os.path.join(ROOT, "models", key_name)]
|
||
yaml_file = os.path.join(ROOT, "extra_model_paths.yaml")
|
||
if os.path.exists(yaml_file):
|
||
with open(yaml_file, "r") as f:
|
||
cfg = yaml.safe_load(f) or {}
|
||
# Each section is a dict with base_path and relative subdirs
|
||
for section in cfg.values():
|
||
base = section.get("base_path", ROOT)
|
||
rel = section.get(key_name)
|
||
if rel:
|
||
dirs.append(os.path.abspath(os.path.join(base, rel)))
|
||
return dirs
|
||
|
||
class PreviewImageNode:
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"ckpt_name": ("STRING",),
|
||
"lora_stack": ("LORA_STACK",),
|
||
}
|
||
}
|
||
RETURN_TYPES = ("IMAGE",)
|
||
RETURN_NAMES = ("preview",)
|
||
FUNCTION = "preview_images"
|
||
CATEGORY = "Mango Node Pack/Loaders"
|
||
|
||
def preview_images(self, ckpt_name, lora_stack):
|
||
# strip off extensions so "foo.ckpt" → "foo"
|
||
base_name = os.path.splitext(ckpt_name)[0]
|
||
|
||
exts = [".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff"]
|
||
|
||
def find_image_in(dirs, name):
|
||
for d in dirs:
|
||
if not os.path.isdir(d):
|
||
continue
|
||
for ext in exts:
|
||
candidate = os.path.join(d, name + ext)
|
||
if os.path.exists(candidate):
|
||
print(f"[PreviewImage] ✔ Found {name+ext} in {d}")
|
||
return candidate
|
||
print(f"[PreviewImage] ✘ No preview for '{name}' in {dirs}")
|
||
return None
|
||
|
||
# build our search lists
|
||
model_dirs = collect_dirs("checkpoints")
|
||
lora_dirs = collect_dirs("loras")
|
||
|
||
imgs = []
|
||
# 1) model preview
|
||
mp = find_image_in(model_dirs, base_name)
|
||
imgs.append(Image.open(mp).convert("RGB") if mp else Image.new("RGB",(512,512),(40,40,40)))
|
||
|
||
# 2) up to 5 LoRAs
|
||
for name, *_ in lora_stack[:5]:
|
||
if not name:
|
||
continue
|
||
lp = find_image_in(lora_dirs, name)
|
||
imgs.append(Image.open(lp).convert("RGB") if lp else Image.new("RGB",(512,512),(80,80,80)))
|
||
|
||
# 3) composite 2×3 grid
|
||
cols = 3
|
||
rows = math.ceil(len(imgs)/cols)
|
||
w, h = imgs[0].size
|
||
grid = Image.new("RGB", (cols*w, rows*h), (20,20,20))
|
||
for i, im in enumerate(imgs):
|
||
grid.paste(im, ((i%cols)*w, (i//cols)*h))
|
||
|
||
return (grid,)
|