diff --git a/modules/api/nudenet.py b/modules/api/nudenet.py index a1243592a..9ad6332dd 100644 --- a/modules/api/nudenet.py +++ b/modules/api/nudenet.py @@ -3,13 +3,14 @@ from modules.api import api def nudenet_censor( - image: str = Body("", title='nudenet input image'), + image: str = Body("", title='NudeNet input image'), score: float = Body(0.2, title='nudenet threshold score'), blocks: int = Body(3, title='nudenet pixelation blocks'), censor: list = Body([], title='nudenet censorship items'), method: str = Body('pixelate', title='nudenet censorship method'), overlay: str = Body('', title='nudenet overlay image path'), ): + """Detect and censor NSFW regions in an image using NudeNet. Returns detections and optionally censored image.""" from scripts.nudenet import nudenet # pylint: disable=no-name-in-module base64image = image image = api.decode_base64_to_image(image) @@ -27,6 +28,7 @@ def prompt_check( lang: str = Body("eng", title='allowed languages'), alphabet: str = Body("latn", title='allowed alphabets'), ): + """Check prompt language and alphabet against allowed values. Returns detected language and pass/fail flags.""" from scripts.nudenet import langdetect # pylint: disable=no-name-in-module res = langdetect.lang_detect(prompt) res = ','.join(res) if isinstance(res, list) else res @@ -41,6 +43,7 @@ def image_guard( image: str = Body("", title='input image'), policy: str = Body("", title='optional policy definition'), ): + """Evaluate an image against a content policy using the ImageGuard classifier.""" from scripts.nudenet import imageguard # pylint: disable=no-name-in-module image = api.decode_base64_to_image(image) res = imageguard.image_guard(image=image, policy=policy) @@ -51,6 +54,7 @@ def banned_words( words: str = Body("", title='comma separated list of banned words'), prompt: str = Body("", title='prompt text'), ): + """Check a prompt against a comma-separated list of banned words. Returns any matches found.""" from scripts.nudenet import bannedwords # pylint: disable=no-name-in-module found = bannedwords.check_banned(words=words, prompt=prompt) return found @@ -58,7 +62,7 @@ def banned_words( def register_api(): from modules.shared import api as api_instance - api_instance.add_api_route("/sdapi/v1/nudenet", nudenet_censor, methods=["POST"], response_model=dict) - api_instance.add_api_route("/sdapi/v1/prompt-lang", prompt_check, methods=["POST"], response_model=dict) - api_instance.add_api_route("/sdapi/v1/image-guard", image_guard, methods=["POST"], response_model=dict) - api_instance.add_api_route("/sdapi/v1/prompt-banned", banned_words, methods=["POST"], response_model=list) + api_instance.add_api_route("/sdapi/v1/nudenet", nudenet_censor, methods=["POST"], response_model=dict, tags=["Processing"]) + api_instance.add_api_route("/sdapi/v1/prompt-lang", prompt_check, methods=["POST"], response_model=dict, tags=["Processing"]) + api_instance.add_api_route("/sdapi/v1/image-guard", image_guard, methods=["POST"], response_model=dict, tags=["Processing"]) + api_instance.add_api_route("/sdapi/v1/prompt-banned", banned_words, methods=["POST"], response_model=list, tags=["Processing"])