84 lines
2.0 KiB
Python
Executable File
84 lines
2.0 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Extract model and LoRA metadata from ComfyUI-generated images.
|
|
|
|
Outputs a JSON mapping of filename → {"model": "...", "loras": ["...", ...]}.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
try:
|
|
from PIL import Image
|
|
except ImportError:
|
|
print("Error: Pillow (PIL) required. Install with: pip install Pillow", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp"}
|
|
|
|
|
|
def extract_ai_meta(filepath):
|
|
"""Return (model, loras) or (None, []) if no ComfyUI metadata found."""
|
|
try:
|
|
img = Image.open(filepath)
|
|
except Exception:
|
|
return None, []
|
|
|
|
if "prompt" not in img.info:
|
|
return None, []
|
|
|
|
try:
|
|
data = json.loads(img.info["prompt"])
|
|
except (json.JSONDecodeError, TypeError):
|
|
return None, []
|
|
|
|
model = None
|
|
loras = []
|
|
|
|
for node in data.values():
|
|
cls = node.get("class_type", "")
|
|
inputs = node.get("inputs", {})
|
|
|
|
if cls == "CheckpointLoaderSimple":
|
|
ckpt = inputs.get("ckpt_name", "")
|
|
if ckpt:
|
|
model = ckpt.rsplit(".", 1)[0]
|
|
|
|
if "lora" in cls.lower():
|
|
lora = inputs.get("lora_name", "")
|
|
if lora:
|
|
loras.append(lora.rsplit(".", 1)[0])
|
|
|
|
loras.sort()
|
|
return model, loras
|
|
|
|
|
|
def main():
|
|
target = sys.argv[1] if len(sys.argv) > 1 else os.path.expanduser("~/Pictures")
|
|
target = os.path.abspath(target)
|
|
|
|
if not os.path.isdir(target):
|
|
print(f"Error: {target} is not a directory", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
meta = {}
|
|
|
|
for entry in sorted(os.listdir(target)):
|
|
fpath = os.path.join(target, entry)
|
|
if not os.path.isfile(fpath):
|
|
continue
|
|
ext = os.path.splitext(entry)[1].lower()
|
|
if ext not in IMAGE_EXTS:
|
|
continue
|
|
|
|
model, loras = extract_ai_meta(fpath)
|
|
if model or loras:
|
|
meta[entry] = {"model": model, "loras": loras}
|
|
|
|
print(json.dumps(meta, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|