| """SafeTensors-only loader for KoharuLayout RF-DETR Seg 2XL.""" | |
| from __future__ import annotations | |
| from pathlib import Path | |
| import warnings | |
| from rfdetr import RFDETRSeg2XLarge | |
| from rfdetr.config import PretrainWeightsCompatibilityWarning | |
| from safetensors.torch import load_file | |
| CLASS_NAMES = ["text", "onomatopoeia", "bubble", "panel"] | |
| def load_model(weights: str | Path) -> RFDETRSeg2XLarge: | |
| with warnings.catch_warnings(): | |
| warnings.simplefilter("ignore", PretrainWeightsCompatibilityWarning) | |
| model = RFDETRSeg2XLarge( | |
| pretrain_weights=None, | |
| resolution=1152, | |
| num_select=160, | |
| num_classes=len(CLASS_NAMES), | |
| ) | |
| incompatible = model.model.model.load_state_dict(load_file(str(weights), device="cpu"), strict=True) | |
| if incompatible.missing_keys or incompatible.unexpected_keys: | |
| raise RuntimeError(f"Incompatible weights: {incompatible}") | |
| model.model.class_names = CLASS_NAMES.copy() | |
| return model | |