re-separate rename/organize, canonical --resort for AI folders
rename-ai-snaps is rename-only again (+--include-renamed, +--only-generic); extract-ai-meta is the shared recursive metadata engine emitting canonical model/lora dest paths; organize-images does all sorting, gains --resort for idempotent tree re-sorts (used to fix ~10K old-scheme folders on the NAS).
This commit is contained in:
309
rename-ai-snaps
309
rename-ai-snaps
@@ -1,19 +1,20 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
rename-ai-snaps — Scan PNGs for AI prompt metadata and reorg into model/lora folders.
|
||||
rename-ai-snaps — Scan PNGs for AI prompt metadata and rename descriptively.
|
||||
|
||||
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
|
||||
Scans for ComfyUI PNGs, reads embedded prompt metadata, extracts keywords,
|
||||
and renames files IN PLACE to:
|
||||
schmeeve-AI-{keywords}.png
|
||||
|
||||
Renaming only — use organize-images to sort files into model/lora folders.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
@@ -85,6 +86,9 @@ NEGATIVE_INDICATORS = {
|
||||
"body out of frame", "bad hands", "bad face", "blurry",
|
||||
}
|
||||
|
||||
# Generic ComfyUI output filenames (ComfyUI_00042_.png, ComfyUI_temp_….png)
|
||||
GENERIC_RE = re.compile(r"^ComfyUI", re.IGNORECASE)
|
||||
|
||||
# ── Helpers ────────────────────────────────────────────────────────────────
|
||||
|
||||
def spinner():
|
||||
@@ -115,82 +119,24 @@ def is_quality_only(text):
|
||||
|
||||
# ── 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."""
|
||||
def load_prompt_data(filepath):
|
||||
"""Return parsed ComfyUI prompt JSON from a PNG, or None."""
|
||||
try:
|
||||
img = Image.open(filepath)
|
||||
except Exception:
|
||||
return None, None
|
||||
return None
|
||||
|
||||
if "prompt" not in img.info:
|
||||
return None, None
|
||||
return None
|
||||
|
||||
try:
|
||||
data = json.loads(img.info["prompt"])
|
||||
return json.loads(img.info["prompt"])
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return None, None
|
||||
return None
|
||||
|
||||
|
||||
def extract_prompts(data):
|
||||
"""Return (positive_prompt, negative_prompt) from parsed prompt JSON."""
|
||||
# Gather all text fields from all nodes
|
||||
candidates = []
|
||||
for node in data.values():
|
||||
@@ -325,69 +271,22 @@ def keywords_to_name(keywords):
|
||||
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 nas_mount_base():
|
||||
"""Return the NAS mount base path for this machine."""
|
||||
candidates = [
|
||||
"/mini.nas/miniShare1",
|
||||
os.path.expanduser("~/mnt/mini.nas/miniShare1"),
|
||||
os.path.expanduser("~/mnt/miniShare1"),
|
||||
]
|
||||
for p in candidates:
|
||||
if os.path.isdir(p):
|
||||
return p
|
||||
return candidates[-1]
|
||||
|
||||
|
||||
def main():
|
||||
import argparse
|
||||
|
||||
DEFAULT_OUTPUT = os.path.join(nas_mount_base(), "Pictures", "Images", "AI")
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Scan PNGs for AI prompt metadata and reorg into model/lora folders.",
|
||||
description="Scan PNGs for AI prompt metadata and rename in place. "
|
||||
"Does not move files — use organize-images for sorting.",
|
||||
)
|
||||
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",
|
||||
help="Auto-rename without prompting",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-r", "--recursive", action="store_true",
|
||||
@@ -396,7 +295,15 @@ def main():
|
||||
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)",
|
||||
help="Show proposed changes without making any",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--include-renamed", action="store_true",
|
||||
help="Also process files already named schmeeve-AI-* (default: skip them)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--only-generic", action="store_true",
|
||||
help="Only rename generic ComfyUI_* filenames",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
@@ -405,8 +312,6 @@ def main():
|
||||
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))
|
||||
@@ -425,20 +330,20 @@ def main():
|
||||
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-"):
|
||||
# Skip already-renamed files (schmeeve-AI-*) unless requested
|
||||
if not args.include_renamed and p.name.startswith("schmeeve-AI-"):
|
||||
continue
|
||||
model, loras, pos, neg = extract_metadata(str(p))
|
||||
if args.only_generic and not GENERIC_RE.match(p.name):
|
||||
continue
|
||||
data = load_prompt_data(str(p))
|
||||
if not data:
|
||||
continue
|
||||
pos, neg = extract_prompts(data)
|
||||
source = pos or neg
|
||||
if source:
|
||||
candidates = extract_all_candidates(source)
|
||||
if candidates:
|
||||
image_data[p] = {
|
||||
"candidates": candidates,
|
||||
"model": model,
|
||||
"loras": loras,
|
||||
"source": source,
|
||||
}
|
||||
image_data[p] = candidates
|
||||
for kw, _ in candidates:
|
||||
if kw not in word_images:
|
||||
word_images[kw] = set()
|
||||
@@ -447,48 +352,35 @@ def main():
|
||||
freq_map = {kw: len(images) for kw, images in word_images.items()}
|
||||
total_images = len(image_data)
|
||||
|
||||
# ── Phase 1b: Propose file paths ──
|
||||
# ── Phase 1b: Propose new names (in place — same directory) ──
|
||||
raw_proposals = {}
|
||||
for p, info in image_data.items():
|
||||
candidates = info["candidates"]
|
||||
model = info["model"]
|
||||
loras = info["loras"]
|
||||
|
||||
for p, candidates in image_data.items():
|
||||
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:
|
||||
if not name or name == p.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
|
||||
raw_proposals[p] = name
|
||||
|
||||
# Clear the spinner line
|
||||
sys.stdout.write("\r" + " " * 60 + "\r")
|
||||
sys.stdout.flush()
|
||||
|
||||
# Deduplicate file names within their target directories
|
||||
# Deduplicate file names within their directories
|
||||
proposals = {}
|
||||
used_paths = {}
|
||||
for p, rel_path in raw_proposals.items():
|
||||
new_path = rel_path
|
||||
counter = 0
|
||||
for p, name in raw_proposals.items():
|
||||
new_path = p.with_name(name)
|
||||
key = str(new_path)
|
||||
counter = 0
|
||||
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}")
|
||||
stem = new_path.stem
|
||||
ext = new_path.suffix
|
||||
new_path = new_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
|
||||
@@ -499,9 +391,9 @@ def main():
|
||||
return
|
||||
|
||||
print(f"\n {'DRY RUN' if args.dry_run else 'PROPOSED'} — "
|
||||
f"{len(proposals)} file(s) to organize into {output_root}/\n")
|
||||
f"{len(proposals)} file(s) to rename in place\n")
|
||||
|
||||
for i, (old_path, rel_path) in enumerate(proposals.items(), 1):
|
||||
for i, (old_path, new_path) in enumerate(proposals.items(), 1):
|
||||
try:
|
||||
src_display = str(old_path.relative_to(scan_dir))
|
||||
except ValueError:
|
||||
@@ -509,7 +401,7 @@ def main():
|
||||
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(f" → {new_path.name}")
|
||||
|
||||
print()
|
||||
|
||||
@@ -517,40 +409,46 @@ def main():
|
||||
print(f" Dry-run complete. No files were changed.\n")
|
||||
return
|
||||
|
||||
# ── Phase 3: Rename/move (interactive or auto) ──
|
||||
def do_rename(old_path, new_path):
|
||||
"""Rename with collision suffix. Returns final path or None on failure."""
|
||||
if old_path == new_path:
|
||||
return new_path
|
||||
if new_path.exists():
|
||||
stem = new_path.stem
|
||||
ext = new_path.suffix
|
||||
counter = 1
|
||||
while new_path.exists():
|
||||
new_path = new_path.with_name(f"{stem}_{counter}{ext}")
|
||||
counter += 1
|
||||
print(f" (name existed, using {new_path.name})")
|
||||
try:
|
||||
old_path.rename(new_path)
|
||||
return new_path
|
||||
except (FileNotFoundError, OSError) as e:
|
||||
print(f" SKIPPED ({e}): {old_path.name}")
|
||||
return None
|
||||
|
||||
# ── Phase 3: Rename (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:
|
||||
shutil.move(old_path, new_path)
|
||||
moved += 1
|
||||
except (FileNotFoundError, OSError) as e:
|
||||
print(f" SKIPPED ({e}): {old_path.name}")
|
||||
print(f" Moved {moved} file(s)." if not args.dry_run else "")
|
||||
renamed = 0
|
||||
for old_path, new_path in proposals.items():
|
||||
if do_rename(old_path, new_path):
|
||||
renamed += 1
|
||||
print(f" Renamed {renamed} file(s).")
|
||||
else:
|
||||
moved = 0
|
||||
renamed = 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)
|
||||
old_path, new_path = items[i]
|
||||
|
||||
print(f"\n [{i+1}/{len(items)}]")
|
||||
print(f" From: {old_path.name}")
|
||||
print(f" To: {output_root / rel_path}")
|
||||
print(f" To: {new_path.name}")
|
||||
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: ")
|
||||
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: ")
|
||||
sys.stdout.flush()
|
||||
choice = sys.stdin.readline().strip().lower()
|
||||
|
||||
@@ -564,15 +462,14 @@ def main():
|
||||
i += 1
|
||||
continue
|
||||
elif choice == "e":
|
||||
sys.stdout.write(f" Edit name (will be 'schmeeve-AI-{rel_path.parent}/...'): ")
|
||||
sys.stdout.write(f" Edit keywords (name becomes 'schmeeve-AI-...'): ")
|
||||
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)
|
||||
new_path = old_path.with_name(f"schmeeve-AI-{custom_desc}.png")
|
||||
else:
|
||||
print(" Invalid name, skipping.")
|
||||
skipped += 1
|
||||
@@ -586,20 +483,9 @@ def main():
|
||||
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:
|
||||
shutil.move(p, np)
|
||||
except (FileNotFoundError, OSError) as e:
|
||||
print(f" SKIPPED ({e}): {p.name}")
|
||||
moved += len(items) - i
|
||||
p, np = items[j]
|
||||
if do_rename(p, np):
|
||||
renamed += 1
|
||||
break
|
||||
else:
|
||||
print(f" Unknown option '{choice}', skipping.")
|
||||
@@ -607,24 +493,11 @@ def main():
|
||||
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:
|
||||
shutil.move(old_path, new_path)
|
||||
moved += 1
|
||||
except (FileNotFoundError, OSError) as e:
|
||||
print(f" SKIPPED ({e}): {old_path.name}")
|
||||
if do_rename(old_path, new_path):
|
||||
renamed += 1
|
||||
i += 1
|
||||
|
||||
print(f"\n Moved: {moved} Skipped: {skipped}")
|
||||
print(f"\n Renamed: {renamed} Skipped: {skipped}")
|
||||
|
||||
# Check for JPEGs with AI metadata
|
||||
jpg_pattern = "**/*.jpg" if args.recursive else "*.jpg"
|
||||
|
||||
Reference in New Issue
Block a user