feat: add ADetailer support (#1785)

scripts/convert_yolov8_to_safetensors.py

@@ -0,0 +1,86 @@
+#!/usr/bin/env python3
+"""Convert an Ultralytics YOLOv8 detection checkpoint for sd.cpp ADetailer."""
+
+import argparse
+import json
+from pathlib import Path
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Convert an Ultralytics YOLOv8 detection .pt checkpoint to safetensors."
+    )
+    parser.add_argument("input", type=Path, help="input YOLOv8 detection checkpoint")
+    parser.add_argument("output", type=Path, help="output safetensors path")
+    parser.add_argument(
+        "--input-size", type=int, default=640, help="detector input size metadata (default: 640)"
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    if args.input_size < 32 or args.input_size % 32 != 0:
+        raise ValueError("--input-size must be a positive multiple of 32")
+    if args.output.suffix.lower() != ".safetensors":
+        raise ValueError("output path must use the .safetensors extension")
+
+    try:
+        import torch
+        from safetensors.torch import save_file
+        from ultralytics import YOLO
+        from ultralytics.nn.modules.head import Detect
+    except ImportError as exc:
+        raise SystemExit("conversion requires ultralytics, torch, and safetensors") from exc
+
+    torch_load = torch.load
+
+    def load_trusted_checkpoint(*load_args, **load_kwargs):
+        load_kwargs.setdefault("weights_only", False)
+        return torch_load(*load_args, **load_kwargs)
+
+    torch.load = load_trusted_checkpoint
+    try:
+        yolo = YOLO(str(args.input))
+    finally:
+        torch.load = torch_load
+    network = yolo.model
+    if not isinstance(network.model[-1], Detect) or network.model[-1].__class__.__name__ != "Detect":
+        raise ValueError("only YOLOv8 detection checkpoints are supported; segmentation is not yet supported")
+
+    network.eval()
+    network.fuse()
+    state_dict = network.state_dict()
+    required = {
+        "model.0.conv.weight",
+        "model.22.cv2.0.2.weight",
+        "model.22.cv3.0.2.weight",
+    }
+    missing = sorted(required.difference(state_dict))
+    if missing:
+        raise ValueError(f"checkpoint does not match the supported YOLOv8 layout; missing {missing}")
+
+    tensors = {}
+    for name, tensor in state_dict.items():
+        if not name.startswith("model.") or ".bn." in name or name.endswith("dfl.conv.weight"):
+            continue
+        if not (name.endswith(".weight") or name.endswith(".bias")):
+            continue
+        dtype = torch.float16 if name.endswith(".weight") else torch.float32
+        tensors[name] = tensor.detach().to(device="cpu", dtype=dtype).contiguous()
+
+    metadata = {
+        "format": "pt",
+        "yolov8.variant": "detect",
+        "yolov8.input_size": str(args.input_size),
+        "yolov8.num_classes": str(int(network.model[-1].nc)),
+        "yolov8.reg_max": str(int(network.model[-1].reg_max)),
+        "yolov8.names": json.dumps(yolo.names, ensure_ascii=False),
+    }
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    save_file(tensors, str(args.output), metadata=metadata)
+    print(f"wrote {args.output}: {len(tensors)} tensors")
+
+
+if __name__ == "__main__":
+    main()