model first

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,
+Scans for ComfyUI PNGs, reads embedded prompt metadata, extracts the
-extracts short keyword descriptions, and interactively renames them to:
+checkpoint model and LoRA names, and moves files into:
-  schmeeve-AI-{keywords}.png
+  {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):
+def extract_metadata(filepath):
-    """Propose 'schmeeve-AI-{keywords}.png' or None.
+    """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
+    if "prompt" not in img.info:
-    keyword selection that favors rare (discriminating) words.
+        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
+    return model, loras, pos, neg
-    if not source:
-        return None
 
-    candidates = extract_all_candidates(source)
-    if not candidates:
-        return None
 
-    if freq_map and total_images > 1:
+def sanitize_dir_name(name):
-        keywords = select_best_keywords(candidates, freq_map, total_images)
+    """Sanitize a string for use as a directory name."""
-    else:
+    name = name.strip().replace(" ", "_")
-        keywords = [kw for kw, _ in candidates[:5]]
+    name = re.sub(r"[^a-zA-Z0-9_.-]", "_", name)
-
+    name = re.sub(r"_+", "_", name).strip("_")
-    return keywords_to_name(keywords)
+    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:
+        if not name:
-            raw_proposals[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
 
     # 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 = {}
+    used_paths = {}
-    for p, name in raw_proposals.items():
+    for p, rel_path in raw_proposals.items():
-        base = name
+        new_path = rel_path
-        if base in name_counts:
+        counter = 0
-            name_counts[base] += 1
+        key = str(new_path)
-        else:
+        while key in used_paths:
-            name_counts[base] = 1
+            counter += 1
-        proposals[p] = name
-
-    for p, name in list(proposals.items()):
-        if name_counts[name] > 1:
             ts = time.strftime("%m%d-%H%M%S", time.localtime(p.stat().st_mtime))
-            stem = name.rsplit(".", 1)[0]
+            stem = rel_path.stem
-            proposals[p] = f"{stem}-{ts}.png"
+            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):
+    print(f"\n  {'DRY RUN' if args.dry_run else 'PROPOSED'} — "
-        old = old_path.name
+          f"{len(proposals)} file(s) to organize into {output_root}/\n")
-        stem = old_path.stem
+
-        ext = old_path.suffix
+    for i, (old_path, rel_path) in enumerate(proposals.items(), 1):
-        # Truncate old name for display
+        try:
-        old_display = old if len(old) < 50 else old[:22] + "…" + old[-25:]
+            src_display = str(old_path.relative_to(scan_dir))
-        print(f"  {i:>3}. {old_display}")
+        except ValueError:
-        print(f"       → {new_name}")
+            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
+        moved = 0
-        for old_path, new_name in proposals.items():
+        for old_path, rel_path in proposals.items():
-            new_path = old_path.with_name(new_name)
+            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_path, rel_path = items[i]
-            old = old_path.name
+            new_path = output_root / rel_path
-            new = new_name
+            new_path_display = str(rel_path)
 
             print(f"\n  [{i+1}/{len(items)}]")
-            print(f"  Current: {old}")
+            print(f"  From: {old_path.name}")
-            print(f"  New:     {new}")
+            print(f"  To:   {output_root / rel_path}")
             remaining = len(items) - i - 1
-            rlabel = f"rename all {remaining}" if remaining else ""
+            rlabel = f"move all {remaining}" if remaining else ""
-            sys.stdout.write(f"  [Enter]=rename  [e]=edit  [s]=skip{'  [a]=' + rlabel if rlabel else ''}  [q]=quit: ")
+            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]
+                    p, rp = items[j]
-                    np = p.with_name(n)
+                    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
+            # Perform move
-            new_path = old_path.with_name(new)
+            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: