diff --git a/CHANGELOG.md b/CHANGELOG.md index 6a050673e..971a8c397 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ now covers both ROCm on Windows and Linux see *main interface -> scripts -> rocm advanced config* - **Features** + - **nudenet** add `LlavaGuard` and `QwenGuard` as image safety evaulation models - installer auto-restart on upgrade - enhanced filename pattern processing allows for any *processing* property name (as defined in `modules/processing_class.py` and saved to `ui-config.json`) diff --git a/modules/api/nudenet.py b/modules/api/nudenet.py index 9c4bfd511..3f0ef5e65 100644 --- a/modules/api/nudenet.py +++ b/modules/api/nudenet.py @@ -42,11 +42,12 @@ def prompt_check( def image_guard( image: str = Body("", title='input image'), policy: str = Body("", title='optional policy definition'), + model: str = Body("", title='optional policy model name'), ): """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) + res = imageguard.image_guard(image=image, policy=policy, model_name=model) return res diff --git a/modules/errors.py b/modules/errors.py index 24dabc808..32bf3f9dd 100644 --- a/modules/errors.py +++ b/modules/errors.py @@ -23,7 +23,7 @@ def display(e: Exception, task: str, suppress=None): suppress = [] if isinstance(e, ErrorLimiterAbort): return - log.critical(f"{task or 'error'}: {type(e).__name__}") + log.error(f"{task or 'error'}: {type(e).__name__}") """ trace = traceback.format_exc() log.error(trace) diff --git a/modules/scripts_manager.py b/modules/scripts_manager.py index 3e5988ae8..45ef6ef0f 100644 --- a/modules/scripts_manager.py +++ b/modules/scripts_manager.py @@ -29,6 +29,9 @@ class PostprocessImageArgs: def __init__(self, image): self.image = image + def __str__(self): + return f'PostprocessImageArgs(image={self.image})' + class PostprocessBatchListArgs: def __init__(self, images): diff --git a/modules/scripts_postprocessing.py b/modules/scripts_postprocessing.py index 70410826a..cd70a601c 100644 --- a/modules/scripts_postprocessing.py +++ b/modules/scripts_postprocessing.py @@ -11,6 +11,9 @@ class PostprocessedImage: self.image = image self.info = info + def __str__(self): + return f'PostprocessedImage(image={self.image} info={self.info})' + class ScriptPostprocessing: filename = None diff --git a/scripts/nudenet/imageguard.py b/scripts/nudenet/imageguard.py index aa3a881be..fbd62e43c 100644 --- a/scripts/nudenet/imageguard.py +++ b/scripts/nudenet/imageguard.py @@ -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)} diff --git a/scripts/nudenet_ext.py b/scripts/nudenet_ext.py index 41a60b162..915776e30 100644 --- a/scripts/nudenet_ext.py +++ b/scripts/nudenet_ext.py @@ -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) diff --git a/scripts/prompt_enhance.py b/scripts/prompt_enhance.py index e919081f9..f74b49c91 100644 --- a/scripts/prompt_enhance.py +++ b/scripts/prompt_enhance.py @@ -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('
') with gr.Group(): with gr.Row():