Bing-su/adetailer
1
0
Fork 0
mirror of https://github.com/Bing-su/adetailer.git synced 2026-03-13 01:10:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: huggingface download exception

Browse Source
This commit is contained in:
Bingsu
2023-05-29 12:41:48 +09:00
parent b3b35ac114
commit 45afcc0a8b
1 changed files with 25 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
adetailer/common.py
View File

@@ -18,6 +18,14 @@ class PredictOutput:
preview: Optional[Image.Image] = None
def hf_download(file: str):
try:
path = hf_hub_download(repo_id, file)
except Exception:
path = "INVALID"
return path
def get_models(
model_dir: Union[str, Path], huggingface: bool = True
) -> OrderedDict[str, Optional[str]]:
@@ -31,29 +39,27 @@ def get_models(
else:
model_paths = []
models = OrderedDict()
if huggingface:
models = OrderedDict(
models.update(
{
"face_yolov8n.pt": hf_hub_download(repo_id, "face_yolov8n.pt"),
"face_yolov8s.pt": hf_hub_download(repo_id, "face_yolov8s.pt"),
"mediapipe_face_full": None,
"mediapipe_face_short": None,
"hand_yolov8n.pt": hf_hub_download(repo_id, "hand_yolov8n.pt"),
"person_yolov8n-seg.pt": hf_hub_download(
repo_id, "person_yolov8n-seg.pt"
),
"person_yolov8s-seg.pt": hf_hub_download(
repo_id, "person_yolov8s-seg.pt"
),
}
)
else:
models = OrderedDict(
{
"mediapipe_face_full": None,
"mediapipe_face_short": None,
"face_yolov8n.pt": hf_download("face_yolov8n.pt"),
"face_yolov8s.pt": hf_download("face_yolov8s.pt"),
"hand_yolov8n.pt": hf_download("hand_yolov8n.pt"),
"person_yolov8n-seg.pt": hf_download("person_yolov8n-seg.pt"),
"person_yolov8s-seg.pt": hf_download("person_yolov8s-seg.pt"),
}
)
models.update(
{
"mediapipe_face_full": None,
"mediapipe_face_short": None,
}
)
invalid_keys = [k for k, v in models.items() if v == "INVALID"]
for key in invalid_keys:
models.pop(key)
for path in model_paths:
if path.name in models: