add flags

This commit is contained in:
2026-06-09 17:59:28 -07:00
parent a808cd2cc7
commit a8a5c448bf
2 changed files with 32 additions and 7 deletions

View File

@@ -18,6 +18,17 @@ except ImportError:
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp"} IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp"}
def is_lora_connected(node_id, data):
"""Check if a LoRA node's output is consumed by any downstream node."""
for nid, node in data.items():
if nid == node_id:
continue
for val in node.get("inputs", {}).values():
if isinstance(val, list) and len(val) == 2 and str(val[0]) == str(node_id):
return True
return False
def extract_ai_meta(filepath): def extract_ai_meta(filepath):
"""Return (model, loras) or (None, []) if no ComfyUI metadata found.""" """Return (model, loras) or (None, []) if no ComfyUI metadata found."""
try: try:
@@ -36,7 +47,7 @@ def extract_ai_meta(filepath):
model = None model = None
loras = [] loras = []
for node in data.values(): for nid, node in data.items():
cls = node.get("class_type", "") cls = node.get("class_type", "")
inputs = node.get("inputs", {}) inputs = node.get("inputs", {})
@@ -48,6 +59,8 @@ def extract_ai_meta(filepath):
if "lora" in cls.lower(): if "lora" in cls.lower():
lora = inputs.get("lora_name", "") lora = inputs.get("lora_name", "")
if lora: if lora:
# Only include LoRAs whose output is actually consumed
if is_lora_connected(nid, data):
loras.append(lora.rsplit(".", 1)[0]) loras.append(lora.rsplit(".", 1)[0])
loras.sort() loras.sort()

View File

@@ -8,6 +8,7 @@ DRY_RUN=true
VERBOSE=false VERBOSE=false
INSTALL=false INSTALL=false
FLAT=false FLAT=false
LIMIT=0
usage() { usage() {
echo "Usage: $(basename "$0") [path] [-e|--execute] [-v|--verbose] [-i|--install] [--flat]" echo "Usage: $(basename "$0") [path] [-e|--execute] [-v|--verbose] [-i|--install] [--flat]"
@@ -16,16 +17,25 @@ usage() {
echo " -v Verbose: show per-file classification" echo " -v Verbose: show per-file classification"
echo " -i Install this script to ~/.local/bin/" echo " -i Install this script to ~/.local/bin/"
echo " --flat Don't sub-categorize by model/LoRA (flat AI/ folder)" echo " --flat Don't sub-categorize by model/LoRA (flat AI/ folder)"
echo " --limit N Only process N files (for batching)"
exit 0 exit 0
} }
ARGS=() ARGS=()
LIMIT_NEXT=false
for arg in "$@"; do for arg in "$@"; do
if [ "$LIMIT_NEXT" = true ]; then
LIMIT="$arg"
LIMIT_NEXT=false
continue
fi
case "$arg" in case "$arg" in
-e|--execute) DRY_RUN=false ;; -e|--execute) DRY_RUN=false ;;
-v|--verbose) VERBOSE=true ;; -v|--verbose) VERBOSE=true ;;
-i|--install) INSTALL=true ;; -i|--install) INSTALL=true ;;
--flat) FLAT=true ;; --flat) FLAT=true ;;
--limit) LIMIT_NEXT=true ;;
--limit=*) LIMIT="${arg#*=}" ;;
-h|--help) usage ;; -h|--help) usage ;;
*) ARGS+=("$arg") ;; *) ARGS+=("$arg") ;;
esac esac
@@ -55,6 +65,7 @@ fi
echo "organize-images — Classifying images in: $TARGET" echo "organize-images — Classifying images in: $TARGET"
[ "$DRY_RUN" = true ] && echo " DRY RUN (use -e to execute)" || echo " EXECUTING" [ "$DRY_RUN" = true ] && echo " DRY RUN (use -e to execute)" || echo " EXECUTING"
[ "$LIMIT" -gt 0 ] && echo " Limit: first $LIMIT files"
# --- Pre-scan: extract AI metadata for model/LoRA routing --- # --- Pre-scan: extract AI metadata for model/LoRA routing ---
AI_META="" AI_META=""
@@ -138,6 +149,10 @@ for f in "$TARGET"/*; do
esac esac
done done
TOTAL=${#FILES[@]} TOTAL=${#FILES[@]}
if [ "$LIMIT" -gt 0 ] && [ "$LIMIT" -lt "$TOTAL" ]; then
FILES=("${FILES[@]:0:$LIMIT}")
TOTAL=$LIMIT
fi
if [ "$TOTAL" -eq 0 ]; then if [ "$TOTAL" -eq 0 ]; then
echo " No image files found." echo " No image files found."
exit 0 exit 0
@@ -192,11 +207,8 @@ for f in "${FILES[@]}"; do
fi fi
if [ "$VERBOSE" = true ]; then if [ "$VERBOSE" = true ]; then
if [ "$cat" = "AI" ] && [ "$FLAT" = false ] && [ -n "$AI_META" ]; then short="${dest#$TARGET/}"
printf " (%d/%d) [%s] %s\n" "$idx" "$TOTAL" "${sub#AI/}" "$(basename "$f")" printf " (%d/%d) %s %s\n" "$idx" "$TOTAL" "$(basename "$f")" "$short"
else
printf " (%d/%d) [%s] %s\n" "$idx" "$TOTAL" "$cat" "$(basename "$f")"
fi
else else
progress_bar "$idx" "$TOTAL" progress_bar "$idx" "$TOTAL"
fi fi