Skip to content

Commit

Permalink
dev(narugo): add wd14 syncer, ci skip
Browse files Browse the repository at this point in the history
  • Loading branch information
narugo1992 committed May 12, 2024
1 parent c535534 commit 090c281
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 4 deletions.
6 changes: 3 additions & 3 deletions imgutils/tagging/wd14.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,12 @@
from functools import lru_cache
from typing import List, Tuple

import huggingface_hub
import numpy as np
import onnxruntime
import pandas as pd
from PIL import Image
from hbutils.testing.requires.version import VersionInfo
from huggingface_hub import hf_hub_download

from .format import remove_underline
from .overlap import drop_overlap_tags
Expand Down Expand Up @@ -64,7 +64,7 @@ def _get_wd14_model(model_name):
:rtype: ONNXModel
"""
_version_support_check(model_name)
return open_onnx_model(huggingface_hub.hf_hub_download(MODEL_NAMES[model_name], MODEL_FILENAME))
return open_onnx_model(hf_hub_download(MODEL_NAMES[model_name], MODEL_FILENAME))


@lru_cache()
Expand All @@ -79,7 +79,7 @@ def _get_wd14_labels(model_name, no_underline: bool = False) -> Tuple[List[str],
:return: A tuple containing the list of tag names, and lists of indexes for rating, general, and character categories.
:rtype: Tuple[List[str], List[int], List[int], List[int]]
"""
path = huggingface_hub.hf_hub_download(MODEL_NAMES[model_name], LABEL_FILENAME)
path = hf_hub_download(MODEL_NAMES[model_name], LABEL_FILENAME)
df = pd.read_csv(path)
name_series = df["name"]
if no_underline:
Expand Down
3 changes: 2 additions & 1 deletion requirements-zoo.txt
Original file line number Diff line number Diff line change
Expand Up @@ -22,4 +22,5 @@ controlnet_aux
lighttuner
natsort
tabulate
hfmirror>=0.0.7
hfmirror>=0.0.7
tabulate

0 comments on commit 090c281

Please sign in to comment.