model first

This commit is contained in:
2026-07-29 17:06:53 -07:00
parent 9752eef0b0
commit e85fce18c6

View File

@@ -1,13 +1,13 @@
#!/usr/bin/env python3
"""
rename-ai-snaps — Scan PNGs for AI prompt metadata and propose descriptive filenames.
rename-ai-snaps — Scan PNGs for AI prompt metadata and reorg into model/lora folders.
Usage:
./rename-ai-snaps [path] [--no-interactive]
./rename-ai-snaps [path] [options]
Scans directory (default ~/Pictures) for ComfyUI PNGs, reads their embedded prompt,
extracts short keyword descriptions, and interactively renames them to:
schmeeve-AI-{keywords}.png
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
@@ -84,7 +84,7 @@ NEGATIVE_INDICATORS = {
"body out of frame", "bad hands", "bad face", "blurry",
}
# ── Helpers ────────────────────────────────────────────────────────────────
def spinner():
chars = "⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏"
@@ -95,7 +95,6 @@ def spinner():
def is_negative_text(text):
"""Heuristic: is this text block a negative prompt?"""
lower = text.lower()
score = 0
for ind in NEGATIVE_INDICATORS:
@@ -105,9 +104,7 @@ def is_negative_text(text):
def is_quality_only(text):
"""Heuristic: does this text block contain only quality/technical tags?"""
lower = text.lower()
# Split into words, strip weighting syntax
words = re.findall(r"[a-z_]+", lower)
if not words:
return False
@@ -115,6 +112,69 @@ def is_quality_only(text):
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:
@@ -134,20 +194,18 @@ def extract_prompts(filepath):
candidates = []
for node in data.values():
inputs = node.get("inputs", {})
# Check all common text-carrying fields
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 (e.g. ["node_id", 0])
# 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], {})
# Check if the referenced node has a 'value' or 'text' field
ref_inputs = ref_node.get("inputs", {})
for rf in ("value", "text", "string"):
rv = ref_inputs.get(rf, "")
@@ -158,11 +216,9 @@ def extract_prompts(filepath):
if not candidates:
return None, None
# Split into positive vs negative
positives = [c for c in candidates if not is_negative_text(c)]
negatives = [c for c in candidates if is_negative_text(c)]
# Further filter: quality-only texts are not useful for naming
positives = [c for c in positives if not is_quality_only(c)]
pos = max(positives, key=len) if positives else None
@@ -176,14 +232,9 @@ def extract_all_candidates(text):
return []
text_lower = text.lower()
# Strip weighting/parenthetical syntax: (word:1.2), [word], {word}, (word)
cleaned = re.sub(r"[\[\(\{][^\]\)\}]*[\]\)\}]", "", text_lower)
# Split on commas, periods, semicolons, exclamation/question marks
segments = re.split(r"[,.;:!?]+", cleaned)
# Collect meaningful keywords, preserving position order
seen = set()
candidates = []
position = 0
@@ -193,12 +244,10 @@ def extract_all_candidates(text):
if not seg or len(seg) < 4:
continue
# Extract individual words from segment
words = re.findall(r"[a-zA-Z_]+", seg)
good = []
for w in words:
wl = w.lower().strip("_")
# Skip short words, stop words, quality tags, negative indicators
if len(wl) < 3:
continue
if wl in STOP_WORDS or wl in GENERIC_WORDS:
@@ -216,7 +265,6 @@ def extract_all_candidates(text):
good.append(wl)
if good:
# Take up to 2 unseen keywords from this segment
taken = 0
for g in good:
if g not in seen and taken < 2:
@@ -229,12 +277,6 @@ def extract_all_candidates(text):
def select_best_keywords(candidates, freq_map, total_images, max_keywords=5):
"""Select most discriminating keywords by rarity and position.
Keywords that appear in fewer images across the batch score higher.
Later-position keywords get a small bonus for tiebreaking.
When total_images <= 1, falls back to first-N selection.
"""
if not candidates:
return []
@@ -244,15 +286,12 @@ def select_best_keywords(candidates, freq_map, total_images, max_keywords=5):
scored = []
for kw, pos in candidates:
freq = freq_map.get(kw, 1)
# Rarity: 1.0 = appears in 1 image, 0.0 = appears in all images
rarity = 1.0 - ((freq - 1) / max(total_images - 1, 1))
# Position bonus: slight preference for later words (up to 0.3)
max_pos = max(len(candidates), 1)
pos_bonus = (pos / max_pos) * 0.3
score = rarity + pos_bonus
scored.append((score, kw, pos))
# Sort by score desc, then by original position asc for stability
scored.sort(key=lambda x: (-x[0], x[2]))
selected = []
@@ -268,7 +307,6 @@ def select_best_keywords(candidates, freq_map, total_images, max_keywords=5):
def keywords_to_name(keywords):
"""Convert keyword list to 'schmeeve-AI-{keywords}.png' name."""
if not keywords:
return None
@@ -276,7 +314,6 @@ def keywords_to_name(keywords):
desc = re.sub(r"[^a-z0-9-]", "-", desc)
desc = re.sub(r"-+", "-", desc).strip("-")
# Truncate at 40 chars, breaking at a word boundary
if len(desc) > 40:
desc = desc[:40].rstrip("-")
if "-" in desc:
@@ -287,27 +324,33 @@ def keywords_to_name(keywords):
return f"schmeeve-AI-{desc}.png" if desc and len(desc) > 3 else None
def propose_name(filepath, freq_map=None, total_images=1):
"""Propose 'schmeeve-AI-{keywords}.png' or 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
When freq_map and total_images > 1 are provided, uses frequency-aware
keyword selection that favors rare (discriminating) words.
"""
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)
source = pos or neg
if not source:
return None
return model, loras, pos, neg
candidates = extract_all_candidates(source)
if not candidates:
return None
if freq_map and total_images > 1:
keywords = select_best_keywords(candidates, freq_map, total_images)
else:
keywords = [kw for kw, _ in candidates[:5]]
return keywords_to_name(keywords)
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 ───────────────────────────────────────────────────────────────────
@@ -315,16 +358,33 @@ def propose_name(filepath, freq_map=None, total_images=1):
def main():
import argparse
DEFAULT_OUTPUT = os.path.expanduser(
"~/mnt/mini.nas/miniShare1/Pictures/Images/AI"
)
parser = argparse.ArgumentParser(
description="Rename AI-generated PNGs based on embedded prompt metadata.",
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-rename without prompting",
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()
@@ -333,7 +393,16 @@ def main():
print(f"Error: {scan_dir} is not a directory")
sys.exit(1)
pngs = sorted(scan_dir.glob("*.png"))
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")
@@ -344,106 +413,132 @@ def main():
for p in pngs:
sys.stdout.write(f"\r {next(spin)} Analyzing PNGs")
sys.stdout.flush()
# Skip already-renamed files
# Skip already-renamed files (schmeeve-AI-*)
if p.name.startswith("schmeeve-AI-"):
continue
pos, neg = extract_prompts(str(p))
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
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)
time.sleep(0.02)
freq_map = {kw: len(images) for kw, images in word_images.items()}
total_images = len(image_data)
# ── Phase 1b: Propose names with frequency-aware selection ──
# ── Phase 1b: Propose file paths ──
raw_proposals = {}
for p, candidates in image_data.items():
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 name:
raw_proposals[p] = name
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 proposed names; use short mtime tag for collisions
# Deduplicate file names within their target directories
proposals = {}
name_counts = {}
for p, name in raw_proposals.items():
base = name
if base in name_counts:
name_counts[base] += 1
else:
name_counts[base] = 1
proposals[p] = name
for p, name in list(proposals.items()):
if name_counts[name] > 1:
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 = name.rsplit(".", 1)[0]
proposals[p] = f"{stem}-{ts}.png"
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 names ──
# ── Phase 2: Display proposed changes ──
if not proposals:
print(" No renamable PNGs found.")
return
for i, (old_path, new_name) in enumerate(proposals.items(), 1):
old = old_path.name
stem = old_path.stem
ext = old_path.suffix
# Truncate old name for display
old_display = old if len(old) < 50 else old[:22] + "…" + old[-25:]
print(f" {i:>3}. {old_display}")
print(f" → {new_name}")
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()
print(f" {len(proposals)} file(s) to rename.\n")
# ── Phase 3: Rename (interactive or auto) ──
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:
renamed = 0
for old_path, new_name in proposals.items():
new_path = old_path.with_name(new_name)
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 = old_path.with_name(f"{stem}_{counter}{old_path.suffix}")
new_path = new_path.with_name(f"{stem}_{counter}{new_path.suffix}")
counter += 1
try:
old_path.rename(new_path)
renamed += 1
moved += 1
except FileNotFoundError:
print(f" SKIPPED (not found): {old_path.name}")
print(f" Renamed {renamed} file(s).")
print(f" Moved {moved} file(s)." if not args.dry_run else "")
else:
renamed = 0
moved = 0
skipped = 0
items = list(proposals.items())
i = 0
while i < len(items):
old_path, new_name = items[i]
old = old_path.name
new = new_name
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" Current: {old}")
print(f" New: {new}")
print(f" From: {old_path.name}")
print(f" To: {output_root / rel_path}")
remaining = len(items) - i - 1
rlabel = f"rename all {remaining}" if remaining else ""
sys.stdout.write(f" [Enter]=rename [e]=edit [s]=skip{' [a]=' + rlabel if rlabel else ''} [q]=quit: ")
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()
@@ -457,44 +552,42 @@ def main():
i += 1
continue
elif choice == "e":
sys.stdout.write(f" Edit name (will be prefixed 'schmeeve-AI-'): ")
sys.stdout.write(f" Edit name (will be 'schmeeve-AI-{rel_path.parent}/...'): ")
sys.stdout.flush()
custom = sys.stdin.readline().strip()
if custom:
# Sanitize
custom_desc = re.sub(r"[^a-z0-9-]", "-", custom.lower())
custom_desc = re.sub(r"-+", "-", custom_desc).strip("-")
if custom_desc:
new = f"schmeeve-AI-{custom_desc}.png"
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:
# empty = skip
skipped += 1
i += 1
continue
# Fall through to rename
elif choice == "":
pass # rename with proposed name
pass
elif choice == "a":
# Rename current and all remaining without further prompts
for j in range(i, len(items)):
p, n = items[j]
np = p.with_name(n)
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 = p.with_name(f"{stem}_{counter}{p.suffix}")
np = np.with_name(f"{stem}_{counter}{np.suffix}")
counter += 1
try:
p.rename(np)
except FileNotFoundError:
print(f" SKIPPED (not found): {p.name}")
renamed += len(items) - i
moved += len(items) - i
break
else:
print(f" Unknown option '{choice}', skipping.")
@@ -502,27 +595,29 @@ def main():
i += 1
continue
# Perform rename
new_path = old_path.with_name(new)
# 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 = old_path.with_name(f"{stem}_{counter}{old_path.suffix}")
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)
renamed += 1
moved += 1
except FileNotFoundError:
print(f" SKIPPED (not found): {old_path.name}")
i += 1
print(f"\n Renamed: {renamed} Skipped: {skipped}")
print(f"\n Moved: {moved} Skipped: {skipped}")
# One more pass: also look at .jpg? (maybe later)
# Check for JPEGs with AI metadata
jpg_pattern = "**/*.jpg" if args.recursive else "*.jpg"
remaining_ai = 0
for f in scan_dir.glob("*.jpg"):
for f in scan_dir.glob(jpg_pattern):
try:
img = Image.open(f)
if "prompt" in img.info: