mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-07-23 11:20:53 -05:00
87 lines
3.1 KiB
Python
87 lines
3.1 KiB
Python
#!/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()
|