"""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