VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-09 06:45:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update tag_images_by_wd14_tagger.py

Browse Source 
add WDV3
This commit is contained in:
青龍聖者@bdsqlsz
2024-03-18 22:29:05 +08:00
parent f9317052ed
commit a7dff592d3
1 changed files with 12 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
finetune/tag_images_by_wd14_tagger.py
View File

@@ -86,9 +86,9 @@ def main(args):
logger.info(f"downloading wd14 tagger model from hf_hub. id: {args.repo_id}") logger.info(f"downloading wd14 tagger model from hf_hub. id: {args.repo_id}")
files = FILES files = FILES
if args.onnx: if args.onnx:
files = ["selected_tags.csv"]
files += FILES_ONNX files += FILES_ONNX
for file in files: else:
hf_hub_download(args.repo_id, file, cache_dir=args.model_dir, force_download=True, force_filename=file)
for file in SUB_DIR_FILES: for file in SUB_DIR_FILES:
hf_hub_download( hf_hub_download(
args.repo_id, args.repo_id,
@@ -98,11 +98,14 @@ def main(args):
force_download=True, force_download=True,
force_filename=file, force_filename=file,
) )
for file in files:
hf_hub_download(args.repo_id, file, cache_dir=args.model_dir, force_download=True, force_filename=file)
else: else:
logger.info("using existing wd14 tagger model") logger.info("using existing wd14 tagger model")
# 画像を読み込む # 画像を読み込む
if args.onnx: if args.onnx:
import torch
import onnx import onnx
import onnxruntime as ort import onnxruntime as ort