diff --git a/adetailer/common.py b/adetailer/common.py index 5b03da5..37c0320 100644 --- a/adetailer/common.py +++ b/adetailer/common.py @@ -20,7 +20,9 @@ class PredictOutput: 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) if model_dir.is_dir(): model_paths = [ @@ -31,17 +33,29 @@ def get_models(model_dir: Union[str, Path]) -> OrderedDict[str, Optional[str]]: else: model_paths = [] - models = OrderedDict( - { - "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"), - } - ) + 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"), + "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, + } + ) for path in model_paths: if path.name in models: diff --git a/preload.py b/preload.py new file mode 100644 index 0000000..10be161 --- /dev/null +++ b/preload.py @@ -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", + ) diff --git a/scripts/!adetailer.py b/scripts/!adetailer.py index d2f626d..7841f2c 100644 --- a/scripts/!adetailer.py +++ b/scripts/!adetailer.py @@ -44,8 +44,9 @@ except Exception: pass AFTER_DETAILER = "After Detailer" +no_huggingface = getattr(cmd_opts, "ad_no_huggingface", False) 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