Keep the served file extension in TrainedModel.download()

CHANGELOG.md

@@ -34,6 +34,13 @@ All notable changes to this project will be documented in this file.
   YOLOv10. Existing models remain available for inference, and YOLO dataset
   formats remain supported ([#572](https://github.com/roboflow/roboflow-python/pull/572)).
 
+### Fixed
+
+- `TrainedModel.download()` saves the file as `weights` plus the extension of
+  the artifact that the API serves, for example `weights.onnx`. Before, it
+  always used `weights.pt`
+  ([#582](https://github.com/roboflow/roboflow-python/pull/582)).
+
 ## 1.7.3
 
 ### Added

roboflow/core/training.py

@@ -10,7 +10,9 @@
 
 import json
 import os
+from pathlib import PurePosixPath
 from typing import List
+from urllib.parse import urlparse
 
 from roboflow.adapters import rfapi
 from roboflow.config import (
@@ -159,12 +161,17 @@ def poll_until_video_results(self, job_id=None) -> dict:
         return self._video_model().poll_until_video_results(job_id)
 
     def download(self, format="pt", location="."):
-        """Download this model's PyTorch weights to ``location/weights.pt``."""
+        """Download this model's weights to ``location`` and return the file path.
+
+        The file is ``weights`` plus the extension of the artifact that the API
+        serves: ``weights.pt`` for most models, ``weights.onnx`` for a NAS child.
+        """
         weights_url = rfapi.get_model_weights_url(
             self.__api_key, self.workspace, self.project, self._weights_id, model_format=format
         )
+        suffix = PurePosixPath(urlparse(weights_url).path).suffix or ".pt"
         os.makedirs(location, exist_ok=True)
-        out_path = os.path.join(location, "weights.pt")
+        out_path = os.path.join(location, f"weights{suffix}")
         response = http.get(weights_url, stream=True)
         response.raise_for_status()
         with open(out_path, "wb") as f:

tests/test_training.py

@@ -1,5 +1,7 @@
+import tempfile
 import unittest
-from unittest.mock import patch
+from pathlib import Path
+from unittest.mock import MagicMock, patch
 
 from roboflow.config import (
     CLASSIFICATION_MODEL,
@@ -88,6 +90,34 @@ def test_poll_reuses_the_predict_video_legacy_model(self):
         self.assertIs(poll.call_args.args[0], model._video_model())
 
 
+class TestTrainedModelDownload(unittest.TestCase):
+    def test_download_keeps_the_extension_of_the_served_artifact(self):
+        signed_query = "?X-Goog-Signature=abc.def"
+        cases = [
+            ("https://storage.googleapis.com/bucket/ws/models/doc/weights.pt", "weights.pt"),
+            ("https://storage.googleapis.com/bucket/ws/proj/3/best.pt", "weights.pt"),
+            ("https://storage.googleapis.com/bucket/ws/models/doc/weights.onnx", "weights.onnx"),
+            ("https://storage.googleapis.com/bucket/ws/models/doc/model.safetensors", "weights.safetensors"),
+            ("https://storage.googleapis.com/bucket/ws/models/doc/weights", "weights.pt"),
+        ]
+
+        for artifact_url, expected_filename in cases:
+            with self.subTest(artifact_url=artifact_url), tempfile.TemporaryDirectory() as location:
+                model = TrainedModel("key", "ws", "proj", "ws/model-slug", model_type="rfdetr-base")
+                response = MagicMock()
+                response.iter_content.return_value = [b"weights-bytes"]
+                weights_url = artifact_url + signed_query
+
+                with (
+                    patch("roboflow.core.training.rfapi.get_model_weights_url", return_value=weights_url),
+                    patch("roboflow.core.training.http.get", return_value=response),
+                ):
+                    out_path = model.download(location=location)
+
+                self.assertEqual(Path(out_path), Path(location) / expected_filename)
+                self.assertEqual(Path(out_path).read_bytes(), b"weights-bytes")
+
+
 class TestTrainingModels(unittest.TestCase):
     def test_models_are_cached_until_refresh(self):
         training = Training("key", "ws", "proj", "1", {"trainingId": "training-1"})