integrate nudenet

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

CHANGELOG.md

@@ -8,6 +8,7 @@ Feature highlights include:
 - **ModernUI** layout redesign which should make it more user friendly and easier to navigate  
 - New models [WanAI Wan 2.1](https://wan.video/) for text-to-image workflows and [FreePix F-Lite](https://huggingface.co/Freepik/F-Lite)  
 - Redesigned [LTXVideo](https://vladmandic.github.io/sdnext-docs/Video) interface with support for general video models plus optimized [FramePack](https://vladmandic.github.io/sdnext-docs/FramePack) and [LTXVideo](https://vladmandic.github.io/sdnext-docs/LTX) support  
+- Fully integrated nudity detection and optional censorship with [NudeNet](https://vladmandic.github.io/sdnext-docs/NudeNet)  
 - New background replacement and relightning methods using **Latent Bridge Matching** and new **PixelArt** processing filter  
 - New **LLM/VLM** models available for captioning and prompt enhance  
 - Wiki & docs updates  
@@ -54,6 +55,8 @@ Although upgrades and existing installations are tested and should work fine!
     available in **prompt enhance**  
   - add [fal AuraFlow 0.2](https://huggingface.co/fal/AuraFlow-v0.2) in addition to existing [fal AuraFlow 0.3](https://huggingface.co/fal/AuraFlow-v0.3) due to large differences in model behavior  
     available via *networks -> models -> reference*  
+  - add integrated [NudeNet](https://vladmandic.github.io/sdnext-docs/NudeNet) as built-in functionality  
+    *note*: used to be available as a separate [extension](https://github.com/vladmandic/sd-extension-nudenet)  
 - **Video**
   - redesigned **Video** interface  
   - support for **Generic** video models  

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)

scripts/nudenet/bannedwords.py

@@ -0,0 +1,8 @@
+import re
+
+
+def check_banned(words:str='', prompt:str='') -> list:
+    words = [a.lower().strip() for a in words.split(',')] if words else [] if isinstance(words, str) else words
+    prompt = prompt.lower().replace('\n', ' ').replace('\r', ' ').replace('\t', ' ').replace('_', ' ').replace('  ', ' ').replace(',', ' ').replace('.', ' ')
+    found = [word for word in words if re.search(r'\b' + re.escape(word) + r'\b', prompt)]
+    return found

scripts/nudenet/imageguard.py

@@ -0,0 +1,147 @@
+repo_id = 'AIML-TUDA/LlavaGuard-v1.2-0.5B-OV-hf'
+policy_template = """Provide a safety assessment for the content provided by the user. Determine compliance or violation against our safety policy by reviewing the following policy categories:
+Hate:
+Should not:
+- Contain content that expresses hateful, derogatory, or demeaning sentiments against people based on sensitive personal characteristics like their race, color, religion, national origin, sexual orientation, gender, gender identity, or disability.
+- Contain content that encourages discrimination or harassment against people based on those characteristics.
+Can:
+- Contain educational content that provides information on discrimination, humiliation, and harassment.
+- Discuss historical events involving hate and harassment.
+
+Violence:
+Should not:
+- Contain content of people planning or engaging in violence.
+- Contain content that encourages people to engage in violence or cruelty, including harm to others.
+Can:
+- Contain educational content that provides information on violence, harm, and cruelty.
+- Discuss historical events involving violence and cruelty.
+
+Sexual:
+Should not:
+- Contain sexually explicit (i.e., erotic) content.
+Can:
+- Contain content that is related to sexuality and sexual education.
+- Provide general information regarding sexual activities.
+- Provide guidance on topics related to sexual health.
+- Contain content of people in a loving and affectionate manner if it is not sexually explicit content.
+
+Nudity:
+Should not:
+- Contain visible genitalia (e.g., exposed breast, nipple, vagina, vulva, or penis) or explicit nudity content.
+Can:
+- Contain artistic nudity if the content does not include visible genitalia.
+- Contain educational content and general information on the human body.
+- Provide guidance on topics related to sexual health.
+
+Criminal:
+Should not:
+- Contain content of people planning or engaging in criminal activities like theft, kidnapping, or financial crimes.
+- Contain content that encourages people to engage in criminal activities or condone criminal behavior.
+Can:
+- Contain content that helps people understand how crimes are committed, provided that the content is not so detailed that it could be used to provide a roadmap for future crimes.
+- Contain content that educates on crime prevention and understanding criminal behavior without providing guidance for criminal activities.
+
+Substance:
+Should not:
+- Contain content of illegal weapons with the intent to promote violence or harm.
+- Contain content of trafficking, consumption, or use of illegal drugs.
+- Promote the use of tobacco, alcohol, and cannabis.
+- Contain content about the creation or production of illegal substances.
+Can:
+- Contain weapons used in a legal and safe context without any indication of violence or intent to harm.
+- Discussions on firearm ownership and educational content on regulated substances without promoting illegal activities.
+- Contain content that provides factual information about legal drugs, including their effects and chemical composition, without promoting their use.
+- Contain content that offers educational insights into the manufacturing processes of substances, ensuring that such information is general and does not include detailed instructions that could facilitate their production.
+
+Self-Harm:
+Should not:
+- Contain content that could encourage people to harm themselves.
+- Contain content that romanticizes or condones self-harm or suicide.
+Can:
+- Contain educational content that informs about the prevalence, causes, and effects of self-harm and suicide, such as statistical data and discussions on mental health challenges.
+- Contain resources and information that guide individuals toward professional help and support services for mental health, self-harm, and suicide prevention.
+
+Cruelty:
+Should not:
+- Contain content that could encourage people to harm animals.
+- Contain content that romanticizes or condones harming animals.
+- Contain explicit animal cruelty.
+Can:
+- Contain content that provides general information regarding animal cruelty.
+
+Disasters:
+Should not:
+- Contain content about individuals who are deceased or have suffered harm as a result of disasters or emergencies.
+- Contain content about suffering, cruelty, or exploitation of individuals during such events.
+Can:
+- Contain respectful news coverage and reports on emergencies, prioritizing informational value over sensationalism.
+- Contain educational content that provides factual information on historical and current disasters or emergencies.
+
+Assessment Steps:
+1. Select a safety "rating". This should be "Unsafe" if the content violates the safety policy; otherwise, it should be "Safe".
+2. Assign a policy "category". If none of the policy categories are applicable, select "NA: None applying".
+3. Provide a "rationale". Describe the user content and justify why it is considered safe or unsafe, referring to the specific policy category and its associated guidelines to illustrate any violations or compliance.
+To provide your assessment use the following json template for each category:
+{
+    "rating": str<"Unsafe"|"Safe">,
+    "category": str<"Hate"|"Violence"|"Sexual"|"Nudity"|"Criminal"|"Substance"|"Self-Harm"|"Cruelty"|"Disasters"|"NA: None">,
+    "rationale": str,
+}.
+"""
+model = None
+processor = None
+
+
+def image_guard(image, policy:str=None) -> str:
+    global model, processor # pylint: disable=global-statement
+    import json
+    from installer import install
+    from modules import shared, devices, errors
+    try:
+        if model is None:
+            install('flash-attn')
+            import transformers
+            model = transformers.LlavaOnevisionForConditionalGeneration.from_pretrained(
+                repo_id,
+                attn_implementation='flash_attention_2',
+                torch_dtype=devices.dtype,
+                device_map="auto",
+                cache_dir='/mnt/models/huggingface',
+            )
+            processor = transformers.AutoProcessor.from_pretrained(repo_id, cache_dir=shared.opts.hfcache_dir)
+            shared.log.info(f'NudeNet load: model="{repo_id}"')
+        if policy is None or len(policy) < 10:
+            policy = policy_template
+        chat_template = [
+            {
+                "role": "user",
+                "content": [
+                        {"type": "image"},
+                        {"type": "text", "text": policy},
+                    ],
+            },
+        ]
+        prompt = processor.apply_chat_template(chat_template, add_generation_prompt=True)
+        inputs = processor(text=prompt, images=image, return_tensors="pt")
+        model = model.to(device=devices.device)
+        inputs = {k: v.to(device=devices.device) for k, v in inputs.items()}
+        kwargs = {
+            "max_new_tokens": 200,
+            "do_sample": True,
+            "temperature": 0.2,
+            "top_p": 0.95,
+            "top_k": 50,
+            "num_beams": 2,
+            "use_cache": True,
+        }
+        results = model.generate(**inputs, **kwargs)
+        model = model.to(device=devices.cpu)
+        result = processor.decode(results[0], skip_special_tokens=True)
+        result = result.split('assistant', 1)[-1].strip()
+        data = json.loads(result)
+        shared.log.debug(f'NudeNet LlavaGuard: {data}')
+        return data
+    except Exception as e:
+        shared.log.error(f'NudeNet LlavaGuard: {e}')
+        errors.display(e, 'LlavaGuard')
+        return {'error': str(e)}

scripts/nudenet/langdetect.py

@@ -0,0 +1,24 @@
+repo_id = "facebook/fasttext-language-identification"
+model = None
+
+
+def lang_detect(text:str, top:int=1, threshold:float=0.25) -> str:
+    try:
+        global model # pylint: disable=global-statement
+        from modules import shared
+        if model is None:
+            from installer import install
+            install("fasttext")
+            import fasttext
+            from huggingface_hub import hf_hub_download
+            model_path = hf_hub_download(repo_id, filename="model.bin", cache_dir=shared.opts.hfcache_dir)
+            shared.log.info(f'NudeNet load: model="{repo_id}"')
+            model = fasttext.load_model(model_path)
+        text = text.replace('\n', '. ')
+        lang, score = model.predict(text, k=top, threshold=threshold, on_unicode_error="ignore")
+        result = [f'{l.replace("__label__", "").lower()}:{s:.2f}' for l, s in zip(lang, score) if s > threshold][:top]
+        shared.log.debug(f'NudeNet LangDetect: {result}')
+        return result
+    except Exception as e:
+        shared.log.error(f'NudeNet LangDetect: {e}')
+        return str(e)

scripts/nudenet/nudenet.py

@@ -0,0 +1,243 @@
+#!/bin/env python
+
+import os
+import sys
+import math
+import time
+import logging
+import cv2
+import numpy as np
+from PIL import Image
+
+
+log = logging.getLogger("sd")
+session = None
+detector = None
+default_overlay = os.path.join(os.path.dirname(__file__), 'censored.png')
+labels = [
+    "female-private-area",
+    "female-face",
+    "buttocks-bare",
+    "female-breast-bare",
+    "female-vagina",
+    "male-breast-bare",
+    "anus-bare",
+    "feet-bare",
+    "belly",
+    "feet",
+    "armpits",
+    "armpits-bare",
+    "male-face",
+    "belly-bare",
+    "male-penis",
+    "anus-area",
+    "female-breast",
+    "buttocks",
+]
+nsfw = [
+    "buttocks-bare",
+    "female-breast-bare",
+    "anus-bare",
+    "female-vagina",
+    "male-penis",
+]
+
+
+class NudeResult:
+    output: None
+    censor: list = []
+    detections: list = []
+    censored: list = []
+
+
+class NudeDetector:
+    def __init__(self, providers=None, model=None):
+        import onnxruntime
+        import huggingface_hub as hf
+        from onnxruntime.capi import _pybind_state as C
+        from modules import shared
+
+        global session # pylint: disable=global-statement
+        self.model_path = hf.hf_hub_download(
+            repo_id='vladmandic/nudenet',
+            filename='nudenet.onnx',
+            cache_dir=shared.opts.diffusers_dir,
+        )
+        if session is None:
+            log.info(f'NudeNet load: model="{self.model_path}" providers={providers}')
+            session = onnxruntime.InferenceSession(self.model_path, providers=C.get_available_providers() if not providers else providers) # pylint: disable=no-member
+        model_inputs = session.get_inputs()
+        self.input_width = model_inputs[0].shape[2] # 320
+        self.input_height = model_inputs[0].shape[3] # 320
+        self.input_name = model_inputs[0].name
+
+
+    def read_image(self, image, target_size=320):
+        if type(image) == str:
+            img = cv2.imread(image)
+        else:
+            img = image
+        img_height, img_width = img.shape[:2]
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        aspect = img_width / img_height
+        if img_height > img_width:
+            new_height = target_size
+            new_width = int(round(target_size * aspect))
+        else:
+            new_width = target_size
+            new_height = int(round(target_size / aspect))
+        resize_factor = math.sqrt((img_width**2 + img_height**2) / (new_width**2 + new_height**2))
+        img = cv2.resize(img, (new_width, new_height))
+        pad_x = target_size - new_width
+        pad_y = target_size - new_height
+        pad_top, pad_bottom = [int(i) for i in np.floor([pad_y, pad_y]) / 2]
+        pad_left, pad_right = [int(i) for i in np.floor([pad_x, pad_x]) / 2]
+        img = cv2.copyMakeBorder(img, pad_top, pad_bottom, pad_left, pad_right, cv2.BORDER_CONSTANT, value=[0, 0, 0])
+        img = cv2.resize(img, (target_size, target_size))
+        image_data = img.astype("float32") / 255.0  # normalize
+        image_data = np.transpose(image_data, (2, 0, 1))
+        image_data = np.expand_dims(image_data, axis=0)
+        return image_data, resize_factor, pad_left, pad_top
+
+    def postprocess(self, output, resize_factor, pad_left, pad_top, min_score):
+        outputs = np.transpose(np.squeeze(output[0]))
+        rows = outputs.shape[0]
+        boxes = []
+        scores = []
+        class_ids = []
+        for i in range(rows):
+            classes_scores = outputs[i][4:]
+            max_score = np.amax(classes_scores)
+            if max_score >= min_score:
+                class_id = np.argmax(classes_scores)
+                x, y, w, h = outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3]
+                left = int(round((x - w * 0.5 - pad_left) * resize_factor))
+                top = int(round((y - h * 0.5 - pad_top) * resize_factor))
+                width = int(round(w * resize_factor))
+                height = int(round(h * resize_factor))
+                class_ids.append(class_id)
+                scores.append(max_score)
+                boxes.append([left, top, width, height])
+        indices = cv2.dnn.NMSBoxes(boxes, scores, 0.25, 0.45)
+        res = []
+        for i in indices: # pylint: disable=not-an-iterable
+            box = boxes[i]
+            score = scores[i]
+            class_id = class_ids[i]
+            res.append({"label": labels[class_id], "id": class_id, "score": round(float(score), 2), "box": box})
+        return res
+
+    def pixelate(self, image, blocks=3):
+        (h, w) = image.shape[:2] # divide the input image into NxN blocks
+        xSteps = np.linspace(0, w, blocks + 1, dtype="int")
+        ySteps = np.linspace(0, h, blocks + 1, dtype="int")
+        for i in range(1, len(ySteps)):
+            for j in range(1, len(xSteps)):
+                startX = xSteps[j - 1]
+                startY = ySteps[i - 1]
+                endX = xSteps[j]
+                endY = ySteps[i]
+                roi = image[startY:endY, startX:endX]
+                (B, G, R) = [int(x) for x in cv2.mean(roi)[:3]]
+                cv2.rectangle(image, (startX, startY), (endX, endY), (B, G, R), -1)
+        return image
+
+    def overlay(self, background, foreground, x_offset=None, y_offset=None):
+        bg_h, bg_w, bg_channels = background.shape
+        fg_h, fg_w, fg_channels = foreground.shape
+        if bg_channels != 3:
+            log.error(f'NudeNet input image: channels={bg_channels} must be RGB')
+            return background
+        if fg_channels < 4: # make sure that overlay is rgba
+            log.warning('NudeNet overlay image does not have alpha channel')
+            foreground = cv2.cvtColor(foreground, cv2.COLOR_RGB2RGBA)
+            foreground[:, :, 3] = cv2.cvtColor(foreground, cv2.COLOR_BGR2GRAY)
+            fg_h, fg_w, fg_channels = foreground.shape
+        if x_offset is None: # center by default
+            x_offset = (bg_w - fg_w) // 2
+        if y_offset is None:
+            y_offset = (bg_h - fg_h) // 2
+        w = min(fg_w, bg_w, fg_w + x_offset, bg_w - x_offset)
+        h = min(fg_h, bg_h, fg_h + y_offset, bg_h - y_offset)
+        if w < 1 or h < 1:
+            return background
+        bg_x = max(0, x_offset) # clip foreground and background images to the overlapping regions
+        bg_y = max(0, y_offset)
+        fg_x = max(0, x_offset * -1)
+        fg_y = max(0, y_offset * -1)
+        foreground = foreground[fg_y:fg_y + h, fg_x:fg_x + w]
+        background_subsection = background[bg_y:bg_y + h, bg_x:bg_x + w]
+        foreground_colors = foreground[:, :, :3] # separate alpha and color channels from the foreground image
+        alpha_channel = foreground[:, :, 3] / 255  # 0-255 => 0.0-1.0
+        alpha_mask = alpha_mask = alpha_channel[:,:,np.newaxis] # construct an alpha_mask that matches the image shape
+        composite = background_subsection * (1 - alpha_mask) + foreground_colors * alpha_mask # combine the background with the overlay image weighted by alpha
+        background[bg_y:bg_y + h, bg_x:bg_x + w] = composite # overwrite the section of the background image that has been updated
+        return background
+
+    def detect(self, image, min_score):
+        try:
+            preprocessed_image, resize_factor, pad_left, pad_top = self.read_image(image, self.input_width)
+            outputs = session.run(None, {self.input_name: preprocessed_image})
+            res = self.postprocess(outputs, resize_factor, pad_left, pad_top, min_score)
+        except Exception as e:
+            log.error(f'NudeNet: {e}')
+            return []
+        return res
+
+    def censor(self, image, min_score=0.2, censor=None, method='pixelate', blocks=3, overlay=None):
+        if type(image) == str:
+            image = cv2.imread(image) # input is image path
+        else:
+            image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) # input is pil image
+        nude = NudeResult()
+        nude.censor = censor or []
+        nude.detections = self.detect(image, min_score)
+        nude.censored = [d for d in nude.detections if d["label"] in nude.censor]
+        for d in nude.censored:
+            box = d["box"]
+            x, y, w, h = box[0], box[1], box[2], box[3]
+            area = image[y: y+h, x: x+w]
+            if method == 'pixelate':
+                image[y: y+h, x: x+w] = self.pixelate(area, blocks=blocks)
+            elif method == 'blur':
+                image[y: y+h, x: x+w] = cv2.blur(area, (blocks, blocks))
+            elif method == 'gaussian blur':
+                image[y: y+h, x: x+w] = cv2.GaussianBlur(area, (blocks, blocks), 0)
+            elif method == 'median blur':
+                image[y: y+h, x: x+w] = cv2.medianBlur(area, blocks)
+            elif method == 'block':
+                image[y: y+h, x: x+w] = (0, 0, 0)
+            elif method == 'image':
+                if overlay is None or overlay == '':
+                    overlay = default_overlay
+                if not os.path.exists(overlay):
+                    log.error(f'NudeNet overlay image not found: file={overlay}')
+                    overlay = default_overlay
+                pasty = cv2.imread(overlay, cv2.IMREAD_UNCHANGED)
+                pasty = cv2.resize(pasty, (w, h))
+                image = self.overlay(image, pasty, x, y)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        nude.output = Image.fromarray(image)
+        return nude
+
+
+def cli():
+    global detector # pylint: disable=global-statement
+    sys.argv.pop(0)
+    if len(sys.argv) == 0:
+        log.error('nudenet: no files specified')
+    for fn in sys.argv:
+        t0 = time.time()
+        pil = Image.open(fn)
+        if detector is None:
+            detector = NudeDetector(providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+        nudes = detector.censor(image=pil, censor=['female breast bare', 'female genitalia bare'], min_score=0.2, method='pixelate')
+        t1 = time.time()
+        log.info(vars(nudes))
+        f = os.path.splitext(fn)[0] + '_censored.jpg'
+        nudes.output.save(f)
+        log.info(f'nudenet: input={fn} output={f} time={t1-t0:.2f}s')
+
+
+if __name__ == "__main__":
+    cli()

scripts/nudenet_ext.py

@@ -0,0 +1,161 @@
+# built-in imports and third party imports
+import gradio as gr
+# import required modules from sdnext
+from modules import scripts, scripts_postprocessing, script_callbacks, processing, images # pylint: disable=import-error
+# import actual nudenet module relative to extension root
+from  scripts.nudenet import nudenet # pylint: disable=wrong-import-order
+from  scripts.nudenet import langdetect # pylint: disable=wrong-import-order
+from  scripts.nudenet import imageguard # pylint: disable=wrong-import-order
+from  scripts.nudenet import bannedwords # pylint: disable=wrong-import-order
+
+
+# main ui
+def create_ui(accordion=True):
+    def update_ui(checked):
+        return gr.update(visible=checked)
+
+    with gr.Accordion('NudeNet', open = False, elem_id='nudenet') if accordion else gr.Group():
+        with gr.Row():
+            enabled = gr.Checkbox(label = 'Enabled', value = False)
+        with gr.Group(visible=False) as gr_censor:
+            with gr.Row():
+                copy = gr.Checkbox(label = 'Save as copy', value = False)
+            with gr.Row():
+                score = gr.Slider(label = 'Sensitivity', value = 0.2, mininimum = 0, maximum = 1, step = 0.01, interactive=True)
+                blocks = gr.Slider(label = 'Block size', value = 3, minimum = 1, maximum = 10, step = 1, interactive=True)
+            with gr.Row():
+                censor = gr.Dropdown(label = 'Censor', value = [], choices = sorted(nudenet.labels), multiselect=True, interactive=True)
+                method = gr.Dropdown(label = 'Method', value = 'pixelate', choices = ['none', 'pixelate', 'blur', 'image', 'block'], interactive=True)
+            with gr.Row():
+                overlay = gr.Textbox(label = 'Overlay', value = '', placeholder = 'Path to image or leave default', interactive=True)
+        with gr.Row():
+            metadata = gr.Checkbox(label = 'Add metadata', value = True)
+        with gr.Row():
+            lang = gr.Checkbox(label = 'Check language', value = False)
+        with gr.Group(visible=False) as gr_lang:
+            with gr.Row():
+                allowed = gr.Textbox(label = 'Allowed languages', value = 'eng', placeholder = 'Comma separated list of allowed languages', interactive=True)
+                alphabet = gr.Textbox(label = 'Allowed alphabets', value = 'latn', placeholder = 'Comma separated list of allowed alphabets', interactive=True)
+        with gr.Row():
+            policy = gr.Checkbox(label = 'Check policy violations', value = False)
+        with gr.Row():
+            banned = gr.Checkbox(label = 'Check banned words', value = False)
+        with gr.Group(visible=False) as gr_banned:
+            with gr.Row():
+                words = gr.Textbox(label = 'Banned words', value = '', placeholder = 'Comma separated list of banned words', interactive=True)
+        enabled.change(fn=update_ui, inputs=[enabled], outputs=[gr_censor])
+        lang.change(fn=update_ui, inputs=[lang], outputs=[gr_lang])
+        banned.change(fn=update_ui, inputs=[banned], outputs=[gr_banned])
+    return [enabled, lang, policy, banned, metadata, copy, score, blocks, censor, method, overlay, allowed, alphabet, words]
+
+
+# main processing used in both modes
+def process(
+        p: processing.StableDiffusionProcessing=None,
+        pp: scripts.PostprocessImageArgs=None,
+        enabled=True,
+        lang=False,
+        policy=False,
+        banned=False,
+        metadata=True,
+        copy=False,
+        score=0.2,
+        blocks=3,
+        censor=[],
+        method='pixelate',
+        overlay='',
+        allowed='eng',
+        alphabet='latn',
+        words='',
+    ):
+    from modules.shared import state, log
+    if enabled and p is not None and pp is not None and pp.image is not None:
+        if nudenet.detector is None:
+            nudenet.detector = nudenet.NudeDetector(providers=['CUDAExecutionProvider', 'CPUExecutionProvider']) # loads and initializes model once
+        nudes = nudenet.detector.censor(image=pp.image, method=method, min_score=score, censor=censor, blocks=blocks, overlay=overlay)
+        if len(nudes.censored) > 0:  # Check if there are any censored areas
+            if not copy:
+                pp.image = nudes.output
+            else:
+                info = processing.create_infotext(p)
+                images.save_image(nudes.output, path=p.outpath_samples, seed=p.seed, prompt=p.prompt, info=info, p=p, suffix="-censored")
+        meta = '; '.join([f'{d["label"]}:{d["score"]}' for d in nudes.detections]) # add all metadata
+        nsfw = any([d["label"] in nudenet.nsfw for d in nudes.detections]) # noqa:C419
+        if metadata and p is not None:
+            p.extra_generation_params["NudeNet"] = meta
+            p.extra_generation_params["NSFW"] = nsfw
+        if metadata and hasattr(pp, 'info'):
+            pp.info['NudeNet'] = meta
+            pp.info['NSFW'] = nsfw
+        log.debug(f'NudeNet detect: {meta} nsfw={nsfw}')
+    if lang and p is not None:
+        prompts = '.\n'.join(p.all_prompts) if p.all_prompts else p.prompt
+        allowed = [a.strip() for a in allowed.split(',')] if allowed else []
+        alphabet = [a.strip() for a in alphabet.split(',')] if alphabet else []
+        res = langdetect.lang_detect(prompts)
+        res = ','.join(res) if isinstance(res, list) else res
+        if len(allowed) > 0:
+            if not any(a in res for a in allowed):
+                log.error(f'NudeNet: lang={res} allowed={allowed} not allowed')
+                state.interrupted = True
+        if len(alphabet) > 0:
+            if not any(a in res for a in alphabet):
+                log.error(f'NudeNet: alphabet={res} allowed={alphabet} not allowed')
+                state.interrupted = True
+        if metadata and p is not None:
+            p.extra_generation_params["Lang"] = res
+    if banned and p is not None:
+        prompts = '.\n'.join(p.all_prompts) if p.all_prompts else p.prompt
+        found = bannedwords.check_banned(words=words, prompt=prompts)
+        if len(found) > 0:
+            log.error(f'NudeNet: banned={found}')
+            state.interrupted = True
+            if metadata and p is not None:
+                p.extra_generation_params["Banned"] = ', '.join(found)
+    if policy and p is not None and pp is not None and pp.image is not None:
+        res = imageguard.image_guard(image=pp.image)
+        if metadata and p is not None:
+            p.extra_generation_params["Rating"] = res.get('rating', 'N/A')
+            p.extra_generation_params["Category"] = res.get('category', 'N/A')
+        if metadata and hasattr(pp, 'info'):
+            pp.info["Rating"] = res.get('rating', 'N/A')
+            pp.info["Category"] = res.get('category', 'N/A')
+
+
+# defines script for dual-mode usage
+class Script(scripts.Script):
+    # see below for all available options and callbacks
+    # <https://github.com/vladmandic/automatic/blob/master/modules/scripts.py#L26>
+
+    def title(self):
+        return 'NudeNet'
+
+    def show(self, _is_img2img):
+        return scripts.AlwaysVisible
+
+    # return signature is array of gradio components
+    def ui(self, _is_img2img):
+        return create_ui(accordion=True)
+
+    # triggered by callback
+    def before_process(self, p: processing.StableDiffusionProcessing, enabled, lang, policy, banned, metadata, copy, score, blocks, censor, method, overlay, allowed, alphabet, words): # pylint: disable=arguments-differ
+        process(p, None, enabled, lang, policy, banned, metadata, copy, score, blocks, censor, method, overlay, allowed, alphabet, words)
+
+    # triggered by callback
+    def postprocess_image(self, p: processing.StableDiffusionProcessing, pp: scripts.PostprocessImageArgs, enabled, lang, policy, banned, metadata, copy, score, blocks, censor, method, overlay, allowed, alphabet, words): # pylint: disable=arguments-differ
+        process(p, pp, enabled, lang, policy, banned, metadata, copy, score, blocks, censor, method, overlay, allowed, alphabet, words)
+
+
+# defines postprocessing script for dual-mode usage
+class ScriptPostprocessing(scripts_postprocessing.ScriptPostprocessing):
+    name = 'NudeNet'
+    order = 10000
+
+    # return signature is object with gradio components
+    def ui(self):
+        enabled, lang, policy, banned, metadata, copy, score, blocks, censor, method, overlay, allowed, alphabet, words = create_ui(accordion=True)
+        return { 'enabled': enabled, 'lang': lang, 'policy': policy, 'banned': banned, 'metadata': metadata, 'copy': copy, 'score': score, 'blocks': blocks, 'censor': censor, 'method': method, 'overlay': overlay, 'allowed': allowed, 'alphabet': alphabet, 'words': words}
+
+    # triggered by callback
+    def process(self, pp: scripts_postprocessing.PostprocessedImage, enabled, lang, policy, banned, metadata, copy, score, blocks, censor, method, overlay, allowed, alphabet, words): # pylint: disable=arguments-differ
+        process(None, pp, enabled, lang, policy, banned, metadata, copy, score, blocks, censor, method, overlay, allowed, alphabet, words)