ComfyUI中PyTorch权重加载失败问题的技术解析与解决方案

问题背景

在ComfyUI项目中使用PyTorch 2.6及以上版本时，用户可能会遇到权重文件加载失败的问题，系统提示"Weights only load failed"。这个问题的根源在于PyTorch 2.6对模型安全性进行了重大改进，默认启用了weights_only=True的安全加载模式。

技术原理分析

PyTorch 2.6版本引入了一个重要的安全变更：torch.load()函数的weights_only参数默认值从False改为True。这一变更旨在防止潜在的不安全代码执行风险，因为传统的PyTorch模型序列化可能包含任意Python代码。

当weights_only=True时，PyTorch会严格限制可以反序列化的对象类型，只允许基本数据类型和特定的安全类。如果模型文件中包含自定义类（如ultralytics.nn.tasks.DetectionModel），系统会拒绝加载并抛出错误。

具体错误分析

错误信息中明确指出两个关键点：

1. 检测到了不被允许的全局类GLOBAL ultralytics.nn.tasks.DetectionModel
2. 提供了两种解决方案：要么关闭weights_only安全检查，要么显式地将自定义类加入安全列表

解决方案

方案一：信任源并关闭安全检查（不推荐）

如果确认模型文件来源可信，可以显式设置weights_only=False：

torch.load(file_path, weights_only=False)

风险提示：这种方法会完全关闭安全检查，可能执行模型中的任意代码，仅应在完全信任模型来源时使用。

方案二：安全地添加自定义类到白名单（推荐）

PyTorch提供了两种方式将自定义类加入安全列表：

1. 使用add_safe_globals函数：

from torch.serialization import add_safe_globals
from ultralytics.nn.tasks import DetectionModel

add_safe_globals([DetectionModel])
model = torch.load(file_path, weights_only=True)

2. 使用上下文管理器：

from torch.serialization import safe_globals
from ultralytics.nn.tasks import DetectionModel

with safe_globals([DetectionModel]):
    model = torch.load(file_path, weights_only=True)

最佳实践建议

1. 优先使用方案二：在保持安全检查的同时解决问题，安全性更高
2. 验证模型来源：确保加载的模型文件来自可信源
3. 考虑模型转换：将模型转换为更安全的格式（如ONNX）
4. 版本兼容性检查：确认使用的PyTorch版本与模型训练版本兼容
5. 错误处理：在代码中添加适当的异常处理，为用户提供友好的错误提示

总结

PyTorch 2.6的安全改进虽然带来了短暂的兼容性问题，但从长远看提升了深度学习应用的安全性。ComfyUI用户遇到此问题时，应理解其背后的安全考量，并根据自身安全需求选择合适的解决方案。对于生产环境，推荐采用方案二的安全加载方式，既解决了兼容性问题，又保持了必要的安全防护。