schmeeve/schmeeve-toolz
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
e85fce18c6a102e6c791fefd78dd6c35c3ccc2bf
schmeeve-toolz/rename-ai-snaps
schmeeve e85fce18c6 model first
2026-07-29 17:06:53 -07:00

633 lines
21 KiB
Python
Executable File
Raw Blame History

#!/usr/bin/env python3
"""
rename-ai-snaps — Scan PNGs for AI prompt metadata and reorg into model/lora folders.
Usage:
./rename-ai-snaps [path] [options]
Scans for ComfyUI PNGs, reads embedded prompt metadata, extracts the
checkpoint model and LoRA names, and moves files into:
{output}/{model}/{lora}/schmeeve-AI-{keywords}.png
"""
import json
import os
import re
import sys
import time
from pathlib import Path
try:
from PIL import Image
except ImportError:
print("Error: Pillow (PIL) is required. Install with: pip install Pillow")
sys.exit(1)
# ── stopwords and filter sets ──────────────────────────────────────────────
QUALITY_TAGS = {
"score_6_up", "score_7_up", "score_8_up", "score_9",
"score_6", "score_7", "score_8",
"masterpiece", "best quality", "good quality", "normal quality",
"high quality", "highly detailed", "very detailed", "extreme detail",
"very_aesthetic", "absurdres", "8k", "4k",
"photorealistic", "photograph",
"depth of field", "solo focus", "cinematic",
"newest", "amazing", "stunning",
}
QUALITY_WORDS = {
"best", "good", "high", "top", "ultra", "super",
"mega", "hyper", "extreme", "extra", "ultimate",
}
TECHNICAL_WORDS = {
"detailed", "focus", "quality", "aesthetic", "realistic",
"cinematic", "lighting", "rendering", "shading", "texture",
"newest", "absurdres",
}
STOP_WORDS = {
"the", "a", "an", "of", "in", "on", "at", "to", "for", "with",
"and", "or", "is", "are", "was", "were", "be", "been", "being",
"have", "has", "had", "do", "does", "did", "will", "would",
"could", "should", "may", "might", "can", "shall", "this",
"that", "these", "those", "it", "its", "by", "from", "as",
"into", "through", "during", "before", "after", "above", "below",
"between", "out", "off", "over", "under", "again", "further",
"then", "once", "here", "there", "when", "where", "why", "how",
"all", "each", "every", "both", "few", "more", "most", "other",
"some", "such", "no", "nor", "not", "only", "own", "same", "so",
"than", "too", "very", "just", "about", "up", "down",
"make", "get", "set", "put", "take", "give", "show", "use",
"like", "look", "see", "want", "need", "let", "close", "full",
"add", "new", "one", "two", "five",
"also", "well", "back", "still", "even", "much",
"you", "your", "my", "me", "we", "our", "they", "them", "their",
}
GENERIC_WORDS = {
"man", "men", "guy", "guys", "boy", "boys", "woman", "women",
"girl", "girls", "people", "person", "human", "figure",
"photo", "image", "picture", "shot", "view", "pose", "posing",
"face", "head", "body", "skin", "hair", "eyes", "hand", "hands",
"dark", "light", "bright", "color", "colour",
}
NEGATIVE_INDICATORS = {
"deformed", "distorted", "disfigured", "poorly drawn", "bad anatomy",
"extra digits", "missing digits", "extra limbs", "missing limbs",
"ugly", "tiling", "low quality", "worst quality", "normal quality",
"lowres", "monochrome", "grayscale", "text", "watermark",
"branding", "border", "cropped", "signature", "username",
"error", "mutation", "mutated", "out of frame", "duplicate", "cloned",
"body out of frame", "bad hands", "bad face", "blurry",
}
# ── Helpers ────────────────────────────────────────────────────────────────
def spinner():
chars = "⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏"
i = 0
while True:
yield chars[i % len(chars)]
i += 1
def is_negative_text(text):
lower = text.lower()
score = 0
for ind in NEGATIVE_INDICATORS:
if ind in lower:
score += 1
return score >= 2
def is_quality_only(text):
lower = text.lower()
words = re.findall(r"[a-z_]+", lower)
if not words:
return False
meaningful = sum(1 for w in words if w not in QUALITY_TAGS and len(w) > 2)
return meaningful == 0
# ── Metadata extraction from ComfyUI prompt JSON ───────────────────────────
CHECKPOINT_CLASS_TYPES = {
"CheckpointLoaderSimple",
"Checkpoint Loader with Name (Image Saver)",
"CheckpointLoader",
}
LORA_CLASS_TYPES = {
"LoraLoader",
"Power Lora Loader (rgthree)",
}
MODEL_CLASS_TYPES = {
"UNETLoader",
}
def _strip_model_name(raw):
"""Strip path prefix and .safetensors extension from a model/LoRA name."""
name = raw.strip()
# Handle backslash/forward-slash paths: "Pony\\model.safetensors" or "Pony/model.safetensors"
name = name.replace("\\", "/")
name = os.path.basename(name)
# Strip .safetensors extension
if name.endswith(".safetensors"):
name = name[:-len(".safetensors")]
return name
def extract_model_name(data):
"""Return the checkpoint/UNET model name from the prompt JSON, or None."""
model_name = None
for node in data.values():
inputs = node.get("inputs", {})
class_type = node.get("class_type", "")
if class_type in CHECKPOINT_CLASS_TYPES and "ckpt_name" in inputs:
model_name = _strip_model_name(inputs["ckpt_name"])
break
if class_type in MODEL_CLASS_TYPES and "unet_name" in inputs:
model_name = _strip_model_name(inputs["unet_name"])
break
return model_name if model_name else None
def extract_lora_names(data):
"""Return a sorted list of unique LoRA names from the prompt JSON."""
names = []
for node in data.values():
inputs = node.get("inputs", {})
class_type = node.get("class_type", "")
if class_type in LORA_CLASS_TYPES and "lora_name" in inputs:
names.append(_strip_model_name(inputs["lora_name"]))
# Some custom nodes use 'lora' or similar — look for any field ending
# in 'lora_name' or 'lora' that contains '.safetensors'
for field, val in inputs.items():
if isinstance(val, str) and val.endswith(".safetensors") and "lora" in field.lower():
name = _strip_model_name(val)
if name and name not in names:
names.append(name)
return sorted(set(names))
def extract_prompts(filepath):
"""Return (positive_prompt, negative_prompt) from a ComfyUI PNG."""
try:
img = Image.open(filepath)
except Exception:
return None, None
if "prompt" not in img.info:
return None, None
try:
data = json.loads(img.info["prompt"])
except (json.JSONDecodeError, TypeError):
return None, None
# Gather all text fields from all nodes
candidates = []
for node in data.values():
inputs = node.get("inputs", {})
for field in ("text", "prompt", "positive", "negative", "value", "string"):
val = inputs.get(field, "")
if isinstance(val, str) and len(val.strip()) > 3:
candidates.append(val.strip())
# Second pass: resolve node references
for node in data.values():
inputs = node.get("inputs", {})
for field in ("text", "prompt", "positive", "negative", "value", "string"):
val = inputs.get(field)
if isinstance(val, list) and len(val) == 2 and isinstance(val[0], str):
ref_node = data.get(val[0], {})
ref_inputs = ref_node.get("inputs", {})
for rf in ("value", "text", "string"):
rv = ref_inputs.get(rf, "")
if isinstance(rv, str) and len(rv.strip()) > 3:
candidates.append(rv.strip())
break
if not candidates:
return None, None
positives = [c for c in candidates if not is_negative_text(c)]
negatives = [c for c in candidates if is_negative_text(c)]
positives = [c for c in positives if not is_quality_only(c)]
pos = max(positives, key=len) if positives else None
neg = max(negatives, key=len) if negatives else None
return pos, neg
def extract_all_candidates(text):
"""Extract all candidate keywords with position index for frequency scoring."""
if not text:
return []
text_lower = text.lower()
cleaned = re.sub(r"[\[\(\{][^\]\)\}]*[\]\)\}]", "", text_lower)
segments = re.split(r"[,.;:!?]+", cleaned)
seen = set()
candidates = []
position = 0
for seg in segments:
seg = seg.strip()
if not seg or len(seg) < 4:
continue
words = re.findall(r"[a-zA-Z_]+", seg)
good = []
for w in words:
wl = w.lower().strip("_")
if len(wl) < 3:
continue
if wl in STOP_WORDS or wl in GENERIC_WORDS:
continue
if wl in QUALITY_TAGS or wl in QUALITY_WORDS:
continue
if wl in TECHNICAL_WORDS:
continue
if wl in NEGATIVE_INDICATORS:
continue
if wl.startswith("score") or wl.startswith("step"):
continue
if wl.isdigit():
continue
good.append(wl)
if good:
taken = 0
for g in good:
if g not in seen and taken < 2:
candidates.append((g, position))
seen.add(g)
taken += 1
position += 1
return candidates
def select_best_keywords(candidates, freq_map, total_images, max_keywords=5):
if not candidates:
return []
if total_images <= 1:
return [kw for kw, _ in candidates[:max_keywords]]
scored = []
for kw, pos in candidates:
freq = freq_map.get(kw, 1)
rarity = 1.0 - ((freq - 1) / max(total_images - 1, 1))
max_pos = max(len(candidates), 1)
pos_bonus = (pos / max_pos) * 0.3
score = rarity + pos_bonus
scored.append((score, kw, pos))
scored.sort(key=lambda x: (-x[0], x[2]))
selected = []
seen = set()
for _, kw, _ in scored:
if kw not in seen:
selected.append(kw)
seen.add(kw)
if len(selected) >= max_keywords:
break
return selected
def keywords_to_name(keywords):
if not keywords:
return None
desc = "-".join(keywords)
desc = re.sub(r"[^a-z0-9-]", "-", desc)
desc = re.sub(r"-+", "-", desc).strip("-")
if len(desc) > 40:
desc = desc[:40].rstrip("-")
if "-" in desc:
truncated = "-".join(desc.split("-")[:-1])
if truncated and len(truncated) > 10:
desc = truncated
return f"schmeeve-AI-{desc}.png" if desc and len(desc) > 3 else None
def extract_metadata(filepath):
"""Return (model_name, lora_names, positive_prompt, negative_prompt)."""
try:
img = Image.open(filepath)
except Exception:
return None, None, None, None
if "prompt" not in img.info:
return None, None, None, None
try:
data = json.loads(img.info["prompt"])
except (json.JSONDecodeError, TypeError):
return None, None, None, None
model = extract_model_name(data)
loras = extract_lora_names(data)
pos, neg = extract_prompts(filepath)
return model, loras, pos, neg
def sanitize_dir_name(name):
"""Sanitize a string for use as a directory name."""
name = name.strip().replace(" ", "_")
name = re.sub(r"[^a-zA-Z0-9_.-]", "_", name)
name = re.sub(r"_+", "_", name).strip("_")
return name
# ── main ───────────────────────────────────────────────────────────────────
def main():
import argparse
DEFAULT_OUTPUT = os.path.expanduser(
"~/mnt/mini.nas/miniShare1/Pictures/Images/AI"
)
parser = argparse.ArgumentParser(
description="Scan PNGs for AI prompt metadata and reorg into model/lora folders.",
)
parser.add_argument(
"path", nargs="?", default=os.path.expanduser("~/Pictures"),
help="Directory to scan for PNG files (default: ~/Pictures)",
)
parser.add_argument(
"-o", "--output", default=DEFAULT_OUTPUT,
help=f"Output root directory (default: {DEFAULT_OUTPUT})",
)
parser.add_argument(
"-n", "--no-interactive", action="store_true",
help="Auto-reorg without prompting",
)
parser.add_argument(
"-r", "--recursive", action="store_true",
help="Search directories recursively (default: off)",
)
parser.add_argument(
"-p", "--preview", "--dry-run", action="store_true",
dest="dry_run",
help="Show proposed changes without making any (combine with -n for full listing)",
)
args = parser.parse_args()
scan_dir = Path(args.path).expanduser().resolve()
if not scan_dir.is_dir():
print(f"Error: {scan_dir} is not a directory")
sys.exit(1)
output_root = Path(args.output).expanduser().resolve()
# Gather PNGs
glob_pattern = "**/*.png" if args.recursive else "*.png"
pngs = sorted(scan_dir.glob(glob_pattern))
if not pngs:
print(f" No PNG files found in {scan_dir}" +
(" (recursive search)" if args.recursive else ""))
return
# ── Phase 1: Collect candidates and build frequency map ──
sys.stdout.write(" Analyzing PNGs")
sys.stdout.flush()
spin = spinner()
image_data = {}
word_images = {}
for p in pngs:
sys.stdout.write(f"\r {next(spin)} Analyzing PNGs")
sys.stdout.flush()
# Skip already-renamed files (schmeeve-AI-*)
if p.name.startswith("schmeeve-AI-"):
continue
model, loras, pos, neg = extract_metadata(str(p))
source = pos or neg
if source:
candidates = extract_all_candidates(source)
if candidates:
image_data[p] = {
"candidates": candidates,
"model": model,
"loras": loras,
"source": source,
}
for kw, _ in candidates:
if kw not in word_images:
word_images[kw] = set()
word_images[kw].add(p)
freq_map = {kw: len(images) for kw, images in word_images.items()}
total_images = len(image_data)
# ── Phase 1b: Propose file paths ──
raw_proposals = {}
for p, info in image_data.items():
candidates = info["candidates"]
model = info["model"]
loras = info["loras"]
if total_images > 1:
keywords = select_best_keywords(candidates, freq_map, total_images)
else:
keywords = [kw for kw, _ in candidates[:5]]
name = keywords_to_name(keywords)
if not name:
continue
# Determine folder structure
model_dir = sanitize_dir_name(model) if model else "unknown"
if loras:
lora_dir = sanitize_dir_name("+".join(loras))
else:
lora_dir = "no-lora"
rel_path = Path(model_dir) / lora_dir / name
raw_proposals[p] = rel_path
# Clear the spinner line
sys.stdout.write("\r" + " " * 60 + "\r")
sys.stdout.flush()
# Deduplicate file names within their target directories
proposals = {}
used_paths = {}
for p, rel_path in raw_proposals.items():
new_path = rel_path
counter = 0
key = str(new_path)
while key in used_paths:
counter += 1
ts = time.strftime("%m%d-%H%M%S", time.localtime(p.stat().st_mtime))
stem = rel_path.stem
ext = rel_path.suffix
new_path = rel_path.with_name(f"{stem}-{ts}{'_' + str(counter) if counter > 1 else ''}{ext}")
key = str(new_path)
used_paths[key] = True
proposals[p] = new_path
# ── Phase 2: Display proposed changes ──
if not proposals:
print(" No renamable PNGs found.")
return
print(f"\n {'DRY RUN' if args.dry_run else 'PROPOSED'} — "
f"{len(proposals)} file(s) to organize into {output_root}/\n")
for i, (old_path, rel_path) in enumerate(proposals.items(), 1):
try:
src_display = str(old_path.relative_to(scan_dir))
except ValueError:
src_display = str(old_path)
if len(src_display) > 55:
src_display = src_display[:25] + "…" + src_display[-27:]
print(f" {i:>3}. {src_display}")
print(f" → {output_root / rel_path}")
print()
if args.dry_run:
print(f" Dry-run complete. No files were changed.\n")
return
# ── Phase 3: Rename/move (interactive or auto) ──
if args.no_interactive:
moved = 0
for old_path, rel_path in proposals.items():
new_path = output_root / rel_path
new_path.parent.mkdir(parents=True, exist_ok=True)
if new_path.exists():
stem = new_path.stem
counter = 1
while new_path.exists():
new_path = new_path.with_name(f"{stem}_{counter}{new_path.suffix}")
counter += 1
try:
old_path.rename(new_path)
moved += 1
except FileNotFoundError:
print(f" SKIPPED (not found): {old_path.name}")
print(f" Moved {moved} file(s)." if not args.dry_run else "")
else:
moved = 0
skipped = 0
items = list(proposals.items())
i = 0
while i < len(items):
old_path, rel_path = items[i]
new_path = output_root / rel_path
new_path_display = str(rel_path)
print(f"\n [{i+1}/{len(items)}]")
print(f" From: {old_path.name}")
print(f" To: {output_root / rel_path}")
remaining = len(items) - i - 1
rlabel = f"move all {remaining}" if remaining else ""
sys.stdout.write(f" [Enter]=move [e]=edit [s]=skip{' [a]=' + rlabel if rlabel else ''} [q]=quit: ")
sys.stdout.flush()
choice = sys.stdin.readline().strip().lower()
if choice == "q":
remaining = len(items) - i - 1
if remaining:
print(f" Skipping remaining {remaining} file(s).")
break
elif choice == "s":
skipped += 1
i += 1
continue
elif choice == "e":
sys.stdout.write(f" Edit name (will be 'schmeeve-AI-{rel_path.parent}/...'): ")
sys.stdout.flush()
custom = sys.stdin.readline().strip()
if custom:
custom_desc = re.sub(r"[^a-z0-9-]", "-", custom.lower())
custom_desc = re.sub(r"-+", "-", custom_desc).strip("-")
if custom_desc:
new_name = f"schmeeve-AI-{custom_desc}.png"
rel_path = rel_path.with_name(new_name)
else:
print(" Invalid name, skipping.")
skipped += 1
i += 1
continue
else:
skipped += 1
i += 1
continue
elif choice == "":
pass
elif choice == "a":
for j in range(i, len(items)):
p, rp = items[j]
np = output_root / rp
np.parent.mkdir(parents=True, exist_ok=True)
if np.exists():
stem = np.stem
counter = 1
while np.exists():
np = np.with_name(f"{stem}_{counter}{np.suffix}")
counter += 1
try:
p.rename(np)
except FileNotFoundError:
print(f" SKIPPED (not found): {p.name}")
moved += len(items) - i
break
else:
print(f" Unknown option '{choice}', skipping.")
skipped += 1
i += 1
continue
# Perform move
new_path = output_root / rel_path
new_path.parent.mkdir(parents=True, exist_ok=True)
if new_path.exists():
stem = new_path.stem
counter = 1
while new_path.exists():
new_path = new_path.with_name(f"{stem}_{counter}{new_path.suffix}")
counter += 1
print(f" (file existed, saved as {new_path.name})")
try:
old_path.rename(new_path)
moved += 1
except FileNotFoundError:
print(f" SKIPPED (not found): {old_path.name}")
i += 1
print(f"\n Moved: {moved} Skipped: {skipped}")
# Check for JPEGs with AI metadata
jpg_pattern = "**/*.jpg" if args.recursive else "*.jpg"
remaining_ai = 0
for f in scan_dir.glob(jpg_pattern):
try:
img = Image.open(f)
if "prompt" in img.info:
remaining_ai += 1
except Exception:
pass
if remaining_ai:
print(f" Note: {remaining_ai} JPEG(s) with AI metadata found (not yet supported).")
if __name__ == "__main__":
main()
Reference in New Issue View Git Blame Copy Permalink