Layout fid

class LayoutFIDArchitecture(StrEnum):
    """Closed set of supported layout FID encoder architectures."""

    fidnet_v3 = auto()
    layoutnet = auto()

class LayoutFIDConfig(PretrainedConfig):
    """Configuration saved with layout FID checkpoints.

    Args:
        dataset_name: Canonical layout dataset name.
        id2label: Dataset-local id-to-label metadata.
        architecture: Encoder architecture selected by the checkpoint.
        source: Released artifact family selected by the checkpoint.
        num_public_labels: Number of public dataset labels.
        num_label_embeddings: Number of model label embeddings.
        max_length: Maximum element count accepted by the checkpoint.
        d_model: Transformer hidden dimension.
        nhead: Number of attention heads.
        num_layers: Number of transformer encoder layers.
        bbox_format_for_model: Internal bbox format consumed by the encoder.
        label_id_offset: Offset applied before model label embedding lookup.
        pad_label_id: Label id used only in padded model tensor positions.
        reference_stats: Relative reference-statistics paths by split.
        kwargs: Extra Hugging Face config fields.

    Raises:
        ValueError: If label counts or enum values are invalid.

    Examples:
        >>> cfg = LayoutFIDConfig(
        ...     dataset_name="publaynet",
        ...     architecture="layoutnet",
        ...     source="layoutflow",
        ...     num_public_labels=5,
        ...     num_label_embeddings=6,
        ...     max_length=20,
        ... )
        >>> cfg.reference_stats["test"]
        'reference_stats/test.npz'
    """

    model_type = "layout-fid"
    has_no_defaults_at_init = True

    def __init__(
        self,
        *,
        dataset_name: DatasetName | str,
        id2label: Mapping[int, str] | Mapping[str, str] | None = None,
        architecture: LayoutFIDArchitecture | str,
        source: LayoutFIDSource | str,
        num_public_labels: int,
        num_label_embeddings: int,
        max_length: int,
        d_model: int = 256,
        nhead: int = 4,
        num_layers: int = 4,
        bbox_format_for_model: BoxFormat | str = "ltrb",
        label_id_offset: int = 0,
        pad_label_id: int = 0,
        reference_stats: dict[str, str] | None = None,
        **kwargs: LayoutFIDConfigValue,
    ) -> None:
        """Initialize a layout FID checkpoint configuration."""
        super().__init__(**kwargs)  # ty: ignore[invalid-argument-type]
        dataset = normalize_dataset_name(dataset_name)
        arch = normalize_architecture(architecture)
        src = normalize_source(source)
        box_format = normalize_box_format(bbox_format_for_model)
        if num_public_labels <= 0:
            raise ValueError("num_public_labels must be positive")

        if num_label_embeddings < num_public_labels:
            raise ValueError("num_label_embeddings must cover public labels")

        if max_length <= 0:
            raise ValueError("max_length must be positive")

        if pad_label_id < 0 or pad_label_id >= num_label_embeddings:
            raise ValueError("pad_label_id must be inside the embedding table")

        raw_id2label = id2label or id2label_for_dataset(dataset)
        self.dataset_name = str(dataset)
        self.id2label = {int(k): v for k, v in raw_id2label.items()}
        self.label2id = {label: idx for idx, label in self.id2label.items()}
        self.architecture = str(arch)
        self.source = str(src)

        self.num_public_labels = num_public_labels
        self.num_label_embeddings = num_label_embeddings
        self.max_length = max_length
        self.d_model = d_model
        self.nhead = nhead
        self.num_layers = num_layers

        self.bbox_format_for_model = str(box_format)
        self.label_id_offset = label_id_offset
        self.pad_label_id = pad_label_id
        self.reference_stats = dict(reference_stats or DEFAULT_REFERENCE_STATS)

    @property
    def feature_dim(self) -> int:
        """Return the layout FID feature dimension."""
        return self.d_model

    def _get_generation_parameters(self) -> dict[str, str]:
        """Return no generation parameters for this evaluator config."""
        return {}

class LayoutFIDSource(StrEnum):
    """Closed set of supported released artifact families."""

    layoutdm = auto()
    layoutflow = auto()

class LayoutFIDStatsSplit(StrEnum):
    """Closed set of bundled reference-statistics splits."""

    val = auto()
    test = auto()

@dataclass(frozen=True)
class LayoutFIDStatistics:
    """Feature distribution statistics used by layout FID."""

    mu: Float[np.ndarray, "channels"]
    sigma: Float[np.ndarray, "channels channels"]
    split: str
    dataset_name: str
    source: str
    feature_dim: int
    num_samples: int | None = None

    @classmethod
    def from_mapping(
        cls,
        values: Mapping[
            str,
            Float[np.ndarray, "..."]
            | list[float]
            | list[list[float]]
            | str
            | int
            | None,
        ],
    ) -> "LayoutFIDStatistics":
        """Create statistics from a mapping."""
        mu = np.asarray(values["mu"], dtype=np.float64)
        sigma = np.asarray(values["sigma"], dtype=np.float64)
        return cls(
            mu=mu,
            sigma=sigma,
            split=str(values.get("split", "candidate")),
            dataset_name=str(values.get("dataset_name", "")),
            source=str(values.get("source", "")),
            feature_dim=int(cast(int | str, values.get("feature_dim", mu.shape[0]))),
            num_samples=(
                None
                if values.get("num_samples") is None
                else int(cast(int | str, values["num_samples"]))
            ),
        )

class LayoutFIDModel(PreTrainedModel):
    """Feature encoder used for layout FID evaluation."""

    config_class = LayoutFIDConfig
    base_model_prefix = "layout_fid"

    def __init__(self, config: LayoutFIDConfig) -> None:
        """Create a layout FID encoder.

        Args:
            config: Explicit layout FID configuration.
        """
        super().__init__(config)
        self.all_tied_weights_keys: dict[str, str] = {}
        self.emb_label = nn.Embedding(config.num_label_embeddings, config.d_model)
        self.fc_bbox = nn.Linear(4, config.d_model)
        self.enc_fc_in = nn.Linear(config.d_model * 2, config.d_model)
        self.enc_transformer = TransformerWithToken(
            d_model=config.d_model,
            dim_feedforward=config.d_model // 2,
            nhead=config.nhead,
            num_layers=config.num_layers,
        )
        self.fc_out_disc = nn.Linear(config.d_model, 1)
        self.pos_token = nn.Parameter(torch.rand(config.max_length, 1, config.d_model))
        self.dec_fc_in = nn.Linear(config.d_model * 2, config.d_model)
        dec_layer = nn.TransformerEncoderLayer(
            d_model=config.d_model,
            nhead=config.nhead,
            dim_feedforward=config.d_model // 2,
        )
        self.dec_transformer = nn.TransformerEncoder(
            dec_layer, num_layers=config.num_layers
        )
        self.fc_out_cls = nn.Linear(config.d_model, config.num_label_embeddings)
        self.fc_out_bbox = nn.Linear(config.d_model, 4)

    def extract_features(
        self,
        *,
        bbox: Float[torch.Tensor, "batch elements 4"],
        labels: Int[torch.Tensor, "batch elements"],
        padding_mask: Bool[torch.Tensor, "batch elements"],
    ) -> Float[torch.Tensor, "batch channels"]:
        """Extract batch-level feature vectors.

        Args:
            bbox: Model-ready boxes.
            labels: Model-ready label ids.
            padding_mask: Boolean mask where ``True`` marks padded elements.

        Returns:
            Feature tensor shaped ``(batch, d_model)``.

        Raises:
            ValueError: If input shapes are inconsistent.

        Examples:
            >>> from layout_fid import LayoutFIDConfig, LayoutFIDModel
            >>> cfg = LayoutFIDConfig(
            ...     dataset_name="publaynet", architecture="layoutnet",
            ...     source="layoutflow", num_public_labels=5,
            ...     num_label_embeddings=6, max_length=2,
            ... )
            >>> model = LayoutFIDModel(cfg)
            >>> out = model.extract_features(
            ...     bbox=torch.zeros(1, 2, 4),
            ...     labels=torch.zeros(1, 2, dtype=torch.long),
            ...     padding_mask=torch.zeros(1, 2, dtype=torch.bool),
            ... )
            >>> tuple(out.shape)
            (1, 256)
        """
        self._validate_inputs(bbox, labels, padding_mask)
        box_features = self.fc_bbox(bbox)
        label_features = self.emb_label(labels)
        hidden = self.enc_fc_in(torch.cat([box_features, label_features], dim=-1))
        hidden = torch.relu(hidden).permute(1, 0, 2)
        encoded = self.enc_transformer(hidden, padding_mask)
        return encoded[0]

    def forward(
        self,
        *,
        bbox: Float[torch.Tensor, "batch elements 4"],
        labels: Int[torch.Tensor, "batch elements"],
        padding_mask: Bool[torch.Tensor, "batch elements"],
        output_reconstruction: bool = False,
        return_dict: bool = True,
    ) -> LayoutFIDOutput | tuple[Shaped[torch.Tensor, "..."], ...]:
        """Run feature extraction and optional reconstruction heads.

        Args:
            bbox: Model-ready boxes.
            labels: Model-ready label ids.
            padding_mask: Boolean mask where ``True`` marks padded elements.
            output_reconstruction: Whether to return class and bbox predictions.
            return_dict: Whether to return ``LayoutFIDOutput``.

        Returns:
            ``LayoutFIDOutput`` or a tuple with the same non-``None`` fields.

        Raises:
            ValueError: If input shapes are inconsistent.
        """
        features = self.extract_features(
            bbox=bbox, labels=labels, padding_mask=padding_mask
        )
        discriminator_logits = self.fc_out_disc(features).squeeze(-1)
        class_logits: Float[torch.Tensor, "... labels"] | None = None
        bbox_pred: Float[torch.Tensor, "... 4"] | None = None
        if output_reconstruction:
            class_logits, bbox_pred = self._decode(features, padding_mask)
        if not return_dict:
            values = (features, discriminator_logits, class_logits, bbox_pred)
            return tuple(value for value in values if value is not None)
        return LayoutFIDOutput(
            features=features,
            discriminator_logits=discriminator_logits,
            class_logits=class_logits,
            bbox_pred=bbox_pred,
        )

    def _decode(
        self,
        features: Float[torch.Tensor, "batch channels"],
        padding_mask: Bool[torch.Tensor, "batch elements"],
    ) -> tuple[Float[torch.Tensor, "... labels"], Float[torch.Tensor, "... 4"]]:
        batch_size, elements = padding_mask.shape
        hidden = features.unsqueeze(0).expand(elements, -1, -1)
        positions = self.pos_token[:elements].expand(-1, batch_size, -1)
        hidden = torch.relu(self.dec_fc_in(torch.cat([hidden, positions], dim=-1)))
        hidden = self.dec_transformer(hidden, src_key_padding_mask=padding_mask)
        hidden = hidden.permute(1, 0, 2)
        class_logits = self.fc_out_cls(hidden)
        bbox_pred = torch.sigmoid(self.fc_out_bbox(hidden))
        if normalize_source(self.config.source) is LayoutFIDSource.layoutflow:
            valid = ~padding_mask
            class_logits = class_logits[valid]
            bbox_pred = bbox_pred[valid]
        return class_logits, bbox_pred

    @staticmethod
    def _validate_inputs(
        bbox: Float[torch.Tensor, "batch elements 4"],
        labels: Int[torch.Tensor, "batch elements"],
        padding_mask: Bool[torch.Tensor, "batch elements"],
    ) -> None:
        if bbox.ndim != 3 or bbox.shape[-1] != 4:
            raise ValueError("bbox must have shape (batch, elements, 4)")

        if labels.shape != bbox.shape[:2]:
            raise ValueError("labels must have shape (batch, elements)")

        if padding_mask.shape != bbox.shape[:2]:
            raise ValueError("padding_mask must have shape (batch, elements)")

@dataclass
class LayoutFIDOutput(ModelOutput):
    """Output returned by ``LayoutFIDModel.forward``.

    Args:
        features: Batch-level layout feature vectors.
        discriminator_logits: Optional discriminator logits.
        class_logits: Optional per-element or valid-element class logits.
        bbox_pred: Optional reconstructed boxes.
        intermediates: Optional diagnostic tensors.
    """

    features: Float[torch.Tensor, "batch channels"]
    discriminator_logits: Float[torch.Tensor, "batch"] | None = None
    class_logits: Float[torch.Tensor, "... labels"] | None = None
    bbox_pred: Float[torch.Tensor, "... 4"] | None = None
    intermediates: dict[str, object] | None = None

class LayoutFIDEvaluator:
    """Compose a layout FID model, processor, and reference statistics."""

    def __init__(
        self,
        *,
        model: LayoutFIDModel,
        processor: LayoutFIDProcessor,
        reference_statistics: Mapping[str, LayoutFIDStatistics] | None = None,
        device: torch.device | str | None = None,
    ) -> None:
        """Create an evaluator."""
        self.model = model
        self.processor = processor
        self.reference_statistics = dict(reference_statistics or {})
        self.device = (
            torch.device(device) if device is not None else torch.device("cpu")
        )
        self.model.to(self.device)
        self.model.eval()

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str | PathLike[str],
        *,
        device: torch.device | str | None = None,
        **kwargs: LayoutFIDLoadKwarg,
    ) -> LayoutFIDEvaluator:
        """Load evaluator components from a local directory or Hub id."""
        model = LayoutFIDModel.from_pretrained(pretrained_model_name_or_path, **kwargs)
        processor = LayoutFIDProcessor.from_pretrained(
            pretrained_model_name_or_path,
            **cast(dict[str, LayoutFIDConfigValue], kwargs),
        )
        stats = cls._load_reference_statistics(pretrained_model_name_or_path, model)
        return cls(
            model=model, processor=processor, reference_statistics=stats, device=device
        )

    def extract_features(
        self,
        *,
        layouts: LayoutGenerationOutput
        | Mapping[
            str,
            Float[torch.Tensor, "batch elements 4"]
            | Float[np.ndarray, "batch elements 4"]
            | Int[torch.Tensor, "batch elements"]
            | Int[np.ndarray, "batch elements"]
            | Bool[torch.Tensor, "batch elements"]
            | Bool[np.ndarray, "batch elements"]
            | Mapping[int, str]
            | Mapping[str, str]
            | None,
        ]
        | None = None,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Float[np.ndarray, "batch elements 4"]
        | Sequence[ArrayLikeInput]
        | None = None,
        labels: Int[torch.Tensor, "batch elements"]
        | Int[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        mask: Bool[torch.Tensor, "batch elements"]
        | Bool[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        id2label: Mapping[int, str] | Mapping[str, str] | None = None,
        box_format: str = "xywh",
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
        batch_size: int = 512,
    ) -> Float[torch.Tensor, "batch channels"]:
        """Extract features from public layout tensors."""
        layout_kwargs = self._layout_kwargs(
            layouts=layouts,
            bbox=bbox,
            labels=labels,
            mask=mask,
            id2label=id2label,
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas_size,
        )
        batch = self.processor(
            bbox=cast(
                Float[torch.Tensor, "batch elements 4"]
                | Float[np.ndarray, "batch elements 4"]
                | Sequence[ArrayLikeInput],
                layout_kwargs["bbox"],
            ),
            labels=cast(
                Int[torch.Tensor, "batch elements"]
                | Int[np.ndarray, "batch elements"]
                | Sequence[ArrayLikeInput],
                layout_kwargs["labels"],
            ),
            mask=cast(
                Bool[torch.Tensor, "batch elements"]
                | Bool[np.ndarray, "batch elements"]
                | Sequence[ArrayLikeInput]
                | None,
                layout_kwargs["mask"],
            ),
            id2label=cast(
                Mapping[int, str] | Mapping[str, str] | None,
                layout_kwargs["id2label"],
            ),
            box_format=cast(str, layout_kwargs["box_format"]),
            normalized=cast(bool, layout_kwargs["normalized"]),
            canvas_size=cast(tuple[int, int] | None, layout_kwargs["canvas_size"]),
            device=self.device,
        )
        outputs: list[Float[torch.Tensor, "batch channels"]] = []
        for start in range(0, batch.bbox.shape[0], batch_size):
            end = start + batch_size
            with torch.no_grad():
                outputs.append(
                    self.model.extract_features(
                        bbox=batch.bbox[start:end],
                        labels=batch.labels[start:end],
                        padding_mask=batch.padding_mask[start:end],
                    ).cpu()
                )
        return torch.cat(outputs, dim=0)

    def compute_statistics(
        self,
        *,
        layouts: LayoutGenerationOutput
        | Mapping[
            str,
            Float[torch.Tensor, "batch elements 4"]
            | Float[np.ndarray, "batch elements 4"]
            | Int[torch.Tensor, "batch elements"]
            | Int[np.ndarray, "batch elements"]
            | Bool[torch.Tensor, "batch elements"]
            | Bool[np.ndarray, "batch elements"]
            | Mapping[int, str]
            | Mapping[str, str]
            | None,
        ]
        | None = None,
        features: Float[torch.Tensor, "batch channels"]
        | Float[np.ndarray, "batch channels"]
        | None = None,
        **layout_kwargs: LayoutFIDLayoutKwarg,
    ) -> LayoutFIDStatistics:
        """Compute candidate feature statistics."""
        if features is not None and (layouts is not None or layout_kwargs):
            raise ValueError("Pass either features or layout inputs, not both")

        if features is None:
            features = self.extract_features(
                layouts=layouts,
                bbox=cast(
                    Float[torch.Tensor, "batch elements 4"]
                    | Float[np.ndarray, "batch elements 4"]
                    | Sequence[ArrayLikeInput]
                    | None,
                    layout_kwargs.get("bbox"),
                ),
                labels=cast(
                    Int[torch.Tensor, "batch elements"]
                    | Int[np.ndarray, "batch elements"]
                    | Sequence[ArrayLikeInput]
                    | None,
                    layout_kwargs.get("labels"),
                ),
                mask=cast(
                    Bool[torch.Tensor, "batch elements"]
                    | Bool[np.ndarray, "batch elements"]
                    | Sequence[ArrayLikeInput]
                    | None,
                    layout_kwargs.get("mask"),
                ),
                id2label=cast(
                    Mapping[int, str] | Mapping[str, str] | None,
                    layout_kwargs.get("id2label"),
                ),
                box_format=cast(str, layout_kwargs.get("box_format", "xywh")),
                normalized=cast(bool, layout_kwargs.get("normalized", True)),
                canvas_size=cast(
                    tuple[int, int] | None, layout_kwargs.get("canvas_size")
                ),
                batch_size=cast(int, layout_kwargs.get("batch_size", 512)),
            )
        return compute_feature_statistics(
            features,
            dataset_name=self.model.config.dataset_name,
            source=self.model.config.source,
        )

    def compute_fid(
        self,
        *,
        layouts: LayoutGenerationOutput
        | Mapping[
            str,
            Float[torch.Tensor, "batch elements 4"]
            | Float[np.ndarray, "batch elements 4"]
            | Int[torch.Tensor, "batch elements"]
            | Int[np.ndarray, "batch elements"]
            | Bool[torch.Tensor, "batch elements"]
            | Bool[np.ndarray, "batch elements"]
            | Mapping[int, str]
            | Mapping[str, str]
            | None,
        ]
        | None = None,
        features: Float[torch.Tensor, "batch channels"]
        | Float[np.ndarray, "batch channels"]
        | None = None,
        statistics: LayoutFIDStatistics
        | Mapping[
            str,
            Float[np.ndarray, ...] | list[float] | list[list[float]] | str | int | None,
        ]
        | None = None,
        reference_statistics: LayoutFIDStatistics
        | Mapping[
            str,
            Float[np.ndarray, ...] | list[float] | list[list[float]] | str | int | None,
        ]
        | None = None,
        reference_split: LayoutFIDStatsSplit | str = "test",
        **layout_kwargs: LayoutFIDLayoutKwarg,
    ) -> float:
        """Compute layout FID against bundled or supplied reference statistics."""
        provided = sum(value is not None for value in (layouts, features, statistics))
        if provided + bool(layout_kwargs) == 0:
            raise ValueError("Pass candidate layouts, features, or statistics")

        if statistics is None:
            statistics = self.compute_statistics(
                layouts=layouts, features=features, **layout_kwargs
            )
        reference = (
            LayoutFIDStatistics.from_mapping(
                cast(
                    Mapping[
                        str,
                        Float[np.ndarray, "..."]
                        | list[float]
                        | list[list[float]]
                        | str
                        | int
                        | None,
                    ],
                    reference_statistics,
                )
            )
            if reference_statistics is not None
            and not isinstance(reference_statistics, LayoutFIDStatistics)
            else reference_statistics
        )
        if reference is None:
            split = str(normalize_stats_split(reference_split))
            try:
                reference = self.reference_statistics[split]
            except KeyError as exc:
                raise ValueError(
                    f"Reference statistics split is not loaded: {split}"
                ) from exc

        candidate = (
            LayoutFIDStatistics.from_mapping(
                cast(
                    Mapping[
                        str,
                        Float[np.ndarray, "..."]
                        | list[float]
                        | list[list[float]]
                        | str
                        | int
                        | None,
                    ],
                    statistics,
                )
            )
            if not isinstance(statistics, LayoutFIDStatistics)
            else statistics
        )
        return compute_layout_fid_from_statistics(candidate, reference)

    @staticmethod
    def _layout_kwargs(
        *,
        layouts: LayoutGenerationOutput
        | Mapping[
            str,
            Float[torch.Tensor, "batch elements 4"]
            | Float[np.ndarray, "batch elements 4"]
            | Int[torch.Tensor, "batch elements"]
            | Int[np.ndarray, "batch elements"]
            | Bool[torch.Tensor, "batch elements"]
            | Bool[np.ndarray, "batch elements"]
            | Mapping[int, str]
            | Mapping[str, str]
            | None,
        ]
        | None,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Float[np.ndarray, "batch elements 4"]
        | Sequence[ArrayLikeInput]
        | None,
        labels: Int[torch.Tensor, "batch elements"]
        | Int[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None,
        mask: Bool[torch.Tensor, "batch elements"]
        | Bool[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None,
        id2label: Mapping[int, str] | Mapping[str, str] | None,
        box_format: str,
        normalized: bool,
        canvas_size: tuple[int, int] | None,
    ) -> dict[
        str,
        Float[torch.Tensor, "batch elements 4"]
        | Float[np.ndarray, "batch elements 4"]
        | Int[torch.Tensor, "batch elements"]
        | Int[np.ndarray, "batch elements"]
        | Bool[torch.Tensor, "batch elements"]
        | Bool[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | Mapping[int, str]
        | Mapping[str, str]
        | str
        | bool
        | tuple[int, int]
        | None,
    ]:
        if layouts is not None and any(
            value is not None for value in (bbox, labels, mask)
        ):
            raise ValueError("Pass either layouts or explicit bbox/labels/mask")

        if layouts is not None:
            if isinstance(layouts, LayoutGenerationOutput):
                bbox = layouts.bbox
                labels = layouts.labels
                mask = layouts.mask
                id2label = layouts.id2label
            else:
                bbox = cast(
                    Float[torch.Tensor, "batch elements 4"]
                    | Float[np.ndarray, "batch elements 4"]
                    | Sequence[ArrayLikeInput],
                    layouts["bbox"],
                )
                labels = cast(
                    Int[torch.Tensor, "batch elements"]
                    | Int[np.ndarray, "batch elements"]
                    | Sequence[ArrayLikeInput],
                    layouts["labels"],
                )
                mask = cast(
                    Bool[torch.Tensor, "batch elements"]
                    | Bool[np.ndarray, "batch elements"]
                    | Sequence[ArrayLikeInput]
                    | None,
                    layouts.get("mask"),
                )
                id2label = cast(
                    Mapping[int, str] | Mapping[str, str] | None,
                    layouts.get("id2label"),
                )
        if bbox is None or labels is None:
            raise ValueError("bbox and labels are required")

        return {
            "bbox": bbox,
            "labels": labels,
            "mask": mask,
            "id2label": id2label,
            "box_format": box_format,
            "normalized": normalized,
            "canvas_size": canvas_size,
        }

    @staticmethod
    def _load_reference_statistics(
        model_path: str | PathLike[str], model: LayoutFIDModel
    ) -> dict[str, LayoutFIDStatistics]:
        path = Path(model_path)
        if not path.exists():
            return {}
        stats: dict[str, LayoutFIDStatistics] = {}
        for split, relative in model.config.reference_stats.items():
            stats_path = path / relative
            if stats_path.exists():
                stats[split] = load_reference_statistics(stats_path)
        return stats