From ac61fce526b87593c2279c4ab67cd25e240a76ed Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sat, 29 Mar 2025 11:43:22 -0400 Subject: [PATCH] prompt enhance add censor detection and debugging Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 1 + modules/prompt_parser_diffusers.py | 2 +- modules/sd_modules.py | 73 ++++++++++++++++++++++++++++++ scripts/prompt_enhance.py | 60 +++++++++++++++++------- wiki | 2 +- 5 files changed, 120 insertions(+), 18 deletions(-) create mode 100644 modules/sd_modules.py diff --git a/CHANGELOG.md b/CHANGELOG.md index b029616ea..80851083d 100644 --- a/CHANGELOG.md +++ b/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* diff --git a/modules/prompt_parser_diffusers.py b/modules/prompt_parser_diffusers.py index fc6b2af52..4aaf49a39 100644 --- a/modules/prompt_parser_diffusers.py +++ b/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 diff --git a/modules/sd_modules.py b/modules/sd_modules.py new file mode 100644 index 000000000..9a619a2f6 --- /dev/null +++ b/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 diff --git a/scripts/prompt_enhance.py b/scripts/prompt_enhance.py index 76c53bc6f..284b9613e 100644 --- a/scripts/prompt_enhance.py +++ b/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): diff --git a/wiki b/wiki index 15afa8e1d..9aff8cd69 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 15afa8e1d865450fa123b5dcdb3a2af2317b65af +Subproject commit 9aff8cd69b01570bd7fd2d52b0f9da6baec9b3be