update nudenet

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

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`)  

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
 
 

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)

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):

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

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)}

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)

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('<br>')
             with gr.Group():
                 with gr.Row():