From d0af5a6157ee1fb36a4eabd065e7791358602872 Mon Sep 17 00:00:00 2001 From: Leandro Rosemberg Date: Wed, 1 Jul 2026 23:13:12 -0300 Subject: [PATCH 01/11] refactor: extract public package_custom_weights helper from deploy flows --- roboflow/core/version.py | 14 +- roboflow/core/workspace.py | 16 +- roboflow/util/model_processor.py | 1271 ++++++++++++++++++++-------- tests/util/test_model_processor.py | 388 +++++++++ 4 files changed, 1299 insertions(+), 390 deletions(-) diff --git a/roboflow/core/version.py b/roboflow/core/version.py index fd4a3f5f..4d8a9af0 100644 --- a/roboflow/core/version.py +++ b/roboflow/core/version.py @@ -34,8 +34,8 @@ from roboflow.models.vlm import VLMModel from roboflow.util.annotations import amend_data_yaml from roboflow.util.general import extract_zip, write_line -from roboflow.util.model_processor import process, validate_model_type_for_project -from roboflow.util.versions import get_model_format, get_wrong_dependencies_versions, normalize_yolo_model_type +from roboflow.util.model_processor import package_custom_weights_interactive, validate_model_type_for_project +from roboflow.util.versions import get_model_format, get_wrong_dependencies_versions if TYPE_CHECKING: import numpy as np @@ -497,14 +497,10 @@ def deploy(self, model_type: str, model_path: str, filename: str = "weights/best model_path (str): File path to the model weights to be uploaded. filename (str, optional): The name of the weights file. Defaults to "weights/best.pt". """ - model_type = normalize_yolo_model_type(model_type) - zip_file_name, model_type = process(model_type, model_path, filename) + bundle = package_custom_weights_interactive(model_type, model_path, filename, build_dir=model_path) - if zip_file_name is None: - raise RuntimeError("Failed to process model") - - self._validate_against_project_type(model_type) - self._upload_zip(model_type, model_path, zip_file_name) + self._validate_against_project_type(bundle.model_type) + self._upload_zip(bundle.model_type, model_path, bundle.archive_path.name) def _validate_against_project_type(self, model_type: str) -> None: validate_model_type_for_project(model_type, self.type, self.project) diff --git a/roboflow/core/workspace.py b/roboflow/core/workspace.py index c1dca153..64a5c430 100644 --- a/roboflow/core/workspace.py +++ b/roboflow/core/workspace.py @@ -787,8 +787,10 @@ def deploy_model( filename (str, optional): The name of the weights file. Defaults to "weights/best.pt". """ - from roboflow.util.model_processor import process, validate_model_type_for_project - from roboflow.util.versions import normalize_yolo_model_type + from roboflow.util.model_processor import ( + package_custom_weights_interactive, + validate_model_type_for_project, + ) if not project_ids: raise ValueError("At least one project ID must be provided") @@ -800,16 +802,12 @@ def deploy_model( if project_id not in projects_by_id: raise ValueError(f"Project {project_id} is not accessible in this workspace") - model_type = normalize_yolo_model_type(model_type) - zip_file_name, model_type = process(model_type, model_path, filename) - - if zip_file_name is None: - raise RuntimeError("Failed to process model") + bundle = package_custom_weights_interactive(model_type, model_path, filename, build_dir=model_path) for project_id in project_ids: - validate_model_type_for_project(model_type, projects_by_id[project_id].get("type", ""), project_id) + validate_model_type_for_project(bundle.model_type, projects_by_id[project_id].get("type", ""), project_id) - self._upload_zip(model_type, model_path, project_ids, model_name, zip_file_name) + self._upload_zip(bundle.model_type, model_path, project_ids, model_name, bundle.archive_path.name) def _upload_zip( self, diff --git a/roboflow/util/model_processor.py b/roboflow/util/model_processor.py index 0022c11a..c6409585 100644 --- a/roboflow/util/model_processor.py +++ b/roboflow/util/model_processor.py @@ -1,8 +1,40 @@ +"""Packaging of custom model weights for Roboflow upload. + +The public, non-interactive entry point is :func:`package_custom_weights`. It +only builds the upload archive; it never prompts, prints, or uploads, so it is +safe to call from servers and other headless environments (for example the +Roboflow MCP server):: + + from roboflow.util.model_processor import package_custom_weights + + bundle = package_custom_weights("yolov8n", "runs/detect/train") + try: + ... # upload bundle.archive_path + finally: + bundle.cleanup() + +Expected, user-correctable failures raise :class:`ModelPackagingError` +subclasses; anything else escaping these helpers is a bug. + +:func:`process` is the legacy entry point kept for backwards compatibility, +and the ``Version.deploy`` / ``Workspace.deploy_model`` flows wrap the packaging +step with :func:`package_custom_weights_interactive`, which preserves the +historical print-and-confirm CLI behavior. +""" + +from __future__ import annotations + import json +import math import os import shutil +import tarfile +import tempfile import zipfile -from typing import Callable, Optional +from dataclasses import dataclass +from importlib import import_module +from pathlib import Path +from typing import Any import yaml @@ -19,7 +51,177 @@ TYPE_OBJECT_DETECTION, TYPE_SEMANTIC_SEGMENTATION, ) -from roboflow.util.versions import print_warn_for_wrong_dependencies_versions +from roboflow.util.versions import get_wrong_dependencies_versions, normalize_yolo_model_type + +SUPPORTED_MODELS = ( + "yolov5", + "yolov7", + "yolov7-seg", + "yolov8", + "yolov9", + "yolov10", + "yolov11", + "yolov12", + "yolo26", + "yolonas", + "paligemma", + "paligemma2", + "florence-2", + "rfdetr", +) + +SUPPORTED_HUGGINGFACE_TYPES = ( + "florence-2-base", + "florence-2-large", + "paligemma-3b-pt-224", + "paligemma-3b-pt-448", + "paligemma-3b-pt-896", + "paligemma2-3b-pt-224", + "paligemma2-3b-pt-448", + "paligemma2-3b-pt-896", + "paligemma2-3b-pt-224-peft", + "paligemma2-3b-pt-448-peft", + "paligemma2-3b-pt-896-peft", +) + +SUPPORTED_RFDETR_TYPES = ( + # Detection models + "rfdetr-base", + "rfdetr-nano", + "rfdetr-small", + "rfdetr-medium", + "rfdetr-large", + "rfdetr-xlarge", + "rfdetr-2xlarge", + # Segmentation models + "rfdetr-seg-nano", + "rfdetr-seg-small", + "rfdetr-seg-medium", + "rfdetr-seg-large", + "rfdetr-seg-xlarge", + "rfdetr-seg-2xlarge", +) + +# YOLO families Roboflow rejects without a size suffix (e.g. `yolov8` must be +# `yolov8n`/`yolov8s`/...). Legacy yolov5/7/9 go through the opt.yaml path and are +# intentionally excluded. +ULTRALYTICS_YOLO_FAMILIES = ("yolov8", "yolov10", "yolov11", "yolov12", "yolo26") + +# Canonical (depth_multiple, width_multiple) -> size letter for the classic YOLO +# scaling (v5/v8/v9/v10). Newer families (v11+) instead store an explicit ``scale`` +# letter in the model yaml, which is read first. +YOLO_DEPTH_WIDTH_TO_SIZE = { + (0.33, 0.25): "n", + (0.33, 0.50): "s", + (0.67, 0.75): "m", + (0.67, 1.00): "b", + (1.00, 1.00): "l", + (1.00, 1.25): "x", +} + +# Position-encoding grid size (DINOv2 tokens per side) each *known* RF-DETR variant +# is built with, mirroring rfdetr/config.py. Roboflow reconstructs the architecture +# from the model_type at the variant's default resolution, so a checkpoint trained at +# that default must match or state_dict loading fails on the backbone +# position_embeddings. Variants absent here (e.g. detection xlarge/2xlarge, which have +# no standard config) are not grid-checked. A checkpoint trained at a custom +# resolution may not match any entry; that case warns rather than blocks. +RFDETR_POSITIONAL_ENCODING_SIZE = { + "rfdetr-nano": 24, + "rfdetr-small": 32, + "rfdetr-medium": 36, + "rfdetr-base": 37, + "rfdetr-large": 44, + "rfdetr-seg-nano": 26, + "rfdetr-seg-small": 32, + "rfdetr-seg-medium": 36, + "rfdetr-seg-large": 42, + "rfdetr-seg-xlarge": 52, + "rfdetr-seg-2xlarge": 64, +} + + +class ModelPackagingError(Exception): + """Custom weights could not be packaged for a user-correctable reason. + + Consumers can treat any instance of this class as an expected input problem + (bad model_type, missing files, mismatched metadata, ...) and surface the + message to the user. Exceptions that are not ModelPackagingError indicate + bugs and are deliberately not wrapped. + """ + + +class UnsupportedModelError(ModelPackagingError, ValueError): + """The model_type is not supported for custom weights upload.""" + + +class TaskMismatchError(ModelPackagingError, ValueError): + """The model_type's task conflicts with the checkpoint or the project type.""" + + +class MissingFileError(ModelPackagingError, FileNotFoundError): + """A file required for packaging was not found.""" + + +class MissingDependencyError(ModelPackagingError, RuntimeError): + """A Python package required to package these weights is not installed.""" + + +class DependencyMismatchError(ModelPackagingError, RuntimeError): + """An installed dependency version differs from the recommended one. + + Retry with ``allow_dependency_mismatch=True`` to package anyway. + """ + + retry_flag = "allow_dependency_mismatch" + + def __init__(self, message: str, *, dependency: str, required: str, installed: str): + super().__init__(message) + self.dependency = dependency + self.required = required + self.installed = installed + + +class SizeMismatchError(ModelPackagingError, ValueError): + """The declared model size/variant conflicts with the checkpoint architecture. + + Retry with ``allow_size_mismatch=True`` to package the requested model_type + as-is. + """ + + retry_flag = "allow_size_mismatch" + + def __init__(self, message: str, *, requested: str, detected: str | None = None): + super().__init__(message) + self.requested = requested + self.detected = detected + + +@dataclass(frozen=True) +class ModelUploadBundle: + """Packaged archive ready to upload through the Roboflow API. + + ``model_type`` is the resolved type (it may differ from the requested one, + e.g. ``yolov8`` filled in as ``yolov8n`` from the checkpoint architecture). + ``owns_build_dir`` is True when :func:`package_custom_weights` created a + temporary build directory; call :meth:`cleanup` once the archive has been + consumed. + """ + + archive_path: Path + build_dir: Path + model_type: str + warnings: tuple[str, ...] = () + owns_build_dir: bool = False + + @property + def size_bytes(self) -> int: + return self.archive_path.stat().st_size + + def cleanup(self) -> None: + """Remove the build directory if this bundle created it (no-op otherwise).""" + if self.owns_build_dir: + shutil.rmtree(self.build_dir, ignore_errors=True) def task_of_model_type(model_type: str) -> str: @@ -36,7 +238,7 @@ def task_of_model_type(model_type: str) -> str: def validate_model_type_for_project(model_type: str, project_type: str, project_id: str) -> None: - """Raise ValueError if model_type's task doesn't match the Roboflow project type.""" + """Raise TaskMismatchError if model_type's task doesn't match the Roboflow project type.""" expected = { TYPE_OBJECT_DETECTION: TASK_DET, TYPE_INSTANCE_SEGMENTATION: TASK_SEG, @@ -48,70 +250,208 @@ def validate_model_type_for_project(model_type: str, project_type: str, project_ return actual = task_of_model_type(model_type) if actual != expected: - raise ValueError( + raise TaskMismatchError( f"Project '{project_id}' is type '{project_type}' (task '{expected}') " f"but model_type '{model_type}' implies task '{actual}'." ) +def package_custom_weights( + model_type: str, + model_path: str, + filename: str = "weights/best.pt", + *, + build_dir: str | Path | None = None, + allow_dependency_mismatch: bool = False, + allow_size_mismatch: bool = False, +) -> ModelUploadBundle: + """Package locally trained custom weights into a Roboflow upload archive. + + This is the public packaging entry point. It is non-interactive and free of + side effects on ``model_path``: it never prompts, prints, or writes into the + source directory. Heavy dependencies (torch, ultralytics) are imported + lazily, only for the model families that need them. + + Args: + model_type: Roboflow model type (e.g. "yolov8n", "rfdetr-base"). + model_path: Directory containing the trained model artifacts. + filename: Weights file path, relative to ``model_path``. + build_dir: Directory to write intermediate artifacts and the final + archive into. Defaults to a fresh temporary directory owned by the + returned bundle; call ``bundle.cleanup()`` when done. + allow_dependency_mismatch: Record a warning instead of raising + DependencyMismatchError when an installed dependency version is not + the recommended one. + allow_size_mismatch: Record a warning instead of raising + SizeMismatchError when the declared model size/variant conflicts + with the checkpoint architecture. + + Returns: + ModelUploadBundle with the archive path, the resolved model_type, and + any warnings collected while packaging. + + Raises: + ModelPackagingError: (or a subclass) for user-correctable problems. + """ + normalized_model_type = normalize_yolo_model_type(model_type.strip()) + source_dir = Path(model_path).expanduser().resolve() + if not source_dir.is_dir(): + raise MissingFileError(f"Model path '{model_path}' does not exist or is not a directory.") + + owns_build_dir = build_dir is None + if build_dir is None: + build_path = Path(tempfile.mkdtemp(prefix="roboflow-package-")) + else: + build_path = Path(build_dir) + build_path.mkdir(parents=True, exist_ok=True) + + try: + archive_path, resolved_model_type, warnings = _process_model( + model_type=normalized_model_type, + model_path=source_dir, + filename=filename, + build_dir=build_path, + allow_dependency_mismatch=allow_dependency_mismatch, + allow_size_mismatch=allow_size_mismatch, + ) + except BaseException: + if owns_build_dir: + shutil.rmtree(build_path, ignore_errors=True) + raise + + return ModelUploadBundle( + archive_path=archive_path, + build_dir=build_path, + model_type=resolved_model_type, + warnings=tuple(warnings), + owns_build_dir=owns_build_dir, + ) + + +def package_custom_weights_interactive( + model_type: str, + model_path: str, + filename: str = "weights/best.pt", + *, + build_dir: str | Path | None = None, +) -> ModelUploadBundle: + """Package weights with the historical interactive SDK behavior. + + Used by ``Version.deploy`` and ``Workspace.deploy_model``: warnings are + printed, and dependency/size mismatches ask for confirmation before + retrying with the corresponding override. Declining re-raises the error. + """ + allow_dependency_mismatch = False + allow_size_mismatch = False + while True: + try: + bundle = package_custom_weights( + model_type, + model_path, + filename, + build_dir=build_dir, + allow_dependency_mismatch=allow_dependency_mismatch, + allow_size_mismatch=allow_size_mismatch, + ) + except (DependencyMismatchError, SizeMismatchError) as error: + print(error) + answer = input("Would you like to continue anyway? y/n: ") + if answer.lower() != "y": + raise + if isinstance(error, DependencyMismatchError): + allow_dependency_mismatch = True + else: + allow_size_mismatch = True + continue + for warning in bundle.warnings: + print(warning) + return bundle + + def process(model_type: str, model_path: str, filename: str) -> tuple[str, str]: - processor = _get_processor_function(model_type) - return processor(model_type, model_path, filename) - - -def _get_processor_function(model_type: str) -> Callable: - supported_models = [ - "yolov5", - "yolov7", - "yolov7-seg", - "yolov8", - "yolov9", - "yolov10", - "yolov11", - "yolov12", - "yolo26", - "yolonas", - "paligemma", - "paligemma2", - "florence-2", - "rfdetr", - ] - - if not any(supported_model in model_type for supported_model in supported_models): - raise (ValueError(f"Model type {model_type} not supported. Supported models are {supported_models}")) + """Legacy packaging entry point, kept for backwards compatibility. - if model_type.startswith(("paligemma", "paligemma2", "florence-2")): - if any(model in model_type for model in ["paligemma", "paligemma2", "florence-2"]): - supported_hf_types = [ - "florence-2-base", - "florence-2-large", - "paligemma-3b-pt-224", - "paligemma-3b-pt-448", - "paligemma-3b-pt-896", - "paligemma2-3b-pt-224", - "paligemma2-3b-pt-448", - "paligemma2-3b-pt-896", - "paligemma2-3b-pt-224-peft", - "paligemma2-3b-pt-448-peft", - "paligemma2-3b-pt-896-peft", - ] - if model_type not in supported_hf_types: - raise RuntimeError( - f"{model_type} not supported for this type of upload." - f"Supported upload types are {supported_hf_types}" - ) - return _process_huggingface + Packages into ``model_path`` (the historical behavior: intermediate + artifacts and the final archive land there), prints any packaging warnings, + and returns ``(archive_file_name, resolved_model_type)``. New code should + call :func:`package_custom_weights` instead. + """ + bundle = package_custom_weights(model_type, model_path, filename, build_dir=model_path) + for warning in bundle.warnings: + print(warning) + return bundle.archive_path.name, bundle.model_type + + +def _process_model( + *, + model_type: str, + model_path: Path, + filename: str, + build_dir: Path, + allow_dependency_mismatch: bool, + allow_size_mismatch: bool, +) -> tuple[Path, str, list[str]]: + if not any(supported in model_type for supported in SUPPORTED_MODELS): + raise UnsupportedModelError( + f"Model type '{model_type}' is not supported for custom weights upload. " + f"Supported families are: {', '.join(SUPPORTED_MODELS)}." + ) + if model_type.startswith(("paligemma", "paligemma2", "florence-2")): + return _process_huggingface(model_type, model_path, build_dir) if "yolonas" in model_type: - return _process_yolonas - + return _process_yolonas(model_type, model_path, filename, build_dir) if "rfdetr" in model_type: - return _process_rfdetr + return _process_rfdetr(model_type, model_path, filename, build_dir, allow_size_mismatch) + return _process_yolo( + model_type, + model_path, + filename, + build_dir, + allow_dependency_mismatch, + allow_size_mismatch, + ) + - return _process_yolo +def _import_required_module(module_name: str, install_hint: str) -> Any: + try: + return import_module(module_name) + except ImportError as exc: + raise MissingDependencyError( + f"The '{module_name}' Python package is required to package these " + f"custom weights. Install it with `{install_hint}`." + ) from exc + + +def _check_dependency_version( + *, + dependency: str, + operator: str, + required_version: str, + allow_mismatch: bool, + warnings: list[str], +) -> None: + mismatches = get_wrong_dependencies_versions([(dependency, operator, required_version)]) + if not mismatches: + return + _, _, _, installed = mismatches[0] + message = ( + f"{dependency}{operator}{required_version} is recommended for this " + f"upload, but {dependency} {installed} is installed." + ) + if allow_mismatch: + warnings.append(message) + return + raise DependencyMismatchError( + f"{message} Retry with allow_dependency_mismatch=True to package with the " + f"installed version, or `pip install {dependency}{operator}{required_version}`.", + dependency=dependency, + required=f"{dependency}{operator}{required_version}", + installed=installed, + ) -def _detect_yolo_task(model_instance) -> Optional[str]: +def _detect_yolo_task(model_instance: Any) -> str | None: """Detect the training task of an Ultralytics model instance via its class name.""" if model_instance is None: return None @@ -125,87 +465,188 @@ def _detect_yolo_task(model_instance) -> Optional[str]: }.get(type(model_instance).__name__) -def _process_yolo(model_type: str, model_path: str, filename: str) -> tuple[str, str]: - if "yolov8" in model_type: - try: - import torch - import ultralytics +def _class_names_from_model_instance(model_instance: Any) -> list[str]: + names = getattr(model_instance, "names", None) + if isinstance(names, list): + return names + if isinstance(names, dict): + return [name for _, name in sorted(names.items(), key=lambda item: item[0])] + raise ModelPackagingError("Could not extract class names from the model checkpoint.") - except ImportError: - raise RuntimeError( - "The ultralytics python package is required to deploy yolov8" - " models. Please install it with `pip install ultralytics`" - ) - print_warn_for_wrong_dependencies_versions([("ultralytics", "==", "8.0.196")], ask_to_continue=True) +def _filtered_args(args: Any) -> dict[str, Any]: + if not isinstance(args, dict): + args = vars(args) + return {k: v for k, v in args.items() if k in {"model", "imgsz", "batch"}} - elif "yolov10" in model_type: - try: - import torch - import ultralytics - except ImportError: - raise RuntimeError( - "The ultralytics python package is required to deploy yolov10" - " models. Please install it with `pip install ultralytics`" - ) +def _load_checkpoint(torch_module: Any, checkpoint_path: Path, *, map_location: str | None = None) -> Any: + kwargs: dict[str, Any] = {"weights_only": False} + if map_location is not None: + kwargs["map_location"] = map_location + return torch_module.load(checkpoint_path, **kwargs) - elif "yolov5" in model_type or "yolov7" in model_type or "yolov9" in model_type: - try: - import torch - except ImportError: - raise RuntimeError( - f"The torch python package is required to deploy {model_type} models." - " Please install it with `pip install torch`" - ) - elif "yolov11" in model_type: - try: - import torch - import ultralytics +def _legacy_yolo_args(opts: dict[str, Any], opt_path: Path) -> dict[str, Any]: + """Return required legacy YOLO upload args from opt.yaml.""" + if "imgsz" in opts: + image_size = opts["imgsz"] + elif "img_size" in opts: + image_size = opts["img_size"] + else: + raise ModelPackagingError(f"{opt_path} is missing required key 'imgsz' or 'img_size'.") + if "batch_size" not in opts: + raise ModelPackagingError(f"{opt_path} is missing required key 'batch_size'.") + return {"imgsz": image_size, "batch": opts["batch_size"]} - except ImportError: - raise RuntimeError( - "The ultralytics python package is required to deploy yolov11" - " models. Please install it with `pip install ultralytics`" - ) - print_warn_for_wrong_dependencies_versions([("ultralytics", ">=", "8.3.0")], ask_to_continue=True) +def _infer_yolo_size(model_instance: Any) -> str | None: + """Infer a YOLO size letter (n/s/m/l/x/...) from a loaded checkpoint. - elif "yolov12" in model_type: - try: - import torch - import ultralytics + Prefers an explicit ``scale`` letter in the model yaml (set by newer + Ultralytics), then maps the ``(depth_multiple, width_multiple)`` pair used by + the classic scaling. Returns None when neither signal is present. + """ + yaml_cfg = getattr(model_instance, "yaml", None) or {} + scale = yaml_cfg.get("scale") + if isinstance(scale, str) and len(scale) == 1 and scale.isalpha(): + return scale.lower() + + depth = yaml_cfg.get("depth_multiple") + width = yaml_cfg.get("width_multiple") + if isinstance(depth, (int, float)) and isinstance(width, (int, float)): + for (ref_depth, ref_width), letter in YOLO_DEPTH_WIDTH_TO_SIZE.items(): + if abs(depth - ref_depth) < 1e-6 and abs(width - ref_width) < 1e-6: + return letter + return None - except ImportError: - raise RuntimeError( - "The ultralytics python package is required to deploy yolov12" - " models. Please install it from `https://github.com/sunsmarterjie/yolov12`" + +def _resolve_yolo_size( + model_type: str, + model_instance: Any, + warnings: list[str], + allow_mismatch: bool = False, +) -> str: + """Fill in or check a YOLO model_type's size suffix against the checkpoint. + + Roboflow rejects bare family names (e.g. ``yolov8``) with an + ``InvalidModelTypeException`` because it needs the model size, and a size that + disagrees with the weights fails conversion. A *missing* size is inferred and + filled in. A *supplied* size that conflicts with the inferred one raises so the + caller can confirm — unless ``allow_mismatch`` is set, in which case the + supplied size is packaged as-is with a warning. A user size is also kept when + the size cannot be inferred. Returns the resolved model_type. + """ + core = model_type.lower().split("-", 1)[0] + family = next((f for f in ULTRALYTICS_YOLO_FAMILIES if core.startswith(f)), None) + if family is None: + return model_type + + inferred = _infer_yolo_size(model_instance) + provided = core[len(family) :] + task_suffix = model_type[len(core) :] + + if inferred is None: + if not provided: + raise SizeMismatchError( + f"model_type '{model_type}' is missing a size suffix and the size " + f"could not be inferred from the checkpoint. Specify it explicitly, " + f"e.g. '{family}n', '{family}s', '{family}m', '{family}l', '{family}x'.", + requested=model_type, ) + return model_type + + if not provided: + warnings.append( + f"Inferred model size '{family}{inferred}' from the checkpoint " + f"architecture (model_type was '{model_type}')." + ) + return f"{family}{inferred}{task_suffix}" + + if provided == inferred: + return model_type - print( - "\n!!! ATTENTION !!!\n" - "Model must be trained and uploaded using ultralytics from https://github.com/sunsmarterjie/yolov12\n" - "or through the Roboflow platform\n" - "!!! ATTENTION !!!\n" + if allow_mismatch: + warnings.append( + f"model_type '{model_type}' declares size '{provided}', but the checkpoint " + f"architecture is '{family}{inferred}'. Packaging as '{model_type}' as requested." ) + return model_type + + raise SizeMismatchError( + f"You specified model_type '{model_type}' (size '{provided}'), but the " + f"checkpoint architecture is '{family}{inferred}'. They don't match, so " + f"Roboflow's weight conversion would fail. Upload as '{family}{inferred}" + f"{task_suffix}', or set allow_size_mismatch=True to upload " + f"'{model_type}' exactly as specified.", + requested=model_type, + detected=f"{family}{inferred}{task_suffix}", + ) + - print_warn_for_wrong_dependencies_versions([("ultralytics", "==", "8.3.63")], ask_to_continue=True) +def _process_yolo( + model_type: str, + model_path: Path, + filename: str, + build_dir: Path, + allow_dependency_mismatch: bool, + allow_size_mismatch: bool, +) -> tuple[Path, str, list[str]]: + warnings: list[str] = [] + torch = _import_required_module("torch", "pip install torch") + ultralytics = None + if "yolov8" in model_type: + ultralytics = _import_required_module("ultralytics", "pip install ultralytics==8.0.196") + _check_dependency_version( + dependency="ultralytics", + operator="==", + required_version="8.0.196", + allow_mismatch=allow_dependency_mismatch, + warnings=warnings, + ) + elif "yolov10" in model_type: + ultralytics = _import_required_module("ultralytics", "pip install ultralytics") + elif "yolov11" in model_type: + ultralytics = _import_required_module("ultralytics", "pip install 'ultralytics>=8.3.0'") + _check_dependency_version( + dependency="ultralytics", + operator=">=", + required_version="8.3.0", + allow_mismatch=allow_dependency_mismatch, + warnings=warnings, + ) + elif "yolov12" in model_type: + ultralytics = _import_required_module( + "ultralytics", + "pip install git+https://github.com/sunsmarterjie/yolov12.git", + ) + warnings.append( + "YOLOv12 uploads must use the Ultralytics fork from " + "https://github.com/sunsmarterjie/yolov12 or a Roboflow-trained model." + ) + _check_dependency_version( + dependency="ultralytics", + operator="==", + required_version="8.3.63", + allow_mismatch=allow_dependency_mismatch, + warnings=warnings, + ) elif "yolo26" in model_type: - try: - import torch - import ultralytics + ultralytics = _import_required_module("ultralytics", "pip install ultralytics") - except ImportError: - raise RuntimeError( - "The ultralytics python package is required to deploy yolo26" - " models. Please install it with `pip install ultralytics`" - ) + checkpoint_path = model_path / filename + if not checkpoint_path.exists(): + raise MissingFileError(f"Model weights file '{checkpoint_path}' was not found.") - model = torch.load(os.path.join(model_path, filename), weights_only=False) + checkpoint = _load_checkpoint(torch, checkpoint_path) + if not isinstance(checkpoint, dict): + raise ModelPackagingError(f"Model weights file '{checkpoint_path}' is not a supported checkpoint dictionary.") + model_instance = checkpoint.get("model") or checkpoint.get("ema") + if model_instance is None: + raise ModelPackagingError("Could not find a 'model' or 'ema' entry in the checkpoint.") - model_instance = model["model"] if "model" in model and model["model"] is not None else model["ema"] + model_type = _resolve_yolo_size(model_type, model_instance, warnings, allow_size_mismatch) detected_task = _detect_yolo_task(model_instance) if detected_task: @@ -213,28 +654,15 @@ def _process_yolo(model_type: str, model_path: str, filename: str) -> tuple[str, if existing_task == TASK_DET and detected_task != TASK_DET: model_type = f"{model_type}-{detected_task}" elif existing_task != detected_task: - raise ValueError( + raise TaskMismatchError( f"model_type '{model_type}' implies task '{existing_task}' but the " f".pt file is a '{detected_task}' checkpoint. Use a matching model_type." ) - if isinstance(model_instance.names, list): - class_names = model_instance.names - else: - class_names = [] - for i, val in enumerate(model_instance.names): - class_names.append((val, model_instance.names[val])) - class_names.sort(key=lambda x: x[0]) - class_names = [x[1] for x in class_names] - - if ( - "yolov8" in model_type - or "yolov10" in model_type - or "yolov11" in model_type - or "yolov12" in model_type - or "yolo26" in model_type - ): - # try except for backwards compatibility with older versions of ultralytics + class_names = _class_names_from_model_instance(model_instance) + if any(name in model_type for name in ULTRALYTICS_YOLO_FAMILIES): + if ultralytics is None: + ultralytics = _import_required_module("ultralytics", "pip install ultralytics") if ( "-cls" in model_type or model_type.startswith("yolov10") @@ -243,78 +671,51 @@ def _process_yolo(model_type: str, model_path: str, filename: str) -> tuple[str, or model_type.startswith("yolo26") ): nc = model_instance.yaml["nc"] - args = model["train_args"] + args = checkpoint["train_args"] else: nc = model_instance.nc args = model_instance.args - try: - model_artifacts = { - "names": class_names, - "yaml": model_instance.yaml, - "nc": nc, - "args": {k: val for k, val in args.items() if ((k == "model") or (k == "imgsz") or (k == "batch"))}, - "ultralytics_version": ultralytics.__version__, - "model_type": model_type, - } - except Exception: - model_artifacts = { - "names": class_names, - "yaml": model_instance.yaml, - "nc": nc, - "args": { - k: val for k, val in args.__dict__.items() if ((k == "model") or (k == "imgsz") or (k == "batch")) - }, - "ultralytics_version": ultralytics.__version__, - "model_type": model_type, - } - elif "yolov5" in model_type or "yolov7" in model_type or "yolov9" in model_type: - # parse from yaml for yolov5 - - with open(os.path.join(model_path, "opt.yaml")) as stream: - opts = yaml.safe_load(stream) - + model_artifacts: dict[str, Any] = { + "names": class_names, + "yaml": model_instance.yaml, + "nc": nc, + "args": _filtered_args(args), + "ultralytics_version": ultralytics.__version__, + "model_type": model_type, + } + else: + # yolov5 / yolov7 / yolov9 read their upload args from opt.yaml + opt_path = model_path / "opt.yaml" + if not opt_path.exists(): + raise MissingFileError(f"You must provide an opt.yaml file at '{opt_path}' for {model_type} uploads.") + with opt_path.open() as stream: + opts = yaml.safe_load(stream) or {} model_artifacts = { "names": class_names, "nc": model_instance.nc, - "args": { - "imgsz": opts["imgsz"] if "imgsz" in opts else opts["img_size"], - "batch": opts["batch_size"], - }, + "args": _legacy_yolo_args(opts, opt_path), "model_type": model_type, } if hasattr(model_instance, "yaml"): model_artifacts["yaml"] = model_instance.yaml - with open(os.path.join(model_path, "model_artifacts.json"), "w") as fp: - json.dump(model_artifacts, fp) - - torch.save(model_instance.state_dict(), os.path.join(model_path, "state_dict.pt")) - - list_files = [ - "results.csv", - "results.png", - "model_artifacts.json", - "state_dict.pt", - ] - - zip_file_name = "roboflow_deploy.zip" - - with zipfile.ZipFile(os.path.join(model_path, zip_file_name), "w") as zipMe: - for file in list_files: - if os.path.exists(os.path.join(model_path, file)): - zipMe.write( - os.path.join(model_path, file), - arcname=file, - compress_type=zipfile.ZIP_DEFLATED, - ) - else: - if file in ["model_artifacts.json", "state_dict.pt"]: - raise (ValueError(f"File {file} not found. Please make sure to provide a valid model path.")) - - return zip_file_name, model_type + (build_dir / "model_artifacts.json").write_text(json.dumps(model_artifacts)) + torch.save(model_instance.state_dict(), build_dir / "state_dict.pt") + + archive_path = build_dir / "roboflow_deploy.zip" + _write_zip( + archive_path, + [ + (model_path / "results.csv", "results.csv", False), + (model_path / "results.png", "results.png", False), + (build_dir / "model_artifacts.json", "model_artifacts.json", True), + (build_dir / "state_dict.pt", "state_dict.pt", True), + ], + ) + return archive_path, model_type, warnings -def _detect_rfdetr_task(checkpoint) -> Optional[str]: +def _detect_rfdetr_task(checkpoint: Any) -> str | None: """Detect the training task of an rf-detr checkpoint. rf-detr currently only supports weight upload for detection and instance @@ -328,83 +729,328 @@ def _detect_rfdetr_task(checkpoint) -> Optional[str]: model_name = checkpoint.get("model_name") if isinstance(model_name, str): return TASK_SEG if TASK_SEG in model_name.lower() else TASK_DET - args = checkpoint.get("args") - if args is None: + raw_args = checkpoint.get("args") + if raw_args is None: return None - seg_head = args.get("segmentation_head") if isinstance(args, dict) else getattr(args, "segmentation_head", None) - if seg_head is True: + args = raw_args if isinstance(raw_args, dict) else vars(raw_args) + segmentation_head = args.get("segmentation_head") + if segmentation_head is True: return TASK_SEG - if seg_head is False: + if segmentation_head is False: return TASK_DET return None -def _process_rfdetr(model_type: str, model_path: str, filename: str) -> tuple[str, str]: - _supported_types = [ - # Detection models - "rfdetr-base", - "rfdetr-nano", - "rfdetr-small", - "rfdetr-medium", - "rfdetr-large", - "rfdetr-xlarge", - "rfdetr-2xlarge", - # Segmentation models - "rfdetr-seg-nano", - "rfdetr-seg-small", - "rfdetr-seg-medium", - "rfdetr-seg-large", - "rfdetr-seg-xlarge", - "rfdetr-seg-2xlarge", - ] - if model_type not in _supported_types: - raise ValueError(f"Model type {model_type} not supported. Supported types are {_supported_types}") - - if not os.path.exists(model_path): - raise FileNotFoundError(f"Model path {model_path} does not exist.") - - model_files = os.listdir(model_path) - pt_file = next((f for f in model_files if f.endswith(".pt") or f.endswith(".pth")), None) - - if pt_file is None: - raise RuntimeError("No .pt or .pth model file found in the provided path") - - import torch - - checkpoint = torch.load(os.path.join(model_path, pt_file), map_location="cpu", weights_only=False) +def _rfdetr_checkpoint_pe_size(checkpoint: Any) -> int | None: + """Return an RF-DETR checkpoint's position-encoding grid size (tokens per side). - detected_task = _detect_rfdetr_task(checkpoint) - if detected_task: - implied_task = task_of_model_type(model_type) - if detected_task != implied_task: - raise ValueError( - f"model_type '{model_type}' implies task '{implied_task}' but the " - f".pt is a '{detected_task}' rfdetr checkpoint. Use a matching model_type." + Prefers the explicit ``positional_encoding_size`` arg, then ``resolution // + patch_size``, then derives it from the backbone ``position_embeddings`` tensor + (``grid² + 1`` tokens). Returns None when the geometry cannot be determined. + """ + if not isinstance(checkpoint, dict): + return None + raw_args = checkpoint.get("args") + if isinstance(raw_args, dict): + args = raw_args + elif raw_args is not None: + args = dict(vars(raw_args)) + else: + args = {} + + pe = args.get("positional_encoding_size") + if isinstance(pe, int) and pe > 0: + return pe + resolution = args.get("resolution") + patch_size = args.get("patch_size") + if isinstance(resolution, int) and isinstance(patch_size, int) and patch_size > 0: + return resolution // patch_size + + state_dict = checkpoint.get("model") + if isinstance(state_dict, dict): + for key, tensor in state_dict.items(): + if not key.endswith("position_embeddings"): + continue + shape = getattr(tensor, "shape", None) + if shape is not None and len(shape) == 3: + grid = math.isqrt(int(shape[1]) - 1) + if grid > 0 and grid * grid == int(shape[1]) - 1: + return grid + return None + + +def _resolve_rfdetr_variant( + model_type: str, + checkpoint: Any, + warnings: list[str], + allow_mismatch: bool = False, +) -> str: + """Check an RF-DETR model_type's size variant against the checkpoint geometry. + + Roboflow rebuilds the architecture from ``model_type`` at the variant's default + resolution before loading the weights, so a variant whose position-encoding grid + differs from the checkpoint fails conversion with a ``position_embeddings`` size + mismatch. Two cases: + + * The checkpoint's grid matches a *different* known variant (e.g. uploaded as + ``rfdetr-seg-nano`` but the grid is ``rfdetr-seg-small``) — a high-confidence + mislabel. Raise, naming the variant that fits, so the caller can confirm. + * The grid matches *no* known variant — likely a custom training resolution. We + cannot tell whether the backend supports it, so warn and proceed rather than + block a possibly-valid upload. + + ``allow_mismatch`` always proceeds with the requested variant (with a warning). + The detection-vs-segmentation task is handled separately and is not changed + here. Returns the resolved model_type. + """ + expected = RFDETR_POSITIONAL_ENCODING_SIZE.get(model_type) + actual = _rfdetr_checkpoint_pe_size(checkpoint) + if actual is None or expected is None or actual == expected: + return model_type + + task = task_of_model_type(model_type) + match = next( + ( + name + for name, grid in RFDETR_POSITIONAL_ENCODING_SIZE.items() + if grid == actual and task_of_model_type(name) == task + ), + None, + ) + + if match is not None and not allow_mismatch: + raise SizeMismatchError( + f"You specified model_type '{model_type}' (a {expected}x{expected} " + f"position-encoding grid), but the checkpoint was trained with " + f"{actual}x{actual}, which matches '{match}'. They don't match, so " + f"Roboflow's weight conversion would fail to load the backbone position " + f"embeddings. Upload as '{match}', or set allow_size_mismatch=True to " + f"upload '{model_type}' exactly as specified.", + requested=model_type, + detected=match, + ) + + if match is None: + warnings.append( + f"model_type '{model_type}' expects a {expected}x{expected} position-encoding " + f"grid, but the checkpoint is {actual}x{actual} and matches no known RF-DETR " + f"variant (it may use a custom training resolution). Packaging as " + f"'{model_type}'; Roboflow's conversion may reject it if it rebuilds at the " + f"variant's default resolution." + ) + else: + warnings.append( + f"model_type '{model_type}' expects a {expected}x{expected} grid, but the " + f"checkpoint is {actual}x{actual} (matches '{match}'). Packaging as " + f"'{model_type}' as requested." + ) + return model_type + + +def _find_rfdetr_checkpoint(model_path: Path, filename: str, warnings: list[str]) -> Path: + """Locate the rf-detr checkpoint, preserving the historical discovery behavior. + + The requested ``filename`` wins when it exists. Otherwise fall back to the + first top-level .pt/.pth file (sorted for determinism), which is how rf-detr + uploads located the checkpoint before ``filename`` was honored for them. + """ + requested_file = model_path / filename + if requested_file.exists(): + return requested_file + + discovered = sorted(path for path in model_path.iterdir() if path.is_file() and path.suffix in {".pt", ".pth"}) + if not discovered: + raise MissingFileError( + f"No .pt or .pth checkpoint found in '{model_path}' (and '{requested_file}' does not exist)." + ) + warnings.append( + f"Weights file '{requested_file}' was not found; using discovered checkpoint '{discovered[0].name}' instead." + ) + return discovered[0] + + +def _write_rfdetr_class_names(model_path: Path, build_dir: Path, checkpoint: Any) -> Path: + class_names_path = model_path / "class_names.txt" + if class_names_path.exists(): + class_names = class_names_path.read_text().splitlines() + else: + raw_args = checkpoint.get("args") if isinstance(checkpoint, dict) else None + if raw_args is None: + args: dict[str, Any] = {} + elif isinstance(raw_args, dict): + args = raw_args + else: + args = dict(vars(raw_args)) + class_names = args.get("class_names") or [] + if not class_names: + raise MissingFileError( + f"No class_names.txt file found in '{model_path}', and the RF-DETR " + "checkpoint does not include args with class_names. This should only " + "happen on rfdetr models trained before version 1.1.0. Create " + "class_names.txt with one class per line or retrain with a newer " + "rfdetr library." ) - get_classnames_txt_for_rfdetr(model_path, pt_file, checkpoint=checkpoint) + if "background_class83422" not in class_names: + class_names = ["background_class83422", *class_names] + output_path = build_dir / "class_names.txt" + output_path.write_text("\n".join(class_names) + "\n") + return output_path + + +def _process_rfdetr( + model_type: str, + model_path: Path, + filename: str, + build_dir: Path, + allow_size_mismatch: bool, +) -> tuple[Path, str, list[str]]: + if model_type not in SUPPORTED_RFDETR_TYPES: + raise UnsupportedModelError( + f"Model type '{model_type}' is not supported for RF-DETR upload. " + f"Supported types are: {', '.join(SUPPORTED_RFDETR_TYPES)}." + ) + torch = _import_required_module("torch", "pip install torch") + warnings: list[str] = [] - # Copy the .pt file to weights.pt if not already named weights.pt - if pt_file != "weights.pt": - shutil.copy(os.path.join(model_path, pt_file), os.path.join(model_path, "weights.pt")) + checkpoint_path = _find_rfdetr_checkpoint(model_path, filename, warnings) + checkpoint = _load_checkpoint(torch, checkpoint_path, map_location="cpu") - required_files = ["weights.pt"] + detected_task = _detect_rfdetr_task(checkpoint) + if detected_task and detected_task != task_of_model_type(model_type): + raise TaskMismatchError( + f"model_type '{model_type}' implies task '{task_of_model_type(model_type)}', " + f"but the checkpoint is a '{detected_task}' RF-DETR model. Use a matching model_type." + ) + + model_type = _resolve_rfdetr_variant(model_type, checkpoint, warnings, allow_size_mismatch) + + shutil.copy(checkpoint_path, build_dir / "weights.pt") + _write_rfdetr_class_names(model_path, build_dir, checkpoint) + archive_path = build_dir / "roboflow_deploy.zip" + _write_zip( + archive_path, + [ + (build_dir / "weights.pt", "weights.pt", True), + (model_path / "results.csv", "results.csv", False), + (model_path / "results.png", "results.png", False), + (model_path / "model_artifacts.json", "model_artifacts.json", False), + (build_dir / "class_names.txt", "class_names.txt", False), + ], + ) + return archive_path, model_type, warnings - optional_files = ["results.csv", "results.png", "model_artifacts.json", "class_names.txt"] - zip_file_name = "roboflow_deploy.zip" - with zipfile.ZipFile(os.path.join(model_path, zip_file_name), "w") as zipMe: - for file in required_files: - zipMe.write(os.path.join(model_path, file), arcname=file, compress_type=zipfile.ZIP_DEFLATED) +def _process_huggingface( + model_type: str, + model_path: Path, + build_dir: Path, +) -> tuple[Path, str, list[str]]: + if model_type not in SUPPORTED_HUGGINGFACE_TYPES: + raise UnsupportedModelError( + f"Model type '{model_type}' is not supported for this type of upload. " + f"Supported types are: {', '.join(SUPPORTED_HUGGINGFACE_TYPES)}." + ) + + model_files = [path for path in model_path.iterdir() if path.is_file()] + safetensors_files = [path for path in model_files if path.suffix == ".safetensors"] + npz_file = next((path for path in model_files if path.suffix == ".npz"), None) + if safetensors_files: + required = { + "preprocessor_config.json", + "special_tokens_map.json", + "tokenizer_config.json", + "tokenizer.json", + } + missing = sorted(required - {path.name for path in model_files}) + if missing: + raise MissingFileError(f"Missing files required for a PyTorch {model_type} upload: {', '.join(missing)}.") + files_to_deploy = model_files + elif npz_file is not None: + files_to_deploy = [npz_file] + else: + raise MissingFileError(f"No .npz or .safetensors model file found in '{model_path}'.") + + archive_path = build_dir / "roboflow_deploy.tar" + with tarfile.open(archive_path, "w") as tar: + for path in files_to_deploy: + tar.add(path, arcname=path.name) + return archive_path, model_type, [] + + +def _process_yolonas( + model_type: str, + model_path: Path, + filename: str, + build_dir: Path, +) -> tuple[Path, str, list[str]]: + torch = _import_required_module("torch", "pip install torch") + weights_path = model_path / filename + if not weights_path.exists(): + raise MissingFileError(f"Model weights file '{weights_path}' was not found.") + + checkpoint = _load_checkpoint(torch, weights_path, map_location="cpu") + class_names = checkpoint["processing_params"]["class_names"] + opt_path = model_path / "opt.yaml" + if not opt_path.exists(): + raise MissingFileError( + f"You must create an opt.yaml file at '{opt_path}' of the format:\n" + f"imgsz: \n" + f"batch_size: \n" + f"architecture: \n" + ) + with opt_path.open() as stream: + opts = yaml.safe_load(stream) or {} + missing = [key for key in ("imgsz", "batch_size", "architecture") if key not in opts] + if missing: + raise ModelPackagingError(f"{opt_path} lacks required keys: {', '.join(missing)}.") + + model_artifacts = { + "names": class_names, + "nc": len(class_names), + "args": { + "imgsz": opts["imgsz"], + "batch": opts["batch_size"], + "architecture": opts["architecture"], + }, + "model_type": model_type, + } + (build_dir / "model_artifacts.json").write_text(json.dumps(model_artifacts)) + shutil.copy(weights_path, build_dir / "state_dict.pt") + + archive_path = build_dir / "roboflow_deploy.zip" + _write_zip( + archive_path, + [ + (model_path / "results.json", "results.json", False), + (model_path / "results.png", "results.png", False), + (build_dir / "model_artifacts.json", "model_artifacts.json", True), + (build_dir / "state_dict.pt", "state_dict.pt", True), + ], + ) + return archive_path, model_type, [] - for file in optional_files: - if os.path.exists(os.path.join(model_path, file)): - zipMe.write(os.path.join(model_path, file), arcname=file, compress_type=zipfile.ZIP_DEFLATED) - return zip_file_name, model_type +def _write_zip( + archive_path: Path, + files: list[tuple[Path, str, bool]], +) -> None: + with zipfile.ZipFile(archive_path, "w") as zip_file: + for path, arcname, required in files: + if path.exists(): + zip_file.write(path, arcname=arcname, compress_type=zipfile.ZIP_DEFLATED) + elif required: + raise MissingFileError(f"Required upload artifact '{path}' was not found.") def get_classnames_txt_for_rfdetr(model_path: str, pt_file: str, checkpoint=None): + """Legacy rf-detr class-names helper, kept for backwards compatibility. + + Writes (and mutates) ``class_names.txt`` inside ``model_path``. The packaging + flow uses :func:`_write_rfdetr_class_names` instead, which leaves the source + directory untouched. + """ class_names_path = os.path.join(model_path, "class_names.txt") if os.path.exists(class_names_path): maybe_prepend_dummy_class(class_names_path) @@ -424,7 +1070,7 @@ def get_classnames_txt_for_rfdetr(model_path: str, pt_file: str, checkpoint=None maybe_prepend_dummy_class(class_names_path) return class_names_path - raise FileNotFoundError( + raise MissingFileError( f"No class_names.txt file found in model path {model_path}.\n" f"This should only happen on rfdetr models trained before version 1.1.0.\n" f"Please re-train your model with the latest version of the rfdetr library, or\n" @@ -442,122 +1088,3 @@ def maybe_prepend_dummy_class(class_name_file: str): class_names.insert(0, dummy_class) with open(class_name_file, "w") as f: f.writelines(class_names) - - -def _process_huggingface( - model_type: str, model_path: str, filename: str = "fine-tuned-paligemma-3b-pt-224.f16.npz" -) -> tuple[str, str]: - # Check if model_path exists - if not os.path.exists(model_path): - raise FileNotFoundError(f"Model path {model_path} does not exist.") - model_files = os.listdir(model_path) - print(f"Model files found in {model_path}: {model_files}") - - files_to_deploy = [] - - # Find first .npz file in model_path - npz_filename = next((file for file in model_files if file.endswith(".npz")), None) - if any([file.endswith(".safetensors") for file in model_files]): - print(f"Found .safetensors file in model path. Deploying PyTorch {model_type} model.") - necessary_files = [ - "preprocessor_config.json", - "special_tokens_map.json", - "tokenizer_config.json", - "tokenizer.json", - ] - for file in necessary_files: - if file not in model_files: - print("Missing necessary file", file) - res = input("Do you want to continue? (y/n)") - if res.lower() != "y": - exit(1) - for file in model_files: - files_to_deploy.append(file) - elif npz_filename is not None: - print(f"Found .npz file {npz_filename} in model path. Deploying JAX PaliGemma model.") - files_to_deploy.append(npz_filename) - else: - raise FileNotFoundError(f"No .npz or .safetensors file found in model path {model_path}.") - - if len(files_to_deploy) == 0: - raise FileNotFoundError(f"No valid files found in model path {model_path}.") - print(f"Zipping files for deploy: {files_to_deploy}") - - import tarfile - - tar_file_name = "roboflow_deploy.tar" - - with tarfile.open(os.path.join(model_path, tar_file_name), "w") as tar: - for file in files_to_deploy: - tar.add(os.path.join(model_path, file), arcname=file) - - print("Uploading to Roboflow... May take several minutes.") - - return tar_file_name, model_type - - -def _process_yolonas(model_type: str, model_path: str, filename: str = "weights/best.pt") -> tuple[str, str]: - try: - import torch - except ImportError: - raise RuntimeError( - "The torch python package is required to deploy yolonas models. Please install it with `pip install torch`" - ) - - model = torch.load(os.path.join(model_path, filename), map_location="cpu") - class_names = model["processing_params"]["class_names"] - - opt_path = os.path.join(model_path, "opt.yaml") - if not os.path.exists(opt_path): - raise RuntimeError( - f"You must create an opt.yaml file at {os.path.join(model_path, '')} of the format:\n" - f"imgsz: \n" - f"batch_size: \n" - f"architecture: \n" - ) - with open(os.path.join(model_path, "opt.yaml")) as stream: - opts = yaml.safe_load(stream) - required_keys = ["imgsz", "batch_size", "architecture"] - for key in required_keys: - if key not in opts: - raise RuntimeError(f"{opt_path} lacks required key {key}. Required keys: {required_keys}") - - model_artifacts = { - "names": class_names, - "nc": len(class_names), - "args": { - "imgsz": opts["imgsz"] if "imgsz" in opts else opts["img_size"], - "batch": opts["batch_size"], - "architecture": opts["architecture"], - }, - "model_type": model_type, - } - - with open(os.path.join(model_path, "model_artifacts.json"), "w") as fp: - json.dump(model_artifacts, fp) - - shutil.copy(os.path.join(model_path, filename), os.path.join(model_path, "state_dict.pt")) - - list_files = [ - "results.json", - "results.png", - "model_artifacts.json", - "state_dict.pt", - ] - - zip_file_name = "roboflow_deploy.zip" - - with zipfile.ZipFile(os.path.join(model_path, zip_file_name), "w") as zipMe: - for file in list_files: - if os.path.exists(os.path.join(model_path, file)): - zipMe.write( - os.path.join(model_path, file), - arcname=file, - compress_type=zipfile.ZIP_DEFLATED, - ) - else: - if file in ["model_artifacts.json", filename]: - raise (ValueError(f"File {file} not found. Please make sure to provide a valid model path.")) - - return zip_file_name, model_type diff --git a/tests/util/test_model_processor.py b/tests/util/test_model_processor.py index 80408602..db848bae 100644 --- a/tests/util/test_model_processor.py +++ b/tests/util/test_model_processor.py @@ -1,14 +1,36 @@ +import json import os +import sys import tempfile +import types import unittest +import zipfile +from pathlib import Path from types import SimpleNamespace +from unittest import mock from roboflow.config import TASK_CLS, TASK_DET, TASK_OBB, TASK_POSE, TASK_SEG, TASK_SEM +from roboflow.util import model_processor from roboflow.util.model_processor import ( + MissingFileError, + ModelPackagingError, + SizeMismatchError, + TaskMismatchError, + UnsupportedModelError, _detect_rfdetr_task, _detect_yolo_task, + _infer_yolo_size, + _legacy_yolo_args, + _resolve_rfdetr_variant, + _resolve_yolo_size, + _rfdetr_checkpoint_pe_size, + _write_rfdetr_class_names, get_classnames_txt_for_rfdetr, + package_custom_weights, + package_custom_weights_interactive, + process, task_of_model_type, + validate_model_type_for_project, ) @@ -110,5 +132,371 @@ def test_namespace_args(self): ) +class ValidateModelTypeForProjectTest(unittest.TestCase): + def test_rejects_detection_for_classification(self): + with self.assertRaises(TaskMismatchError) as ctx: + validate_model_type_for_project("yolov8", "classification", "widgets") + self.assertIn("classification", str(ctx.exception)) + self.assertIn("task 'cls'", str(ctx.exception)) + + def test_task_mismatch_is_a_value_error(self): + # Callers that caught the historical ValueError keep working. + with self.assertRaises(ValueError): + validate_model_type_for_project("yolov8", "classification", "widgets") + + def test_unknown_project_type_is_ignored(self): + validate_model_type_for_project("yolov8", "some-new-type", "widgets") + + +class LegacyYoloArgsTest(unittest.TestCase): + def test_reports_missing_batch_size(self): + with self.assertRaises(ModelPackagingError) as ctx: + _legacy_yolo_args({"imgsz": 640}, Path("opt.yaml")) + self.assertIn("batch_size", str(ctx.exception)) + + def test_reports_missing_image_size(self): + with self.assertRaises(ModelPackagingError) as ctx: + _legacy_yolo_args({"batch_size": 8}, Path("opt.yaml")) + self.assertIn("imgsz", str(ctx.exception)) + + def test_accepts_either_image_size_key(self): + self.assertEqual(_legacy_yolo_args({"imgsz": 640, "batch_size": 8}, Path("x")), {"imgsz": 640, "batch": 8}) + self.assertEqual( + _legacy_yolo_args({"img_size": 416, "batch_size": 4}, Path("x")), {"imgsz": 416, "batch": 4} + ) + + +class _FakeYoloModel: + def __init__(self, yaml): + self.yaml = yaml + + +class InferYoloSizeTest(unittest.TestCase): + def test_from_depth_width_multiples(self): + model = _FakeYoloModel({"depth_multiple": 0.33, "width_multiple": 0.25}) + self.assertEqual(_infer_yolo_size(model), "n") + + def test_explicit_scale_letter_wins(self): + model = _FakeYoloModel({"scale": "m", "depth_multiple": 0.67, "width_multiple": 0.75}) + self.assertEqual(_infer_yolo_size(model), "m") + + def test_unknown_returns_none(self): + self.assertIsNone(_infer_yolo_size(_FakeYoloModel({}))) + + +class ResolveYoloSizeTest(unittest.TestCase): + def test_fills_bare_family_from_architecture(self): + warnings: list = [] + model = _FakeYoloModel({"depth_multiple": 0.33, "width_multiple": 0.25}) + self.assertEqual(_resolve_yolo_size("yolov8", model, warnings), "yolov8n") + self.assertTrue(warnings and "Inferred model size 'yolov8n'" in warnings[0]) + + def test_preserves_task_suffix_when_filling(self): + model = _FakeYoloModel({"depth_multiple": 0.33, "width_multiple": 0.50}) + self.assertEqual(_resolve_yolo_size("yolov8-seg", model, []), "yolov8s-seg") + + def test_raises_when_size_cannot_be_inferred(self): + with self.assertRaises(SizeMismatchError) as ctx: + _resolve_yolo_size("yolov8", _FakeYoloModel({}), []) + self.assertIn("could not be inferred", str(ctx.exception)) + self.assertIn("yolov8n", str(ctx.exception)) + + def test_raises_on_declared_size_conflict(self): + model = _FakeYoloModel({"depth_multiple": 0.33, "width_multiple": 0.25}) + with self.assertRaises(SizeMismatchError) as ctx: + _resolve_yolo_size("yolov8m", model, []) + self.assertIn("yolov8n", str(ctx.exception)) + self.assertIn("allow_size_mismatch=True", str(ctx.exception)) + self.assertEqual(ctx.exception.requested, "yolov8m") + self.assertEqual(ctx.exception.detected, "yolov8n") + + def test_allow_mismatch_keeps_declared_size(self): + warnings: list = [] + model = _FakeYoloModel({"depth_multiple": 0.33, "width_multiple": 0.25}) + self.assertEqual(_resolve_yolo_size("yolov8m", model, warnings, allow_mismatch=True), "yolov8m") + self.assertTrue(warnings and "as requested" in warnings[0]) + + def test_keeps_user_size_when_not_inferable(self): + warnings: list = [] + self.assertEqual(_resolve_yolo_size("yolov8m", _FakeYoloModel({}), warnings), "yolov8m") + self.assertEqual(warnings, []) + + def test_keeps_matching_sized_type_without_warning(self): + warnings: list = [] + model = _FakeYoloModel({"depth_multiple": 0.33, "width_multiple": 0.25}) + self.assertEqual(_resolve_yolo_size("yolov8n", model, warnings), "yolov8n") + self.assertEqual(warnings, []) + + +class ResolveRfdetrVariantTest(unittest.TestCase): + def test_raises_on_size_conflict_naming_the_fit(self): + checkpoint = {"args": {"resolution": 384, "patch_size": 12}} + with self.assertRaises(SizeMismatchError) as ctx: + _resolve_rfdetr_variant("rfdetr-seg-nano", checkpoint, []) + self.assertIn("rfdetr-seg-small", str(ctx.exception)) + self.assertIn("32x32", str(ctx.exception)) + self.assertIn("allow_size_mismatch=True", str(ctx.exception)) + + def test_allow_mismatch_keeps_requested_variant(self): + warnings: list = [] + checkpoint = {"args": {"resolution": 384, "patch_size": 12}} + resolved = _resolve_rfdetr_variant("rfdetr-seg-nano", checkpoint, warnings, allow_mismatch=True) + self.assertEqual(resolved, "rfdetr-seg-nano") + self.assertTrue(warnings and "as requested" in warnings[0]) + + def test_keeps_matching_grid_without_warning(self): + warnings: list = [] + checkpoint = {"args": {"positional_encoding_size": 32}} + self.assertEqual(_resolve_rfdetr_variant("rfdetr-seg-small", checkpoint, warnings), "rfdetr-seg-small") + self.assertEqual(warnings, []) + + def test_warns_but_allows_custom_resolution(self): + warnings: list = [] + resolved = _resolve_rfdetr_variant("rfdetr-seg-nano", {"args": {"positional_encoding_size": 99}}, warnings) + self.assertEqual(resolved, "rfdetr-seg-nano") + self.assertTrue(warnings and "matches no known RF-DETR variant" in warnings[0]) + + def test_pe_size_derived_from_position_embeddings_tensor(self): + class FakeTensor: + shape = (1, 1025, 384) + + checkpoint = {"model": {"backbone.0.encoder.encoder.embeddings.position_embeddings": FakeTensor()}} + self.assertEqual(_rfdetr_checkpoint_pe_size(checkpoint), 32) + + +class WriteRfdetrClassNamesTest(unittest.TestCase): + def test_fails_cleanly_without_checkpoint_args(self): + with tempfile.TemporaryDirectory() as tmp: + with self.assertRaises(MissingFileError) as ctx: + _write_rfdetr_class_names(Path(tmp), Path(tmp), checkpoint={}) + self.assertIn("does not include args with class_names", str(ctx.exception)) + + def test_does_not_mutate_existing_class_names_file(self): + with tempfile.TemporaryDirectory() as tmp, tempfile.TemporaryDirectory() as build: + source = Path(tmp) / "class_names.txt" + source.write_text("cat\ndog\n") + output = _write_rfdetr_class_names(Path(tmp), Path(build), checkpoint={}) + self.assertEqual(source.read_text(), "cat\ndog\n") + self.assertEqual( + output.read_text().splitlines(), + ["background_class83422", "cat", "dog"], + ) + + +def _fake_torch(load_result, calls=None): + module = types.ModuleType("torch") + + def load(path, **kwargs): + if calls is not None: + calls.append((Path(path), kwargs)) + return load_result + + def save(obj, path): + Path(path).write_bytes(b"fake-state-dict") + + module.load = load + module.save = save + return module + + +def _import_patch(modules): + def _import(module_name, install_hint): + return modules[module_name] + + return mock.patch.object(model_processor, "_import_required_module", side_effect=_import) + + +def _write_yolonas_inputs(model_dir: Path): + weights = model_dir / "weights" / "best.pt" + weights.parent.mkdir() + weights.write_bytes(b"checkpoint") + (model_dir / "opt.yaml").write_text("imgsz: 640\nbatch_size: 8\narchitecture: yolo_nas_s\n") + + +class PackageCustomWeightsTest(unittest.TestCase): + """Contract tests for the public non-interactive helper.""" + + def _package_yolonas(self, model_dir: Path, **kwargs): + calls: list = [] + torch = _fake_torch({"processing_params": {"class_names": ["widget"]}}, calls) + with _import_patch({"torch": torch}): + bundle = package_custom_weights("yolonas", str(model_dir), **kwargs) + return bundle, calls + + def test_never_prompts_or_exits(self): + prompt_guard = mock.patch( + "builtins.input", side_effect=AssertionError("package_custom_weights must not prompt") + ) + exit_guard = mock.patch.object( + sys, "exit", side_effect=AssertionError("package_custom_weights must not exit") + ) + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + _write_yolonas_inputs(model_dir) + with prompt_guard, exit_guard: + bundle, calls = self._package_yolonas(model_dir) + try: + self.assertTrue(bundle.archive_path.exists()) + self.assertEqual(calls[0][1], {"weights_only": False, "map_location": "cpu"}) + finally: + bundle.cleanup() + + def test_does_not_write_into_model_path(self): + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + _write_yolonas_inputs(model_dir) + before = sorted(path for path in model_dir.rglob("*")) + bundle, _ = self._package_yolonas(model_dir) + try: + self.assertEqual(sorted(path for path in model_dir.rglob("*")), before) + self.assertNotEqual(bundle.build_dir, model_dir) + self.assertTrue(bundle.owns_build_dir) + finally: + bundle.cleanup() + self.assertFalse(bundle.build_dir.exists()) + + def test_explicit_build_dir_is_used_and_not_cleaned_up(self): + with tempfile.TemporaryDirectory() as tmp, tempfile.TemporaryDirectory() as build: + model_dir = Path(tmp) + _write_yolonas_inputs(model_dir) + bundle, _ = self._package_yolonas(model_dir, build_dir=build) + self.assertEqual(bundle.build_dir, Path(build)) + self.assertFalse(bundle.owns_build_dir) + bundle.cleanup() + self.assertTrue(bundle.archive_path.exists()) + + def test_owned_build_dir_is_removed_on_failure(self): + with tempfile.TemporaryDirectory() as tmp, tempfile.TemporaryDirectory() as fake_build: + with mock.patch.object(model_processor.tempfile, "mkdtemp", return_value=fake_build): + with self.assertRaises(UnsupportedModelError): + package_custom_weights("not-a-model", tmp) + self.assertFalse(Path(fake_build).exists()) + + def test_missing_model_path_raises_missing_file(self): + with self.assertRaises(MissingFileError): + package_custom_weights("yolonas", "/nonexistent/path/for/test") + + def test_rfdetr_falls_back_to_discovered_checkpoint(self): + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + (model_dir / "other.pt").write_bytes(b"checkpoint") + torch = _fake_torch({"args": {"class_names": ["widget"]}}) + with _import_patch({"torch": torch}): + bundle = package_custom_weights("rfdetr-base", str(model_dir)) + try: + self.assertTrue(any("other.pt" in warning for warning in bundle.warnings)) + with zipfile.ZipFile(bundle.archive_path) as archive: + self.assertIn("weights.pt", archive.namelist()) + self.assertIn("class_names.txt", archive.namelist()) + finally: + bundle.cleanup() + + def test_rfdetr_without_any_checkpoint_raises(self): + with tempfile.TemporaryDirectory() as tmp: + torch = _fake_torch({}) + with _import_patch({"torch": torch}): + with self.assertRaises(MissingFileError): + package_custom_weights("rfdetr-base", tmp) + + def test_yolov8_full_flow_builds_artifacts(self): + checkpoint_names = {1: "dog", 0: "cat"} + + class DetectionModel: + names = checkpoint_names + nc = 2 + yaml = {"nc": 2, "depth_multiple": 0.33, "width_multiple": 0.25} + args = {"model": "yolov8n.yaml", "imgsz": 640, "batch": 16, "lr0": 0.01} + + def state_dict(self): + return {"weight": b"w"} + + fake_ultralytics = types.ModuleType("ultralytics") + fake_ultralytics.__version__ = "8.0.196" + fake_torch = _fake_torch({"model": DetectionModel()}) + + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + weights = model_dir / "weights" / "best.pt" + weights.parent.mkdir() + weights.write_bytes(b"checkpoint") + + with ( + _import_patch({"torch": fake_torch, "ultralytics": fake_ultralytics}), + mock.patch.dict(sys.modules, {"ultralytics": fake_ultralytics}), + ): + bundle = package_custom_weights("yolov8", str(model_dir)) + try: + self.assertEqual(bundle.model_type, "yolov8n") + self.assertTrue(any("Inferred model size 'yolov8n'" in warning for warning in bundle.warnings)) + with zipfile.ZipFile(bundle.archive_path) as archive: + artifacts = json.loads(archive.read("model_artifacts.json")) + self.assertIn("state_dict.pt", archive.namelist()) + self.assertEqual(artifacts["names"], ["cat", "dog"]) + self.assertEqual(artifacts["model_type"], "yolov8n") + self.assertEqual(artifacts["ultralytics_version"], "8.0.196") + self.assertEqual(artifacts["args"], {"model": "yolov8n.yaml", "imgsz": 640, "batch": 16}) + finally: + bundle.cleanup() + + +class ProcessCompatTest(unittest.TestCase): + """The legacy process() entry point keeps its historical contract.""" + + def test_packages_into_model_path_and_returns_tuple(self): + calls: list = [] + torch = _fake_torch({"processing_params": {"class_names": ["widget"]}}, calls) + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + _write_yolonas_inputs(model_dir) + with _import_patch({"torch": torch}): + zip_file_name, model_type = process("yolonas", str(model_dir), "weights/best.pt") + + self.assertEqual(zip_file_name, "roboflow_deploy.zip") + self.assertEqual(model_type, "yolonas") + # Historical side effects: artifacts and archive land in model_path. + self.assertTrue((model_dir / "roboflow_deploy.zip").exists()) + self.assertTrue((model_dir / "model_artifacts.json").exists()) + self.assertTrue((model_dir / "state_dict.pt").exists()) + + +class PackageCustomWeightsInteractiveTest(unittest.TestCase): + def _bundle(self): + return model_processor.ModelUploadBundle( + archive_path=Path("roboflow_deploy.zip"), + build_dir=Path("."), + model_type="yolov8n", + warnings=("some warning",), + ) + + def test_retries_with_size_override_on_confirmation(self): + error = SizeMismatchError("size conflict", requested="yolov8m", detected="yolov8n") + outcomes = [error, self._bundle()] + + def fake_package(*args, **kwargs): + outcome = outcomes.pop(0) + if isinstance(outcome, Exception): + raise outcome + self.assertTrue(kwargs["allow_size_mismatch"]) + return outcome + + with ( + mock.patch.object(model_processor, "package_custom_weights", side_effect=fake_package), + mock.patch("builtins.input", return_value="y"), + mock.patch("builtins.print"), + ): + bundle = package_custom_weights_interactive("yolov8m", "/models") + self.assertEqual(bundle.model_type, "yolov8n") + + def test_reraises_when_user_declines(self): + error = SizeMismatchError("size conflict", requested="yolov8m") + with ( + mock.patch.object(model_processor, "package_custom_weights", side_effect=error), + mock.patch("builtins.input", return_value="n"), + mock.patch("builtins.print"), + ): + with self.assertRaises(SizeMismatchError): + package_custom_weights_interactive("yolov8m", "/models") + + if __name__ == "__main__": unittest.main() From 7d6db31ce7330c0ff269d5f208d6185a63671796 Mon Sep 17 00:00:00 2001 From: Leandro Rosemberg Date: Thu, 2 Jul 2026 00:16:55 -0300 Subject: [PATCH 02/11] refactor: keep legacy interactive confirm behavior in process() --- roboflow/util/model_processor.py | 21 ++++++++++---------- tests/util/test_model_processor.py | 31 ++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 11 deletions(-) diff --git a/roboflow/util/model_processor.py b/roboflow/util/model_processor.py index c6409585..096bbbab 100644 --- a/roboflow/util/model_processor.py +++ b/roboflow/util/model_processor.py @@ -16,10 +16,10 @@ Expected, user-correctable failures raise :class:`ModelPackagingError` subclasses; anything else escaping these helpers is a bug. -:func:`process` is the legacy entry point kept for backwards compatibility, -and the ``Version.deploy`` / ``Workspace.deploy_model`` flows wrap the packaging -step with :func:`package_custom_weights_interactive`, which preserves the -historical print-and-confirm CLI behavior. +The legacy :func:`process` entry point and the ``Version.deploy`` / +``Workspace.deploy_model`` flows wrap the packaging step with +:func:`package_custom_weights_interactive`, which preserves the historical +print-and-confirm CLI behavior. """ from __future__ import annotations @@ -371,14 +371,13 @@ def package_custom_weights_interactive( def process(model_type: str, model_path: str, filename: str) -> tuple[str, str]: """Legacy packaging entry point, kept for backwards compatibility. - Packages into ``model_path`` (the historical behavior: intermediate - artifacts and the final archive land there), prints any packaging warnings, - and returns ``(archive_file_name, resolved_model_type)``. New code should - call :func:`package_custom_weights` instead. + Preserves the historical contract end to end: packages into ``model_path`` + (intermediate artifacts and the final archive land there), prints packaging + warnings, asks for confirmation on dependency/size mismatches, and returns + ``(archive_file_name, resolved_model_type)``. Headless code should call + :func:`package_custom_weights` instead. """ - bundle = package_custom_weights(model_type, model_path, filename, build_dir=model_path) - for warning in bundle.warnings: - print(warning) + bundle = package_custom_weights_interactive(model_type, model_path, filename, build_dir=model_path) return bundle.archive_path.name, bundle.model_type diff --git a/tests/util/test_model_processor.py b/tests/util/test_model_processor.py index db848bae..44620ea4 100644 --- a/tests/util/test_model_processor.py +++ b/tests/util/test_model_processor.py @@ -458,6 +458,37 @@ def test_packages_into_model_path_and_returns_tuple(self): self.assertTrue((model_dir / "model_artifacts.json").exists()) self.assertTrue((model_dir / "state_dict.pt").exists()) + def test_prompts_and_retries_on_mismatch_like_before(self): + error = model_processor.DependencyMismatchError( + "wrong ultralytics", + dependency="ultralytics", + required="ultralytics==8.0.196", + installed="8.3.0", + ) + bundle = model_processor.ModelUploadBundle( + archive_path=Path("roboflow_deploy.zip"), + build_dir=Path("."), + model_type="yolov8n", + ) + outcomes = [error, bundle] + + def fake_package(*args, **kwargs): + outcome = outcomes.pop(0) + if isinstance(outcome, Exception): + raise outcome + self.assertTrue(kwargs["allow_dependency_mismatch"]) + return outcome + + with ( + mock.patch.object(model_processor, "package_custom_weights", side_effect=fake_package), + mock.patch("builtins.input", return_value="y"), + mock.patch("builtins.print"), + ): + zip_file_name, model_type = process("yolov8m", "/models", "weights/best.pt") + + self.assertEqual(zip_file_name, "roboflow_deploy.zip") + self.assertEqual(model_type, "yolov8n") + class PackageCustomWeightsInteractiveTest(unittest.TestCase): def _bundle(self): From f4a378afc0ec7030102114ee1e174549b753b85a Mon Sep 17 00:00:00 2001 From: Leandro Rosemberg Date: Thu, 2 Jul 2026 00:35:57 -0300 Subject: [PATCH 03/11] fix: validate model_type family by prefix instead of substring --- roboflow/util/model_processor.py | 8 ++++---- tests/util/test_model_processor.py | 7 +++++++ 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/roboflow/util/model_processor.py b/roboflow/util/model_processor.py index 096bbbab..03cfe8bf 100644 --- a/roboflow/util/model_processor.py +++ b/roboflow/util/model_processor.py @@ -390,17 +390,17 @@ def _process_model( allow_dependency_mismatch: bool, allow_size_mismatch: bool, ) -> tuple[Path, str, list[str]]: - if not any(supported in model_type for supported in SUPPORTED_MODELS): + if not model_type.startswith(SUPPORTED_MODELS): raise UnsupportedModelError( f"Model type '{model_type}' is not supported for custom weights upload. " - f"Supported families are: {', '.join(SUPPORTED_MODELS)}." + f"It must start with a supported family: {', '.join(SUPPORTED_MODELS)}." ) if model_type.startswith(("paligemma", "paligemma2", "florence-2")): return _process_huggingface(model_type, model_path, build_dir) - if "yolonas" in model_type: + if model_type.startswith("yolonas"): return _process_yolonas(model_type, model_path, filename, build_dir) - if "rfdetr" in model_type: + if model_type.startswith("rfdetr"): return _process_rfdetr(model_type, model_path, filename, build_dir, allow_size_mismatch) return _process_yolo( model_type, diff --git a/tests/util/test_model_processor.py b/tests/util/test_model_processor.py index 44620ea4..2a8c1dce 100644 --- a/tests/util/test_model_processor.py +++ b/tests/util/test_model_processor.py @@ -376,6 +376,13 @@ def test_missing_model_path_raises_missing_file(self): with self.assertRaises(MissingFileError): package_custom_weights("yolonas", "/nonexistent/path/for/test") + def test_family_must_be_a_prefix_not_a_substring(self): + # 'foo-yolov8n' merely contains a family token; the backend would + # reject it after upload, so the gate must reject it up front. + with tempfile.TemporaryDirectory() as tmp: + with self.assertRaises(UnsupportedModelError): + package_custom_weights("foo-yolov8n", tmp) + def test_rfdetr_falls_back_to_discovered_checkpoint(self): with tempfile.TemporaryDirectory() as tmp: model_dir = Path(tmp) From 0f536de0f72e8a727dfe1026b81a6cfb118ce675 Mon Sep 17 00:00:00 2001 From: Leandro Rosemberg Date: Mon, 6 Jul 2026 12:01:23 -0300 Subject: [PATCH 04/11] fix: harden custom-weights packaging edge cases from PR review --- roboflow/util/model_processor.py | 67 ++++++++++++++++++++---------- tests/util/test_model_processor.py | 62 ++++++++++++++++++++++++--- 2 files changed, 101 insertions(+), 28 deletions(-) diff --git a/roboflow/util/model_processor.py b/roboflow/util/model_processor.py index 03cfe8bf..4983564f 100644 --- a/roboflow/util/model_processor.py +++ b/roboflow/util/model_processor.py @@ -102,6 +102,8 @@ "rfdetr-seg-2xlarge", ) +DEFAULT_WEIGHTS_FILENAME = "weights/best.pt" + # YOLO families Roboflow rejects without a size suffix (e.g. `yolov8` must be # `yolov8n`/`yolov8s`/...). Legacy yolov5/7/9 go through the opt.yaml path and are # intentionally excluded. @@ -237,6 +239,20 @@ def task_of_model_type(model_type: str) -> str: return TASK_DET +def _checkpoint_args_as_dict(raw_args: Any) -> dict[str, Any]: + """Coerce a checkpoint's ``args`` (dict, argparse.Namespace, or None) to a dict. + + A corrupt checkpoint may store ``args`` as a scalar or list; those have no + meaningful attributes, so treat them as no-args rather than letting + ``vars()`` raise a bare ``TypeError`` outside the ModelPackagingError contract. + """ + if isinstance(raw_args, dict): + return raw_args + if hasattr(raw_args, "__dict__"): + return dict(vars(raw_args)) + return {} + + def validate_model_type_for_project(model_type: str, project_type: str, project_id: str) -> None: """Raise TaskMismatchError if model_type's task doesn't match the Roboflow project type.""" expected = { @@ -259,7 +275,7 @@ def validate_model_type_for_project(model_type: str, project_type: str, project_ def package_custom_weights( model_type: str, model_path: str, - filename: str = "weights/best.pt", + filename: str = DEFAULT_WEIGHTS_FILENAME, *, build_dir: str | Path | None = None, allow_dependency_mismatch: bool = False, @@ -331,7 +347,7 @@ def package_custom_weights( def package_custom_weights_interactive( model_type: str, model_path: str, - filename: str = "weights/best.pt", + filename: str = DEFAULT_WEIGHTS_FILENAME, *, build_dir: str | Path | None = None, ) -> ModelUploadBundle: @@ -547,6 +563,13 @@ def _resolve_yolo_size( if inferred is None: if not provided: + if allow_mismatch: + warnings.append( + f"Could not infer a size for '{model_type}' from the checkpoint; " + f"uploading the bare family name as requested. Roboflow may reject it " + f"if it requires an explicit size." + ) + return model_type raise SizeMismatchError( f"model_type '{model_type}' is missing a size suffix and the size " f"could not be inferred from the checkpoint. Specify it explicitly, " @@ -731,7 +754,7 @@ def _detect_rfdetr_task(checkpoint: Any) -> str | None: raw_args = checkpoint.get("args") if raw_args is None: return None - args = raw_args if isinstance(raw_args, dict) else vars(raw_args) + args = _checkpoint_args_as_dict(raw_args) segmentation_head = args.get("segmentation_head") if segmentation_head is True: return TASK_SEG @@ -749,13 +772,7 @@ def _rfdetr_checkpoint_pe_size(checkpoint: Any) -> int | None: """ if not isinstance(checkpoint, dict): return None - raw_args = checkpoint.get("args") - if isinstance(raw_args, dict): - args = raw_args - elif raw_args is not None: - args = dict(vars(raw_args)) - else: - args = {} + args = _checkpoint_args_as_dict(checkpoint.get("args")) pe = args.get("positional_encoding_size") if isinstance(pe, int) and pe > 0: @@ -847,16 +864,24 @@ def _resolve_rfdetr_variant( def _find_rfdetr_checkpoint(model_path: Path, filename: str, warnings: list[str]) -> Path: - """Locate the rf-detr checkpoint, preserving the historical discovery behavior. + """Locate the rf-detr checkpoint. - The requested ``filename`` wins when it exists. Otherwise fall back to the - first top-level .pt/.pth file (sorted for determinism), which is how rf-detr - uploads located the checkpoint before ``filename`` was honored for them. + An explicitly-requested ``filename`` (anything other than the default) must + exist: falling back to a different checkpoint on a typo would silently + package the wrong weights. Only the default path falls back to discovering + the first top-level .pt/.pth file (sorted for determinism), preserving how + rf-detr uploads located the checkpoint before ``filename`` was honored. """ requested_file = model_path / filename if requested_file.exists(): return requested_file + if filename != DEFAULT_WEIGHTS_FILENAME: + raise MissingFileError( + f"RF-DETR weights file '{requested_file}' was not found. Set filename to the " + f"checkpoint's exact .pt or .pth path relative to model_path." + ) + discovered = sorted(path for path in model_path.iterdir() if path.is_file() and path.suffix in {".pt", ".pth"}) if not discovered: raise MissingFileError( @@ -874,13 +899,7 @@ def _write_rfdetr_class_names(model_path: Path, build_dir: Path, checkpoint: Any class_names = class_names_path.read_text().splitlines() else: raw_args = checkpoint.get("args") if isinstance(checkpoint, dict) else None - if raw_args is None: - args: dict[str, Any] = {} - elif isinstance(raw_args, dict): - args = raw_args - else: - args = dict(vars(raw_args)) - class_names = args.get("class_names") or [] + class_names = _checkpoint_args_as_dict(raw_args).get("class_names") or [] if not class_names: raise MissingFileError( f"No class_names.txt file found in '{model_path}', and the RF-DETR " @@ -924,7 +943,11 @@ def _process_rfdetr( model_type = _resolve_rfdetr_variant(model_type, checkpoint, warnings, allow_size_mismatch) - shutil.copy(checkpoint_path, build_dir / "weights.pt") + weights_dest = build_dir / "weights.pt" + # In the legacy deploy flow build_dir is model_path, so a checkpoint already + # named weights.pt is its own destination; copying would raise SameFileError. + if checkpoint_path.resolve() != weights_dest.resolve(): + shutil.copy(checkpoint_path, weights_dest) _write_rfdetr_class_names(model_path, build_dir, checkpoint) archive_path = build_dir / "roboflow_deploy.zip" _write_zip( diff --git a/tests/util/test_model_processor.py b/tests/util/test_model_processor.py index 2a8c1dce..4a902849 100644 --- a/tests/util/test_model_processor.py +++ b/tests/util/test_model_processor.py @@ -17,6 +17,7 @@ SizeMismatchError, TaskMismatchError, UnsupportedModelError, + _checkpoint_args_as_dict, _detect_rfdetr_task, _detect_yolo_task, _infer_yolo_size, @@ -114,6 +115,20 @@ def test_unrecognized_returns_none(self): self.assertIsNone(_detect_rfdetr_task({"model_name": None})) self.assertIsNone(_detect_rfdetr_task({"args": SimpleNamespace(other=1)})) + def test_scalar_args_do_not_raise_typeerror(self): + # A corrupt checkpoint storing args as a bare scalar must not escape the + # ModelPackagingError contract with a raw vars() TypeError. + self.assertIsNone(_detect_rfdetr_task({"args": 640})) + + +class CheckpointArgsAsDictTest(unittest.TestCase): + def test_coerces_dict_namespace_none_and_scalar(self): + self.assertEqual(_checkpoint_args_as_dict({"a": 1}), {"a": 1}) + self.assertEqual(_checkpoint_args_as_dict(SimpleNamespace(a=1)), {"a": 1}) + self.assertEqual(_checkpoint_args_as_dict(None), {}) + self.assertEqual(_checkpoint_args_as_dict(640), {}) + self.assertEqual(_checkpoint_args_as_dict(["not", "a", "dict"]), {}) + class GetClassnamesTxtForRfdetrTest(unittest.TestCase): def _classnames(self, args): @@ -161,9 +176,7 @@ def test_reports_missing_image_size(self): def test_accepts_either_image_size_key(self): self.assertEqual(_legacy_yolo_args({"imgsz": 640, "batch_size": 8}, Path("x")), {"imgsz": 640, "batch": 8}) - self.assertEqual( - _legacy_yolo_args({"img_size": 416, "batch_size": 4}, Path("x")), {"imgsz": 416, "batch": 4} - ) + self.assertEqual(_legacy_yolo_args({"img_size": 416, "batch_size": 4}, Path("x")), {"imgsz": 416, "batch": 4}) class _FakeYoloModel: @@ -227,6 +240,20 @@ def test_keeps_matching_sized_type_without_warning(self): self.assertEqual(_resolve_yolo_size("yolov8n", model, warnings), "yolov8n") self.assertEqual(warnings, []) + def test_bare_family_uninferable_raises_by_default(self): + with self.assertRaises(SizeMismatchError): + _resolve_yolo_size("yolov10", _FakeYoloModel({}), []) + + def test_bare_family_uninferable_proceeds_under_allow_mismatch(self): + # The interactive retry sets allow_mismatch=True after the user confirms; + # this branch must then converge (return) rather than raise forever. + warnings: list = [] + self.assertEqual( + _resolve_yolo_size("yolov10", _FakeYoloModel({}), warnings, allow_mismatch=True), + "yolov10", + ) + self.assertTrue(warnings and "bare family name" in warnings[0]) + class ResolveRfdetrVariantTest(unittest.TestCase): def test_raises_on_size_conflict_naming_the_fit(self): @@ -327,9 +354,7 @@ def test_never_prompts_or_exits(self): prompt_guard = mock.patch( "builtins.input", side_effect=AssertionError("package_custom_weights must not prompt") ) - exit_guard = mock.patch.object( - sys, "exit", side_effect=AssertionError("package_custom_weights must not exit") - ) + exit_guard = mock.patch.object(sys, "exit", side_effect=AssertionError("package_custom_weights must not exit")) with tempfile.TemporaryDirectory() as tmp: model_dir = Path(tmp) _write_yolonas_inputs(model_dir) @@ -405,6 +430,31 @@ def test_rfdetr_without_any_checkpoint_raises(self): with self.assertRaises(MissingFileError): package_custom_weights("rfdetr-base", tmp) + def test_rfdetr_explicit_missing_filename_does_not_fall_back(self): + # A typo'd explicit filename must fail loudly, not silently upload a + # different checkpoint that happens to be in the directory. + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + (model_dir / "some_other_checkpoint.pth").write_bytes(b"checkpoint") + torch = _fake_torch({"args": {"class_names": ["widget"]}}) + with _import_patch({"torch": torch}): + with self.assertRaises(MissingFileError): + package_custom_weights("rfdetr-base", str(model_dir), filename="checkpoint_epoch50.pth") + + def test_rfdetr_legacy_deploy_layout_does_not_self_copy(self): + # Legacy deploy passes build_dir=model_path with a top-level weights.pt; + # copying weights.pt onto itself would raise shutil.SameFileError. + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + (model_dir / "weights.pt").write_bytes(b"checkpoint") + torch = _fake_torch({"args": {"class_names": ["widget"]}}) + with _import_patch({"torch": torch}): + bundle = package_custom_weights( + "rfdetr-base", str(model_dir), filename="weights.pt", build_dir=model_dir + ) + with zipfile.ZipFile(bundle.archive_path) as archive: + self.assertIn("weights.pt", archive.namelist()) + def test_yolov8_full_flow_builds_artifacts(self): checkpoint_names = {1: "dog", 0: "cat"} From 08a1bcca98f8917f1ef0ef3db771fa2b4b64058e Mon Sep 17 00:00:00 2001 From: Leandro Rosemberg Date: Mon, 6 Jul 2026 13:46:20 -0300 Subject: [PATCH 05/11] fix: reject bare/incomplete checkpoints before upload with actionable errors --- roboflow/util/model_processor.py | 32 +++++++++++++++++++--- tests/util/test_model_processor.py | 43 ++++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+), 4 deletions(-) diff --git a/roboflow/util/model_processor.py b/roboflow/util/model_processor.py index 4983564f..c5043122 100644 --- a/roboflow/util/model_processor.py +++ b/roboflow/util/model_processor.py @@ -490,9 +490,9 @@ def _class_names_from_model_instance(model_instance: Any) -> list[str]: def _filtered_args(args: Any) -> dict[str, Any]: - if not isinstance(args, dict): - args = vars(args) - return {k: v for k, v in args.items() if k in {"model", "imgsz", "batch"}} + # A corrupt checkpoint may store args as a scalar/None; coerce via the shared + # helper so it becomes {} rather than raising a raw TypeError from vars(). + return {k: v for k, v in _checkpoint_args_as_dict(args).items() if k in {"model", "imgsz", "batch"}} def _load_checkpoint(torch_module: Any, checkpoint_path: Path, *, map_location: str | None = None) -> Any: @@ -934,6 +934,20 @@ def _process_rfdetr( checkpoint_path = _find_rfdetr_checkpoint(model_path, filename, warnings) checkpoint = _load_checkpoint(torch, checkpoint_path, map_location="cpu") + # Roboflow's server-side RF-DETR conversion reads checkpoint["args"] (the + # class names, class count, and model config). A bare inference state_dict — + # e.g. {"model": } with nothing else — would otherwise package and + # upload fine, then fail conversion with an opaque KeyError: 'args'. Catch it + # here so the caller gets an actionable error before uploading. + if not isinstance(checkpoint, dict) or checkpoint.get("args") is None: + raise ModelPackagingError( + f"The RF-DETR checkpoint '{checkpoint_path.name}' is missing its 'args' " + "metadata; it looks like a bare inference state_dict. Roboflow's weight " + "conversion needs the full training checkpoint (args carries the class " + "names, class count, and model config). Re-export the checkpoint from your " + "training run, or download the deploy checkpoint from Roboflow." + ) + detected_task = _detect_rfdetr_task(checkpoint) if detected_task and detected_task != task_of_model_type(model_type): raise TaskMismatchError( @@ -1012,7 +1026,17 @@ def _process_yolonas( raise MissingFileError(f"Model weights file '{weights_path}' was not found.") checkpoint = _load_checkpoint(torch, weights_path, map_location="cpu") - class_names = checkpoint["processing_params"]["class_names"] + # A SuperGradients YOLO-NAS checkpoint carries processing_params.class_names. + # A bare state_dict (e.g. torch.save(net.state_dict())) lacks it and would + # otherwise raise a raw KeyError/TypeError instead of an actionable error. + processing_params = checkpoint.get("processing_params") if isinstance(checkpoint, dict) else None + class_names = processing_params.get("class_names") if isinstance(processing_params, dict) else None + if not class_names: + raise ModelPackagingError( + f"The YOLO-NAS checkpoint '{weights_path.name}' is missing " + "processing_params.class_names; it looks like a bare state_dict. Provide the " + "full training checkpoint saved by SuperGradients." + ) opt_path = model_path / "opt.yaml" if not opt_path.exists(): raise MissingFileError( diff --git a/tests/util/test_model_processor.py b/tests/util/test_model_processor.py index 4a902849..5a44cc4d 100644 --- a/tests/util/test_model_processor.py +++ b/tests/util/test_model_processor.py @@ -20,6 +20,7 @@ _checkpoint_args_as_dict, _detect_rfdetr_task, _detect_yolo_task, + _filtered_args, _infer_yolo_size, _legacy_yolo_args, _resolve_rfdetr_variant, @@ -130,6 +131,20 @@ def test_coerces_dict_namespace_none_and_scalar(self): self.assertEqual(_checkpoint_args_as_dict(["not", "a", "dict"]), {}) +class FilteredArgsTest(unittest.TestCase): + def test_keeps_only_upload_keys_from_dict_or_namespace(self): + self.assertEqual( + _filtered_args({"model": "m", "imgsz": 640, "batch": 8, "lr0": 0.01}), + {"model": "m", "imgsz": 640, "batch": 8}, + ) + self.assertEqual(_filtered_args(SimpleNamespace(imgsz=320, batch=4, extra=1)), {"imgsz": 320, "batch": 4}) + + def test_scalar_or_none_args_do_not_raise(self): + # A corrupt .args must coerce to {} instead of raising a raw TypeError. + self.assertEqual(_filtered_args(None), {}) + self.assertEqual(_filtered_args(640), {}) + + class GetClassnamesTxtForRfdetrTest(unittest.TestCase): def _classnames(self, args): with tempfile.TemporaryDirectory() as model_path: @@ -430,6 +445,34 @@ def test_rfdetr_without_any_checkpoint_raises(self): with self.assertRaises(MissingFileError): package_custom_weights("rfdetr-base", tmp) + def test_rfdetr_bare_state_dict_without_args_fails_before_upload(self): + # A stripped inference checkpoint ({"model": state_dict} with no "args") + # would package fine but fail Roboflow's server-side conversion with an + # opaque KeyError: 'args'. Packaging must reject it up front. + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + (model_dir / "weights.pt").write_bytes(b"checkpoint") + torch = _fake_torch({"model": {"backbone.weight": object()}}) + with _import_patch({"torch": torch}): + with self.assertRaises(ModelPackagingError) as ctx: + package_custom_weights("rfdetr-base", str(model_dir), filename="weights.pt") + self.assertIn("args", str(ctx.exception)) + self.assertIn("state_dict", str(ctx.exception)) + + def test_yolonas_bare_state_dict_without_class_names_fails_before_upload(self): + # A bare YOLO-NAS state_dict lacks processing_params.class_names; it must + # raise an actionable error instead of a raw KeyError/TypeError. + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + (model_dir / "weights").mkdir() + (model_dir / "weights" / "best.pt").write_bytes(b"checkpoint") + torch = _fake_torch({"backbone.weight": object()}) # bare state_dict, no processing_params + with _import_patch({"torch": torch}): + with self.assertRaises(ModelPackagingError) as ctx: + package_custom_weights("yolonas", str(model_dir)) + self.assertIn("class_names", str(ctx.exception)) + self.assertIn("state_dict", str(ctx.exception)) + def test_rfdetr_explicit_missing_filename_does_not_fall_back(self): # A typo'd explicit filename must fail loudly, not silently upload a # different checkpoint that happens to be in the directory. From 29d1aec3fcf7d8f75ebfb97275030f86e88ea135 Mon Sep 17 00:00:00 2001 From: Leandro Rosemberg Date: Mon, 6 Jul 2026 20:25:09 -0300 Subject: [PATCH 06/11] fix: reject non-exact yolonas model types before upload --- roboflow/util/model_processor.py | 6 ++++++ tests/util/test_model_processor.py | 7 +++++++ 2 files changed, 13 insertions(+) diff --git a/roboflow/util/model_processor.py b/roboflow/util/model_processor.py index c5043122..72e81823 100644 --- a/roboflow/util/model_processor.py +++ b/roboflow/util/model_processor.py @@ -1020,6 +1020,12 @@ def _process_yolonas( filename: str, build_dir: Path, ) -> tuple[Path, str, list[str]]: + if model_type != "yolonas": + raise UnsupportedModelError( + f"Model type '{model_type}' is not supported for YOLO-NAS upload. The only " + "supported YOLO-NAS type is 'yolonas'; the architecture size goes in opt.yaml " + "as 'architecture: yolo_nas_s' (or _m / _l)." + ) torch = _import_required_module("torch", "pip install torch") weights_path = model_path / filename if not weights_path.exists(): diff --git a/tests/util/test_model_processor.py b/tests/util/test_model_processor.py index 5a44cc4d..0f4bf622 100644 --- a/tests/util/test_model_processor.py +++ b/tests/util/test_model_processor.py @@ -423,6 +423,13 @@ def test_family_must_be_a_prefix_not_a_substring(self): with self.assertRaises(UnsupportedModelError): package_custom_weights("foo-yolov8n", tmp) + def test_yolonas_requires_exact_model_type(self): + # Only 'yolonas' is valid; a suffixed typo like 'yolonas-foo' passes the + # family prefix gate but must be rejected before upload, not by the backend. + with tempfile.TemporaryDirectory() as tmp: + with self.assertRaises(UnsupportedModelError): + package_custom_weights("yolonas-foo", tmp) + def test_rfdetr_falls_back_to_discovered_checkpoint(self): with tempfile.TemporaryDirectory() as tmp: model_dir = Path(tmp) From 64111d9c13ba1e72d1acd9b53894389b287a31d0 Mon Sep 17 00:00:00 2001 From: Leandro Rosemberg Date: Tue, 7 Jul 2026 16:27:25 -0300 Subject: [PATCH 07/11] feat: adapt raw PyTorch-Lightning rf-detr checkpoints on packaging --- roboflow/util/model_processor.py | 138 ++++++++++++++++++++++--------- 1 file changed, 99 insertions(+), 39 deletions(-) diff --git a/roboflow/util/model_processor.py b/roboflow/util/model_processor.py index 72e81823..6c52a8f8 100644 --- a/roboflow/util/model_processor.py +++ b/roboflow/util/model_processor.py @@ -84,23 +84,32 @@ "paligemma2-3b-pt-896-peft", ) -SUPPORTED_RFDETR_TYPES = ( - # Detection models - "rfdetr-base", - "rfdetr-nano", - "rfdetr-small", - "rfdetr-medium", - "rfdetr-large", - "rfdetr-xlarge", - "rfdetr-2xlarge", - # Segmentation models - "rfdetr-seg-nano", - "rfdetr-seg-small", - "rfdetr-seg-medium", - "rfdetr-seg-large", - "rfdetr-seg-xlarge", - "rfdetr-seg-2xlarge", -) +# Minimum rf-detr release shipping ``RFDETR.export_for_roboflow`` (used to rebuild +# an upload bundle from a raw PyTorch-Lightning checkpoint). +RFDETR_MIN_VERSION = "1.8.0" + +# rf-detr model_type -> RFDETR subclass name. Single source of truth for both the +# supported-type check and the ``from_checkpoint`` fallback (used when a raw +# checkpoint lacks the metadata rf-detr needs to infer its own class). +_RFDETR_MODEL_TYPE_TO_CLASS = { + # Detection + "rfdetr-base": "RFDETRBase", + "rfdetr-nano": "RFDETRNano", + "rfdetr-small": "RFDETRSmall", + "rfdetr-medium": "RFDETRMedium", + "rfdetr-large": "RFDETRLarge", + "rfdetr-xlarge": "RFDETRXLarge", + "rfdetr-2xlarge": "RFDETR2XLarge", + # Segmentation + "rfdetr-seg-nano": "RFDETRSegNano", + "rfdetr-seg-small": "RFDETRSegSmall", + "rfdetr-seg-medium": "RFDETRSegMedium", + "rfdetr-seg-large": "RFDETRSegLarge", + "rfdetr-seg-xlarge": "RFDETRSegXLarge", + "rfdetr-seg-2xlarge": "RFDETRSeg2XLarge", +} + +SUPPORTED_RFDETR_TYPES = tuple(_RFDETR_MODEL_TYPE_TO_CLASS) DEFAULT_WEIGHTS_FILENAME = "weights/best.pt" @@ -750,7 +759,13 @@ def _detect_rfdetr_task(checkpoint: Any) -> str | None: return None model_name = checkpoint.get("model_name") if isinstance(model_name, str): - return TASK_SEG if TASK_SEG in model_name.lower() else TASK_DET + name = model_name.lower() + # Keypoint rf-detr checkpoints (e.g. 'RFDETRKeypointPreview') are not a + # supported upload type; classify them as pose so the task check rejects + # them instead of silently uploading a keypoint model as detection. + if "keypoint" in name: + return TASK_POSE + return TASK_SEG if TASK_SEG in name else TASK_DET raw_args = checkpoint.get("args") if raw_args is None: return None @@ -916,6 +931,34 @@ def _write_rfdetr_class_names(model_path: Path, build_dir: Path, checkpoint: Any return output_path +def _is_ptl_checkpoint(checkpoint: Any) -> bool: + """True if the checkpoint is a raw PyTorch-Lightning rf-detr checkpoint dict.""" + return isinstance(checkpoint, dict) and "pytorch-lightning_version" in checkpoint + + +def _require_rfdetr() -> Any: + """Import ``rfdetr`` and verify it ships the upload-bundle helpers. + + Raises :class:`MissingDependencyError` (a ModelPackagingError, so callers see + an actionable 400 rather than an opaque server error) when ``rfdetr`` is + missing or too old to export a Roboflow upload bundle. + """ + try: + import rfdetr + except ImportError as exc: + raise MissingDependencyError( + "The 'rfdetr' package is required to package raw PyTorch-Lightning rf-detr " + f"checkpoints. Install it with `pip install 'rfdetr>={RFDETR_MIN_VERSION}'`." + ) from exc + + if not hasattr(rfdetr.RFDETR, "export_for_roboflow"): + raise MissingDependencyError( + "The installed 'rfdetr' is too old to package raw PyTorch-Lightning rf-detr " + f"checkpoints. Upgrade it with `pip install --upgrade 'rfdetr>={RFDETR_MIN_VERSION}'`." + ) + return rfdetr + + def _process_rfdetr( model_type: str, model_path: Path, @@ -934,20 +977,8 @@ def _process_rfdetr( checkpoint_path = _find_rfdetr_checkpoint(model_path, filename, warnings) checkpoint = _load_checkpoint(torch, checkpoint_path, map_location="cpu") - # Roboflow's server-side RF-DETR conversion reads checkpoint["args"] (the - # class names, class count, and model config). A bare inference state_dict — - # e.g. {"model": } with nothing else — would otherwise package and - # upload fine, then fail conversion with an opaque KeyError: 'args'. Catch it - # here so the caller gets an actionable error before uploading. - if not isinstance(checkpoint, dict) or checkpoint.get("args") is None: - raise ModelPackagingError( - f"The RF-DETR checkpoint '{checkpoint_path.name}' is missing its 'args' " - "metadata; it looks like a bare inference state_dict. Roboflow's weight " - "conversion needs the full training checkpoint (args carries the class " - "names, class count, and model config). Re-export the checkpoint from your " - "training run, or download the deploy checkpoint from Roboflow." - ) - + # Task detection + mismatch runs for every checkpoint shape (it also rejects + # keypoint rf-detr, which is not a supported upload type). detected_task = _detect_rfdetr_task(checkpoint) if detected_task and detected_task != task_of_model_type(model_type): raise TaskMismatchError( @@ -955,14 +986,43 @@ def _process_rfdetr( f"but the checkpoint is a '{detected_task}' RF-DETR model. Use a matching model_type." ) - model_type = _resolve_rfdetr_variant(model_type, checkpoint, warnings, allow_size_mismatch) + if _is_ptl_checkpoint(checkpoint): + # Raw PyTorch-Lightning checkpoint: let rf-detr rebuild a proper upload + # bundle (weights.pt with args.resolution + class_names.txt) into build_dir, + # so the caller's model_path is never mutated. + rfdetr = _require_rfdetr() + try: + model = rfdetr.RFDETR.from_checkpoint(str(checkpoint_path)) + except ValueError: + # Checkpoint lacks model_name/pretrain_weights signals; fall back to the + # already-validated model_type to pick the RFDETR subclass. + model_cls = getattr(rfdetr, _RFDETR_MODEL_TYPE_TO_CLASS[model_type]) + model = model_cls(pretrain_weights=str(checkpoint_path)) + model.export_for_roboflow(str(build_dir)) # writes weights.pt + class_names.txt + else: + # Roboflow's server-side RF-DETR conversion reads checkpoint["args"] (the + # class names, class count, and model config). A bare inference state_dict — + # e.g. {"model": } with nothing else — would otherwise package and + # upload fine, then fail conversion with an opaque KeyError: 'args'. Catch it + # here so the caller gets an actionable error before uploading. + if not isinstance(checkpoint, dict) or checkpoint.get("args") is None: + raise ModelPackagingError( + f"The RF-DETR checkpoint '{checkpoint_path.name}' is missing its 'args' " + "metadata; it looks like a bare inference state_dict. Roboflow's weight " + "conversion needs the full training checkpoint (args carries the class " + "names, class count, and model config). Re-export the checkpoint from your " + "training run, or download the deploy checkpoint from Roboflow." + ) + + model_type = _resolve_rfdetr_variant(model_type, checkpoint, warnings, allow_size_mismatch) + + weights_dest = build_dir / "weights.pt" + # In the legacy deploy flow build_dir is model_path, so a checkpoint already + # named weights.pt is its own destination; copying would raise SameFileError. + if checkpoint_path.resolve() != weights_dest.resolve(): + shutil.copy(checkpoint_path, weights_dest) + _write_rfdetr_class_names(model_path, build_dir, checkpoint) - weights_dest = build_dir / "weights.pt" - # In the legacy deploy flow build_dir is model_path, so a checkpoint already - # named weights.pt is its own destination; copying would raise SameFileError. - if checkpoint_path.resolve() != weights_dest.resolve(): - shutil.copy(checkpoint_path, weights_dest) - _write_rfdetr_class_names(model_path, build_dir, checkpoint) archive_path = build_dir / "roboflow_deploy.zip" _write_zip( archive_path, From 7ae585600609c5e32db0afafe56351b084a6fac1 Mon Sep 17 00:00:00 2001 From: Leandro Rosemberg Date: Wed, 8 Jul 2026 16:54:04 -0300 Subject: [PATCH 08/11] fix: constrain upload filename to model_path and guard incomplete YOLO checkpoints --- roboflow/util/model_processor.py | 83 ++++++++++++-- tests/util/test_model_processor.py | 167 ++++++++++++++++++++++++++++- 2 files changed, 238 insertions(+), 12 deletions(-) diff --git a/roboflow/util/model_processor.py b/roboflow/util/model_processor.py index 6c52a8f8..61afbbdf 100644 --- a/roboflow/util/model_processor.py +++ b/roboflow/util/model_processor.py @@ -262,6 +262,27 @@ def _checkpoint_args_as_dict(raw_args: Any) -> dict[str, Any]: return {} +def _resolve_within_source(source_dir: Path, filename: str) -> Path: + """Resolve ``filename`` against ``source_dir``, refusing to escape it. + + ``filename`` is documented as relative to ``model_path`` and is forwarded + verbatim by hosted callers (the MCP server). An absolute path or ``..`` + segments would let a caller read weights from outside the model directory, + so reject both instead of silently packaging a file the caller never + intended. ``source_dir`` is already resolved; resolving the join collapses + ``..`` and follows symlinks before the containment check. + """ + if Path(filename).is_absolute(): + raise ModelPackagingError(f"filename '{filename}' must be a path relative to model_path, not an absolute path.") + resolved = (source_dir / filename).resolve() + if resolved != source_dir and source_dir not in resolved.parents: + raise ModelPackagingError( + f"filename '{filename}' resolves outside model_path '{source_dir}'. " + "It must point to a checkpoint inside the model directory." + ) + return resolved + + def validate_model_type_for_project(model_type: str, project_type: str, project_id: str) -> None: """Raise TaskMismatchError if model_type's task doesn't match the Roboflow project type.""" expected = { @@ -322,12 +343,13 @@ def package_custom_weights( source_dir = Path(model_path).expanduser().resolve() if not source_dir.is_dir(): raise MissingFileError(f"Model path '{model_path}' does not exist or is not a directory.") + _resolve_within_source(source_dir, filename) owns_build_dir = build_dir is None if build_dir is None: build_path = Path(tempfile.mkdtemp(prefix="roboflow-package-")) else: - build_path = Path(build_dir) + build_path = Path(build_dir).expanduser().resolve() build_path.mkdir(parents=True, exist_ok=True) try: @@ -615,6 +637,32 @@ def _resolve_yolo_size( ) +def _require_model_attr(model_instance: Any, attr: str, model_type: str) -> Any: + """Return ``model_instance.`` or raise a ModelPackagingError. + + Roboflow's server-side conversion needs these fields; a stripped checkpoint + missing one would otherwise raise a raw ``AttributeError`` outside the + ModelPackagingError contract (an opaque 500 for hosted callers). + """ + value = getattr(model_instance, attr, None) + if value is None: + raise ModelPackagingError( + f"The {model_type} checkpoint's model is missing '{attr}'; it does not look " + "like a complete Ultralytics training checkpoint. Re-export it from your training run." + ) + return value + + +def _require_checkpoint_field(checkpoint: Any, key: str, model_type: str) -> Any: + """Return ``checkpoint[key]`` or raise a ModelPackagingError (see _require_model_attr).""" + if not isinstance(checkpoint, dict) or key not in checkpoint: + raise ModelPackagingError( + f"The {model_type} checkpoint is missing '{key}'; it does not look like a " + "complete Ultralytics training checkpoint. Re-export it from your training run." + ) + return checkpoint[key] + + def _process_yolo( model_type: str, model_path: Path, @@ -694,6 +742,7 @@ def _process_yolo( if any(name in model_type for name in ULTRALYTICS_YOLO_FAMILIES): if ultralytics is None: ultralytics = _import_required_module("ultralytics", "pip install ultralytics") + model_yaml = _require_model_attr(model_instance, "yaml", model_type) if ( "-cls" in model_type or model_type.startswith("yolov10") @@ -701,14 +750,19 @@ def _process_yolo( or model_type.startswith("yolov12") or model_type.startswith("yolo26") ): - nc = model_instance.yaml["nc"] - args = checkpoint["train_args"] + if not isinstance(model_yaml, dict) or "nc" not in model_yaml: + raise ModelPackagingError( + f"The {model_type} checkpoint's model config (model.yaml) is missing 'nc'; " + "it does not look like a complete Ultralytics training checkpoint." + ) + nc = model_yaml["nc"] + args = _require_checkpoint_field(checkpoint, "train_args", model_type) else: - nc = model_instance.nc - args = model_instance.args + nc = _require_model_attr(model_instance, "nc", model_type) + args = _require_model_attr(model_instance, "args", model_type) model_artifacts: dict[str, Any] = { "names": class_names, - "yaml": model_instance.yaml, + "yaml": model_yaml, "nc": nc, "args": _filtered_args(args), "ultralytics_version": ultralytics.__version__, @@ -723,7 +777,7 @@ def _process_yolo( opts = yaml.safe_load(stream) or {} model_artifacts = { "names": class_names, - "nc": model_instance.nc, + "nc": _require_model_attr(model_instance, "nc", model_type), "args": _legacy_yolo_args(opts, opt_path), "model_type": model_type, } @@ -902,9 +956,18 @@ def _find_rfdetr_checkpoint(model_path: Path, filename: str, warnings: list[str] raise MissingFileError( f"No .pt or .pth checkpoint found in '{model_path}' (and '{requested_file}' does not exist)." ) - warnings.append( - f"Weights file '{requested_file}' was not found; using discovered checkpoint '{discovered[0].name}' instead." - ) + if len(discovered) > 1: + others = ", ".join(path.name for path in discovered) + warnings.append( + f"Weights file '{requested_file}' was not found and '{model_path}' holds multiple " + f"checkpoints ({others}); packaging '{discovered[0].name}'. Set filename to pick a " + "specific checkpoint if that is not the one you want." + ) + else: + warnings.append( + f"Weights file '{requested_file}' was not found; using discovered checkpoint " + f"'{discovered[0].name}' instead." + ) return discovered[0] diff --git a/tests/util/test_model_processor.py b/tests/util/test_model_processor.py index c4765782..1da01f6c 100644 --- a/tests/util/test_model_processor.py +++ b/tests/util/test_model_processor.py @@ -1,6 +1,7 @@ import json import os import sys +import tarfile import tempfile import types import unittest @@ -12,12 +13,12 @@ from roboflow.config import TASK_CLS, TASK_DET, TASK_OBB, TASK_POSE, TASK_SEG, TASK_SEM from roboflow.util import model_processor from roboflow.util.model_processor import ( + _RFDETR_MODEL_TYPE_TO_CLASS, MissingFileError, ModelPackagingError, SizeMismatchError, TaskMismatchError, UnsupportedModelError, - _RFDETR_MODEL_TYPE_TO_CLASS, _checkpoint_args_as_dict, _detect_rfdetr_task, _detect_yolo_task, @@ -408,7 +409,7 @@ def test_explicit_build_dir_is_used_and_not_cleaned_up(self): model_dir = Path(tmp) _write_yolonas_inputs(model_dir) bundle, _ = self._package_yolonas(model_dir, build_dir=build) - self.assertEqual(bundle.build_dir, Path(build)) + self.assertEqual(bundle.build_dir, Path(build).resolve()) self.assertFalse(bundle.owns_build_dir) bundle.cleanup() self.assertTrue(bundle.archive_path.exists()) @@ -553,6 +554,168 @@ def state_dict(self): finally: bundle.cleanup() + def test_rejects_absolute_filename(self): + # A hosted caller (the MCP server) forwards filename verbatim; an absolute + # path would read weights from outside the model directory. + with tempfile.TemporaryDirectory() as tmp: + with self.assertRaises(ModelPackagingError) as ctx: + package_custom_weights("yolonas", tmp, filename="/etc/passwd") + self.assertIn("absolute", str(ctx.exception)) + + def test_rejects_filename_escaping_model_path(self): + # '..' segments must not let the caller escape the model directory. + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) / "model" + model_dir.mkdir() + (Path(tmp) / "secret.pt").write_bytes(b"outside") + with self.assertRaises(ModelPackagingError) as ctx: + package_custom_weights("yolonas", str(model_dir), filename="../secret.pt") + self.assertIn("outside model_path", str(ctx.exception)) + + def _attempt_ultralytics_yolo(self, model_type, checkpoint): + fake_ultralytics = types.ModuleType("ultralytics") + fake_ultralytics.__version__ = "8.3.0" + fake_torch = _fake_torch(checkpoint) + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + weights = model_dir / "weights" / "best.pt" + weights.parent.mkdir() + weights.write_bytes(b"checkpoint") + with ( + _import_patch({"torch": fake_torch, "ultralytics": fake_ultralytics}), + mock.patch.dict(sys.modules, {"ultralytics": fake_ultralytics}), + ): + # allow_dependency_mismatch keeps a version mismatch from masking the + # completeness error under test. + return package_custom_weights(model_type, str(model_dir), allow_dependency_mismatch=True) + + def test_yolov8_missing_nc_attr_raises_packaging_error(self): + # A stripped Ultralytics checkpoint that lacks model.nc must raise a + # ModelPackagingError, not a raw AttributeError (an opaque 500 in the MCP). + class DetectionModel: + names = {0: "cat"} + yaml = {"nc": 1, "depth_multiple": 0.33, "width_multiple": 0.25} + args = {"imgsz": 640} + + def state_dict(self): + return {} + + with self.assertRaises(ModelPackagingError) as ctx: + self._attempt_ultralytics_yolo("yolov8n", {"model": DetectionModel()}) + self.assertIn("nc", str(ctx.exception)) + + def test_yolov8_missing_args_attr_raises_packaging_error(self): + class DetectionModel: + names = {0: "cat"} + nc = 1 + yaml = {"nc": 1, "depth_multiple": 0.33, "width_multiple": 0.25} + + def state_dict(self): + return {} + + with self.assertRaises(ModelPackagingError) as ctx: + self._attempt_ultralytics_yolo("yolov8n", {"model": DetectionModel()}) + self.assertIn("args", str(ctx.exception)) + + def test_yolov8_missing_yaml_attr_raises_packaging_error(self): + class DetectionModel: + names = {0: "cat"} + nc = 1 + args = {"imgsz": 640} + + def state_dict(self): + return {} + + with self.assertRaises(ModelPackagingError) as ctx: + self._attempt_ultralytics_yolo("yolov8n", {"model": DetectionModel()}) + self.assertIn("yaml", str(ctx.exception)) + + def test_yolov11_missing_train_args_raises_packaging_error(self): + # yolov10/11/12/26 read args from checkpoint["train_args"]; a checkpoint + # without it must fail with the ModelPackagingError contract. + class DetectionModel: + names = {0: "cat"} + yaml = {"nc": 1, "scale": "n"} + + def state_dict(self): + return {} + + with self.assertRaises(ModelPackagingError) as ctx: + self._attempt_ultralytics_yolo("yolov11n", {"model": DetectionModel()}) + self.assertIn("train_args", str(ctx.exception)) + + def test_yolov11_missing_nc_in_model_yaml_raises_packaging_error(self): + class DetectionModel: + names = {0: "cat"} + yaml = {"scale": "n"} + + def state_dict(self): + return {} + + with self.assertRaises(ModelPackagingError) as ctx: + self._attempt_ultralytics_yolo("yolov11n", {"model": DetectionModel(), "train_args": {"imgsz": 640}}) + self.assertIn("nc", str(ctx.exception)) + + def test_rfdetr_multiple_discovered_checkpoints_warns_which_used(self): + # With no explicit filename and several checkpoints present, discovery is + # ambiguous; the warning must name the candidates and which one was used. + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + (model_dir / "alpha.pt").write_bytes(b"checkpoint") + (model_dir / "beta.pth").write_bytes(b"checkpoint") + torch = _fake_torch({"args": {"class_names": ["widget"]}}) + with _import_patch({"torch": torch}): + bundle = package_custom_weights("rfdetr-base", str(model_dir)) + try: + warning = "\n".join(bundle.warnings) + self.assertIn("multiple", warning) + self.assertIn("alpha.pt", warning) + self.assertIn("beta.pth", warning) + finally: + bundle.cleanup() + + +class ProcessHuggingfaceTest(unittest.TestCase): + """Packaging for HuggingFace-backed models (paligemma / florence-2).""" + + def test_missing_companion_files_raise_missing_file(self): + # A safetensors checkpoint without its tokenizer/preprocessor companions + # now raises MissingFileError instead of prompting for the files. + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + (model_dir / "model.safetensors").write_bytes(b"weights") + with mock.patch("builtins.input", side_effect=AssertionError("must not prompt")): + with self.assertRaises(MissingFileError) as ctx: + package_custom_weights("florence-2-base", str(model_dir)) + message = str(ctx.exception) + self.assertIn("tokenizer.json", message) + self.assertIn("preprocessor_config.json", message) + + def test_no_model_file_raises_missing_file(self): + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + (model_dir / "readme.txt").write_text("no weights here") + with self.assertRaises(MissingFileError) as ctx: + package_custom_weights("paligemma-3b-pt-224", str(model_dir)) + self.assertIn("safetensors", str(ctx.exception)) + + def test_npz_checkpoint_packages_into_tar(self): + with tempfile.TemporaryDirectory() as tmp: + model_dir = Path(tmp) + (model_dir / "model.npz").write_bytes(b"npz-bytes") + bundle = package_custom_weights("paligemma-3b-pt-224", str(model_dir)) + try: + self.assertTrue(bundle.archive_path.name.endswith(".tar")) + with tarfile.open(bundle.archive_path) as tar: + self.assertIn("model.npz", tar.getnames()) + finally: + bundle.cleanup() + + def test_unsupported_huggingface_type_raises(self): + with tempfile.TemporaryDirectory() as tmp: + with self.assertRaises(UnsupportedModelError): + package_custom_weights("florence-2-tiny", str(tmp)) + class ProcessCompatTest(unittest.TestCase): """The legacy process() entry point keeps its historical contract.""" From 7f27b8b44ba62b8888a48866ea486c4583cb2460 Mon Sep 17 00:00:00 2001 From: Leandro Rosemberg Date: Wed, 8 Jul 2026 17:02:09 -0300 Subject: [PATCH 09/11] fix: reject filenames absolute under either OS convention (Windows CI) --- roboflow/util/model_processor.py | 9 +++++++-- tests/util/test_model_processor.py | 10 ++++++---- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/roboflow/util/model_processor.py b/roboflow/util/model_processor.py index 61afbbdf..1ecb581a 100644 --- a/roboflow/util/model_processor.py +++ b/roboflow/util/model_processor.py @@ -33,7 +33,7 @@ import zipfile from dataclasses import dataclass from importlib import import_module -from pathlib import Path +from pathlib import Path, PurePosixPath, PureWindowsPath from typing import Any import yaml @@ -271,8 +271,13 @@ def _resolve_within_source(source_dir: Path, filename: str) -> Path: so reject both instead of silently packaging a file the caller never intended. ``source_dir`` is already resolved; resolving the join collapses ``..`` and follows symlinks before the containment check. + + Absoluteness is tested under both OS conventions, not just the host's: the + hosted MCP forwards caller-supplied paths, so a POSIX absolute like + ``/etc/passwd`` must be rejected even when packaging happens to run on + Windows (where ``Path.is_absolute()`` alone would miss it, and vice versa). """ - if Path(filename).is_absolute(): + if PurePosixPath(filename).is_absolute() or PureWindowsPath(filename).is_absolute(): raise ModelPackagingError(f"filename '{filename}' must be a path relative to model_path, not an absolute path.") resolved = (source_dir / filename).resolve() if resolved != source_dir and source_dir not in resolved.parents: diff --git a/tests/util/test_model_processor.py b/tests/util/test_model_processor.py index 1da01f6c..caa73157 100644 --- a/tests/util/test_model_processor.py +++ b/tests/util/test_model_processor.py @@ -556,11 +556,13 @@ def state_dict(self): def test_rejects_absolute_filename(self): # A hosted caller (the MCP server) forwards filename verbatim; an absolute - # path would read weights from outside the model directory. + # path — POSIX or Windows-style, regardless of the packaging host's OS — + # must be rejected rather than reading weights from outside model_path. with tempfile.TemporaryDirectory() as tmp: - with self.assertRaises(ModelPackagingError) as ctx: - package_custom_weights("yolonas", tmp, filename="/etc/passwd") - self.assertIn("absolute", str(ctx.exception)) + for bad in ("/etc/passwd", "C:\\Windows\\System32\\config"): + with self.assertRaises(ModelPackagingError) as ctx: + package_custom_weights("yolonas", tmp, filename=bad) + self.assertIn("absolute", str(ctx.exception)) def test_rejects_filename_escaping_model_path(self): # '..' segments must not let the caller escape the model directory. From 5dcb2e7ef62d5e25bbbf75e4e63b0e2d72b801a2 Mon Sep 17 00:00:00 2001 From: Leandro Rosemberg Date: Wed, 8 Jul 2026 17:59:41 -0300 Subject: [PATCH 10/11] fix: reject filenames that resolve to a directory instead of a checkpoint --- roboflow/util/model_processor.py | 5 +++++ tests/util/test_model_processor.py | 11 +++++++++++ 2 files changed, 16 insertions(+) diff --git a/roboflow/util/model_processor.py b/roboflow/util/model_processor.py index 1ecb581a..46389b35 100644 --- a/roboflow/util/model_processor.py +++ b/roboflow/util/model_processor.py @@ -285,6 +285,11 @@ def _resolve_within_source(source_dir: Path, filename: str) -> Path: f"filename '{filename}' resolves outside model_path '{source_dir}'. " "It must point to a checkpoint inside the model directory." ) + # '' / '.' resolve to model_path itself and a subdirectory passes the + # containment check; all would reach torch.load() and leak a raw + # IsADirectoryError outside the ModelPackagingError contract. + if resolved.is_dir(): + raise ModelPackagingError(f"filename '{filename}' must point to a checkpoint file, not a directory.") return resolved diff --git a/tests/util/test_model_processor.py b/tests/util/test_model_processor.py index caa73157..c28cfdd0 100644 --- a/tests/util/test_model_processor.py +++ b/tests/util/test_model_processor.py @@ -574,6 +574,17 @@ def test_rejects_filename_escaping_model_path(self): package_custom_weights("yolonas", str(model_dir), filename="../secret.pt") self.assertIn("outside model_path", str(ctx.exception)) + def test_rejects_filename_pointing_at_a_directory(self): + # '' / '.' resolve to model_path itself and a subdirectory stays inside + # it; all three would otherwise reach torch.load() and leak a raw + # IsADirectoryError outside the ModelPackagingError contract. + with tempfile.TemporaryDirectory() as tmp: + (Path(tmp) / "weights").mkdir() + for bad in ("", ".", "weights"): + with self.assertRaises(ModelPackagingError) as ctx: + package_custom_weights("yolonas", tmp, filename=bad) + self.assertIn("not a directory", str(ctx.exception)) + def _attempt_ultralytics_yolo(self, model_type, checkpoint): fake_ultralytics = types.ModuleType("ultralytics") fake_ultralytics.__version__ = "8.3.0" From d246f4f11e50e3dc8b6f7cc45b3d0fb6659ce1e0 Mon Sep 17 00:00:00 2001 From: Leandro Rosemberg Date: Wed, 8 Jul 2026 21:19:03 -0300 Subject: [PATCH 11/11] chore: bump package version to 1.3.13 --- roboflow/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/roboflow/__init__.py b/roboflow/__init__.py index 056ae873..260c1261 100644 --- a/roboflow/__init__.py +++ b/roboflow/__init__.py @@ -21,7 +21,7 @@ CLIPModel = None # type: ignore[assignment,misc] GazeModel = None # type: ignore[assignment,misc] -__version__ = "1.3.12" +__version__ = "1.3.13" def check_key(api_key, model, notebook, num_retries=0):