解决 PyTorch 2.6.0 中 weights_only 安全检查导致的模型加载失败问题
- 记录一次在 Stable Diffusion Forge 项目中解决 ADetailer 插件模型加载问题的过程*
问题背景
在 PyTorch 2.6.0 版本中,为了增强安全性,torch.load()
函数的默认行为发生了重大变化:weights_only
参数的默认值从 False
改为 True
。这意味着默认情况下,PyTorch 现在只允许加载纯权重数据,拒绝加载任何可能包含可执行代码的模型文件。
这导致在使用 ADetailer 插件时出现以下错误:
│ ValueError: 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. │
│ │
│ Check the documentation of torch.load to learn more about types accepted by default with weights_only │
│ https://pytorch.org/docs/stable/generated/torch.load.html. │
│ File corrupted: /home/epors/29095/models/adetailer/FacesV1.pt │
│ Forge has tried to move the corrupted file to /home/epors/29095/models/adetailer/FacesV1.pt.corrupted │
│ You may try again now and Forge will download models again.
环境信息
- 项目: Stable Diffusion Forge
- 版本: f2.0.1v1.10.1-1.10.1
- Python: 3.10.16 (Conda)
- PyTorch: 2.6.0+cu124
- 问题插件: ADetailer
根本原因分析
ADetailer 使用的模型文件(例如使用这款Furry人脸识别 FacesV1.pt)不仅包含模型权重,还包含对 ultralytics.nn.tasks.DetectionModel
类的引用。在 PyTorch 2.6.0 的默认安全设置下,这种包含代码引用的模型文件被拒绝加载。
解决方案
经过代码分析,发现需要在两个关键位置修改 torch.load
的调用,强制使用 weights_only=False
来恢复旧版行为。
修改 1:Forge 项目后端
文件位置: backend/utils.py
修改前:
if not 'weights_only' in torch.load.__code__.co_varnames:
print("Warning torch.load doesn't support weights_only on this pytorch version, loading unsafely.")
pl_sd = torch.load(ckpt, map_location=device, weights_only=True)
修改后:
# 强制使用 weights_only=False 以兼容 PyTorch 2.6.0 的安全限制
pl_sd = torch.load(ckpt, map_location=device, weights_only=False)
说明: 移除了版本检查逻辑,直接强制使用 weights_only=False
,确保在 PyTorch 2.6.0+ 中也能正确加载包含代码引用的模型文件。
修改 2:Ultralytics 库
因为我使用的是Conda环境,在出现问题时后台页面日志已告知了 torch.load
调用位置:
│ │ /home/epors/.conda/envs/forge/lib/python3.10/site-packages/ultralytics/nn/tasks.py:832 in │ │
│ │ torch_safe_load │ │
│ │ │ │
│ │ 831 │ │ │ else: │ │
│ │ ❱ 832 │ │ │ │ ckpt = torch.load(file, map_location="cpu") │ │
│ │ 833 │ │
│ │ │ │
│ │ /home/epors/29095/modules_forge/patch_basic.py:54 in loader │ │
│ │ │ │
│ │ 53 │ │ │ │ │ exp += f'You may try again now and Forge will download models ag │ │
│ │ ❱ 54 │ │ │ raise ValueError(exp) │ │
│ │ 55 │ │ return result │ │
文件位置: /home/epors/.conda/envs/forge/lib/python3.10/site-packages/ultralytics/nn/tasks.py
修改前 (约第832行):
ckpt = torch.load(file, map_location="cpu")
修改后:
ckpt = torch.load(file, map_location="cpu", weights_only=False)
说明: 在 ultralytics 库的模型加载函数中显式添加 weights_only=False
参数,确保该库加载模型时也能绕过安全检查。
技术细节
搜索相关代码的方法
在定位问题时,使用了以下命令行工具搜索项目中的相关代码:
# 使用 ripgrep 搜索 weights_only 相关代码
rg "weights_only"
# 搜索所有 torch.load 调用
rg "torch.load" --type py
修改的影响
- 安全性考虑: 此修改降低了安全性,但鉴于我们信任 ADetailer 和 ultralytics 的模型来源,这个风险是可接受的。
- 兼容性: 确保与现有模型文件的兼容性,避免需要重新训练或转换所有模型文件。
- 临时性: 这是一个临时解决方案,最终应该迁移到使用
torch.serialization.add_safe_globals()
来明确允许特定的类。
替代方案考虑
在实施此修复前,考虑了以下替代方案:
- 更新模型文件: 尝试下载新格式的模型文件,但服务器尚未更新
- 降级 PyTorch: 降级到 2.5.x 版本,但可能影响其他功能
- 使用环境变量: 设置
TORCH_LOAD_WEIGHTS_ONLY_FALSE=1
,但测试发现不够彻底
验证结果
修改后重启 Forge 项目,ADetailer 插件能够正常加载模型文件,面部检测功能工作正常,不再出现 weights_only
相关的错误。