diff --git a/CHANGELOG.md b/CHANGELOG.md index f8ca70969..9a01d73a8 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/modules/api/api.py b/modules/api/api.py index 5e4d65198..dfc4cacb3 100644 --- a/modules/api/api.py +++ b/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: diff --git a/modules/api/loras.py b/modules/api/loras.py index 7e65a709d..4fbae29c5 100644 --- a/modules/api/loras.py +++ b/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"]) diff --git a/modules/api/nudenet.py b/modules/api/nudenet.py new file mode 100644 index 000000000..391a581fa --- /dev/null +++ b/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) diff --git a/scripts/nudenet/bannedwords.py b/scripts/nudenet/bannedwords.py new file mode 100644 index 000000000..38028276c --- /dev/null +++ b/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 diff --git a/scripts/nudenet/imageguard.py b/scripts/nudenet/imageguard.py new file mode 100644 index 000000000..55bcfaf86 --- /dev/null +++ b/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)} diff --git a/scripts/nudenet/langdetect.py b/scripts/nudenet/langdetect.py new file mode 100644 index 000000000..6c439726d --- /dev/null +++ b/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) diff --git a/scripts/nudenet/nudenet.py b/scripts/nudenet/nudenet.py new file mode 100755 index 000000000..1a59e568c --- /dev/null +++ b/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() diff --git a/scripts/nudenet_ext.py b/scripts/nudenet_ext.py new file mode 100644 index 000000000..f60bd2647 --- /dev/null +++ b/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 + # + + 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) diff --git a/wiki b/wiki index daacd1e61..d069faf08 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit daacd1e61d4ee86d275970fefaa269bb8670c761 +Subproject commit d069faf08f7a7b46f844488bd973de8790015ed1