""" WD14 tagger: imagenes -> prompts Danbooru, en lote y por consola. Descarga el modelo de HuggingFace la primera vez (unos 250 MB, cacheado despues). python Engine/wd14_tagger.py python Engine/wd14_tagger.py --threshold 0.45 --limit 20 python Engine/wd14_tagger.py --model eva02 --batch 16 python Engine/wd14_tagger.py --rethreshold 0.3 --from-json Escribe un .txt por imagen en Archivos/wd14/tags/ y, con --json, un .json con la confianza de cada tag para poder reprocesar sin volver a correr la inferencia. Los defaults usan SmilingWolf/wd-swinv2-tagger-v3, el unico de la gama que identifica personajes (67% de acierto medido). Usa --model vittagger para reproducir el Space deepghs/wd14_tagging_online, que va mas enfocado pero no sabe de personajes. --model eva02 da mas tags (43/img) a cambio del doble de tiempo. """ import argparse import csv import json import os import sys import time from pathlib import Path BASE_DIR = Path(__file__).resolve().parent.parent DEFAULT_INBOX = BASE_DIR / "Archivos" / "wd14" / "inbox" DEFAULT_TAGS = BASE_DIR / "Archivos" / "wd14" / "tags" MODELS = { "swinv2v3": "SmilingWolf/wd-swinv2-tagger-v3", "eva02": "SmilingWolf/wd-eva02-large-tagger-v3", "convnextv3": "SmilingWolf/wd-convnext-tagger-v3", "convnextv2": "SmilingWolf/wd-v1-4-convnext-tagger-v2", "swinv2v2": "SmilingWolf/wd-v1-4-swinv2-tagger-v2", "vittagger": "SmilingWolf/wd-v1-4-vit-tagger", "convnextv14": "SmilingWolf/wd-v1-4-convnext-tagger", } IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".bmp") SIZE = 448 RATING_CATEGORY = "9" def load_model(repo, threads=None, device="auto"): import onnxruntime as ort from huggingface_hub import hf_hub_download print(f"[-] Modelo: {repo}") model_path = hf_hub_download(repo_id=repo, filename="model.onnx") csv_path = hf_hub_download(repo_id=repo, filename="selected_tags.csv") names, ratings = [], [] with open(csv_path, newline="", encoding="utf-8") as f: for row in csv.DictReader(f): names.append(row["name"]) ratings.append(row["category"] == RATING_CATEGORY) print(f"[-] {len(names)} tags en el modelo") providers = ort.get_available_providers() if device == "cpu": use = ["CPUExecutionProvider"] elif device == "gpu": if "CUDAExecutionProvider" not in providers: raise SystemExit("[!] --device gpu pero CUDA no disponible " f"({providers}). Requiere onnxruntime-gpu + CUDA.") use = ["CUDAExecutionProvider", "CPUExecutionProvider"] else: use = ["CUDAExecutionProvider", "CPUExecutionProvider"] \ if "CUDAExecutionProvider" in providers else ["CPUExecutionProvider"] print(f"[-] Ejecucion: {use[0]}") opts = None if threads: opts = ort.SessionOptions() opts.intra_op_num_threads = threads print(f"[-] Hilos ORT: intra_op={threads}") session = ort.InferenceSession(model_path, sess_options=opts, providers=use) shape = session.get_inputs()[0].shape fixed_batch = isinstance(shape[0], int) if fixed_batch: print("[-] El modelo fija el batch a 1: se procesa imagen a imagen") return session, names, ratings, fixed_batch def preprocess(path, pad=True): """Cuadra con fondo blanco y luego escala. Mantiene la proporcion, cosa que el resize directo NO hace: aplasta una vertical 832x1216 a 448x448 y el modelo deja de entender el cuerpo (inventa 'pauldrons', 'headphones'...). Con pad=False se comporta como el preprocesado de referencia.""" from PIL import Image import numpy as np img = Image.open(path).convert("RGB") if pad: side = max(max(img.size), SIZE) canvas = Image.new("RGB", (side, side), (255, 255, 255)) canvas.paste(img, ((side - img.width) // 2, (side - img.height) // 2)) img = canvas img = img.resize((SIZE, SIZE), Image.Resampling.LANCZOS) arr = np.asarray(img)[:, :, ::-1].astype(np.float32) return arr[np.newaxis, ...] def tag_batch(session, paths, pad=True): import numpy as np batch = np.concatenate([preprocess(p, pad) for p in paths], axis=0) name = session.get_inputs()[0].name return session.run(None, {name: batch})[0] def rank(probs, names, ratings): pairs = ((n, s) for n, s, r in zip(names, probs, ratings) if not r) return sorted(pairs, key=lambda t: -float(t[1])) def render(ranked, threshold, spaces=False): out = [] for tag, score in ranked: if score < threshold: continue out.append(tag.replace("_", " ") if spaces else tag) return ", ".join(out) def collect(folder, limit=None, recursive=True): folder = Path(folder) if not folder.exists(): sys.exit(f"[!] No existe la carpeta: {folder}") it = folder.rglob("*") if recursive else folder.iterdir() files = sorted(p for p in it if p.is_file() and p.suffix.lower() in IMAGE_EXTS) if limit: files = files[:limit] if not files: sys.exit(f"[!] No hay imagenes en {folder} ({', '.join(IMAGE_EXTS)})") return files def main(): p = argparse.ArgumentParser(description="WD14: imagenes -> prompts Danbooru") p.add_argument("--folder", default=str(DEFAULT_INBOX), help="carpeta de imagenes") p.add_argument("--out", default=str(DEFAULT_TAGS), help="carpeta de salida .txt") p.add_argument("--model", default="swinv2v3", choices=sorted(MODELS), help="modelo WD14 (default swinv2v3; el unico que identifica personajes)") p.add_argument("--threshold", type=float, default=0.50, help="confianza minima (default 0.50; 0.35 recupera mas tags)") p.add_argument("--batch", type=int, default=8, help="imagenes por inferencia") p.add_argument("--limit", type=int, help="procesa solo las N primeras") p.add_argument("--spaces", action="store_true", help="convierte guiones bajos en espacios (default: los conserva)") p.add_argument("--stretch", action="store_true", help="resize directo sin cuadrar (degrada las verticales)") p.add_argument("--json", action="store_true", help="guarda tambien las confianzas en un .json") p.add_argument("--force", action="store_true", help="reescribe .txt existentes") p.add_argument("--no-recursive", action="store_true", help="no baja a subcarpetas (SFW/, NSFW/...)") p.add_argument("--threads", type=int, help="hilos ORT intra_op (default: todos; p.ej. 4 para no doggear la CPU)") p.add_argument("--device", choices=("auto", "cpu", "gpu"), default="auto", help="auto (CUDA si hay) / cpu / gpu (onnxruntime-gpu)") args = p.parse_args() files = collect(args.folder, args.limit, not args.no_recursive) src_root = Path(args.folder) def rel(p): return p.relative_to(src_root).with_suffix("") todo = [f for f in files if args.force or not (Path(args.out) / rel(f).with_suffix(".txt")).exists()] skipped = len(files) - len(todo) if skipped: print(f"[-] {skipped} ya etiquetadas (usa --force para rehacerlas)") out_dir = Path(args.out) out_dir.mkdir(parents=True, exist_ok=True) session, names, ratings, fixed_batch = load_model( MODELS[args.model], threads=args.threads, device=args.device) batch_size = 1 if fixed_batch else max(1, args.batch) scores_dump = {} t0 = time.time() n_ok = 0 for i in range(0, len(todo), batch_size): chunk = todo[i:i + batch_size] try: probs = tag_batch(session, chunk, not args.stretch) except Exception as e: if len(chunk) == 1: print(f"[!] {chunk[0].name} fallo: {e}") continue batch_size = 1 print(f"[-] Batch no soportado ({e}); bajo a imagen a imagen") break for path, row in zip(chunk, probs): ranked = rank(row, names, ratings) scores_dump[str(rel(path)).replace("\\", "/")] = \ [[t, round(float(s), 4)] for t, s in ranked] text = render(ranked, args.threshold, args.spaces) dest = out_dir / rel(path).with_suffix(".txt") dest.parent.mkdir(parents=True, exist_ok=True) dest.write_text(text, encoding="utf-8") n_ok += 1 done = min(i + len(chunk), len(todo)) rate = done / max(time.time() - t0, 1e-6) print(f"[-] {done}/{len(todo)} ({rate:.1f} img/s)") if args.json and scores_dump: (out_dir / "scores.json").write_text( json.dumps(scores_dump, ensure_ascii=False, indent=1), encoding="utf-8") print(f"[-] confianzas -> {out_dir / 'scores.json'}") dt = time.time() - t0 print(f"[OK] {n_ok} etiquetadas en {dt:.1f}s -> {out_dir}") if __name__ == "__main__": main()