From 6a10fca1d09d793d654fdb3ed28004975faf22b1 Mon Sep 17 00:00:00 2001
From: vladmandic <mandic00@live.com>
Date: Thu, 9 Apr 2026 14:33:35 +0200
Subject: [PATCH] update image guard

Signed-off-by: vladmandic <mandic00@live.com>
---
 scripts/nudenet/imageguard.py | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/scripts/nudenet/imageguard.py b/scripts/nudenet/imageguard.py
index fbd62e43c..c707402b2 100644
--- a/scripts/nudenet/imageguard.py
+++ b/scripts/nudenet/imageguard.py
@@ -96,7 +96,8 @@ 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:
+
+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">,
@@ -113,14 +114,12 @@ processor: AutoProcessor | None = None
 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, 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 or repo != repo_id:
-            # install('flash-attn')
             import transformers
             if 'LlavaGuard' in repo_id:
                 cls = transformers.LlavaOnevisionForConditionalGeneration
@@ -129,7 +128,6 @@ def image_guard(image, policy:str | None=None, model_name:str=''):
             quant_args = model_quant.create_config(module='LLM')
             model = cls.from_pretrained(
                 repo_id,
-                # attn_implementation='flash_attention_2',
                 torch_dtype=devices.dtype,
                 device_map="auto",
                 cache_dir=shared.opts.hfcache_dir,
@@ -151,6 +149,7 @@ def image_guard(image, policy:str | None=None, model_name:str=''):
         ]
         prompt = processor.apply_chat_template(chat_template, add_generation_prompt=True)
         inputs = processor(text=prompt, images=image, return_tensors="pt")
+        input_ids = inputs.input_ids
         model = model.to(device=devices.device)
         inputs = {k: v.to(device=devices.device) for k, v in inputs.items()}
         kwargs = {
@@ -162,12 +161,19 @@ def image_guard(image, policy:str | None=None, model_name:str=''):
             "num_beams": 2,
             "use_cache": True,
         }
-        results = model.generate(**inputs, **kwargs)
+        generated_ids = model.generate(**inputs, **kwargs)
         model = model.to(device=devices.cpu)
-        result = processor.decode(results[0], skip_special_tokens=True)
+
+        trimmed_ids = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(input_ids, generated_ids)]
+        # output_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+
+        result = processor.decode(trimmed_ids[0], skip_special_tokens=True)
         result = result.split('assistant', 1)[-1].strip()
-        data = json.loads(result)
-        log.debug(f'NudeNet LlavaGuard: {data}')
+        try:
+            data = json.loads(result)
+        except Exception:
+            data = { 'text': result }
+        log.debug(f'NudeNet Safety: {data}')
         return data
     except Exception as e:
         log.error(f'NudeNet Safety: {e}')