Error loading YOLO model with PyTorch >= 2.6: Unsupported global: ultralytics.nn.tasks.DetectionModel #13
Description
Okay, here is a template for a GitHub Issue you can post to the ComfyUI_Sapiens repository. Remember to replace the bracketed placeholders [Your ... Version] with your specific environment details.
Title: Error loading YOLO model with PyTorch >= 2.6: Unsupported global: ultralytics.nn.tasks.DetectionModel
Body:
Environment:
Operating System: Windows 11
Python Version: 3.12.8
PyTorch Version: 2.6.0
CUDA Version (if applicable): 12.6
ultralytics Version: 8.3.127 (Updated, issue persists)
ComfyUI_Sapiens Version: f48b7c7
Bug Description:
When using the SapiensLoader node with the use_yolo option enabled, an error occurs during the initialization of the YOLO detector. This seems to be caused by an incompatibility between the way the ultralytics library loads its standard .pt model files (which contain pickled Python objects like DetectionModel) and the default setting of weights_only=True for torch.load introduced in PyTorch 2.2 and made the default behavior more broadly noticeable since PyTorch 2.6 changes.
Steps to Reproduce:
Set up an environment with PyTorch version 2.6 or higher.
Install or update ComfyUI_Sapiens.
Ensure the YOLO model (e.g., yolov8m.pt) is present in the ComfyUI/models/sapiens/ directory (or let the node download it).
In ComfyUI, create a workflow including the SapiensLoader node.
Enable the use_yolo checkbox on the SapiensLoader node.
Attempt to queue the prompt / run the workflow.
The error occurs during node execution.
Error Message:
Weights only load failed. This file can still be loaded, to do so you have two options, do those steps only if you trust the source of the checkpoint.
(1) In PyTorch 2.6, we changed the default value of the weights_only argument in torch.load from False to True. Re-running torch.load with weights_only set to False will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
(2) Alternatively, to load with weights_only=True please check the recommended steps in the following error message.
WeightsUnpickler error: Unsupported global: GLOBAL ultralytics.nn.tasks.DetectionModel was not an allowed global by default. Please use torch.serialization.add_safe_globals([DetectionModel]) or the torch.serialization.safe_globals([DetectionModel]) context manager to allowlist this global if you trust this class/function
[Full traceback leading to the error, if available]
(Please paste the full traceback here if you have it)
Troubleshooting Done:
Updated the ultralytics library to the latest version (8.3.127), but the issue persists.
Possible Cause / Suggestion:
The issue seems to originate from the call self.model = YOLO(model_path) in Sapiens_Pytorch/detector.py (around line 23). This uses the ultralytics library's standard loading mechanism for .pt files.
Could the node be modified to handle this incompatibility? Potential solutions might include:
Loading a different model format like TorchScript (.torchscript) or ONNX (.onnx) which are generally safer and don't rely on Python pickling in the same way. This would require exporting the .pt model first.
Exploring if it's possible to pass loading arguments or use a context manager (like torch.serialization.safe_globals) when the YOLO() class loads the model, although this might be difficult as the torch.load call is likely internal to the ultralytics library.
Thank you for looking into this!