BooruDatasetTagManager项目中离线加载AutoTagger模型的方法

BooruDatasetTagManager项目中离线加载AutoTagger模型的方法

BooruDatasetTagManager 项目地址: https://gitcode.com/gh_mirrors/bo/BooruDatasetTagManager

背景介绍

在BooruDatasetTagManager项目中，AutoTagger功能通常需要从Hugging Face平台下载预训练模型。然而在实际使用中，用户可能会遇到网络连接问题，或者希望使用自定义模型。本文将详细介绍如何通过模拟Hugging Face缓存结构的方式，实现离线加载AutoTagger模型。

缓存目录结构解析

Hugging Face的模型缓存采用特定的目录结构。理解这一结构是成功离线加载模型的关键。标准的缓存目录结构如下：

<CACHE_DIR>
└─ models--<组织名>--<模型名>
   ├─ blobs
   ├─ refs
   │  └─ main
   └─ snapshots
      └─ <40位哈希值>
         ├─ model.onnx
         └─ selected_tags.csv

离线加载标准模型的步骤

1. 创建基础目录结构：

   • 在缓存目录下创建models--SmilingWolf--wd-vit-tagger-v3文件夹
   • 在该文件夹下创建blobs、refs和snapshots子目录

2. 设置模型文件：

   • 在snapshots下创建一个以40位随机哈希值命名的文件夹
   • 将下载好的model.onnx和selected_tags.csv文件放入该文件夹

3. 配置引用文件：

   • 在refs目录下创建名为main的文本文件
   • 文件中写入刚才创建的40位哈希值文件夹名（不含换行符）

4. 启用离线模式：

   • 在BooruDatasetTagManager的AutoTagger设置中勾选"Skip internet requests"选项

加载自定义模型的特殊处理

对于非标准模型（如Z3D-E621-Convnext），除了上述步骤外，还需要修改项目代码：

1. 修改interrogator.py：

   • 在WD_TAGGER_NAMES列表中添加新模型名称
   • 在WD_TAGGER_THRESHOLDS列表中添加对应的阈值

2. 创建对应的缓存目录：

   • 目录名格式为models--SmilingWolf--Z3D-E621-Convnext
   • 其余结构与标准模型相同

技术原理

这种方法利用了Hugging Face Hub的缓存机制。当启用"Skip internet requests"选项时，系统会直接查找本地缓存而非联网下载。通过模拟标准的缓存目录结构，我们可以让系统认为模型已经正确下载并缓存。

注意事项

1. 40位哈希值理论上可以随机生成，但建议保持一致性以避免潜在问题
2. 不同模型可能需要不同的预处理参数，添加新模型时应确保了解其特性
3. 模型文件必须与原始结构完全一致，包括文件名和扩展名
4. 对于大型模型，确保有足够的磁盘空间

总结

通过模拟Hugging Face的缓存结构，我们可以灵活地在BooruDatasetTagManager中使用离线模型。这种方法不仅解决了网络连接问题，还为用户提供了使用自定义模型的途径。理解这一机制后，用户可以更自由地管理模型文件，提高工作效率。

BooruDatasetTagManager 项目地址: https://gitcode.com/gh_mirrors/bo/BooruDatasetTagManager