Compare commits
2 Commits
e85fce18c6
...
bf1361262a
| Author | SHA1 | Date | |
|---|---|---|---|
| bf1361262a | |||
| 1ee7d7210a |
92
comfyui-sync
92
comfyui-sync
@@ -1,92 +0,0 @@
|
|||||||
#!/bin/bash
|
|
||||||
set -euo pipefail
|
|
||||||
|
|
||||||
OUTPUT_SOURCE="$HOME/ComfyUI/output"
|
|
||||||
VIDEO_SOURCE="$HOME/ComfyUI/videos"
|
|
||||||
AUDIO_SOURCE="$HOME/ComfyUI/output/audio"
|
|
||||||
NAS_MOUNT="/mini.nas/miniShare1"
|
|
||||||
SHARE_DIR="$NAS_MOUNT/Pictures"
|
|
||||||
STATE_FILE="$HOME/.local/state/comfyui-sync.state"
|
|
||||||
NEATCLI="$(command -v /home/linuxbrew/.linuxbrew/bin/neatcli || true)"
|
|
||||||
RENAME_SCRIPT="$HOME/git/schmeeve-toolz/rename-ai-snaps"
|
|
||||||
ORGANIZE_SCRIPT="$HOME/git/schmeeve-toolz/organize-images"
|
|
||||||
|
|
||||||
mkdir -p "$(dirname "$STATE_FILE")"
|
|
||||||
|
|
||||||
if [ ! -d "$SHARE_DIR" ]; then
|
|
||||||
echo "[comfyui-sync] ERROR: Share dir $SHARE_DIR does not exist"
|
|
||||||
exit 1
|
|
||||||
fi
|
|
||||||
|
|
||||||
if ! mountpoint -q "$NAS_MOUNT"; then
|
|
||||||
echo "[comfyui-sync] $NAS_MOUNT is not mounted, attempting to mount..."
|
|
||||||
if ! sudo -n mount "$NAS_MOUNT" 2>&1; then
|
|
||||||
echo "[comfyui-sync] ERROR: mount attempt failed"
|
|
||||||
exit 1
|
|
||||||
fi
|
|
||||||
if ! mountpoint -q "$NAS_MOUNT"; then
|
|
||||||
echo "[comfyui-sync] ERROR: $NAS_MOUNT still not mounted after mount attempt, refusing to copy locally"
|
|
||||||
exit 1
|
|
||||||
fi
|
|
||||||
echo "[comfyui-sync] Mount succeeded."
|
|
||||||
fi
|
|
||||||
|
|
||||||
echo "[comfyui-sync] Starting at $(date)"
|
|
||||||
|
|
||||||
current_files=$( {
|
|
||||||
find "$OUTPUT_SOURCE" -maxdepth 1 -type f -printf '%f\n' 2>/dev/null
|
|
||||||
find "$VIDEO_SOURCE" -maxdepth 1 -type f -printf '%f\n' 2>/dev/null
|
|
||||||
find "$AUDIO_SOURCE" -maxdepth 1 -type f -printf '%f\n' 2>/dev/null
|
|
||||||
} | sort)
|
|
||||||
|
|
||||||
if [ -f "$STATE_FILE" ]; then
|
|
||||||
new_files=$(comm -13 <(sort "$STATE_FILE") <(echo "$current_files"))
|
|
||||||
else
|
|
||||||
echo "[comfyui-sync] First run — all files considered new."
|
|
||||||
new_files="$current_files"
|
|
||||||
fi
|
|
||||||
|
|
||||||
new_count=$(echo "$new_files" | grep -c '[^[:space:]]' || true)
|
|
||||||
|
|
||||||
if [ "$new_count" -gt 0 ]; then
|
|
||||||
echo "[comfyui-sync] Copying $new_count new file(s)..."
|
|
||||||
while IFS= read -r f; do
|
|
||||||
[ -z "$f" ] && continue
|
|
||||||
if [ -f "$OUTPUT_SOURCE/$f" ]; then
|
|
||||||
cp "$OUTPUT_SOURCE/$f" "$SHARE_DIR/"
|
|
||||||
elif [ -f "$VIDEO_SOURCE/$f" ]; then
|
|
||||||
cp "$VIDEO_SOURCE/$f" "$SHARE_DIR/"
|
|
||||||
elif [ -f "$AUDIO_SOURCE/$f" ]; then
|
|
||||||
cp "$AUDIO_SOURCE/$f" "$SHARE_DIR/"
|
|
||||||
fi
|
|
||||||
done <<<"$new_files"
|
|
||||||
|
|
||||||
echo "$current_files" >"$STATE_FILE"
|
|
||||||
echo "[comfyui-sync] State file updated."
|
|
||||||
else
|
|
||||||
echo "[comfyui-sync] No new files to copy."
|
|
||||||
fi
|
|
||||||
|
|
||||||
if [ -n "$NEATCLI" ]; then
|
|
||||||
echo "[comfyui-sync] Running neatcli organize --by-type -e $SHARE_DIR"
|
|
||||||
"$NEATCLI" organize --by-type -e "$SHARE_DIR"
|
|
||||||
else
|
|
||||||
echo "[comfyui-sync] WARNING: neatcli not found, skipping organize step"
|
|
||||||
fi
|
|
||||||
|
|
||||||
IMAGES_DIR="$SHARE_DIR/Images"
|
|
||||||
if [ -d "$IMAGES_DIR" ] && [ -x "$RENAME_SCRIPT" ]; then
|
|
||||||
echo "[comfyui-sync] Running rename-ai-snaps on $IMAGES_DIR"
|
|
||||||
cd "$IMAGES_DIR" && "$RENAME_SCRIPT" . --no-interactive
|
|
||||||
elif [ ! -x "$RENAME_SCRIPT" ]; then
|
|
||||||
echo "[comfyui-sync] WARNING: rename-ai-snaps not found, skipping rename step"
|
|
||||||
fi
|
|
||||||
|
|
||||||
if [ -d "$IMAGES_DIR" ] && [ -x "$ORGANIZE_SCRIPT" ]; then
|
|
||||||
echo "[comfyui-sync] Running organize-images on $IMAGES_DIR"
|
|
||||||
"$ORGANIZE_SCRIPT" "$IMAGES_DIR" --execute
|
|
||||||
elif [ ! -x "$ORGANIZE_SCRIPT" ]; then
|
|
||||||
echo "[comfyui-sync] WARNING: organize-images not found, skipping organize step"
|
|
||||||
fi
|
|
||||||
|
|
||||||
echo "[comfyui-sync] Done at $(date)"
|
|
||||||
632
rename-ai-snaps
632
rename-ai-snaps
@@ -1,632 +0,0 @@
|
|||||||
#!/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
Block a user