Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 09:14:35 +02:00
Code Issues Packages Projects Releases Wiki Activity

update nudenet

Browse Source 
Signed-off-by: vladmandic <mandic00@live.com>
This commit is contained in:
vladmandic
2026-04-09 13:48:54 +02:00
parent f1d2888f10
commit 14ba5e4ce7
8 changed files with 85 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/nudenet/imageguard.py
+28 -10
View File
@@ -6,7 +6,14 @@ if TYPE_CHECKING:
from transformers import AutoProcessor, LlavaOnevisionForConditionalGeneration
repo_id = 'AIML-TUDA/LlavaGuard-v1.2-0.5B-OV-hf'
safety_models = [
'AIML-TUDA/QwenGuard-v1.2-3B',
'AIML-TUDA/QwenGuard-v1.2-7B',
'AIML-TUDA/LlavaGuard-v1.2-0.5B-OV-hf',
'AIML-TUDA/LlavaGuard-v1.2-7B-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:
@@ -98,28 +105,39 @@ To provide your assessment use the following json template for each category:
"""
repo: str = ''
model: LlavaOnevisionForConditionalGeneration | None = None
processor: AutoProcessor | None = None
def image_guard(image, policy:str | None=None):
global model, processor # pylint: disable=global-statement
def image_guard(image, policy:str | None=None, model_name:str=''):
global repo, model, processor # pylint: disable=global-statement
import json
from installer import install
from modules import shared, devices, errors
from modules import shared, devices, errors, model_quant
repo_id = model_name if model_name in safety_models else safety_models[0]
if policy is None or len(policy) < 10:
policy = policy_template
try:
if model is None:
install('flash-attn')
if model is None or repo != repo_id:
# install('flash-attn')
import transformers
model = transformers.LlavaOnevisionForConditionalGeneration.from_pretrained(
if 'LlavaGuard' in repo_id:
cls = transformers.LlavaOnevisionForConditionalGeneration
else:
cls = transformers.Qwen2_5_VLForConditionalGeneration
quant_args = model_quant.create_config(module='LLM')
model = cls.from_pretrained(
repo_id,
attn_implementation='flash_attention_2',
# attn_implementation='flash_attention_2',
torch_dtype=devices.dtype,
device_map="auto",
cache_dir=shared.opts.hfcache_dir,
**quant_args,
)
processor = transformers.AutoProcessor.from_pretrained(repo_id, cache_dir=shared.opts.hfcache_dir)
log.info(f'NudeNet load: model="{repo_id}"')
repo = repo_id
if policy is None or len(policy) < 10:
policy = policy_template
chat_template = [
@@ -152,6 +170,6 @@ def image_guard(image, policy:str | None=None):
log.debug(f'NudeNet LlavaGuard: {data}')
return data
except Exception as e:
log.error(f'NudeNet LlavaGuard: {e}')
errors.display(e, 'LlavaGuard')
log.error(f'NudeNet Safety: {e}')
errors.display(e, 'NudeNet')
return {'error': str(e)}
scripts/nudenet_ext.py
+37 -25
View File
@@ -14,10 +14,11 @@ def create_ui(accordion=True):
with gr.Accordion('NudeNet', open = False, elem_id='postprocess_nudenet_accordion') if accordion else gr.Group():
with gr.Row():
enabled = gr.Checkbox(label = 'Enabled', value = False)
copy = gr.Checkbox(label = 'Save copy', value = False)
metadata = gr.Checkbox(label = 'Update metadata', value = True)
with gr.Row():
enabled = gr.Checkbox(label = 'NudeNet 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)
@@ -26,8 +27,6 @@ def create_ui(accordion=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:
@@ -36,15 +35,21 @@ def create_ui(accordion=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.Group(visible=False) as gr_policy:
with gr.Row():
policy_model = gr.Dropdown(label = 'Policy model', value = imageguard.safety_models[0], choices = imageguard.safety_models, interactive=True)
with gr.Row():
policy_text = gr.Textbox(label = 'Policy template', value = '', placeholder = 'Custom policy template', interactive=True, lines=2)
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)
words = gr.Textbox(label = 'Banned words', value = '', placeholder = 'Comma separated list of banned words', interactive=True, lines=2)
enabled.change(fn=update_ui, inputs=[enabled], outputs=[gr_censor])
lang.change(fn=update_ui, inputs=[lang], outputs=[gr_lang])
policy.change(fn=update_ui, inputs=[policy], outputs=[gr_policy])
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]
return [enabled, lang, policy, banned, metadata, copy, score, blocks, censor, method, overlay, allowed, alphabet, words, policy_model, policy_text]
# main processing used in both modes
@@ -65,18 +70,20 @@ def process(
allowed='eng',
alphabet='latn',
words='',
policy_model='',
policy_text='',
):
from modules.shared import state, log
if enabled and pp is not None and pp.image is not None:
if enabled 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
t0 = time.time()
nudes = nudenet.detector.censor(image=pp.image, method=method, min_score=score, censor=censor, blocks=blocks, overlay=overlay)
t1 = time.time()
if len(nudes.censored) > 0: # Check if there are any censored areas
if p is None:
pp.image = nudes.output
else:
pp.image = nudes.output
if copy and p is not None:
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")
dct = {d["label"]: d["score"] for d in nudes.detections}
@@ -89,7 +96,8 @@ def process(
pp.info['NudeNet'] = meta
pp.info['NSFW'] = nsfw
log.debug(f'NudeNet detect: {dct} nsfw={nsfw} time={(t1 - t0):.2f}')
if lang and p is not None:
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 []
@@ -103,9 +111,10 @@ def process(
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:
if metadata and (p is not None):
p.extra_generation_params["Lang"] = res
if banned and p is not None:
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:
@@ -113,9 +122,12 @@ def process(
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:
if (not copy) and (pp is not None) and (pp.image is not None):
pp.image = None
if policy and (pp is not None) and (pp.image is not None):
res = imageguard.image_guard(image=pp.image, policy=policy_text, model_name=policy_model)
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 isinstance(pp, scripts_postprocessing.PostprocessedImage):
@@ -139,12 +151,12 @@ class ScriptNudeNet(scripts.Script):
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)
def before_process(self, p: processing.StableDiffusionProcessing, enabled, lang, policy, banned, metadata, copy, score, blocks, censor, method, overlay, allowed, alphabet, words, policy_model, policy_text): # pylint: disable=arguments-differ
process(p, None, enabled, lang, policy, banned, metadata, copy, score, blocks, censor, method, overlay, allowed, alphabet, words, policy_model, policy_text)
# 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)
def postprocess_image(self, p: processing.StableDiffusionProcessing, pp: scripts.PostprocessImageArgs, enabled, lang, policy, banned, metadata, copy, score, blocks, censor, method, overlay, allowed, alphabet, words, policy_model, policy_text): # pylint: disable=arguments-differ
process(p, pp, enabled, lang, policy, banned, metadata, copy, score, blocks, censor, method, overlay, allowed, alphabet, words, policy_model, policy_text)
# defines postprocessing script for dual-mode usage
@@ -154,9 +166,9 @@ class ScriptPostprocessingNudeNet(scripts_postprocessing.ScriptPostprocessing):
# 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}
enabled, lang, policy, banned, metadata, copy, score, blocks, censor, method, overlay, allowed, alphabet, words, policy_model, policy_text = 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, 'policy_model': policy_model, 'policy_text': policy_text}
# 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)
def process(self, pp: scripts_postprocessing.PostprocessedImage, enabled, lang, policy, banned, metadata, copy, score, blocks, censor, method, overlay, allowed, alphabet, words, policy_model, policy_text): # pylint: disable=arguments-differ
process(None, pp, enabled, lang, policy, banned, metadata, copy, score, blocks, censor, method, overlay, allowed, alphabet, words, policy_model, policy_text)
scripts/prompt_enhance.py
+10 -6
View File
@@ -420,7 +420,8 @@ class PromptEnhanceScript(scripts_manager.Script):
self.compile()
except Exception as e:
log.error(f'Prompt enhance: load {e}')
errors.display(e, 'Prompt enhance')
if debug_enabled:
errors.display(e, 'Prompt enhance')
devices.torch_gc()
self.busy = False
@@ -763,7 +764,8 @@ class PromptEnhanceScript(scripts_manager.Script):
debug_log(f'Prompt enhance: len={input_len} shape={inputs["input_ids"].shape} sample={sample} temp={temperature} penalty={penalty} max={tokens}')
except Exception as e:
log.error(f'Prompt enhance tokenize: {e}')
errors.display(e, 'Prompt enhance')
if debug_enabled:
errors.display(e, 'Prompt enhance')
self.busy = False
return prompt_text # Return original text part on error
try:
@@ -792,7 +794,8 @@ class PromptEnhanceScript(scripts_manager.Script):
except Exception as e:
outputs = None
log.error(f'Prompt enhance generate: {e}')
errors.display(e, 'Prompt enhance')
if debug_enabled:
errors.display(e, 'Prompt enhance')
self.busy = False
response = f'Error: {str(e)}'
finally:
@@ -853,8 +856,9 @@ class PromptEnhanceScript(scripts_manager.Script):
"""Update vision toggle interactivity and value based on model selection."""
repo_name = get_model_repo_from_display(model_name)
is_vl = is_vision_model(repo_name)
# When non-VL model: disable and uncheck. When VL model: enable and check.
return gr.update(interactive=is_vl, value=is_vl)
if not is_vl:
return gr.update(interactive=False, value=False)
return gr.update(interactive=is_vl)
def ui(self, _is_img2img):
with gr.Accordion('Prompt enhance', open=False, elem_id='prompt_enhance'):
@@ -867,7 +871,7 @@ class PromptEnhanceScript(scripts_manager.Script):
with gr.Row():
# Set initial state based on whether default model supports vision
default_is_vl = is_vision_model(Options.default)
use_vision = gr.Checkbox(label='Use vision', value=default_is_vl, interactive=default_is_vl, elem_id='prompt_enhance_use_vision')
use_vision = gr.Checkbox(label='Use vision', value=False, interactive=default_is_vl, elem_id='prompt_enhance_use_vision')
gr.HTML('<br>')
with gr.Group():
with gr.Row():