diff --git a/rename-ai-snaps b/rename-ai-snaps index cafff55..5786998 100755 --- a/rename-ai-snaps +++ b/rename-ai-snaps @@ -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: