| |
| """Booru Lighting Tagger: lighting tags for anime-style illustrations. |
| |
| Command line (run from anywhere; the model files are read from this script's folder): |
| |
| python inference.py image.png other.jpg # print the tags |
| python inference.py dataset/ --recursive --append-to-captions |
| python inference.py dataset/ --jsonl results.jsonl # every score, for your own thresholds |
| python inference.py dataset/ -r --append-to-captions --preset recall # more tags, a little less precise |
| |
| Python: |
| |
| from inference import LightingTagger |
| tagger = LightingTagger() # model files next to this script |
| for result in tagger.tag_files(["image.png"]): |
| print(result["tags"], result["scores"]["backlighting"]) |
| |
| The model is three fine-tuned copies ("seeds") of SmilingWolf/wd-eva02-large-tagger-v3. Each weight file |
| holds only the parameters that were trained (last 16 transformer blocks, final norm, new head); everything |
| else is loaded from the base model, pinned to the revision used for training. Scores are averaged over the |
| seeds. See README.md for the labels and their reliability. |
| |
| Threshold presets: "default" emits as many tags as the teacher did on the validation images; "recall" emits |
| about 1.5 times as many (see README.md, "Threshold presets"). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import contextlib |
| import copy |
| import json |
| import sys |
| from concurrent.futures import ThreadPoolExecutor |
| from pathlib import Path |
| from typing import Iterable, Iterator |
|
|
| import numpy as np |
| import torch |
| from PIL import Image |
|
|
| CONFIG_FILE = "tagger_config.json" |
| PRESETS = ("default", "recall") |
| IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif", ".tif", ".tiff"} |
| WHITE = (255, 255, 255) |
|
|
|
|
| |
|
|
| def to_rgb(image: Image.Image) -> Image.Image: |
| """RGB image; transparency is flattened onto white.""" |
| if image.mode not in ("RGB", "RGBA"): |
| image = image.convert("RGBA" if "transparency" in image.info else "RGB") |
| if image.mode == "RGBA": |
| background = Image.new("RGBA", image.size, (255, 255, 255, 255)) |
| image = Image.alpha_composite(background, image).convert("RGB") |
| return image |
|
|
|
|
| def pad_square(image: Image.Image) -> Image.Image: |
| """Pads to a square on white, centered, without cropping: the whole lighting stays visible.""" |
| side = max(image.size) |
| square = Image.new("RGB", (side, side), WHITE) |
| square.paste(image, ((side - image.width) // 2, (side - image.height) // 2)) |
| return square |
|
|
|
|
| def prepare(image: Image.Image, transform) -> torch.Tensor: |
| """WD preprocessing: white square padding, the base model's timm transform, RGB -> BGR.""" |
| return transform(pad_square(to_rgb(image)))[[2, 1, 0]] |
|
|
|
|
| |
|
|
| class _Student(torch.nn.Module): |
| """WD backbone with a new linear head (parameter names match the weight files).""" |
|
|
| def __init__(self, backbone, n_outputs: int): |
| super().__init__() |
| self.backbone = backbone |
| self.head = torch.nn.Linear(backbone.num_features, n_outputs) |
|
|
| def forward(self, x): |
| return self.head(self.backbone.forward_head(self.backbone.forward_features(x), pre_logits=True)) |
|
|
|
|
| def _load_base(repo: str, revision: str): |
| """Base WD tagger at the pinned revision, on the CPU, and its timm eval transform.""" |
| import timm |
|
|
| ref = f"{repo}@{revision}" |
| model = timm.create_model(f"hf-hub:{ref}").eval() |
| model.load_state_dict(timm.models.load_state_dict_from_hf(ref)) |
| transform = timm.data.create_transform(**timm.data.resolve_data_config(model.pretrained_cfg, model=model)) |
| return model, transform |
|
|
|
|
| def _trained_parameters(model: _Student, blocks: int) -> set[str]: |
| depth = len(model.backbone.blocks) |
| prefixes = tuple(f"backbone.blocks.{i}." for i in range(depth - blocks, depth)) + ("backbone.fc_norm.", "head.") |
| return {name for name, _ in model.named_parameters() if name.startswith(prefixes)} |
|
|
|
|
| def _load_seed(base, path: Path, n_outputs: int, blocks: int) -> _Student: |
| """A copy of the base model with the trained parameters of one seed loaded on top.""" |
| from safetensors.torch import load_file |
|
|
| model = _Student(copy.deepcopy(base), n_outputs) |
| state = load_file(str(path), device="cpu") |
| expected = _trained_parameters(model, blocks) |
| if set(state) != expected: |
| missing, extra = sorted(expected - set(state)), sorted(set(state) - expected) |
| raise ValueError(f"{path.name}: unexpected parameters (missing {missing[:5]}, extra {extra[:5]})") |
| model.load_state_dict(state, strict=False) |
| return model.eval() |
|
|
|
|
| def _decode(logits: np.ndarray, n_labels: int, steps: int) -> tuple[np.ndarray, np.ndarray]: |
| """Raw outputs -> (scores 0-10, tonal key probabilities). |
| |
| Ordinal target: for each label, `steps` outputs give P(score >= t) for t = 1..steps; the predicted score |
| is their sum. The last outputs are the tonal key logits (softmax).""" |
| width = n_labels * steps |
| sigmoid = 1 / (1 + np.exp(-logits[:, :width].astype(np.float64))) |
| scores = sigmoid.reshape(-1, n_labels, steps).sum(-1) |
| keys = logits[:, width:].astype(np.float64) |
| keys = np.exp(keys - keys.max(1, keepdims=True)) |
| return scores.astype(np.float32), (keys / keys.sum(1, keepdims=True)).astype(np.float32) |
|
|
|
|
| def _resolve_device(device: str) -> torch.device: |
| if device == "auto": |
| if torch.cuda.is_available(): |
| device = "cuda" |
| elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available(): |
| device = "mps" |
| else: |
| device = "cpu" |
| return torch.device(device) |
|
|
|
|
| def _autocast_dtype(device: torch.device, precision: str): |
| """Mixed precision on CUDA only (bf16, as in training; fp16 on GPUs without bf16).""" |
| if device.type != "cuda" or precision == "fp32": |
| return None |
| if precision == "bf16": |
| return torch.bfloat16 |
| if precision == "fp16": |
| return torch.float16 |
| return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 |
|
|
|
|
| class LightingTagger: |
| """Loads the seeds once; `tag_images` and `tag_files` return one result dict per image.""" |
|
|
| def __init__(self, model_dir: str | Path | None = None, device: str = "auto", seeds: int | None = None, |
| precision: str = "auto", preset: str = "default"): |
| self.model_dir = Path(model_dir) if model_dir else Path(__file__).resolve().parent |
| config_path = self.model_dir / CONFIG_FILE |
| if not config_path.exists(): |
| raise FileNotFoundError(f"{config_path} not found: pass the folder that holds the model files") |
| self.config = json.loads(config_path.read_text(encoding="utf-8")) |
| self.version = self.config["version"] |
| self.labels = [label["name"] for label in self.config["labels"]] |
| |
| self.presets = { |
| "default": np.array([label["threshold"] for label in self.config["labels"]], dtype=np.float32), |
| "recall": np.array([label.get("threshold_recall", label["threshold"]) for label in self.config["labels"]], |
| dtype=np.float32), |
| } |
| if preset not in self.presets: |
| raise ValueError(f"preset must be one of {PRESETS}") |
| self.preset = preset |
| self.thresholds = self.presets[preset] |
| self.status = {label["name"]: label["reliability"]["status"] for label in self.config["labels"]} |
| self.keys = [key["name"] for key in self.config["tonal_keys"]] |
| self.steps = int(self.config["ordinal_steps"]) |
| self.device = _resolve_device(device) |
| self.dtype = _autocast_dtype(self.device, precision) |
|
|
| all_seeds = self.config["seeds"] |
| if seeds is not None and not 1 <= seeds <= len(all_seeds): |
| raise ValueError(f"seeds must be between 1 and {len(all_seeds)}") |
| base_model = self.config["base_model"] |
| base, self.transform = _load_base(base_model["repo"], base_model["revision"]) |
| n_outputs = len(self.labels) * self.steps + len(self.keys) |
| self.models = [_load_seed(base, self.model_dir / seed["file"], n_outputs, self.config["trained_blocks"]) |
| .to(self.device) for seed in all_seeds[:seeds]] |
| del base |
|
|
| @classmethod |
| def from_pretrained(cls, repo_id: str, revision: str | None = None, **kwargs) -> LightingTagger: |
| """Downloads the model from the Hugging Face Hub (cached), then loads it.""" |
| from huggingface_hub import snapshot_download |
|
|
| return cls(snapshot_download(repo_id, revision=revision), **kwargs) |
|
|
| |
|
|
| @torch.inference_mode() |
| def predict(self, batch: torch.Tensor) -> tuple[np.ndarray, np.ndarray]: |
| """Preprocessed batch -> (scores: images x labels, tonal key probabilities: images x 4), seed average.""" |
| batch = batch.to(self.device) |
| scores, keys = [], [] |
| for model in self.models: |
| autocast = (torch.autocast(device_type="cuda", dtype=self.dtype) if self.dtype is not None |
| else contextlib.nullcontext()) |
| with autocast: |
| logits = model(batch).float().cpu().numpy() |
| s, k = _decode(logits, len(self.labels), self.steps) |
| scores.append(s) |
| keys.append(k) |
| return np.mean(scores, axis=0), np.mean(keys, axis=0) |
|
|
| def format(self, scores: np.ndarray, key_probs: np.ndarray, validated_only: bool = False, |
| tonal_key_tag: bool = True, preset: str | None = None) -> dict: |
| """One image: tags (by descending score, then the tonal key), all scores and the tonal key. |
| `preset` overrides the tagger's threshold preset for this call ("default" or "recall").""" |
| thresholds = self.thresholds if preset is None else self.presets[preset] |
| tagged = [k for k in np.flatnonzero(scores >= thresholds) |
| if not validated_only or self.status[self.labels[k]] == "validated"] |
| tags = [self.labels[k] for k in sorted(tagged, key=lambda k: -scores[k])] |
| key = self.keys[int(np.argmax(key_probs))] |
| if tonal_key_tag and key != "normal": |
| tags.append(key) |
| return { |
| "tags": tags, |
| "scores": {name: round(float(s), 3) for name, s in zip(self.labels, scores)}, |
| "tonal_key": key, |
| "tonal_key_probabilities": {name: round(float(p), 4) for name, p in zip(self.keys, key_probs)}, |
| } |
|
|
| def tag_images(self, images: list[Image.Image], **options) -> list[dict]: |
| """PIL images -> one result per image (see `format`).""" |
| if not images: |
| return [] |
| scores, keys = self.predict(torch.stack([prepare(image, self.transform) for image in images])) |
| return [self.format(s, k, **options) for s, k in zip(scores, keys)] |
|
|
| def tag_files(self, paths: Iterable[str | Path], batch_size: int = 8, workers: int = 4, |
| **options) -> Iterator[dict]: |
| """Image files -> one result per file, in order, as they are computed. An unreadable file gets an |
| `error` entry and does not stop the run.""" |
| def load(path: Path): |
| try: |
| with Image.open(path) as image: |
| return prepare(image, self.transform) |
| except Exception as exc: |
| return f"{type(exc).__name__}: {exc}" |
|
|
| paths = [Path(p) for p in paths] |
| with ThreadPoolExecutor(max_workers=max(1, workers)) as pool: |
| for start in range(0, len(paths), batch_size): |
| chunk = paths[start:start + batch_size] |
| loaded = list(pool.map(load, chunk)) |
| ok = [k for k, item in enumerate(loaded) if not isinstance(item, str)] |
| results = {} |
| if ok: |
| scores, keys = self.predict(torch.stack([loaded[k] for k in ok])) |
| results = {k: self.format(s, p, **options) for k, s, p in zip(ok, scores, keys)} |
| for k, path in enumerate(chunk): |
| if k in results: |
| yield {"file": str(path), **results[k], "error": None} |
| else: |
| yield {"file": str(path), "tags": [], "scores": None, "tonal_key": None, |
| "tonal_key_probabilities": None, "error": loaded[k]} |
|
|
|
|
| |
|
|
| def _normalize(tag: str) -> str: |
| return tag.strip().lower().replace("_", " ") |
|
|
|
|
| def caption_path(image: Path, extension: str) -> Path: |
| return image.parent / (image.stem + extension) |
|
|
|
|
| def append_to_caption(path: Path, tags: list[str]) -> int: |
| """Adds the tags missing from a comma-separated caption file (created if needed). Returns the count.""" |
| text = path.read_text(encoding="utf-8") if path.exists() else "" |
| present = {_normalize(tag) for tag in text.split(",") if tag.strip()} |
| new = [tag for tag in tags if _normalize(tag) not in present] |
| if new: |
| text = text.rstrip() |
| path.write_text((text + ", " if text else "") + ", ".join(new), encoding="utf-8") |
| return len(new) |
|
|
|
|
| |
|
|
| def find_images(inputs: list[str], recursive: bool) -> list[Path]: |
| paths = [] |
| for item in map(Path, inputs): |
| if item.is_dir(): |
| pattern = "**/*" if recursive else "*" |
| paths += sorted(p for p in item.glob(pattern) if p.is_file() and p.suffix.lower() in IMAGE_EXTENSIONS) |
| elif item.is_file(): |
| paths.append(item) |
| else: |
| print(f"not found: {item}", file=sys.stderr) |
| return paths |
|
|
|
|
| def main(argv: list[str] | None = None) -> int: |
| parser = argparse.ArgumentParser( |
| description="Booru Lighting Tagger: adds lighting tags (backlighting, rim_light, sunset...) to images.") |
| parser.add_argument("inputs", nargs="+", help="image files and/or folders") |
| parser.add_argument("-r", "--recursive", action="store_true", help="look into subfolders") |
| parser.add_argument("--jsonl", type=Path, help="write every result (all scores) to this JSON Lines file") |
| parser.add_argument("--append-to-captions", nargs="?", const=".txt", metavar="EXT", |
| help="add the missing lighting tags to each image's caption file " |
| "(<image name>EXT, default .txt; created if missing)") |
| parser.add_argument("--write-captions", nargs="?", const=".lighting.txt", metavar="EXT", |
| help="write the lighting tags alone to <image name>EXT (default .lighting.txt), " |
| "replacing that file") |
| parser.add_argument("--spaces", action="store_true", |
| help="write tags with spaces instead of underscores (rim light, not rim_light)") |
| parser.add_argument("--validated-only", action="store_true", |
| help="only emit the labels validated against human answers (see README)") |
| parser.add_argument("--preset", choices=PRESETS, default="default", |
| help="threshold preset: default (as many tags as the teacher) or recall (about 1.5 times " |
| "as many tags, a little less precise; see README)") |
| parser.add_argument("--no-tonal-key", action="store_true", |
| help="do not emit the tonal key tag (dark, dim_lighting, high_key)") |
| parser.add_argument("--seeds", type=int, help="number of seeds to use, 1 to 3 (default 3; fewer is faster " |
| "and slightly less accurate)") |
| parser.add_argument("--batch-size", type=int, default=8) |
| parser.add_argument("--workers", type=int, default=4, help="image decoding threads") |
| parser.add_argument("--device", default="auto", help="auto, cuda, cuda:1, mps or cpu") |
| parser.add_argument("--precision", default="auto", choices=("auto", "bf16", "fp16", "fp32"), |
| help="CUDA only (auto: bf16 when supported, as in training)") |
| parser.add_argument("--model-dir", type=Path, help="folder with the model files (default: this script's)") |
| parser.add_argument("--repo", help="download the model from this Hugging Face repo instead") |
| parser.add_argument("--revision", help="with --repo: a tag such as v1.0.0, a branch or a commit") |
| parser.add_argument("-q", "--quiet", action="store_true", help="do not print a line per image") |
| args = parser.parse_args(argv) |
| for stream in (sys.stdout, sys.stderr): |
| with contextlib.suppress(AttributeError, ValueError): |
| stream.reconfigure(encoding="utf-8", errors="replace") |
|
|
| paths = find_images(args.inputs, args.recursive) |
| if not paths: |
| print("no image found", file=sys.stderr) |
| return 1 |
| options = {"device": args.device, "seeds": args.seeds, "precision": args.precision, "preset": args.preset} |
| tagger = (LightingTagger.from_pretrained(args.repo, revision=args.revision, **options) if args.repo |
| else LightingTagger(args.model_dir, **options)) |
| print(f"booru-lighting-tagger {tagger.version}: {len(tagger.models)} seed(s) on {tagger.device}, " |
| f"{tagger.preset} thresholds, {len(paths)} image(s)", file=sys.stderr) |
|
|
| jsonl = args.jsonl.open("w", encoding="utf-8") if args.jsonl else None |
| done = errors = tagged_images = added = 0 |
| try: |
| for result in tagger.tag_files(paths, batch_size=args.batch_size, workers=args.workers, |
| validated_only=args.validated_only, tonal_key_tag=not args.no_tonal_key): |
| done += 1 |
| if jsonl: |
| jsonl.write(json.dumps(result, ensure_ascii=False) + "\n") |
| if result["error"]: |
| errors += 1 |
| print(f"{result['file']}: error: {result['error']}", file=sys.stderr) |
| continue |
| tags = [tag.replace("_", " ") for tag in result["tags"]] if args.spaces else result["tags"] |
| tagged_images += bool(tags) |
| image = Path(result["file"]) |
| if args.append_to_captions: |
| added += append_to_caption(caption_path(image, args.append_to_captions), tags) |
| if args.write_captions: |
| caption_path(image, args.write_captions).write_text(", ".join(tags), encoding="utf-8") |
| if not args.quiet: |
| print(f"{result['file']}: {', '.join(tags) if tags else '(no lighting tag)'}") |
| finally: |
| if jsonl: |
| jsonl.close() |
| summary = f"{done} image(s), {tagged_images} with at least one tag, {errors} error(s)" |
| if args.append_to_captions: |
| summary += f", {added} tag(s) added to captions" |
| print(summary, file=sys.stderr) |
| return 0 if errors == 0 else 2 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
