diff --git a/scripts/reactor_sfw.py b/scripts/reactor_sfw.py index f91cc66..2db04a7 100644 --- a/scripts/reactor_sfw.py +++ b/scripts/reactor_sfw.py @@ -4,19 +4,18 @@ import logging import os from reactor_utils import download -def ensure_nsfw_model(model_path): +def ensure_nsfw_model(nsfwdet_model_path): """Download NSFW detection model if it doesn't exist""" - if not os.path.exists(model_path): - os.makedirs(model_path) + if not os.path.exists(nsfwdet_model_path): + os.makedirs(nsfwdet_model_path) nd_urls = [ "https://huggingface.co/AdamCodd/vit-base-nsfw-detector/resolve/main/config.json", - "https://huggingface.co/AdamCodd/vit-base-nsfw-detector/resolve/main/confusion_matrix.png", "https://huggingface.co/AdamCodd/vit-base-nsfw-detector/resolve/main/model.safetensors", "https://huggingface.co/AdamCodd/vit-base-nsfw-detector/resolve/main/preprocessor_config.json", ] for model_url in nd_urls: model_name = os.path.basename(model_url) - model_path = os.path.join(model_path, model_name) + model_path = os.path.join(nsfwdet_model_path, model_name) download(model_url, model_path, model_name) SCORE = 0.965 # 0.965 and less - is safety content