解决 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

修改的影响

  1. 安全性考虑: 此修改降低了安全性,但鉴于我们信任 ADetailer 和 ultralytics 的模型来源,这个风险是可接受的。
  2. 兼容性: 确保与现有模型文件的兼容性,避免需要重新训练或转换所有模型文件。
  3. 临时性: 这是一个临时解决方案,最终应该迁移到使用 torch.serialization.add_safe_globals() 来明确允许特定的类。

替代方案考虑

在实施此修复前,考虑了以下替代方案:

  1. 更新模型文件: 尝试下载新格式的模型文件,但服务器尚未更新
  2. 降级 PyTorch: 降级到 2.5.x 版本,但可能影响其他功能
  3. 使用环境变量: 设置 TORCH_LOAD_WEIGHTS_ONLY_FALSE=1,但测试发现不够彻底

验证结果

修改后重启 Forge 项目,ADetailer 插件能够正常加载模型文件,面部检测功能工作正常,不再出现 weights_only 相关的错误。

最后修改:2025 年 10 月 09 日
喜欢就请我喝一杯奶茶吧~