add flags
This commit is contained in:
@@ -18,6 +18,17 @@ except ImportError:
|
||||
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):
|
||||
"""Return (model, loras) or (None, []) if no ComfyUI metadata found."""
|
||||
try:
|
||||
@@ -36,7 +47,7 @@ def extract_ai_meta(filepath):
|
||||
model = None
|
||||
loras = []
|
||||
|
||||
for node in data.values():
|
||||
for nid, node in data.items():
|
||||
cls = node.get("class_type", "")
|
||||
inputs = node.get("inputs", {})
|
||||
|
||||
@@ -48,6 +59,8 @@ def extract_ai_meta(filepath):
|
||||
if "lora" in cls.lower():
|
||||
lora = inputs.get("lora_name", "")
|
||||
if lora:
|
||||
# Only include LoRAs whose output is actually consumed
|
||||
if is_lora_connected(nid, data):
|
||||
loras.append(lora.rsplit(".", 1)[0])
|
||||
|
||||
loras.sort()
|
||||
|
||||
@@ -8,6 +8,7 @@ DRY_RUN=true
|
||||
VERBOSE=false
|
||||
INSTALL=false
|
||||
FLAT=false
|
||||
LIMIT=0
|
||||
|
||||
usage() {
|
||||
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 " -i Install this script to ~/.local/bin/"
|
||||
echo " --flat Don't sub-categorize by model/LoRA (flat AI/ folder)"
|
||||
echo " --limit N Only process N files (for batching)"
|
||||
exit 0
|
||||
}
|
||||
|
||||
ARGS=()
|
||||
LIMIT_NEXT=false
|
||||
for arg in "$@"; do
|
||||
if [ "$LIMIT_NEXT" = true ]; then
|
||||
LIMIT="$arg"
|
||||
LIMIT_NEXT=false
|
||||
continue
|
||||
fi
|
||||
case "$arg" in
|
||||
-e|--execute) DRY_RUN=false ;;
|
||||
-v|--verbose) VERBOSE=true ;;
|
||||
-i|--install) INSTALL=true ;;
|
||||
--flat) FLAT=true ;;
|
||||
--limit) LIMIT_NEXT=true ;;
|
||||
--limit=*) LIMIT="${arg#*=}" ;;
|
||||
-h|--help) usage ;;
|
||||
*) ARGS+=("$arg") ;;
|
||||
esac
|
||||
@@ -55,6 +65,7 @@ fi
|
||||
|
||||
echo "organize-images — Classifying images in: $TARGET"
|
||||
[ "$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 ---
|
||||
AI_META=""
|
||||
@@ -138,6 +149,10 @@ for f in "$TARGET"/*; do
|
||||
esac
|
||||
done
|
||||
TOTAL=${#FILES[@]}
|
||||
if [ "$LIMIT" -gt 0 ] && [ "$LIMIT" -lt "$TOTAL" ]; then
|
||||
FILES=("${FILES[@]:0:$LIMIT}")
|
||||
TOTAL=$LIMIT
|
||||
fi
|
||||
if [ "$TOTAL" -eq 0 ]; then
|
||||
echo " No image files found."
|
||||
exit 0
|
||||
@@ -192,11 +207,8 @@ for f in "${FILES[@]}"; do
|
||||
fi
|
||||
|
||||
if [ "$VERBOSE" = true ]; then
|
||||
if [ "$cat" = "AI" ] && [ "$FLAT" = false ] && [ -n "$AI_META" ]; then
|
||||
printf " (%d/%d) [%s] %s\n" "$idx" "$TOTAL" "${sub#AI/}" "$(basename "$f")"
|
||||
else
|
||||
printf " (%d/%d) [%s] %s\n" "$idx" "$TOTAL" "$cat" "$(basename "$f")"
|
||||
fi
|
||||
short="${dest#$TARGET/}"
|
||||
printf " (%d/%d) %s → %s\n" "$idx" "$TOTAL" "$(basename "$f")" "$short"
|
||||
else
|
||||
progress_bar "$idx" "$TOTAL"
|
||||
fi
|
||||
|
||||
Reference in New Issue
Block a user