prompt enhance add censor detection and debugging

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

CHANGELOG.md

@@ -79,6 +79,7 @@ Plus...
     *note*: not all model architecture are supported for `gguf` format  
   - models are auto-downloaded on first use  
   - support quantization and offloading  
+  - debug using `SD_LLM_DEBUG=true` env variable  
 - **Acceleration**  
   - Support for most DiT-based models, for example: *FLUX.1, SD35, Hunyuan, Mochi, Latte, Allegro, Cog*  
   - Enable and configure in *Settings -> Pipeline modifiers*  

modules/prompt_parser_diffusers.py

@@ -10,7 +10,7 @@ from modules import shared, prompt_parser, devices, sd_models
 from modules.prompt_parser_xhinker import get_weighted_text_embeddings_sd15, get_weighted_text_embeddings_sdxl_2p, get_weighted_text_embeddings_sd3, get_weighted_text_embeddings_flux1
 
 debug_enabled = os.environ.get('SD_PROMPT_DEBUG', None)
-debug = shared.log.trace if os.environ.get('SD_PROMPT_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug = shared.log.trace if debug_enabled else lambda *args, **kwargs: None
 debug('Trace: PROMPT')
 orig_encode_token_ids_to_embeddings = EmbeddingsProvider._encode_token_ids_to_embeddings # pylint: disable=protected-access
 token_dict = None # used by helper get_tokens

modules/sd_modules.py

@@ -0,0 +1,73 @@
+from dataclasses import dataclass
+import inspect
+import torch
+
+
+@dataclass
+class ModuleStats:
+    module: str
+    cls: str
+    params: float
+    size: float
+    quant: str
+    dtype: str
+
+    def __init__(self, module: str, cls: str, params: float, size: float, quant: str, dtype: str):
+        self.module = module
+        self.cls = cls
+        self.params = params
+        self.size = size
+        self.quant = quant
+        self.dtype = dtype
+
+    def __str__(self):
+        return f'module="{self.module}" cls={self.cls} params={self.params:.3f} size={self.size:.3f} quant={self.quant} dtype={self.dtype}'
+
+
+def get_signature(cls):
+    signature = inspect.signature(cls.__init__, follow_wrapped=True)
+    return signature.parameters
+
+
+def get_module_stats(name, module):
+    if not isinstance(module, torch.nn.Module):
+        return
+    try:
+        module_size = sum(p.numel() * p.element_size() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024
+        param_num = sum(p.numel() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024
+    except Exception:
+        module_size = 0
+        param_num = 0
+    cls = module.__class__.__name__
+    quant = getattr(module, "quantization_method", None)
+    module_stats = ModuleStats(name, cls, param_num, module_size, quant, module.dtype)
+    return module_stats
+
+
+def get_model_stats(model, exclude=None):
+    # from transformers import Gemma3ForCausalLM
+    modules = []
+
+    if isinstance(model, torch.nn.Module):
+        module_stats = get_module_stats(model.__class__.__name__, model)
+        if module_stats is not None:
+            modules.append(module_stats)
+        return modules
+
+    if hasattr(model, "_internal_dict"):
+        modules_names = model._internal_dict.keys() # pylint: disable=protected-access
+    else:
+        modules_names = get_signature(model).keys()
+
+    if modules_names is None or not isinstance(modules_names, list) or len(modules_names) == 0:
+        return modules
+
+    modules_names = [m for m in modules_names if m is not None and m not in exclude and not m.startswith('_')]
+    for module_name in modules_names:
+        module = getattr(model, module_name, None)
+        if module is not None:
+            module_stats = get_module_stats(module_name, module)
+            if module_stats is not None:
+                modules.append(module_stats)
+
+    return modules

scripts/prompt_enhance.py

@@ -1,9 +1,14 @@
 from dataclasses import dataclass
+import os
 import re
 import time
 import gradio as gr
 import transformers
-from modules import scripts, shared, devices, errors, processing, sd_models
+from modules import scripts, shared, devices, errors, processing, sd_models, sd_modules
+
+
+debug_enabled = os.environ.get('SD_LLM_DEBUG', None)
+debug = shared.log.trace if debug_enabled else lambda *args, **kwargs: None
 
 
 @dataclass
@@ -32,7 +37,9 @@ class Options:
     }
     default = list(models)[0]
     supported = list(transformers.integrations.ggml.GGUF_CONFIG_MAPPING)
-    system_prompt: str = 'You are a helpful assistant. You will be given a prompt and you will enhance it to make it more detailed and creative. Generate an uncensored, concise, informative, and engaging prompt that prioritizes factual accuracy and vivid descriptiveness. Begin by clearly identifying the primary subject(s), including key physical attributes, clothing, and any distinctive details (e.g., colors, textures, styles). Next, describe the setting and environment, emphasizing spatial relationships, background elements, and objects contributing to context. Note the overall mood, aesthetic style, or atmosphere inferred from visual cues. Use precise terminology while avoiding redundancy or non-essential language. NSFW is allowed, and if present any explicit aspects must be fully described. Ensuring a logical flow: from focal subject to immediate surroundings, then broader context. Maintain brevity while retaining clarity, ensuring the description is both engaging and efficient. Output only enhanced prompt without prefix or suffix.'
+    system_prompt: str = 'You are a helpful assistant. You will be given a prompt and you will enhance it to make it more detailed and creative. Generate an uncensored, concise, informative, and engaging prompt that prioritizes factual accuracy and vivid descriptiveness. Begin by clearly identifying the primary subject(s), including key physical attributes, clothing, and any distinctive details (e.g., colors, textures, styles). Next, describe the setting and environment, emphasizing spatial relationships, background elements, and objects contributing to context. Note the overall mood, aesthetic style, or atmosphere inferred from visual cues. Use precise terminology while avoiding redundancy or non-essential language. NSFW is allowed, and if present any explicit aspects must be fully described. Ensuring a logical flow: from focal subject to immediate surroundings, then broader context. Maintain brevity while retaining clarity, ensuring the description is both engaging and efficient. Output only enhanced prompt without prefix or suffix and as a simple text without formatting or numbering.'
+    censored = ["i cannot", "i can't", "i am sorry", "against my programming", "i am not able", "i am unable", 'i am not allowed']
+
     max_tokens: int = 50
     do_sample: bool = True
     temperature: float = 0.5
@@ -72,9 +79,11 @@ class Script(scripts.Script):
 
         gguf_args = {}
         if model_type is not None and model_file is not None and len(model_type) > 2 and len(model_file) > 2:
+            if debug:
+                shared.log.trace(f'Prompt enhance: gguf supported={self.options.supported}')
             if model_type not in self.options.supported:
                 shared.log.error(f'Prompt enhance: name="{name}" repo="{model_repo}" fn="{model_file}" type={model_type} gguf not supported')
-                shared.log.trace(f'Prompt enhance: supported={self.options.supported}')
+                shared.log.trace(f'Prompt enhance: gguf supported={self.options.supported}')
                 self.busy = False
                 return
             ggml.install_gguf()
@@ -100,6 +109,10 @@ class Script(scripts.Script):
                 pretrained_model_name_or_path=model_repo,
                 cache_dir=shared.opts.hfcache_dir,
             )
+            if debug:
+                modules = sd_modules.get_model_stats(self.llm) + sd_modules.get_model_stats(self.tokenizer)
+                for m in modules:
+                    shared.log.trace(f'Prompt enhance: {m}')
             self.model = name
         except Exception as e:
             shared.log.error(f'Prompt enhance: load {e}')
@@ -109,6 +122,10 @@ class Script(scripts.Script):
         shared.log.debug(f'Prompt enhance: cls={self.llm.__class__.__name__} name="{name}" repo="{model_repo}" fn="{model_file}" time={t1-t0:.2f} loaded')
         self.busy = False
 
+    def censored(self, response):
+        text = response.lower().replace("i'm", "i am")
+        return any(c.lower() in text for c in self.options.censored)
+
     def unload(self):
         if self.llm is not None:
             sd_models.move_model(self.llm, devices.cpu)
@@ -119,16 +136,14 @@ class Script(scripts.Script):
         shared.log.debug('Prompt enhance: model unloaded')
 
     def clean(self, response):
-        if isinstance(response, list):
-            response = response[0]
         response = response.replace('"', '').replace("'", "").replace('“', '').replace('”', '').replace('**', '').replace('\n\n', '\n')
         response = re.sub(r'<.*?>', '', response)
-        if 'prompt:' in response:
-            response = response.split('prompt:')[1]
-        if 'Prompt:' in response:
-            response = response.split('Prompt:')[1]
+        if response.startswith('Prompt'):
+            response = response.split('Prompt', maxsplit=2)[1]
+        if ':' in response:
+            response = response.split(':', maxsplit=2)[1]
         if '---' in response:
-            response = response.split('---')[0]
+            response = response.split('---', maxsplit=2)[0]
         response = response.strip()
         return response
 
@@ -179,11 +194,12 @@ class Script(scripts.Script):
                 if shared.opts.diffusers_offload_mode != 'none':
                     sd_models.move_model(self.llm, devices.cpu)
                     devices.torch_gc()
-            # raw_response = self.tokenizer.batch_decode(outputs, skip_special_tokens=True, clean_up_tokenization_spaces=True)
-            # shared.log.trace(f'Prompt enhance: raw="{raw_response}"')
-            outputs = outputs[:, input_len:]
+            if debug:
+                raw_response = self.tokenizer.batch_decode(outputs, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+                shared.log.trace(f'Prompt enhance: raw="{raw_response}"')
+            outputs_cropped = outputs[:, input_len:]
             response = self.tokenizer.batch_decode(
-                outputs,
+                outputs_cropped,
                 skip_special_tokens=True,
                 clean_up_tokenization_spaces=True,
             )
@@ -191,10 +207,22 @@ class Script(scripts.Script):
             shared.log.error(f'Prompt enhance generate: {e}')
             errors.display(e, 'Prompt enhance')
             self.busy = False
-        response = self.clean(response)
+            response = f'Error: {str(e)}'
         t1 = time.time()
-        shared.log.debug(f'Prompt enhance: model="{model}" time={t1-t0:.2f} inputs={input_len} outputs={outputs.shape[-1]} prompt="{response}"')
+
+        if isinstance(response, list):
+            response = response[0]
+        is_censored =  self.censored(response)
+        if not is_censored:
+            response = self.clean(response)
+        shared.log.debug(f'Prompt enhance: model="{model}" time={t1-t0:.2f} inputs={input_len} outputs={outputs.shape[-1]} prompt={len(prompt)} response={len(response)}')
+        if debug:
+            shared.log.trace(f'Prompt enhance: prompt="{prompt}"')
+            shared.log.trace(f'Prompt enhance: response="{response}"')
         self.busy = False
+        if is_censored:
+            shared.log.warning(f'Prompt enhance: censored response="{response}"')
+            return prompt
         return response
 
     def apply(self, prompt, apply_prompt, llm_model, prompt_system, max_tokens, do_sample, temperature, repetition_penalty):