From b7e30eea5e9f03389153c4f21bd78b495e22ee41 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gourieff=20=7C=20=E5=8F=A4=E4=BB=81?= Date: Sat, 1 Mar 2025 18:51:48 +0700 Subject: [PATCH] FIX: Black Image Issue Thanks @mschuettlerTNG for pointing out https://github.com/Gourieff/ComfyUI-ReActor/issues/62#issuecomment-2689122264 Issues #15 #28 #57 #62 --- scripts/reactor_sfw.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/scripts/reactor_sfw.py b/scripts/reactor_sfw.py index b2dfd0c..7c730c1 100644 --- a/scripts/reactor_sfw.py +++ b/scripts/reactor_sfw.py @@ -3,6 +3,7 @@ from PIL import Image import logging import os from reactor_utils import download +from scripts.reactor_logger import logger def ensure_nsfw_model(nsfwdet_model_path): """Download NSFW detection model if it doesn't exist""" @@ -18,13 +19,16 @@ def ensure_nsfw_model(nsfwdet_model_path): model_path = os.path.join(nsfwdet_model_path, model_name) download(model_url, model_path, model_name) -SCORE = 0.972 +SCORE = 0.96 -logging.getLogger('transformers').setLevel(logging.ERROR) +logging.getLogger("transformers").setLevel(logging.ERROR) def nsfw_image(img_path: str, model_path: str): ensure_nsfw_model(model_path) with Image.open(img_path) as img: predict = pipeline("image-classification", model=model_path) result = predict(img) - return True if result[0]["score"] > SCORE else False + if result[0]["label"] == "nsfw" and result[0]["score"] > SCORE: + logger.status(f"NSFW content detected, skipping...") + return True + return False