integrate nudenet

Signed-off-by: Vladimir Mandic <mandic00@live.com>

modules/api/api.py

@@ -5,7 +5,7 @@ from fastapi import FastAPI, APIRouter, Depends, Request
 from fastapi.security import HTTPBasic, HTTPBasicCredentials
 from fastapi.exceptions import HTTPException
 from modules import errors, shared, postprocessing
-from modules.api import models, endpoints, script, helpers, server, nvml, generate, process, control, gallery, loras, docs
+from modules.api import models, endpoints, script, helpers, server, nvml, generate, process, control, docs
 
 
 errors.install()
@@ -100,13 +100,17 @@ class Api:
         self.add_api_route("/sdapi/v1/latents", endpoints.post_latent_history, methods=["POST"], response_model=int)
 
         # lora api
-        self.add_api_route("/sdapi/v1/lora", loras.get_lora, methods=["GET"], response_model=dict)
-        self.add_api_route("/sdapi/v1/loras", loras.get_loras, methods=["GET"], response_model=List[dict])
-        self.add_api_route("/sdapi/v1/refresh-loras", loras.post_refresh_loras, methods=["POST"])
+        from modules.api import loras
+        loras.register_api()
 
         # gallery api
+        from modules.api import gallery
         gallery.register_api(self.app)
 
+        # nudenet api
+        from modules.api import nudenet
+        nudenet.register_api()
+
 
     def add_api_route(self, path: str, endpoint, **kwargs):
         if (shared.cmd_opts.auth or shared.cmd_opts.auth_file) and shared.cmd_opts.api_only:

modules/api/loras.py

@@ -1,3 +1,4 @@
+from typing import List
 from fastapi.exceptions import HTTPException
 
 
@@ -19,3 +20,10 @@ def get_loras():
 def post_refresh_loras():
     from modules.lora import lora_load
     return lora_load.list_available_networks()
+
+
+def register_api():
+    from modules.shared import api
+    api.add_api_route("/sdapi/v1/lora", get_lora, methods=["GET"], response_model=dict)
+    api.add_api_route("/sdapi/v1/loras", get_loras, methods=["GET"], response_model=List[dict])
+    api.add_api_route("/sdapi/v1/refresh-loras", post_refresh_loras, methods=["POST"])

modules/api/nudenet.py

@@ -0,0 +1,64 @@
+from fastapi import Body
+from modules.api import api
+
+
+def nudenet_censor(
+    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'),
+):
+    from scripts.nudenet import nudenet
+    base64image = image
+    image = api.decode_base64_to_image(image)
+    if nudenet.detector is None:
+        nudenet.detector = nudenet.NudeDetector() # loads and initializes model once
+    nudes = nudenet.detector.censor(image=image, method=method, min_score=score, censor=censor, blocks=blocks, overlay=overlay)
+    if len(censor) > 0: # replace image if anything is censored
+        base64image = api.encode_pil_to_base64(nudes.output).decode("utf-8")
+    detections_dict = { d["label"]: d["score"] for d in nudes.detections }
+    return { "image": base64image, "detections": detections_dict }
+
+
+def prompt_check(
+    prompt: str = Body("", title='prompt text'),
+    lang: str = Body("eng", title='allowed languages'),
+    alphabet: str = Body("latn", title='allowed alphabets'),
+):
+    from scripts.nudenet import langdetect
+    res = langdetect.lang_detect(prompt)
+    res = ','.join(res) if isinstance(res, list) else res
+    lang = [a.strip() for a in lang.split(',')] if lang else []
+    alphabet = [a.strip() for a in alphabet.split(',')] if alphabet else []
+    lang_ok = any(a in res for a in lang) if len(lang) > 0 else True
+    alph_ok = any(a in res for a in alphabet) if len(alphabet) > 0 else True
+    return { "lang": res, "lang_ok": lang_ok, "alph_ok": alph_ok }
+
+
+def image_guard(
+    image: str = Body("", title='input image'),
+    policy: str = Body("", title='optional policy definition'),
+):
+    from scripts.nudenet import imageguard
+    image = api.decode_base64_to_image(image)
+    res = imageguard.image_guard(image=image, policy=policy)
+    return res
+
+
+def banned_words(
+    words: str = Body("", title='comma separated list of banned words'),
+    prompt: str = Body("", title='prompt text'),
+):
+    from scripts.nudenet import bannedwords
+    found = bannedwords.check_banned(words=words, prompt=prompt)
+    return found
+
+
+def register_api():
+    from modules.shared import api
+    api.add_api_route("/sdapi/v1//nudenet", nudenet_censor, methods=["POST"], response_model=dict)
+    api.add_api_route("/sdapi/v1//prompt-lang", prompt_check, methods=["POST"], response_model=dict)
+    api.add_api_route("/sdapi/v1//image-guard", image_guard, methods=["POST"], response_model=dict)
+    api.add_api_route("/sdapi/v1//prompt-banned", banned_words, methods=["POST"], response_model=list)