chore: import upstream snapshot with attribution

tests/inference/__init__.py

@@ -0,0 +1,5 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------

tests/inference/helpers.py

@@ -0,0 +1,106 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Shared test helpers for the inference test suite.
+
+Plain classes and functions (not pytest fixtures) shared across multiple test modules to avoid verbatim duplication.
+Import with a relative import::
+
+    from .helpers import _BaseFakeRFDETR, _DummyModel, _DummyRFDETR
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Any
+
+import torch
+
+from rfdetr.detr import RFDETR
+
+
+class _BaseFakeRFDETR(RFDETR):
+    """RFDETR test double that skips weight downloads and returns a minimal model config.
+
+    Subclasses must override ``get_model`` to supply the model context appropriate for
+    the scenario under test.
+
+    Examples:
+        This class is imported directly by test modules that need a weight-free RFDETR.
+    """
+
+    def maybe_download_pretrain_weights(self) -> None:
+        """Skip weight download in tests."""
+        return None
+
+    def get_model_config(self, **kwargs: object) -> SimpleNamespace:
+        """Return a minimal config sufficient for most test scenarios."""
+        return SimpleNamespace(num_channels=3)
+
+
+class _DummyModel:
+    """Minimal model stub that returns deterministic postprocessed results.
+
+    Examples:
+        >>> m = _DummyModel(labels=[0, 1])
+        >>> len(m._labels)
+        2
+    """
+
+    def __init__(
+        self,
+        class_names: list[str] | None = None,
+        labels: list[int] | None = None,
+        include_keypoints: bool = False,
+        num_keypoints: int = 17,
+    ) -> None:
+        """Initialise stub with optional class names, label list, and keypoint flag."""
+        self.device = torch.device("cpu")
+        self.resolution = 28
+        self.model = torch.nn.Identity()
+        self.class_names = class_names
+        self._labels = labels if labels is not None else [1]
+        self._include_keypoints = include_keypoints
+        self._num_keypoints = num_keypoints
+
+    def postprocess(self, predictions: Any, target_sizes: torch.Tensor) -> list[dict[str, torch.Tensor]]:
+        """Return fixed scores/boxes (and optional keypoints) for every image in the batch."""
+        batch = target_sizes.shape[0]
+        results = []
+        for _ in range(batch):
+            result: dict[str, torch.Tensor] = {
+                "scores": torch.tensor([0.9] * len(self._labels)),
+                "labels": torch.tensor(self._labels),
+                "boxes": torch.tensor([[0.0, 0.0, 1.0, 1.0]] * len(self._labels)),
+            }
+            if self._include_keypoints:
+                result["keypoints"] = torch.full((len(self._labels), self._num_keypoints, 3), 0.5, dtype=torch.float32)
+                result["keypoint_precision_cholesky"] = torch.full(
+                    (len(self._labels), self._num_keypoints, 3), 0.25, dtype=torch.float32
+                )
+            results.append(result)
+        return results
+
+
+class _DummyRFDETR(RFDETR):
+    """Weight-free RFDETR that delegates to ``_DummyModel`` for all inference.
+
+    Examples:
+        >>> m = _DummyRFDETR()
+        >>> isinstance(m.model, _DummyModel)
+        True
+    """
+
+    def maybe_download_pretrain_weights(self) -> None:
+        """Skip weight download in tests."""
+        return None
+
+    def get_model_config(self, **kwargs: object) -> SimpleNamespace:
+        """Return a minimal namespace with just ``num_channels``."""
+        return SimpleNamespace(num_channels=3)
+
+    def get_model(self, config: SimpleNamespace) -> _DummyModel:
+        """Return a fresh ``_DummyModel`` instance."""
+        return _DummyModel()

tests/inference/test_from_checkpoint.py

@@ -0,0 +1,857 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for RFDETR.from_checkpoint classmethod.
+
+The inference logic is isolated by patching ``torch.load`` and the target model class inside ``rfdetr.variants`` (or
+``rfdetr.platform.models`` for plus models).  No model weights are downloaded or GPU memory allocated.
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import warnings
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from rfdetr.config import PretrainWeightsCompatibilityWarning
+from rfdetr.detr import RFDETR
+from rfdetr.detr import logger as detr_logger
+from rfdetr.platform import _IS_RFDETR_PLUS_AVAILABLE
+from rfdetr.variants import RFDETRSmall
+
+
+class _CustomObj:
+    """Module-level class that weights_only=True rejects (not in safe globals).
+
+    Must be at module scope so pickle can resolve the fully-qualified name during torch.save.  Local/nested classes
+    cannot be pickled.
+    """
+
+
+def _ns(pretrain_weights: str, num_classes: int = 80) -> dict:
+    """Fake legacy checkpoint with argparse.Namespace args."""
+    return {"args": argparse.Namespace(pretrain_weights=pretrain_weights, num_classes=num_classes)}
+
+
+def _dict(pretrain_weights: str, num_classes: int = 80) -> dict:
+    """Fake PTL-style checkpoint with dict args."""
+    return {"args": {"pretrain_weights": pretrain_weights, "num_classes": num_classes}}
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _call_from_checkpoint(ckpt: dict, path: Path, cls_patch_target: str, **kwargs):
+    """Invoke RFDETR.from_checkpoint with torch.load mocked to return *ckpt* and the model class at *cls_patch_target*
+    replaced by a MagicMock.
+
+    Returns:
+        Tuple of (result, mock_class).
+    """
+    mock_instance = MagicMock()
+    with (
+        patch("rfdetr.detr.torch.load", return_value=ckpt),
+        patch(cls_patch_target) as mock_cls,
+    ):
+        mock_cls.return_value = mock_instance
+        result = RFDETR.from_checkpoint(path, **kwargs)
+    return result, mock_cls
+
+
+# ---------------------------------------------------------------------------
+# Namespace args (legacy .pth checkpoints)
+# ---------------------------------------------------------------------------
+
+
+class TestFromCheckpointNamespaceArgs:
+    """from_checkpoint with argparse.Namespace args (legacy engine.py format)."""
+
+    @pytest.mark.parametrize(
+        ("pretrain_weights, patch_target"),
+        [
+            ("rf-detr-nano.pth", "RFDETRNano"),
+            ("rf-detr-small.pth", "RFDETRSmall"),
+            ("rf-detr-medium.pth", "RFDETRMedium"),
+            ("rf-detr-large.pth", "RFDETRLarge"),
+            ("rf-detr-keypoint-preview-xlarge.pth", "RFDETRKeypointPreview"),
+            ("rf-detr-base.pth", "RFDETRBase"),
+            ("rf-detr-seg-nano.pt", "RFDETRSegNano"),
+            ("rf-detr-seg-small.pt", "RFDETRSegSmall"),
+            ("rf-detr-seg-medium.pt", "RFDETRSegMedium"),
+            ("rf-detr-seg-large.pt", "RFDETRSegLarge"),
+            ("rf-detr-seg-xlarge.pt", "RFDETRSegXLarge"),
+            ("rf-detr-seg-xxlarge.pt", "RFDETRSeg2XLarge"),
+            ("rf-detr-seg-preview.pt", "RFDETRSegPreview"),
+        ],
+    )
+    def test_characterization_infers_correct_class_namespace(
+        self,
+        tmp_path: Path,
+        pretrain_weights: str,
+        patch_target: str,
+    ) -> None:
+        """Namespace-style args: correct subclass is called for each model size."""
+        result, mock_cls = _call_from_checkpoint(
+            _ns(pretrain_weights), tmp_path / "ckpt.pth", f"rfdetr.variants.{patch_target}"
+        )
+
+        mock_cls.assert_called_once()
+        call_kwargs = mock_cls.call_args.kwargs
+        assert call_kwargs.get("num_classes") == 80
+        assert call_kwargs.get("pretrain_weights") == str(tmp_path / "ckpt.pth")
+        assert result is mock_cls.return_value
+
+    @pytest.mark.parametrize(
+        "missing_value",
+        [
+            pytest.param("none", id="bare-none"),
+            pytest.param("null", id="bare-null"),
+            pytest.param("", id="empty"),
+            pytest.param("  None  ", id="whitespace-None"),
+            pytest.param("  ", id="whitespace-only"),
+            pytest.param(" null ", id="whitespace-null"),
+            pytest.param(None, id="python-None"),
+        ],
+    )
+    def test_namespace_args_falls_back_to_checkpoint_filename_when_pretrain_weights_missing(
+        self, tmp_path: Path, missing_value: str | None
+    ) -> None:
+        """Namespace args: filename fallback fires when pretrain_weights is unset-like."""
+        ckpt = _ns(missing_value)  # type: ignore[arg-type]
+        _, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "rf-detr-small.pth", "rfdetr.variants.RFDETRSmall")
+        mock_cls.assert_called_once()
+        assert mock_cls.call_args.kwargs["num_classes"] == 80
+
+
+# ---------------------------------------------------------------------------
+# Dict args (PTL / converted checkpoints)
+# ---------------------------------------------------------------------------
+
+
+class TestFromCheckpointDictArgs:
+    """from_checkpoint with dict-style args (PTL or convert_legacy_checkpoint output)."""
+
+    @pytest.mark.parametrize(
+        ("pretrain_weights, patch_target"),
+        [
+            ("rf-detr-small.pth", "RFDETRSmall"),
+            ("rf-detr-base.pth", "RFDETRBase"),
+        ],
+    )
+    def test_characterization_infers_correct_class_dict(
+        self,
+        tmp_path: Path,
+        pretrain_weights: str,
+        patch_target: str,
+    ) -> None:
+        """Dict-style args: correct subclass is called without AttributeError."""
+        _, mock_cls = _call_from_checkpoint(
+            _dict(pretrain_weights), tmp_path / "ckpt.pth", f"rfdetr.variants.{patch_target}"
+        )
+
+        mock_cls.assert_called_once()
+        call_kwargs = mock_cls.call_args.kwargs
+        assert call_kwargs.get("num_classes") == 80
+
+    def test_characterization_dict_args_missing_num_classes_uses_default(self, tmp_path: Path) -> None:
+        """Dict args without num_classes: constructor is called without num_classes kwarg."""
+        ckpt = {"args": {"pretrain_weights": "rf-detr-small.pth"}}
+        _, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "ckpt.pth", "rfdetr.variants.RFDETRSmall")
+
+        call_kwargs = mock_cls.call_args.kwargs
+        assert "num_classes" not in call_kwargs
+
+
+# ---------------------------------------------------------------------------
+# Edge cases
+# ---------------------------------------------------------------------------
+
+
+class TestFromCheckpointEdgeCases:
+    """Edge-case handling in from_checkpoint."""
+
+    def test_nonexistent_path_raises_file_not_found(self, tmp_path: Path) -> None:
+        """from_checkpoint raises FileNotFoundError when path does not exist."""
+        with pytest.raises(FileNotFoundError):
+            RFDETR.from_checkpoint(tmp_path / "nope.pth")
+
+    def test_directory_path_raises_os_error(self, tmp_path: Path) -> None:
+        """from_checkpoint raises OSError when path is a directory, not a file."""
+        with pytest.raises((OSError, IsADirectoryError)):
+            RFDETR.from_checkpoint(tmp_path)
+
+    def test_characterization_unknown_pretrain_weights_raises_value_error(self, tmp_path: Path) -> None:
+        """Unrecognised pretrain_weights name raises a descriptive ValueError."""
+        ckpt = _ns("/my/custom/finetuned.pth")
+        with patch("rfdetr.detr.torch.load", return_value=ckpt):
+            with pytest.raises(ValueError, match="Could not infer model class"):
+                RFDETR.from_checkpoint(tmp_path / "ckpt.pth")
+
+    def test_filename_fallback_unrecognized_name_raises_value_error(self, tmp_path: Path) -> None:
+        """ValueError fires via filename-fallback path when filename has no known model token."""
+        ckpt = {"args": {"pretrain_weights": "none", "num_classes": 80}}
+        with patch("rfdetr.detr.torch.load", return_value=ckpt):
+            with pytest.raises(ValueError, match="Could not infer model class"):
+                RFDETR.from_checkpoint(tmp_path / "finetuned.pth")
+
+    @pytest.mark.skipif(_IS_RFDETR_PLUS_AVAILABLE, reason="rfdetr_plus is installed — guard not active")
+    def test_filename_fallback_xlarge_without_plus_raises_import_error(self, tmp_path: Path) -> None:
+        """ImportError fires via filename-fallback path when rfdetr_plus is absent."""
+        ckpt = {"args": {"pretrain_weights": "none", "num_classes": 80}}
+        with patch("rfdetr.detr.torch.load", return_value=ckpt):
+            with pytest.raises(ImportError):
+                RFDETR.from_checkpoint(tmp_path / "rf-detr-xlarge-starter.pth")
+
+    def test_characterization_missing_args_key_raises_key_error(self, tmp_path: Path) -> None:
+        """Checkpoint without 'args' key raises KeyError."""
+        ckpt = {"model": {}}
+        with patch("rfdetr.detr.torch.load", return_value=ckpt):
+            with pytest.raises(KeyError):
+                RFDETR.from_checkpoint(tmp_path / "ckpt.pth")
+
+    def test_characterization_callable_on_subclass(self, tmp_path: Path) -> None:
+        """from_checkpoint can be called on a concrete subclass (RFDETRSmall)."""
+        mock_instance = MagicMock()
+        with (
+            patch("rfdetr.detr.torch.load", return_value=_ns("rf-detr-small.pth")),
+            patch("rfdetr.variants.RFDETRSmall") as mock_cls,
+        ):
+            mock_cls.return_value = mock_instance
+            result = RFDETRSmall.from_checkpoint(tmp_path / "ckpt.pth")
+
+        assert result is mock_instance
+        mock_cls.assert_called_once()
+
+    def test_characterization_extra_kwargs_forwarded(self, tmp_path: Path) -> None:
+        """Extra **kwargs are forwarded to the model constructor."""
+        _, mock_cls = _call_from_checkpoint(
+            _ns("rf-detr-small.pth"),
+            tmp_path / "ckpt.pth",
+            "rfdetr.variants.RFDETRSmall",
+            resolution=640,
+        )
+        call_kwargs = mock_cls.call_args.kwargs
+        assert call_kwargs.get("resolution") == 640
+
+    def test_characterization_pretrain_weights_in_kwargs_is_overridden(self, tmp_path: Path) -> None:
+        """pretrain_weights passed in **kwargs is silently overridden by the checkpoint path."""
+        _, mock_cls = _call_from_checkpoint(
+            _ns("rf-detr-small.pth"),
+            tmp_path / "ckpt.pth",
+            "rfdetr.variants.RFDETRSmall",
+            pretrain_weights="/should/be/overridden.pth",
+        )
+        call_kwargs = mock_cls.call_args.kwargs
+        assert call_kwargs["pretrain_weights"] == str(tmp_path / "ckpt.pth")
+
+    def test_characterization_caller_num_classes_overrides_checkpoint(self, tmp_path: Path) -> None:
+        """Caller-supplied num_classes takes precedence over the checkpoint's stored value."""
+        _, mock_cls = _call_from_checkpoint(
+            _ns("rf-detr-small.pth", num_classes=80),
+            tmp_path / "ckpt.pth",
+            "rfdetr.variants.RFDETRSmall",
+            num_classes=5,
+        )
+        call_kwargs = mock_cls.call_args.kwargs
+        assert call_kwargs["num_classes"] == 5
+
+    def test_checkpoint_model_config_forwarded_to_constructor(self, tmp_path: Path) -> None:
+        """Reload should preserve schema-dependent model config from PTL ``.pth`` checkpoints."""
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-keypoint-preview-xlarge.pth", "num_classes": 1},
+            "model_name": "RFDETRKeypointPreview",
+            "model_config": {
+                "num_keypoints_per_class": [0, 17],
+                "use_grouppose_keypoints": True,
+                "dual_projector": True,
+                "pretrain_weights": "/old/path.pth",
+            },
+        }
+        _, mock_cls = _call_from_checkpoint(
+            ckpt,
+            tmp_path / "checkpoint_best_total.pth",
+            "rfdetr.variants.RFDETRKeypointPreview",
+        )
+
+        call_kwargs = mock_cls.call_args.kwargs
+        assert call_kwargs["num_keypoints_per_class"] == [0, 17]
+        assert call_kwargs["use_grouppose_keypoints"] is True
+        assert call_kwargs["dual_projector"] is True
+        assert call_kwargs["num_classes"] == 1
+        assert call_kwargs["pretrain_weights"] == str(tmp_path / "checkpoint_best_total.pth")
+
+    @pytest.mark.skipif(_IS_RFDETR_PLUS_AVAILABLE, reason="rfdetr_plus is installed — guard not active")
+    def test_characterization_xlarge_without_plus_raises_import_error(self, tmp_path: Path) -> None:
+        """Xlarge checkpoint without rfdetr_plus raises ImportError instead of wrong class."""
+        for weights in ("rf-detr-xlarge.pth", "rf-detr-xxlarge.pth"):
+            ckpt = _ns(weights)
+            with patch("rfdetr.detr.torch.load", return_value=ckpt):
+                with pytest.raises(ImportError):
+                    RFDETR.from_checkpoint(tmp_path / "ckpt.pth")
+
+    def test_trust_gate_rejects_custom_class_by_default(self, tmp_path: Path) -> None:
+        """from_checkpoint raises RuntimeError for custom-class checkpoints without trust_checkpoint=True.
+
+        Scenario: user calls from_checkpoint on a file containing an unrecognised Python class.
+        The default trust_checkpoint=False must reject it to prevent arbitrary code execution.
+        """
+        ckpt_path = tmp_path / "custom_obj.pth"
+        torch.save({"model": {}, "args": _CustomObj(), "model_name": "RFDETRSmall"}, ckpt_path)
+
+        with pytest.raises(RuntimeError, match="trust_checkpoint=True"):
+            RFDETR.from_checkpoint(ckpt_path)
+
+
+# ---------------------------------------------------------------------------
+# Deprecated class instantiation
+# ---------------------------------------------------------------------------
+
+
+class TestDeprecatedClassInstantiation:
+    """Deprecated model classes emit deprecation warnings on instantiation."""
+
+    @pytest.mark.parametrize(
+        ("cls_name, import_path"),
+        [
+            ("RFDETRBase", "rfdetr.variants.RFDETRBase"),
+            ("RFDETRLargeDeprecated", "rfdetr.variants.RFDETRLargeDeprecated"),
+            ("RFDETRSegPreview", "rfdetr.variants.RFDETRSegPreview"),
+        ],
+    )
+    def test_direct_instantiation_is_allowed(self, cls_name: str, import_path: str) -> None:
+        """Direct instantiation of a deprecated class does not raise RuntimeError."""
+        import importlib
+
+        module_path, attr = import_path.rsplit(".", 1)
+        module = importlib.import_module(module_path)
+        cls = getattr(module, attr)
+        with patch("rfdetr.detr.RFDETR.__init__", return_value=None):
+            model = cls()
+        assert model.__class__.__name__ == cls_name
+
+    @pytest.mark.parametrize("pretrain_weights", ["rf-detr-base.pth", "rf-detr-seg-preview.pt"])
+    def test_from_checkpoint_resolves_deprecated_class(
+        self,
+        tmp_path: Path,
+        pretrain_weights: str,
+    ) -> None:
+        """from_checkpoint still resolves deprecated classes without KeyError on minimal mocked checkpoints."""
+        ckpt = _ns(pretrain_weights)
+        with (
+            patch("rfdetr.detr.torch.load", return_value=ckpt),
+            patch("rfdetr.detr.RFDETR.__init__", return_value=None),
+        ):
+            model = RFDETR.from_checkpoint(tmp_path / "ckpt.pth")
+        assert model.__class__.__name__ in {"RFDETRBase", "RFDETRSegPreview"}
+
+
+# ---------------------------------------------------------------------------
+# model_name in checkpoint (#887)
+# ---------------------------------------------------------------------------
+
+
+def _ckpt_with_model_name(model_name: str, num_classes: int = 80) -> dict:
+    """Fake checkpoint with model_name key (new format)."""
+    return {
+        "args": {"pretrain_weights": "rf-detr-small.pth", "num_classes": num_classes},
+        "model_name": model_name,
+    }
+
+
+class TestFromCheckpointModelName:
+    """from_checkpoint uses model_name when present in checkpoint."""
+
+    @pytest.mark.parametrize(
+        ("model_name, patch_target"),
+        [
+            ("RFDETRNano", "RFDETRNano"),
+            ("RFDETRSmall", "RFDETRSmall"),
+            ("RFDETRMedium", "RFDETRMedium"),
+            ("RFDETRLarge", "RFDETRLarge"),
+            ("RFDETRKeypointPreview", "RFDETRKeypointPreview"),
+            ("RFDETRBase", "RFDETRBase"),
+            ("RFDETRSegNano", "RFDETRSegNano"),
+            ("RFDETRSegPreview", "RFDETRSegPreview"),
+            ("RFDETRSegSmall", "RFDETRSegSmall"),
+            ("RFDETRSegMedium", "RFDETRSegMedium"),
+            ("RFDETRSegLarge", "RFDETRSegLarge"),
+            ("RFDETRSegXLarge", "RFDETRSegXLarge"),
+            ("RFDETRSeg2XLarge", "RFDETRSeg2XLarge"),
+        ],
+    )
+    def test_model_name_resolves_correct_class(self, tmp_path: Path, model_name: str, patch_target: str) -> None:
+        """model_name in checkpoint maps directly to the correct subclass."""
+        result, mock_cls = _call_from_checkpoint(
+            _ckpt_with_model_name(model_name), tmp_path / "ckpt.pth", f"rfdetr.variants.{patch_target}"
+        )
+        mock_cls.assert_called_once()
+        assert result is mock_cls.return_value
+
+    def test_model_name_takes_priority_over_pretrain_weights(self, tmp_path: Path) -> None:
+        """model_name is used even when pretrain_weights points to a different size."""
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-nano.pth", "num_classes": 80},
+            "model_name": "RFDETRLarge",
+        }
+        _, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "ckpt.pth", "rfdetr.variants.RFDETRLarge")
+        mock_cls.assert_called_once()
+
+    def test_falls_back_to_pretrain_weights_without_model_name(self, tmp_path: Path) -> None:
+        """Old checkpoints without model_name still work via pretrain_weights parsing."""
+        ckpt = _dict("rf-detr-small.pth")
+        assert "model_name" not in ckpt
+        _, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "ckpt.pth", "rfdetr.variants.RFDETRSmall")
+        mock_cls.assert_called_once()
+
+    @pytest.mark.parametrize(
+        "missing_value",
+        [
+            pytest.param("none", id="bare-none"),
+            pytest.param("null", id="bare-null"),
+            pytest.param("", id="empty"),
+            pytest.param("  None  ", id="whitespace-None"),
+            pytest.param("  ", id="whitespace-only"),
+            pytest.param(" null ", id="whitespace-null"),
+            pytest.param(None, id="python-None"),
+        ],
+    )
+    def test_falls_back_to_checkpoint_filename_when_pretrain_weights_missing(
+        self, tmp_path: Path, missing_value: str | None
+    ) -> None:
+        """When pretrain_weights is missing-like, from_checkpoint infers class from checkpoint filename."""
+        ckpt = {"args": {"pretrain_weights": missing_value, "num_classes": 80}}
+        _, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "rf-detr-small.pth", "rfdetr.variants.RFDETRSmall")
+        mock_cls.assert_called_once()
+        assert mock_cls.call_args.kwargs["num_classes"] == 80
+
+    def test_unknown_model_name_falls_back_to_pretrain_weights(self, tmp_path: Path) -> None:
+        """Unrecognised model_name falls back to pretrain_weights parsing."""
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-small.pth", "num_classes": 80},
+            "model_name": "UnknownModel",
+        }
+        _, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "ckpt.pth", "rfdetr.variants.RFDETRSmall")
+        mock_cls.assert_called_once()
+
+    def test_model_name_with_whitespace_is_stripped(self, tmp_path: Path) -> None:
+        """Leading/trailing whitespace in model_name is stripped before class resolution."""
+        ckpt = _ckpt_with_model_name("  RFDETRSmall  ")
+        _, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "ckpt.pth", "rfdetr.variants.RFDETRSmall")
+        mock_cls.assert_called_once()
+
+    @pytest.mark.parametrize(
+        "model_name, expected_class",
+        [
+            ("RFDETRBase", "RFDETRBase"),
+            ("RFDETRSegPreview", "RFDETRSegPreview"),
+        ],
+    )
+    def test_model_name_deprecated_class_resolves_and_instantiates(
+        self, tmp_path: Path, model_name: str, expected_class: str
+    ) -> None:
+        """from_checkpoint resolves deprecated model_name values and instantiates the resolved class."""
+        ckpt = _ckpt_with_model_name(model_name)
+        with (
+            patch("rfdetr.detr.torch.load", return_value=ckpt),
+            patch("rfdetr.detr.RFDETR.__init__", return_value=None),
+        ):
+            model = RFDETR.from_checkpoint(tmp_path / "ckpt.pth")
+        assert model.__class__.__name__ == expected_class
+
+    def test_large_deprecated_model_name_resolves_to_deprecated_class(self, tmp_path: Path) -> None:
+        """Checkpoints saved with model_name='RFDETRLargeDeprecated' must load as RFDETRLargeDeprecated.
+
+        Before the fix, RFDETRLargeDeprecated was absent from _name_map; the substring matcher would pick RFDETRLarge,
+        which fails with a pydantic literal_error when the saved model_config carries encoder='dinov2_windowed_base'
+        (only valid for the deprecated Large configuration).
+        """
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-large.pth", "num_classes": 80},
+            "model_name": "RFDETRLargeDeprecated",
+            "model_config": {
+                "encoder": "dinov2_windowed_base",
+                "projector_scale": "P4",
+            },
+        }
+        result, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "ckpt.pth", "rfdetr.variants.RFDETRLargeDeprecated")
+        mock_cls.assert_called_once()
+        assert result is mock_cls.return_value
+
+    @pytest.mark.skipif(_IS_RFDETR_PLUS_AVAILABLE, reason="rfdetr_plus is installed — guard not active")
+    @pytest.mark.parametrize("model_name", ["RFDETRXLarge", "RFDETR2XLarge"])
+    def test_plus_model_name_without_plus_raises_import_error(self, tmp_path: Path, model_name: str) -> None:
+        """Plus checkpoints using model_name raise install guidance without rfdetr_plus."""
+        ckpt = {
+            "args": {"pretrain_weights": "", "num_classes": 80},
+            "model_name": model_name,
+        }
+        with patch("rfdetr.detr.torch.load", return_value=ckpt):
+            with pytest.raises(ImportError, match="rfdetr_plus package"):
+                RFDETR.from_checkpoint(tmp_path / "ckpt.pth")
+
+
+# ---------------------------------------------------------------------------
+# num_classes provenance (fine-tuning a from_checkpoint model on a new dataset)
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def args_only_checkpoint(tmp_path: Path) -> Path:
+    """Minimal checkpoint with num_classes in args only; model_config carries no num_classes key.
+
+    Covers the legacy checkpoint format where num_classes is embedded in the args dict rather
+    than in model_config.  from_checkpoint extracts it via the args path (detr.py:454-457) and
+    injects it into constructor_kwargs — this fixture verifies that path also clears the
+    Pydantic provenance marker.  Only exercises the args-injection path; model_config path
+    covered by ``two_class_checkpoint``.
+
+    Args:
+        tmp_path: Pytest temporary directory.
+
+    Returns:
+        Path to the saved checkpoint file.
+    """
+    path = tmp_path / "small_two_class_args_only.pth"
+    torch.save(
+        {
+            "model": {"class_embed.bias": torch.zeros(3)},
+            "model_name": "RFDETRSmall",
+            "model_config": {},
+            "args": {"class_names": ["cat", "dog"], "num_classes": 2},
+        },
+        path,
+    )
+    return path
+
+
+@pytest.fixture
+def two_class_checkpoint(tmp_path: Path) -> Path:
+    """Save a minimal synthetic 2-class checkpoint to disk (no downloads, no real weights).
+
+    Follows the lightweight checkpoint pattern used elsewhere in the suite (``test_detr_shim``,
+    ``test_load_pretrain_weights``): write only what ``from_checkpoint``/``load_pretrain_weights`` actually inspect —
+    the ``class_embed.bias`` tensor sized for 2 classes + background, plus the metadata used to resolve the model
+    (``model_name``) and the class count (``model_config`` carrying ``num_classes=2``).  A *non-default*
+    ``num_classes`` is what trips the user-override guards, so it is written explicitly rather than relying on a
+    published checkpoint (whose default 90 would not trip them).  ``from_checkpoint`` still builds a real model from
+    this, which is what the provenance and head-shape assertions exercise.
+
+    Args:
+        tmp_path: Pytest temporary directory.
+
+    Returns:
+        Path to the saved checkpoint file.
+    """
+    path = tmp_path / "small_two_class.pth"
+    torch.save(
+        {
+            "model": {"class_embed.bias": torch.zeros(3)},
+            "model_name": "RFDETRSmall",
+            "model_config": {"num_classes": 2},
+            "args": {"class_names": ["cat", "dog"]},
+        },
+        path,
+    )
+    return path
+
+
+class TestFromCheckpointNumClassesProvenance:
+    """Checkpoint-derived num_classes must not be treated as a user override.
+
+    Regression tests for https://github.com/roboflow/rf-detr/issues/1092: ``from_checkpoint`` copies ``num_classes``
+    out of the checkpoint into the constructor kwargs, which used to mark the field as explicitly user-set.  Both
+    provenance guards (``RFDETR._align_num_classes_from_dataset`` and the head re-init logic in
+    ``rfdetr.models.weights.load_pretrain_weights``) then refused to adapt the detection head to a new dataset's
+    class count, breaking fine-tuning from a checkpoint.
+    """
+
+    def test_checkpoint_num_classes_is_not_marked_user_set(self, two_class_checkpoint: Path) -> None:
+        """from_checkpoint adopts the checkpoint class count without warning about pretrained weights."""
+        with warnings.catch_warnings():
+            warnings.filterwarnings("error", category=PretrainWeightsCompatibilityWarning)
+            model = RFDETR.from_checkpoint(two_class_checkpoint)
+
+        assert model.model_config.num_classes == 2
+        assert model.model.model.class_embed.bias.shape[0] == 3, "Head must match checkpoint (2 classes + background)."
+        assert "num_classes" not in model.model_config.model_fields_set, (
+            "Checkpoint-derived num_classes must not be recorded as explicitly user-set; "
+            "otherwise train() refuses to align the head to a new dataset's class count."
+        )
+
+    def test_train_alignment_adapts_head_to_new_dataset(
+        self, two_class_checkpoint: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Fine-tuning a from_checkpoint model on a dataset with a different class count adapts the head."""
+        model = RFDETR.from_checkpoint(two_class_checkpoint)
+        monkeypatch.setattr(RFDETR, "_detect_num_classes_for_training", staticmethod(lambda *a, **k: 5))
+
+        model._align_num_classes_from_dataset("<five-class-dataset>")
+
+        assert model.model_config.num_classes == 5
+        assert model.model.args.num_classes == 5
+        # train() rebuilds the model from model_config (inside RFDETRModelModule), reloading the checkpoint
+        # weights with the aligned class count; the rebuilt head must adopt the dataset class count.
+        rebuilt = model.get_model(model.model_config)
+        assert rebuilt.model.class_embed.bias.shape[0] == 6, "Rebuilt head must have 5 classes + background."
+
+    def test_explicit_num_classes_kwarg_still_wins(
+        self,
+        two_class_checkpoint: Path,
+        monkeypatch: pytest.MonkeyPatch,
+        caplog: pytest.LogCaptureFixture,
+    ) -> None:
+        """An explicit num_classes kwarg to from_checkpoint stays authoritative over the dataset."""
+        model = RFDETR.from_checkpoint(two_class_checkpoint, num_classes=7)
+
+        assert model.model_config.num_classes == 7
+        assert "num_classes" in model.model_config.model_fields_set
+        assert model.model.model.class_embed.bias.shape[0] == 8, "Head must expand to 7 classes + background."
+
+        monkeypatch.setattr(RFDETR, "_detect_num_classes_for_training", staticmethod(lambda *a, **k: 5))
+        monkeypatch.setattr(detr_logger, "propagate", True)
+        with caplog.at_level(logging.WARNING, logger="rf-detr"):
+            model._align_num_classes_from_dataset("<five-class-dataset>")
+
+        assert model.model_config.num_classes == 7, "Explicit user num_classes must be preserved."
+        assert any("Using the model's configured value" in record.message for record in caplog.records)
+
+    def test_checkpoint_num_classes_from_args_not_marked_user_set(self, args_only_checkpoint: Path) -> None:
+        """num_classes injected from checkpoint args (not model_config) is cleared from model_fields_set."""
+        model = RFDETR.from_checkpoint(args_only_checkpoint)
+
+        assert model.model_config.num_classes == 2
+        assert "num_classes" not in model.model_config.model_fields_set, (
+            "num_classes from checkpoint args must not be recorded as explicitly user-set; "
+            "otherwise train() refuses to adapt the head to a new dataset's class count."
+        )
+
+    def test_explicit_default_num_classes_pins_head(
+        self,
+        two_class_checkpoint: Path,
+        monkeypatch: pytest.MonkeyPatch,
+        caplog: pytest.LogCaptureFixture,
+    ) -> None:
+        """Passing num_classes equal to the ModelConfig default still pins the detection head.
+
+        An explicit num_classes is honored regardless of whether it equals the class default:
+        ``_align_num_classes_from_dataset`` keys off whether the field was set, not whether the
+        value differs from the default, so the dataset count cannot silently override it.  This
+        guards against re-introducing the ``value != default`` clause, whose asymmetric behavior
+        (default silently aligned, non-default preserved) was the bug this test now pins.
+        """
+        model = RFDETR.from_checkpoint(two_class_checkpoint)
+        default_nc = type(model.model_config).model_fields["num_classes"].default
+        # Simulate calling from_checkpoint(path, num_classes=<default>):
+        # assigning the field adds "num_classes" to model_fields_set automatically (Pydantic v2).
+        model.model_config.num_classes = default_nc
+
+        assert "num_classes" in model.model_config.model_fields_set
+        monkeypatch.setattr(RFDETR, "_detect_num_classes_for_training", staticmethod(lambda *a, **k: 5))
+        monkeypatch.setattr(detr_logger, "propagate", True)
+        with caplog.at_level(logging.WARNING, logger="rf-detr"):
+            model._align_num_classes_from_dataset("<five-class-dataset>")
+
+        assert model.model_config.num_classes == default_nc, (
+            "Explicitly passing the ModelConfig default for num_classes must pin the head; "
+            "the dataset class count must not silently override an explicit user setting."
+        )
+        assert any("Using the model's configured value" in record.message for record in caplog.records)
+
+    def test_explicit_default_num_classes_via_from_checkpoint_integrated(
+        self,
+        two_class_checkpoint: Path,
+        monkeypatch: pytest.MonkeyPatch,
+        caplog: pytest.LogCaptureFixture,
+    ) -> None:
+        """from_checkpoint(path, num_classes=<default>) pins head via the integrated code path.
+
+        Unlike test_explicit_default_num_classes_pins_head which simulates the explicit-default scenario via post-
+        construction assignment, this test calls from_checkpoint directly with num_classes=default_nc.  A regression in
+        how from_checkpoint passes num_classes into the constructor would be caught here but not by the proxy-based
+        test.
+        """
+        default_nc = RFDETRSmall._model_config_class.model_fields["num_classes"].default
+        model = RFDETR.from_checkpoint(two_class_checkpoint, num_classes=default_nc)
+
+        assert model.model_config.num_classes == default_nc
+        assert "num_classes" in model.model_config.model_fields_set, (
+            "from_checkpoint with explicit num_classes must keep it in model_fields_set; "
+            "only checkpoint-derived num_classes should be cleared."
+        )
+
+        monkeypatch.setattr(RFDETR, "_detect_num_classes_for_training", staticmethod(lambda *a, **k: 5))
+        monkeypatch.setattr(detr_logger, "propagate", True)
+        with caplog.at_level(logging.WARNING, logger="rf-detr"):
+            model._align_num_classes_from_dataset("<five-class-dataset>")
+
+        assert model.model_config.num_classes == default_nc, (
+            "Head must remain pinned at default_nc after alignment; "
+            "from_checkpoint-supplied num_classes must not be silently overridden."
+        )
+        assert any("Using the model's configured value" in record.message for record in caplog.records)
+
+    def test_equal_class_count_does_not_rebuild_head(
+        self, two_class_checkpoint: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Checkpoint and dataset sharing the same class count leaves the head unchanged."""
+        model = RFDETR.from_checkpoint(two_class_checkpoint)
+        original_bias_shape = model.model.model.class_embed.bias.shape
+        monkeypatch.setattr(RFDETR, "_detect_num_classes_for_training", staticmethod(lambda *a, **k: 2))
+
+        model._align_num_classes_from_dataset("<two-class-dataset>")
+
+        assert model.model_config.num_classes == 2
+        assert model.model.model.class_embed.bias.shape == original_bias_shape, (
+            "Head must not be rebuilt when dataset class count matches checkpoint class count."
+        )
+
+
+# ---------------------------------------------------------------------------
+# Weight-based schema inference
+# ---------------------------------------------------------------------------
+
+
+def _make_kp_active_mask(schema: list[int]) -> torch.Tensor:
+    """Build a bool _kp_active_mask tensor encoding *schema* (mirrors LwDetr._create_kp_active_mask).
+
+    Args:
+        schema: Keypoints-per-class list, e.g. ``[0, 33]`` for background + 33-kp class.
+
+    Returns:
+        Bool tensor of shape ``[len(schema), max(schema)]`` with True in active keypoint slots.
+    """
+    if not schema or max(schema) == 0:
+        return torch.zeros(0, 0, dtype=torch.bool)
+    max_kp = max(schema)
+    mask = torch.zeros(len(schema), max_kp, dtype=torch.bool)
+    for idx, n_kp in enumerate(schema):
+        mask[idx, :n_kp] = True
+    return mask
+
+
+class TestFromCheckpointWeightInference:
+    """from_checkpoint infers schema from checkpoint weights when model_config is absent or stale.
+
+    Regression tests for the bug where a fine-tuned 33-kp keypoint model loaded with the COCO default [0, 17] schema
+    because model_config["num_keypoints_per_class"] was never updated from the default before the checkpoint was saved.
+    The authoritative schema is embedded in the checkpoint weights via the _kp_active_mask buffer; from_checkpoint now
+    reads it directly.
+    """
+
+    def test_infers_keypoint_schema_from_kp_active_mask(self, tmp_path: Path) -> None:
+        """Stale model_config kp schema [0, 17] is overridden by weight-inferred [0, 33]."""
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-keypoint-preview-xlarge.pth"},
+            "model_name": "RFDETRKeypointPreview",
+            "model_config": {"num_keypoints_per_class": [0, 17]},
+            "model": {"_kp_active_mask": _make_kp_active_mask([0, 33])},
+        }
+        _, mock_cls = _call_from_checkpoint(
+            ckpt, tmp_path / "checkpoint_best_total.pth", "rfdetr.variants.RFDETRKeypointPreview"
+        )
+
+        assert mock_cls.call_args.kwargs["num_keypoints_per_class"] == [0, 33]
+
+    def test_infers_keypoint_schema_when_model_config_absent(self, tmp_path: Path) -> None:
+        """num_keypoints_per_class is inferred from _kp_active_mask when model_config is missing."""
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-keypoint-preview-xlarge.pth"},
+            "model_name": "RFDETRKeypointPreview",
+            "model": {"_kp_active_mask": _make_kp_active_mask([0, 33])},
+        }
+        _, mock_cls = _call_from_checkpoint(
+            ckpt, tmp_path / "checkpoint_best_total.pth", "rfdetr.variants.RFDETRKeypointPreview"
+        )
+
+        assert mock_cls.call_args.kwargs["num_keypoints_per_class"] == [0, 33]
+
+    def test_user_kwarg_wins_over_weight_inferred_keypoint_schema(self, tmp_path: Path) -> None:
+        """Explicit num_keypoints_per_class kwarg overrides weight-inferred [0, 33] schema."""
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-keypoint-preview-xlarge.pth"},
+            "model_name": "RFDETRKeypointPreview",
+            "model": {"_kp_active_mask": _make_kp_active_mask([0, 33])},
+        }
+        _, mock_cls = _call_from_checkpoint(
+            ckpt,
+            tmp_path / "checkpoint_best_total.pth",
+            "rfdetr.variants.RFDETRKeypointPreview",
+            num_keypoints_per_class=[0, 17],
+        )
+
+        assert mock_cls.call_args.kwargs["num_keypoints_per_class"] == [0, 17]
+
+    def test_infers_num_classes_from_class_embed_weight(self, tmp_path: Path) -> None:
+        """Stale model_config num_classes=90 is overridden by class_embed.weight shape inference."""
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-small.pth"},
+            "model_name": "RFDETRSmall",
+            "model_config": {"num_classes": 90},
+            "model": {"class_embed.weight": torch.zeros(3, 256)},
+        }
+        _, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "checkpoint_best_total.pth", "rfdetr.variants.RFDETRSmall")
+
+        assert mock_cls.call_args.kwargs["num_classes"] == 2
+
+    def test_user_kwarg_wins_over_weight_inferred_num_classes(self, tmp_path: Path) -> None:
+        """Explicit num_classes kwarg overrides weight-inferred value from class_embed.weight."""
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-small.pth"},
+            "model_name": "RFDETRSmall",
+            "model": {"class_embed.weight": torch.zeros(3, 256)},
+        }
+        _, mock_cls = _call_from_checkpoint(
+            ckpt,
+            tmp_path / "checkpoint_best_total.pth",
+            "rfdetr.variants.RFDETRSmall",
+            num_classes=90,
+        )
+
+        assert mock_cls.call_args.kwargs["num_classes"] == 90
+
+    def test_infers_schema_from_ptl_ckpt_state_dict_format(self, tmp_path: Path) -> None:
+        """Weight inference works for PTL-native .ckpt format (state_dict with model.
+
+        prefix).
+        """
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-keypoint-preview-xlarge.pth"},
+            "model_name": "RFDETRKeypointPreview",
+            "state_dict": {
+                "model._kp_active_mask": _make_kp_active_mask([0, 33]),
+                "model.class_embed.weight": torch.zeros(3, 256),
+            },
+        }
+        _, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "checkpoint.ckpt", "rfdetr.variants.RFDETRKeypointPreview")
+
+        call_kwargs = mock_cls.call_args.kwargs
+        assert call_kwargs["num_keypoints_per_class"] == [0, 33]
+        assert call_kwargs["num_classes"] == 2
+
+    def test_consistent_checkpoint_produces_no_override(self, tmp_path: Path) -> None:
+        """When model_config and weights agree, weight inference leaves constructor_kwargs unchanged."""
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-keypoint-preview-xlarge.pth"},
+            "model_name": "RFDETRKeypointPreview",
+            "model_config": {"num_keypoints_per_class": [0, 33], "num_classes": 2},
+            "model": {
+                "_kp_active_mask": _make_kp_active_mask([0, 33]),
+                "class_embed.weight": torch.zeros(3, 256),
+            },
+        }
+        _, mock_cls = _call_from_checkpoint(
+            ckpt, tmp_path / "checkpoint_best_total.pth", "rfdetr.variants.RFDETRKeypointPreview"
+        )
+
+        call_kwargs = mock_cls.call_args.kwargs
+        assert call_kwargs["num_keypoints_per_class"] == [0, 33]
+        assert call_kwargs["num_classes"] == 2

tests/inference/test_inference.py

@@ -0,0 +1,54 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for rfdetr.inference weight-adaptation helpers."""
+
+import pytest
+import torch
+
+from rfdetr.inference import _adapt_input_conv
+
+
+@pytest.fixture(autouse=True)
+def reset_random_seeds():
+    """Ensure reproducible random state for every test in this module."""
+    torch.manual_seed(0)
+
+
+class TestAdaptInputConv:
+    @pytest.mark.parametrize(
+        ("num_channels", "expected_shape", "expected_builder"),
+        [
+            pytest.param(3, (8, 3, 3, 3), lambda weight: weight, id="identity_3ch"),
+            pytest.param(1, (8, 1, 3, 3), lambda weight: weight.mean(dim=1, keepdim=True), id="mean_1ch"),
+            pytest.param(
+                4,
+                (8, 4, 3, 3),
+                lambda weight: torch.cat([weight, weight], dim=1)[:, :4] * (3.0 / 4.0),
+                id="tile_4ch",
+            ),
+            pytest.param(
+                6,
+                (8, 6, 3, 3),
+                lambda weight: torch.cat([weight, weight], dim=1)[:, :6] * (3.0 / 6.0),
+                id="tile_6ch",
+            ),
+            pytest.param(
+                2,
+                (8, 2, 3, 3),
+                lambda weight: weight[:, :2] * (3.0 / 2.0),
+                id="tile_2ch",
+            ),
+        ],
+    )
+    def test_adapt_input_conv(self, num_channels, expected_shape, expected_builder):
+        """Verify shape and values for each _adapt_input_conv branch."""
+        conv_weight = torch.randn(8, 3, 3, 3)
+
+        adapted_weight = _adapt_input_conv(num_channels, conv_weight)
+        expected_weight = expected_builder(conv_weight)
+
+        assert adapted_weight.shape == expected_shape
+        torch.testing.assert_close(adapted_weight, expected_weight)

tests/inference/test_inference_keypoints.py

@@ -0,0 +1,50 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Focused predict() contract tests for keypoint and non-keypoint outputs."""
+
+import numpy as np
+import PIL.Image
+import supervision as sv
+
+from .helpers import _DummyModel, _DummyRFDETR
+
+
+def test_predict_returns_supervision_keypoints() -> None:
+    """Keypoint model predictions return ``sv.KeyPoints`` with detection details."""
+    image = PIL.Image.new("RGB", (64, 48), color=(128, 128, 128))
+    model = _DummyRFDETR()
+    model.model = _DummyModel(labels=[0, 1], include_keypoints=True)
+
+    key_points = model.predict(image)
+
+    assert isinstance(key_points, sv.KeyPoints)
+    assert key_points.xy.shape == (2, 17, 2)
+    assert key_points.keypoint_confidence.shape == (2, 17)
+    assert key_points.data["xyxy"].shape == (2, 4)
+    assert key_points.detection_confidence.shape == (2,)
+    assert np.isfinite(key_points.xy).all()
+    assert np.isfinite(key_points.keypoint_confidence).all()
+    assert "keypoint_precision_cholesky" in key_points.data
+    keypoint_precision = key_points.data["keypoint_precision_cholesky"]
+    assert isinstance(keypoint_precision, np.ndarray)
+    assert keypoint_precision.shape == (2, 17, 3)
+    assert np.isfinite(keypoint_precision).all()
+    assert "source_image" in key_points.data
+    assert len(key_points.data["source_image"]) == 2
+
+
+def test_predict_default_detection_without_keypoints_unchanged() -> None:
+    """Default detection prediction keeps legacy output structure."""
+    image = PIL.Image.new("RGB", (64, 48), color=(128, 128, 128))
+    model = _DummyRFDETR()
+
+    detections = model.predict(image)
+
+    assert "keypoints" not in detections.data
+    assert not hasattr(detections, "keypoints")
+    assert "class_name" in detections.data
+    assert "source_shape" in detections.data
+    assert detections.data["source_shape"].shape[1] == 2

tests/inference/test_model_device_move.py

@@ -0,0 +1,68 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for the lazy device move running under ``torch.inference_mode()``.
+
+``predict()`` stacks ``@torch.inference_mode()`` on top of ``@_ensure_model_on_device``, so the deferred CPU-to-
+accelerator move happens while inference mode is active.  Tensors materialised under inference mode are *inference
+tensors*: they can never require gradients, so a later ``train()`` / auto-batch probe silently produces no gradients.
+The move itself must therefore always run with inference mode disabled.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Any
+
+import torch
+from torch import nn
+
+from rfdetr.detr import _move_model_context_to_device
+
+
+class _RecordingModule(nn.Module):
+    """Module whose ``to()`` records whether inference mode was active at move time."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.linear = nn.Linear(2, 2)
+        self.inference_mode_at_move: bool | None = None
+
+    def to(self, *args: Any, **kwargs: Any) -> "_RecordingModule":
+        """Record the inference-mode state instead of performing a real device move."""
+        self.inference_mode_at_move = torch.is_inference_mode_enabled()
+        return self
+
+
+class TestMoveModelContextUnderInferenceMode:
+    """The deferred device move must never materialise parameters as inference tensors."""
+
+    def test_moved_params_are_not_inference_tensors(self) -> None:
+        """A real ``.to()`` move inside ``torch.inference_mode()`` must not create inference-tensor parameters."""
+        ctx = SimpleNamespace(device=torch.device("meta"), model=nn.Linear(2, 2))
+
+        with torch.inference_mode():
+            _move_model_context_to_device(ctx)
+
+        assert not any(p.is_inference() for p in ctx.model.parameters())
+
+    def test_move_still_materializes_on_target_device(self) -> None:
+        """The inference-mode guard must not suppress the device move itself."""
+        ctx = SimpleNamespace(device=torch.device("meta"), model=nn.Linear(2, 2))
+
+        with torch.inference_mode():
+            _move_model_context_to_device(ctx)
+
+        assert all(p.device.type == "meta" for p in ctx.model.parameters())
+
+    def test_move_runs_with_inference_mode_disabled(self) -> None:
+        """The ``.to()`` call itself must observe inference mode as disabled."""
+        module = _RecordingModule()
+        ctx = SimpleNamespace(device=torch.device("meta"), model=module)
+
+        with torch.inference_mode():
+            _move_model_context_to_device(ctx)
+
+        assert module.inference_mode_at_move is False

tests/inference/test_optimize_for_inference.py

@@ -0,0 +1,532 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for RFDETR.optimize_for_inference()."""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from rfdetr.detr import RFDETR
+
+
+class _FakeModel(torch.nn.Module):
+    """Minimal nn.Module that satisfies the optimize_for_inference contract."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.linear = torch.nn.Linear(1, 1)
+
+    def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
+        return {"pred_boxes": self.linear(x[:, :1, :1, :1].squeeze(-1).squeeze(-1))}
+
+    def export(self) -> None:
+        pass
+
+
+class _FakeModelContext:
+    def __init__(self, device: torch.device | str = torch.device("cpu"), resolution: int = 28) -> None:
+        self.device = torch.device(device) if not isinstance(device, torch.device) else device
+        self.resolution = resolution
+        self.model = _FakeModel()
+        self.inference_model = None
+
+
+class _FakeRFDETR(RFDETR):
+    def maybe_download_pretrain_weights(self) -> None:
+        return None
+
+    def get_model_config(self, **kwargs) -> SimpleNamespace:
+        return SimpleNamespace(num_channels=3)
+
+    def get_model(self, config: SimpleNamespace) -> _FakeModelContext:
+        return _FakeModelContext()
+
+
+class TestOptimizeForInferenceDtype:
+    """Dtype coercion and validation tests."""
+
+    def test_string_dtype_float32_is_accepted(self) -> None:
+        """Passing dtype='float32' (str) should be coerced to torch.float32."""
+        rfdetr = _FakeRFDETR()
+
+        with patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model):
+            rfdetr.optimize_for_inference(compile=False, dtype="float32")
+
+        assert rfdetr._optimized_dtype == torch.float32
+
+    def test_string_dtype_float16_is_accepted(self) -> None:
+        """Passing dtype='float16' (str) should be coerced to torch.float16."""
+        rfdetr = _FakeRFDETR()
+
+        with patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model):
+            rfdetr.optimize_for_inference(compile=False, dtype="float16")
+
+        assert rfdetr._optimized_dtype == torch.float16
+
+    def test_torch_dtype_is_passed_through(self) -> None:
+        """Passing dtype=torch.float32 directly should work as before."""
+        rfdetr = _FakeRFDETR()
+
+        with patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model):
+            rfdetr.optimize_for_inference(compile=False, dtype=torch.float32)
+
+        assert rfdetr._optimized_dtype == torch.float32
+
+    def test_invalid_dtype_type_raises_type_error(self) -> None:
+        """Passing an invalid dtype type (e.g. int) should raise TypeError."""
+        rfdetr = _FakeRFDETR()
+
+        with pytest.raises(TypeError, match="dtype must be a torch.dtype or a string name of a dtype"):
+            rfdetr.optimize_for_inference(compile=False, dtype=42)  # type: ignore[arg-type]
+
+    def test_invalid_dtype_string_raises_type_error(self) -> None:
+        """Passing a non-existent dtype string should raise TypeError with a descriptive message."""
+        rfdetr = _FakeRFDETR()
+
+        with pytest.raises(TypeError, match="dtype must be a torch.dtype or a string name of a dtype"):
+            rfdetr.optimize_for_inference(compile=False, dtype="not_a_dtype")
+
+    def test_valid_torch_attr_that_is_not_dtype_raises_type_error(self) -> None:
+        """'Tensor' is a valid torch attribute but not a torch.dtype — should raise TypeError."""
+        rfdetr = _FakeRFDETR()
+
+        with pytest.raises(TypeError, match="dtype must be a torch.dtype or a string name of a dtype"):
+            rfdetr.optimize_for_inference(compile=False, dtype="Tensor")  # type: ignore[arg-type]
+
+    @pytest.mark.parametrize("dtype_str", ["float32", "float16", "bfloat16"])
+    def test_string_dtype_variants_are_accepted(self, dtype_str: str) -> None:
+        """Common dtype string names should be accepted and coerced to the matching torch.dtype."""
+        rfdetr = _FakeRFDETR()
+        expected = getattr(torch, dtype_str)
+
+        with patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model):
+            rfdetr.optimize_for_inference(compile=False, dtype=dtype_str)
+
+        assert rfdetr._optimized_dtype == expected
+
+
+class TestOptimizeForInferenceCudaDeviceContext:
+    """Verify that optimize_for_inference wraps operations in the correct device context."""
+
+    @pytest.mark.gpu
+    @patch("rfdetr.detr._move_model_context_to_device")
+    @patch("rfdetr.detr.deepcopy")
+    @patch("torch.cuda.device")
+    def test_cuda_device_context_manager_is_used_for_cuda_device(
+        self,
+        mock_cuda_device,
+        mock_deepcopy,
+        _mock_move_model_context_to_device,
+    ) -> None:
+        """torch.cuda.device() context should be entered when model is on CUDA."""
+        rfdetr = _FakeRFDETR()
+        # Simulate a CUDA device without actually requiring CUDA hardware
+        rfdetr.model.device = torch.device("cuda", 0)
+        mock_deepcopy.return_value = rfdetr.model.model
+
+        entered_devices: list[torch.device] = []
+
+        class _CapturingDeviceCtx:
+            def __init__(self, captured_device):
+                entered_devices.append(captured_device)
+
+            def __enter__(self):
+                return self
+
+            def __exit__(self, *args):
+                pass
+
+        mock_cuda_device.side_effect = _CapturingDeviceCtx
+        rfdetr.optimize_for_inference(compile=False, dtype=torch.float32)
+
+        assert len(entered_devices) == 1
+        assert entered_devices[0] == torch.device("cuda", 0)
+
+    def test_nullcontext_used_for_cpu_device(self) -> None:
+        """contextlib.nullcontext() should be used when model is on CPU (no CUDA init)."""
+        rfdetr = _FakeRFDETR()
+        rfdetr.model.device = torch.device("cpu")
+
+        # torch.cuda.device should NOT be called for CPU devices
+        with (
+            patch("torch.cuda.device") as mock_cuda_device,
+            patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model),
+        ):
+            rfdetr.optimize_for_inference(compile=False, dtype=torch.float32)
+
+        mock_cuda_device.assert_not_called()
+
+    @pytest.mark.gpu
+    @patch("rfdetr.detr._move_model_context_to_device")
+    @patch("rfdetr.detr.deepcopy")
+    @patch("torch.cuda.device")
+    def test_cuda_device_context_uses_model_device(
+        self,
+        mock_cuda_device,
+        mock_deepcopy,
+        _mock_move_model_context_to_device,
+    ) -> None:
+        """The device passed to torch.cuda.device() should match self.model.device."""
+        rfdetr = _FakeRFDETR()
+        expected_device = torch.device("cuda", 2)
+        rfdetr.model.device = expected_device
+        mock_deepcopy.return_value = rfdetr.model.model
+
+        captured: dict[str, torch.device] = {}
+
+        class _CapturingCtx:
+            def __init__(self, captured_device):
+                captured["device"] = captured_device
+
+            def __enter__(self):
+                return self
+
+            def __exit__(self, *args):
+                pass
+
+        mock_cuda_device.side_effect = _CapturingCtx
+        rfdetr.optimize_for_inference(compile=False)
+
+        assert captured.get("device") == expected_device
+
+
+class TestOptimizeForInferenceCompile:
+    """Tests for the compile=True path (JIT trace)."""
+
+    def test_compile_true_calls_jit_trace(self) -> None:
+        """torch.jit.trace should be called with the model and a correctly-shaped dummy input."""
+        rfdetr = _FakeRFDETR()
+        mock_traced = rfdetr.model.model
+
+        with (
+            patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model),
+            patch("torch.jit.trace", return_value=mock_traced) as mock_trace,
+        ):
+            rfdetr.optimize_for_inference(compile=True, batch_size=2)
+
+        assert mock_trace.called
+        dummy_input: torch.Tensor = mock_trace.call_args.args[1]
+        resolution = rfdetr.model.resolution
+        assert dummy_input.shape == (2, 3, resolution, resolution)
+
+    def test_compile_true_sets_compiled_flags(self) -> None:
+        """_optimized_has_been_compiled=True and _optimized_batch_size should be set after compile=True."""
+        rfdetr = _FakeRFDETR()
+
+        with (
+            patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model),
+            patch("torch.jit.trace", return_value=rfdetr.model.model),
+        ):
+            rfdetr.optimize_for_inference(compile=True, batch_size=4)
+
+        assert rfdetr._optimized_has_been_compiled is True
+        assert rfdetr._optimized_batch_size == 4
+
+    def test_compile_false_skips_jit_trace(self) -> None:
+        """torch.jit.trace should NOT be called when compile=False."""
+        rfdetr = _FakeRFDETR()
+
+        with (
+            patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model),
+            patch("torch.jit.trace") as mock_trace,
+        ):
+            rfdetr.optimize_for_inference(compile=False)
+
+        mock_trace.assert_not_called()
+        assert rfdetr._optimized_has_been_compiled is False
+        assert rfdetr._optimized_batch_size is None
+
+
+class TestOptimizeForInferenceState:
+    """Verify that optimize_for_inference correctly sets internal state flags."""
+
+    def test_is_optimized_inplace_false_before_optimization(self) -> None:
+        """is_optimized_inplace is False before any optimization is applied."""
+        rfdetr = _FakeRFDETR()
+        assert rfdetr.is_optimized_inplace is False
+
+    def test_is_optimized_flag_set(self) -> None:
+        """_is_optimized_for_inference should be True after optimization."""
+        rfdetr = _FakeRFDETR()
+
+        with patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model):
+            rfdetr.optimize_for_inference(compile=False)
+
+        assert rfdetr._is_optimized_for_inference is True
+
+    def test_inference_model_set(self) -> None:
+        """model.inference_model should be set after optimization."""
+        rfdetr = _FakeRFDETR()
+
+        with patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model):
+            rfdetr.optimize_for_inference(compile=False)
+
+        assert rfdetr.model.inference_model is not None
+
+    def test_remove_optimized_model_clears_state(self) -> None:
+        """remove_optimized_model() should clear all optimization flags."""
+        rfdetr = _FakeRFDETR()
+
+        with patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model):
+            rfdetr.optimize_for_inference(compile=False)
+
+        rfdetr.remove_optimized_model()
+
+        assert rfdetr._is_optimized_for_inference is False
+        assert rfdetr.model.inference_model is None
+        assert rfdetr._optimized_dtype is None
+        assert rfdetr._optimized_resolution is None
+        assert rfdetr._optimized_has_been_compiled is False
+        assert rfdetr._optimized_batch_size is None
+        assert rfdetr.is_optimized_inplace is False
+
+
+class TestOptimizeForInferenceInplace:
+    """Tests for the low-memory in-place optimization path."""
+
+    def test_inplace_false_keeps_deepcopy_behavior(self) -> None:
+        """The default path should still deep-copy the loaded module."""
+        rfdetr = _FakeRFDETR()
+        original_model = rfdetr.model.model
+        copied_model = _FakeModel()
+
+        with patch("rfdetr.detr.deepcopy", return_value=copied_model) as mock_deepcopy:
+            rfdetr.optimize_for_inference(compile=False)
+
+        mock_deepcopy.assert_called_once_with(original_model)
+        assert rfdetr.model.model is original_model
+        assert rfdetr.model.inference_model is copied_model
+        assert rfdetr._is_optimized_for_inference is True
+        assert rfdetr.is_optimized_inplace is False
+
+    def test_inplace_true_compile_false_does_not_deepcopy(self) -> None:
+        """Inplace=True with compile=False should use the loaded module directly."""
+        rfdetr = _FakeRFDETR()
+        original_model = rfdetr.model.model
+
+        with patch("rfdetr.detr.deepcopy") as mock_deepcopy:
+            rfdetr.optimize_for_inference(compile=False, inplace=True)
+
+        mock_deepcopy.assert_not_called()
+        assert rfdetr.model.model is None
+        assert rfdetr.model.inference_model is original_model
+        assert rfdetr._is_optimized_for_inference is True
+        assert rfdetr.is_optimized_inplace is True
+
+    def test_remove_optimized_model_after_inplace_warns_and_preserves_state(self) -> None:
+        """remove_optimized_model() after inplace optimization issues UserWarning and no-ops."""
+        rfdetr = _FakeRFDETR()
+        original_model = rfdetr.model.model
+
+        rfdetr.optimize_for_inference(compile=False, inplace=True)
+
+        with pytest.warns(UserWarning, match="no effect after inplace optimization"):
+            rfdetr.remove_optimized_model()
+
+        assert rfdetr.model.model is None
+        assert rfdetr.model.inference_model is original_model
+        assert rfdetr._is_optimized_for_inference is True
+        assert rfdetr.is_optimized_inplace is True
+
+    def test_second_optimize_after_inplace_raises_runtime_error(self) -> None:
+        """Calling optimize_for_inference() again after inplace=True raises RuntimeError."""
+        rfdetr = _FakeRFDETR()
+
+        rfdetr.optimize_for_inference(compile=False, inplace=True)
+
+        with pytest.raises(RuntimeError, match="base model has been cleared"):
+            rfdetr.optimize_for_inference(compile=False)
+
+    def test_inplace_true_default_dtype_float32_does_not_cast(self) -> None:
+        """Inplace=True with default dtype (float32) leaves weights unchanged — no casting occurs."""
+        rfdetr = _FakeRFDETR()
+        original_model = rfdetr.model.model
+        original_dtype = original_model.linear.weight.dtype
+
+        rfdetr.optimize_for_inference(compile=False, inplace=True)
+
+        assert rfdetr.model.inference_model is original_model
+        assert original_model.linear.weight.dtype == original_dtype
+        assert rfdetr._optimized_dtype == torch.float32
+        assert rfdetr._optimized_has_been_compiled is False
+        assert rfdetr._optimized_batch_size is None
+
+    @pytest.mark.parametrize(
+        "dtype",
+        [
+            pytest.param(torch.float16, id="float16"),
+            pytest.param(torch.bfloat16, id="bfloat16"),
+        ],
+    )
+    def test_inplace_true_allows_destructive_dtype_casting(self, dtype: torch.dtype) -> None:
+        """In-place optimization may cast the original module to the target dtype."""
+        rfdetr = _FakeRFDETR()
+        original_model = rfdetr.model.model
+
+        rfdetr.optimize_for_inference(compile=False, dtype=dtype, inplace=True)
+
+        assert rfdetr.model.model is None
+        assert rfdetr.model.inference_model is original_model
+        assert original_model.linear.weight.dtype == dtype
+        assert rfdetr._optimized_dtype == dtype
+
+    def test_inplace_export_failure_keeps_base_model(self) -> None:
+        """Export failure in the in-place path should not clear model.model."""
+        rfdetr = _FakeRFDETR()
+        original_model = rfdetr.model.model
+
+        with (
+            patch("rfdetr.detr.deepcopy") as mock_deepcopy,
+            patch.object(original_model, "export", side_effect=RuntimeError("export failed")),
+            pytest.raises(RuntimeError, match="export failed"),
+        ):
+            rfdetr.optimize_for_inference(compile=False, inplace=True)
+
+        mock_deepcopy.assert_not_called()
+        assert rfdetr.model.model is original_model
+        assert rfdetr.model.inference_model is None
+        assert rfdetr._is_optimized_for_inference is False
+        assert rfdetr.is_optimized_inplace is False
+
+    @pytest.mark.parametrize(
+        "dtype",
+        [
+            pytest.param(torch.int8, id="torch-int8"),
+            pytest.param("int8", id="string-int8"),
+        ],
+    )
+    def test_inplace_non_floating_dtype_raises_before_export(self, dtype: torch.dtype | str) -> None:
+        """In-place optimization rejects non-floating dtypes before mutating the base model."""
+        rfdetr = _FakeRFDETR()
+        original_model = rfdetr.model.model
+
+        with (
+            patch("rfdetr.detr.deepcopy") as mock_deepcopy,
+            patch.object(original_model, "export") as mock_export,
+            pytest.raises(ValueError, match="floating-point torch.dtype"),
+        ):
+            rfdetr.optimize_for_inference(compile=False, dtype=dtype, inplace=True)
+
+        mock_deepcopy.assert_not_called()
+        mock_export.assert_not_called()
+        assert rfdetr.model.model is original_model
+        assert rfdetr.model.inference_model is None
+        assert rfdetr._is_optimized_for_inference is False
+        assert rfdetr.is_optimized_inplace is False
+
+    def test_inplace_compile_true_raises_before_export_or_trace(self) -> None:
+        """In-place optimization rejects compile=True before mutating the base model."""
+        rfdetr = _FakeRFDETR()
+        original_model = rfdetr.model.model
+
+        with (
+            patch("rfdetr.detr.deepcopy") as mock_deepcopy,
+            patch.object(original_model, "export") as mock_export,
+            patch("torch.jit.trace") as mock_trace,
+            pytest.raises(ValueError, match="inplace=True.*compile=False"),
+        ):
+            rfdetr.optimize_for_inference(compile=True, inplace=True)
+
+        mock_deepcopy.assert_not_called()
+        mock_export.assert_not_called()
+        mock_trace.assert_not_called()
+        assert rfdetr.model.model is original_model
+        assert rfdetr.model.inference_model is None
+        assert rfdetr._is_optimized_for_inference is False
+        assert rfdetr._optimized_has_been_compiled is False
+        assert rfdetr._optimized_batch_size is None
+        assert rfdetr.is_optimized_inplace is False
+
+
+class TestOptimizeForInferenceExceptionRecovery:
+    """Verify state consistency when optimization fails mid-execution."""
+
+    def test_deepcopy_failure_leaves_clean_state(self) -> None:
+        """If deepcopy raises, inference_model should be None and _is_optimized_for_inference False."""
+        rfdetr = _FakeRFDETR()
+        # Simulate a previously-optimized state to confirm remove_optimized_model ran
+        rfdetr._is_optimized_for_inference = True
+        rfdetr.model.inference_model = rfdetr.model.model
+
+        with (
+            patch("rfdetr.detr.deepcopy", side_effect=RuntimeError("deepcopy failed")),
+            pytest.raises(RuntimeError, match="deepcopy failed"),
+        ):
+            rfdetr.optimize_for_inference(compile=False)
+
+        assert rfdetr.model.inference_model is None
+        assert rfdetr._is_optimized_for_inference is False
+
+    def test_export_failure_leaves_is_optimized_false(self) -> None:
+        """If export() raises after deepcopy succeeds, _is_optimized_for_inference stays False."""
+        rfdetr = _FakeRFDETR()
+        fake_copy = _FakeModel()
+
+        with (
+            patch("rfdetr.detr.deepcopy", return_value=fake_copy),
+            patch.object(fake_copy, "export", side_effect=RuntimeError("export failed")),
+            pytest.raises(RuntimeError, match="export failed"),
+        ):
+            rfdetr.optimize_for_inference(compile=False)
+
+        assert rfdetr._is_optimized_for_inference is False
+
+    def test_jit_trace_failure_leaves_compiled_flags_false(self) -> None:
+        """If jit.trace raises, _optimized_has_been_compiled and _optimized_batch_size stay unset."""
+        rfdetr = _FakeRFDETR()
+
+        with (
+            patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model),
+            patch("torch.jit.trace", side_effect=RuntimeError("trace failed")),
+            pytest.raises(RuntimeError, match="trace failed"),
+        ):
+            rfdetr.optimize_for_inference(compile=True, batch_size=2)
+
+        assert rfdetr._optimized_has_been_compiled is False
+        assert rfdetr._optimized_batch_size is None
+
+    def test_jit_trace_failure_leaves_model_fully_unoptimized(self) -> None:
+        """jit.trace failure leaves both _is_optimized_for_inference=False and inference_model=None."""
+        rfdetr = _FakeRFDETR()
+
+        with (
+            patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model),
+            patch("torch.jit.trace", side_effect=RuntimeError("trace failed")),
+            pytest.raises(RuntimeError, match="trace failed"),
+        ):
+            rfdetr.optimize_for_inference(compile=True)
+
+        assert rfdetr._is_optimized_for_inference is False
+        assert rfdetr.model.inference_model is None
+
+    def test_inplace_export_failure_module_mutations_are_not_undone(self) -> None:
+        """RFDETR resets flags on export failure but cannot undo module-level mutations.
+
+        Production export() may mutate the module (e.g. forward->forward_export) before raising; those changes are not
+        reversed by the exception-recovery path.
+        """
+        rfdetr = _FakeRFDETR()
+        original_model = rfdetr.model.model
+        mutated: dict[str, bool] = {"happened": False}
+
+        def _mutating_export() -> None:
+            mutated["happened"] = True
+            raise RuntimeError("export failed mid-mutation")
+
+        with (
+            patch("rfdetr.detr.deepcopy"),
+            patch.object(original_model, "export", side_effect=_mutating_export),
+            pytest.raises(RuntimeError, match="export failed mid-mutation"),
+        ):
+            rfdetr.optimize_for_inference(compile=False, inplace=True)
+
+        assert rfdetr._is_optimized_for_inference is False
+        assert rfdetr.is_optimized_inplace is False
+        assert rfdetr.model.model is original_model
+        # The mutation happened and cannot be undone by RFDETR's recovery path
+        assert mutated["happened"] is True

tests/inference/test_predict_eval_mode.py

@@ -0,0 +1,136 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests that unoptimized inference always runs the module in eval mode."""
+
+from types import SimpleNamespace
+
+import PIL.Image
+import pytest
+import torch
+
+from rfdetr import detr as detr_module
+
+from .helpers import _BaseFakeRFDETR
+
+
+class _FakeModelWithDropout(torch.nn.Module):
+    """Minimal module whose behavior differs between train and eval mode."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.dropout = torch.nn.Dropout(p=0.5)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Pass input through dropout, active only in train mode."""
+        return self.dropout(x)
+
+
+class _FakeModelContext:
+    """Minimal model context supplying the attributes predict() and train() need."""
+
+    def __init__(self) -> None:
+        self.device = torch.device("cpu")
+        self.resolution = 28
+        self.model = _FakeModelWithDropout()
+        self.inference_model = None
+
+
+class _FakeRFDETR(_BaseFakeRFDETR):
+    """Concrete test double: provides a dropout-bearing model for eval-mode tests."""
+
+    def get_model(self, config: SimpleNamespace) -> _FakeModelContext:
+        """Return a minimal model context with a dropout-bearing module."""
+        return _FakeModelContext()
+
+
+class TestUnoptimizedInferenceEvalMode:
+    """`_ensure_eval_mode_for_unoptimized_inference` must keep the module in eval mode."""
+
+    def test_eval_mode_reasserted_after_train_round_trip(self) -> None:
+        """Eval mode must be applied to whatever self.model.model currently points to.
+
+        ``train()`` rebinds ``self.model.model`` to a brand-new module left in training mode, so eval must be re-applied
+        to the *current* object on every call — not to a cached reference captured at init.
+        """
+        rfdetr = _FakeRFDETR()
+
+        # First inference call: warns once and switches to eval mode.
+        rfdetr._ensure_eval_mode_for_unoptimized_inference()
+        assert rfdetr.model.model.training is False
+
+        # Simulate train() rebinding self.model.model to a fresh training-mode module.
+        rfdetr.model.model = _FakeModelWithDropout()
+        assert rfdetr.model.model.training is True  # new object starts in train mode
+
+        # Every subsequent inference call must re-assert eval on the *new* object.
+        rfdetr._ensure_eval_mode_for_unoptimized_inference()
+        assert rfdetr.model.model.training is False
+
+    def test_optimized_model_skips_eval_assertion(self) -> None:
+        """When _is_optimized_for_inference is True, the method must be a no-op.
+
+        The compiled inference_model snapshot is already in eval mode; calling eval() on the stale self.model.model
+        would target the wrong object.
+        """
+        rfdetr = _FakeRFDETR()
+        rfdetr._is_optimized_for_inference = True
+        rfdetr.model.model.train()
+        assert rfdetr.model.model.training is True  # confirm starting state
+
+        rfdetr._ensure_eval_mode_for_unoptimized_inference()
+
+        assert rfdetr.model.model.training is True  # must remain unchanged
+
+    def test_not_optimized_warning_emitted_only_once(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        """The not-optimized warning is logged at most once across repeated calls."""
+        warnings: list[str] = []
+        monkeypatch.setattr(detr_module.logger, "warning", lambda msg, *a, **k: warnings.append(msg))
+
+        rfdetr = _FakeRFDETR()
+        rfdetr._ensure_eval_mode_for_unoptimized_inference()
+        rfdetr.model.model.train()
+        rfdetr._ensure_eval_mode_for_unoptimized_inference()
+        rfdetr._ensure_eval_mode_for_unoptimized_inference()
+
+        assert len(warnings) == 1
+
+    def test_eval_mode_applied_on_every_call(self) -> None:
+        """Eval() must run on every call, not just when the warning fires.
+
+        Simulate the code path where the warning has already been emitted
+        (``_has_warned_about_not_being_optimized_for_inference=True``) and verify
+        that ``eval()`` is still applied to the current module.
+        """
+        rfdetr = _FakeRFDETR()
+        rfdetr._has_warned_about_not_being_optimized_for_inference = True
+        rfdetr.model.model.train()
+
+        rfdetr._ensure_eval_mode_for_unoptimized_inference()
+
+        assert rfdetr.model.model.training is False
+
+    def test_predict_puts_module_in_eval_mode(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        """Predict() must delegate to _ensure_eval_mode_for_unoptimized_inference, leaving module in eval mode."""
+        rfdetr = _FakeRFDETR()
+        img = PIL.Image.new("RGB", (640, 640), color=(128, 128, 128))
+
+        monkeypatch.setattr(
+            rfdetr.model.model,
+            "forward",
+            lambda batch: {"pred_logits": torch.zeros(1, 10, 81), "pred_boxes": torch.zeros(1, 10, 4)},
+        )
+        monkeypatch.setattr(
+            rfdetr.model,
+            "postprocess",
+            lambda preds, target_sizes: [
+                {"scores": torch.zeros(0), "labels": torch.zeros(0, dtype=torch.long), "boxes": torch.zeros(0, 4)}
+            ],
+            raising=False,
+        )
+
+        rfdetr.predict(img)
+
+        assert rfdetr.model.model.training is False

tests/inference/test_public_api_keypoints.py

@@ -0,0 +1,34 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Public API tests for the keypoint preview variant."""
+
+from rfdetr import RFDETRKeypointPreview
+from rfdetr.config import KeypointTrainConfig, RFDETRKeypointPreviewConfig
+from rfdetr.detr import RFDETRKeypointPreview as RFDETRKeypointPreviewFromDetr
+from rfdetr.variants import RFDETRKeypointPreview as RFDETRKeypointPreviewFromVariants
+
+
+def test_keypoint_preview_top_level_import() -> None:
+    """RFDETRKeypointPreview must be importable from top-level package and keep shared identity."""
+    assert RFDETRKeypointPreview is RFDETRKeypointPreviewFromVariants
+    assert RFDETRKeypointPreview is RFDETRKeypointPreviewFromDetr
+
+
+def test_keypoint_preview_variant_metadata() -> None:
+    """RFDETRKeypointPreview exposes the expected variant metadata and config class."""
+    assert RFDETRKeypointPreview.size == "rfdetr-keypoint-preview"
+    assert RFDETRKeypointPreview._model_config_class is RFDETRKeypointPreviewConfig
+    assert RFDETRKeypointPreview._train_config_class is KeypointTrainConfig
+    assert RFDETRKeypointPreviewConfig.model_fields["pretrain_weights"].default == "rf-detr-keypoint-preview-xlarge.pth"
+
+
+def test_predict_docstring_mentions_one_class_keypoint_mapping() -> None:
+    """The public predict() contract must document active-first keypoint class-id mapping."""
+    doc = RFDETRKeypointPreview.predict.__doc__ or ""
+    assert "one-class preview keypoint setup" in doc
+    assert "class_id=0" in doc
+    assert "class_id=1" in doc
+    assert "__background__" in doc

tests/inference/test_supervision_editable_keypoints.py

@@ -0,0 +1,24 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Smoke test for RF-DETR keypoints carried by Supervision KeyPoints."""
+
+import numpy as np
+import supervision as sv
+
+
+def test_rfdetr_keypoints_include_detection_details() -> None:
+    """RF-DETR-style keypoints preserve detection boxes and scores."""
+    key_points = sv.KeyPoints(
+        xy=np.array([[[1.0, 2.0], [3.0, 4.0]]], dtype=np.float32),
+        keypoint_confidence=np.array([[0.9, 0.8]], dtype=np.float32),
+        detection_confidence=np.array([0.95], dtype=np.float32),
+        class_id=np.array([1], dtype=int),
+        data={"xyxy": np.array([[0, 0, 10, 10]], dtype=np.float32)},
+    )
+
+    assert key_points.xy.shape == (1, 2, 2)
+    np.testing.assert_array_equal(key_points.data["xyxy"], np.array([[0, 0, 10, 10]], dtype=np.float32))
+    np.testing.assert_array_equal(key_points.detection_confidence, np.array([0.95], dtype=np.float32))

tests/inference/test_trt_inference.py

@@ -0,0 +1,110 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+from PIL import Image
+
+from rfdetr.export.benchmark import TRTInference, infer_transforms
+
+
+class TestTRTInference:
+    def test_synchronize_sync_mode_does_not_require_stream(self, monkeypatch) -> None:
+        """`synchronize()` should not access stream in sync mode."""
+        inference = TRTInference.__new__(TRTInference)
+        inference.sync_mode = True
+
+        mock_is_available = Mock(return_value=True)
+        mock_cuda_sync = Mock()
+        monkeypatch.setattr("torch.cuda.is_available", mock_is_available)
+        monkeypatch.setattr("torch.cuda.synchronize", mock_cuda_sync)
+
+        inference.synchronize()
+
+        mock_is_available.assert_called_once()
+        mock_cuda_sync.assert_called_once()
+
+    def test_synchronize_async_mode_uses_stream_sync(self, monkeypatch) -> None:
+        """`synchronize()` should use stream synchronization in async mode."""
+        inference = TRTInference.__new__(TRTInference)
+        inference.sync_mode = False
+        inference.stream = Mock()
+
+        mock_cuda_sync = Mock()
+        monkeypatch.setattr("torch.cuda.synchronize", mock_cuda_sync)
+
+        inference.synchronize()
+
+        inference.stream.synchronize.assert_called_once()
+        mock_cuda_sync.assert_not_called()
+
+    def test_infer_transforms_accepts_none_target(self) -> None:
+        """Benchmark inference preprocessing should support image-only input."""
+        image = Image.new("RGB", (320, 240))
+
+        image_tensor, target = infer_transforms()(image, None)
+
+        assert isinstance(image_tensor, torch.Tensor)
+        assert image_tensor.shape == (3, 640, 640)
+        assert image_tensor.dtype == torch.float32
+        assert target is None
+
+
+class TestBenchmarkShapeParameterization:
+    """Benchmark preprocessing/postprocessing read input size and query count instead of hardcoding 640/300."""
+
+    def test_infer_transforms_uses_requested_size(self) -> None:
+        """infer_transforms resizes to the caller-supplied (height, width)."""
+        image = Image.new("RGB", (320, 240))
+
+        image_tensor, _ = infer_transforms((512, 384))(image, None)
+
+        assert image_tensor.shape == (3, 512, 384)
+
+    def test_infer_transforms_defaults_to_640(self) -> None:
+        """The default input size stays 640x640 for callers that do not pass a size."""
+        image = Image.new("RGB", (320, 240))
+
+        image_tensor, _ = infer_transforms()(image, None)
+
+        assert image_tensor.shape == (3, 640, 640)
+
+    def test_static_dim_returns_concrete_int(self) -> None:
+        """A concrete positive dimension is returned unchanged."""
+        from rfdetr.export.benchmark import _static_dim
+
+        assert _static_dim(384, 640) == 384
+
+    @pytest.mark.parametrize(
+        "value",
+        [
+            pytest.param("height", id="dynamic-string"),
+            pytest.param(None, id="none"),
+            pytest.param(-1, id="negative"),
+        ],
+    )
+    def test_static_dim_falls_back_for_dynamic_axis(self, value) -> None:
+        """Dynamic/unknown axes fall back to the provided default."""
+        from rfdetr.export.benchmark import _static_dim
+
+        assert _static_dim(value, 640) == 640
+
+    def test_post_process_respects_num_queries(self) -> None:
+        """post_process selects exactly num_queries detections per image."""
+        from rfdetr.export.benchmark import post_process
+
+        num_queries = 5
+        outputs = {
+            "labels": torch.rand(1, 20, 3),
+            "dets": torch.rand(1, 20, 4),
+        }
+        target_sizes = torch.tensor([[480, 640]])
+
+        results = post_process(outputs, target_sizes, num_queries=num_queries)
+
+        assert results[0]["scores"].shape == (num_queries,)