feat: no huggingface option

This commit is contained in:
Bingsu
2023-05-13 15:17:08 +09:00
parent a0657caf3e
commit 3d892728c6
3 changed files with 37 additions and 13 deletions

View File

@@ -20,7 +20,9 @@ class PredictOutput:
preview: Optional[Image.Image] = None preview: Optional[Image.Image] = None
def get_models(model_dir: Union[str, Path]) -> OrderedDict[str, Optional[str]]: def get_models(
model_dir: Union[str, Path], huggingface: bool = True
) -> OrderedDict[str, Optional[str]]:
model_dir = Path(model_dir) model_dir = Path(model_dir)
if model_dir.is_dir(): if model_dir.is_dir():
model_paths = [ model_paths = [
@@ -31,17 +33,29 @@ def get_models(model_dir: Union[str, Path]) -> OrderedDict[str, Optional[str]]:
else: else:
model_paths = [] model_paths = []
models = OrderedDict( if huggingface:
{ models = OrderedDict(
"face_yolov8n.pt": hf_hub_download(repo_id, "face_yolov8n.pt"), {
"face_yolov8s.pt": hf_hub_download(repo_id, "face_yolov8s.pt"), "face_yolov8n.pt": hf_hub_download(repo_id, "face_yolov8n.pt"),
"mediapipe_face_full": None, "face_yolov8s.pt": hf_hub_download(repo_id, "face_yolov8s.pt"),
"mediapipe_face_short": None, "mediapipe_face_full": None,
"hand_yolov8n.pt": hf_hub_download(repo_id, "hand_yolov8n.pt"), "mediapipe_face_short": None,
"person_yolov8n-seg.pt": hf_hub_download(repo_id, "person_yolov8n-seg.pt"), "hand_yolov8n.pt": hf_hub_download(repo_id, "hand_yolov8n.pt"),
"person_yolov8s-seg.pt": hf_hub_download(repo_id, "person_yolov8s-seg.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,
}
)
for path in model_paths: for path in model_paths:
if path.name in models: if path.name in models:

9
preload.py Normal file
View File

@@ -0,0 +1,9 @@
import argparse
def preload(parser: argparse.ArgumentParser):
parser.add_argument(
"--ad-no-huggingface",
action="store_true",
help="Don't use adetailer models from huggingface",
)

View File

@@ -44,8 +44,9 @@ except Exception:
pass pass
AFTER_DETAILER = "After Detailer" AFTER_DETAILER = "After Detailer"
no_huggingface = getattr(cmd_opts, "ad_no_huggingface", False)
adetailer_dir = Path(models_path, "adetailer") adetailer_dir = Path(models_path, "adetailer")
model_mapping = get_models(adetailer_dir) model_mapping = get_models(adetailer_dir, huggingface=not no_huggingface)
txt2img_submit_button = img2img_submit_button = None txt2img_submit_button = img2img_submit_button = None