mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 09:14:35 +02:00
prompt enhance custom model loader
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
+9
-7
@@ -67,6 +67,15 @@ Plus...
|
||||
download text encoders into folder set in settings -> system paths -> text encoders (default is *models/Text-encoder*)
|
||||
load using *settings -> text encoder*
|
||||
*tip*: add *sd_text_encoder* to your *settings -> user interface -> quicksettings* list to have it appear at the top of the ui
|
||||
- **Prompt Enhance**
|
||||
- new built-in extension available in text/image/control tabs
|
||||
- can be used to manually or automatically enhance prompts using LLM
|
||||
- built-in presets for **Gemma-3, Qwen-2.5, Phi-4, Llama-3.2, SmolLM2, Dolphin-3**
|
||||
- support for custom models
|
||||
load any models hosted on huggingface
|
||||
load either model in huggingface format or `gguf` format
|
||||
- models are auto-downloaded on first use
|
||||
- support quantization and offloading
|
||||
- **Acceleration**
|
||||
- Support for most DiT-based models, for example: *FLUX.1, SD35, Hunyuan, Mochi, Latte, Allegro, Cog*
|
||||
- Enable and configure in *Settings -> Pipeline modifiers*
|
||||
@@ -84,13 +93,6 @@ Plus...
|
||||
- [ByteDance/Sa2VA](https://huggingface.co/ByteDance/Sa2VA-1B) 1B, 4B
|
||||
simply select from list of available models in caption tab
|
||||
- add option to set system prompt for vlm models that support it: *Gemma, Smol, Qwen*
|
||||
- **Prompt Enhance**
|
||||
- new built-in extension available in text/image/control tabs
|
||||
- can be used to manually or automatically enhance prompts using LLM
|
||||
- supports **Gemma-3, Qwen-2.5, Phi-4, Llama-3.2, SmolLM2**
|
||||
models are auto-downloaded on first use
|
||||
also supports custom models that are compatible with `transformers/AutoModelForCausalLM`
|
||||
- support quantization and offloading
|
||||
- [NudeNet](https://github.com/vladmandic/sd-extension-nudenet/) extension updates
|
||||
- add detection of prompt language and alphabet and filter based on those values
|
||||
- add image policy checks using `LlavaGuard` VLM to detect policy violations (and reasons)
|
||||
|
||||
@@ -116,7 +116,7 @@ button.custom-button { border-radius: var(--button-large-radius); padding: var(-
|
||||
#txt2img_seed, #img2img_seed, #control_seed, #video_seed { min-width: 90px !important }
|
||||
#video_generate_box>button { max-width: unset; }
|
||||
#interrogate_output_prompt>textarea { resize: vertical; }
|
||||
#prompt_enhance_apply, #prompt_enhance_model { max-width: unset; }
|
||||
#prompt_enhance_apply, #prompt_enhance_model, #prompt_enhance_custom_load { max-width: unset; min-width: 100% !important; }
|
||||
#prompt_enhance_system textarea { color: var(--body-text-color-subdued) !important }
|
||||
|
||||
.interrogate { position: absolute; right: 2.8em; top: 0.2em; max-width: fit-content; background: none !important; z-index: 50; font-size: 1.5em !important; }
|
||||
|
||||
+123
-45
@@ -3,25 +3,34 @@ import re
|
||||
import time
|
||||
import gradio as gr
|
||||
import transformers
|
||||
from modules import scripts, shared, devices, processing, sd_models
|
||||
from modules import scripts, shared, devices, errors, processing, sd_models
|
||||
|
||||
|
||||
@dataclass
|
||||
class Options:
|
||||
models = [
|
||||
'Qwen/Qwen2.5-0.5B-Instruct',
|
||||
'Qwen/Qwen2.5-1.5B-Instruct',
|
||||
'Qwen/Qwen2.5-3B-Instruct',
|
||||
'google/gemma-3-1b-it',
|
||||
'google/gemma-3-4b-it',
|
||||
'microsoft/Phi-4-mini-instruct',
|
||||
'HuggingFaceTB/SmolLM2-135M-Instruct',
|
||||
'HuggingFaceTB/SmolLM2-360M-Instruct',
|
||||
'HuggingFaceTB/SmolLM2-1.7B-Instruct',
|
||||
'meta-llama/Llama-3.2-1B-Instruct',
|
||||
'meta-llama/Llama-3.2-3B-Instruct',
|
||||
]
|
||||
default = models[3]
|
||||
models = {
|
||||
'google/gemma-3-1b-it': {},
|
||||
'google/gemma-3-4b-it': {},
|
||||
'Qwen/Qwen2.5-0.5B-Instruct': {},
|
||||
'Qwen/Qwen2.5-1.5B-Instruct': {},
|
||||
'Qwen/Qwen2.5-3B-Instruct': {},
|
||||
'microsoft/Phi-4-mini-instruct': {},
|
||||
'HuggingFaceTB/SmolLM2-135M-Instruct': {},
|
||||
'HuggingFaceTB/SmolLM2-360M-Instruct': {},
|
||||
'HuggingFaceTB/SmolLM2-1.7B-Instruct': {},
|
||||
'meta-llama/Llama-3.2-1B-Instruct': {},
|
||||
'meta-llama/Llama-3.2-3B-Instruct': {},
|
||||
'cognitivecomputations/Dolphin3.0-Llama3.2-1B': {},
|
||||
'cognitivecomputations/Dolphin3.0-Llama3.2-3B': {},
|
||||
'mradermacher/Llama-3.2-1B-Instruct-Uncensored-i1-GGUF': {
|
||||
'repo': 'meta-llama/Llama-3.2-1B-Instruct', # original repo so we can load missing components
|
||||
'type': 'llama', # required so gguf loader knows what to do
|
||||
'gguf': 'mradermacher/Llama-3.2-1B-Instruct-Uncensored-i1-GGUF', # gguf repo
|
||||
'file': 'Llama-3.2-1B-Instruct-Uncensored.i1-Q4_0.gguf', # gguf file inside repo
|
||||
},
|
||||
}
|
||||
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.'
|
||||
max_tokens: int = 50
|
||||
do_sample: bool = True
|
||||
@@ -34,6 +43,7 @@ class Script(scripts.Script):
|
||||
model: str = None
|
||||
llm: transformers.AutoModelForCausalLM = None
|
||||
tokenizer: transformers.AutoProcessor = None
|
||||
busy: bool = False
|
||||
options = Options()
|
||||
|
||||
def title(self):
|
||||
@@ -42,31 +52,61 @@ class Script(scripts.Script):
|
||||
def show(self, _is_img2img):
|
||||
return scripts.AlwaysVisible
|
||||
|
||||
def load(self, model:str=None):
|
||||
model = model or self.options.default
|
||||
if self.model is None or self.model != model:
|
||||
t0 = time.time()
|
||||
from modules import modelloader, model_quant
|
||||
modelloader.hf_login()
|
||||
quant_args = model_quant.create_config(module='LLM')
|
||||
def load(self, name:str=None, model_repo:str=None, model_gguf:str=None, model_type:str=None, model_file:str=None):
|
||||
name = name or self.options.default
|
||||
if self.busy:
|
||||
shared.log.debug('Prompt enhance: busy')
|
||||
return
|
||||
self.busy = True
|
||||
if self.model is not None and self.model == name:
|
||||
return
|
||||
|
||||
t0 = time.time()
|
||||
from modules import modelloader, model_quant, ggml
|
||||
modelloader.hf_login()
|
||||
model_repo = model_repo or self.options.models.get(name, {}).get('repo', None) or name
|
||||
model_gguf = model_gguf or self.options.models.get(name, {}).get('gguf', None) or model_repo
|
||||
model_type = model_type or self.options.models.get(name, {}).get('type', None)
|
||||
model_file = model_file or self.options.models.get(name, {}).get('file', None)
|
||||
|
||||
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 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}')
|
||||
self.busy = False
|
||||
return
|
||||
ggml.install_gguf()
|
||||
gguf_args['model_type'] = model_type
|
||||
gguf_args['gguf_file'] = model_file
|
||||
|
||||
quant_args = model_quant.create_config(module='LLM') if not gguf_args else {}
|
||||
|
||||
try:
|
||||
self.model = None
|
||||
self.llm = None
|
||||
self.llm = transformers.AutoModelForCausalLM.from_pretrained(
|
||||
model,
|
||||
pretrained_model_name_or_path=model_repo if not gguf_args else model_gguf,
|
||||
trust_remote_code=True,
|
||||
torch_dtype=devices.dtype,
|
||||
cache_dir=shared.opts.hfcache_dir,
|
||||
_attn_implementation="eager",
|
||||
**gguf_args,
|
||||
**quant_args,
|
||||
)
|
||||
self.llm.eval()
|
||||
self.tokenizer = transformers.AutoTokenizer.from_pretrained(
|
||||
model,
|
||||
pretrained_model_name_or_path=model_repo,
|
||||
cache_dir=shared.opts.hfcache_dir,
|
||||
)
|
||||
self.model = model
|
||||
devices.torch_gc()
|
||||
t1 = time.time()
|
||||
shared.log.debug(f'Prompt enhance: model="{model}" cls={self.llm.__class__.__name__} time={t1-t0:.2f} loaded')
|
||||
self.model = name
|
||||
except Exception as e:
|
||||
shared.log.error(f'Prompt enhance: load {e}')
|
||||
errors.display(e, 'Prompt enhance')
|
||||
devices.torch_gc()
|
||||
t1 = time.time()
|
||||
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 unload(self):
|
||||
if self.llm is not None:
|
||||
@@ -99,6 +139,8 @@ class Script(scripts.Script):
|
||||
penalty = penalty or self.options.repetition_penalty
|
||||
temperature = temperature or self.options.temperature
|
||||
sample = sample if sample is not None else self.options.do_sample
|
||||
while self.busy:
|
||||
time.sleep(0.1)
|
||||
self.load(model)
|
||||
if self.llm is None:
|
||||
shared.log.error('Prompt enhance: model not loaded')
|
||||
@@ -108,6 +150,7 @@ class Script(scripts.Script):
|
||||
{ "role": "user", "content": prompt },
|
||||
]
|
||||
t0 = time.time()
|
||||
self.busy = True
|
||||
try:
|
||||
inputs = self.tokenizer.apply_chat_template(
|
||||
chat_template,
|
||||
@@ -119,6 +162,8 @@ class Script(scripts.Script):
|
||||
input_len = inputs['input_ids'].shape[1]
|
||||
except Exception as e:
|
||||
shared.log.error(f'Prompt enhance tokenize: {e}')
|
||||
errors.display(e, 'Prompt enhance')
|
||||
self.busy = False
|
||||
return prompt
|
||||
try:
|
||||
with devices.inference_context():
|
||||
@@ -136,12 +181,19 @@ class Script(scripts.Script):
|
||||
# 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:]
|
||||
response = self.tokenizer.batch_decode(outputs, skip_special_tokens=True, clean_up_tokenization_spaces=True)
|
||||
response = self.tokenizer.batch_decode(
|
||||
outputs,
|
||||
skip_special_tokens=True,
|
||||
clean_up_tokenization_spaces=True,
|
||||
)
|
||||
except Exception as e:
|
||||
shared.log.error(f'Prompt enhance generate: {e}')
|
||||
errors.display(e, 'Prompt enhance')
|
||||
self.busy = False
|
||||
response = self.clean(response)
|
||||
t1 = time.time()
|
||||
shared.log.debug(f'Prompt enhance: model="{model}" time={t1-t0:.2f} inputs={input_len} outputs={outputs.shape[-1]} prompt="{response}"')
|
||||
self.busy = False
|
||||
return response
|
||||
|
||||
def apply(self, prompt, apply_prompt, llm_model, prompt_system, max_tokens, do_sample, temperature, repetition_penalty):
|
||||
@@ -158,6 +210,13 @@ class Script(scripts.Script):
|
||||
return [response, response]
|
||||
return [response, gr.update()]
|
||||
|
||||
def get_custom(self, name):
|
||||
model_repo = self.options.models.get(name, {}).get('repo', None) or name
|
||||
model_gguf = self.options.models.get(name, {}).get('gguf', None)
|
||||
model_type = self.options.models.get(name, {}).get('type', None)
|
||||
model_file = self.options.models.get(name, {}).get('file', None)
|
||||
return [model_repo, model_gguf, model_type, model_file]
|
||||
|
||||
def ui(self, _is_img2img):
|
||||
with gr.Accordion('Prompt enhance', open=False, elem_id='prompt_enhance'):
|
||||
with gr.Row():
|
||||
@@ -165,29 +224,48 @@ class Script(scripts.Script):
|
||||
with gr.Row():
|
||||
apply_prompt = gr.Checkbox(label='Apply to prompt', value=False)
|
||||
apply_auto = gr.Checkbox(label='Auto enhance', value=False)
|
||||
gr.HTML('<br>')
|
||||
with gr.Group():
|
||||
with gr.Row():
|
||||
llm_model = gr.Dropdown(label='LLM model', choices=self.options.models, value=self.options.default, interactive=True, allow_custom_value=True, elem_id='prompt_enhance_model')
|
||||
llm_model = gr.Dropdown(label='LLM model', choices=list(self.options.models), value=self.options.default, interactive=True, allow_custom_value=True, elem_id='prompt_enhance_model')
|
||||
with gr.Row():
|
||||
load_btn = gr.Button(value='Load model', elem_id='prompt_enhance_load', variant='secondary')
|
||||
load_btn.click(fn=self.load, inputs=[llm_model], outputs=[])
|
||||
unload_btn = gr.Button(value='Unload model', elem_id='prompt_enhance_unload', variant='secondary')
|
||||
unload_btn.click(fn=self.unload, inputs=[], outputs=[])
|
||||
with gr.Row():
|
||||
prompt_system = gr.Textbox(label='System prompt', value=self.options.system_prompt, interactive=True, lines=4, elem_id='prompt_enhance_system')
|
||||
with gr.Row():
|
||||
max_tokens = gr.Slider(label='Max tokens', value=self.options.max_tokens, minimum=10, maximum=1024, step=1, interactive=True)
|
||||
do_sample = gr.Checkbox(label='Do sample', value=self.options.do_sample, interactive=True)
|
||||
with gr.Row():
|
||||
temperature = gr.Slider(label='Temperature', value=self.options.temperature, minimum=0.0, maximum=1.0, step=0.01, interactive=True)
|
||||
repetition_penalty = gr.Slider(label='Repetition penalty', value=self.options.repetition_penalty, minimum=0.0, maximum=2.0, step=0.01, interactive=True)
|
||||
with gr.Row():
|
||||
prompt_output = gr.Textbox(label='Output', value='', interactive=True, lines=4)
|
||||
with gr.Row():
|
||||
clear_btn = gr.Button(value='Clear', elem_id='prompt_enhance_clear', variant='secondary')
|
||||
clear_btn.click(fn=lambda: '', inputs=[], outputs=[prompt_output])
|
||||
copy_btn = gr.Button(value='Set prompt', elem_id='prompt_enhance_copy', variant='secondary')
|
||||
copy_btn.click(fn=lambda x: x, inputs=[prompt_output], outputs=[self.prompt])
|
||||
with gr.Accordion('Custom model', open=False, elem_id='prompt_enhance_custom'):
|
||||
with gr.Row():
|
||||
model_repo = gr.Textbox(label='Model repo', value=None, interactive=True, elem_id='prompt_enhance_model_repo', placeholder='Original model repo on huggingface')
|
||||
with gr.Row():
|
||||
model_gguf = gr.Textbox(label='Model gguf', value=None, interactive=True, elem_id='prompt_enhance_model_gguf', placeholder='Optional GGUF model repo on huggingface')
|
||||
with gr.Row():
|
||||
model_type = gr.Textbox(label='Model type', value=None, interactive=True, elem_id='prompt_enhance_model_type', placeholder='Optional GGUF model type')
|
||||
with gr.Row():
|
||||
model_file = gr.Textbox(label='Model file', value=None, interactive=True, elem_id='prompt_enhance_model_file', placeholder='Optional GGUF model file inside GGUF model repo')
|
||||
with gr.Row():
|
||||
custom_btn = gr.Button(value='Load custom model', elem_id='prompt_enhance_custom_load', variant='secondary')
|
||||
custom_btn.click(fn=self.load, inputs=[model_file, model_repo, model_gguf, model_type, model_file], outputs=[])
|
||||
llm_model.change(fn=self.get_custom, inputs=[llm_model], outputs=[model_repo, model_gguf, model_type, model_file])
|
||||
gr.HTML('<br>')
|
||||
with gr.Accordion('Options', open=False, elem_id='prompt_enhance_options'):
|
||||
with gr.Row():
|
||||
max_tokens = gr.Slider(label='Max tokens', value=self.options.max_tokens, minimum=10, maximum=1024, step=1, interactive=True)
|
||||
do_sample = gr.Checkbox(label='Do sample', value=self.options.do_sample, interactive=True)
|
||||
with gr.Row():
|
||||
temperature = gr.Slider(label='Temperature', value=self.options.temperature, minimum=0.0, maximum=1.0, step=0.01, interactive=True)
|
||||
repetition_penalty = gr.Slider(label='Repetition penalty', value=self.options.repetition_penalty, minimum=0.0, maximum=2.0, step=0.01, interactive=True)
|
||||
gr.HTML('<br>')
|
||||
with gr.Accordion('Input', open=False, elem_id='prompt_enhance_system_prompt'):
|
||||
with gr.Row():
|
||||
prompt_system = gr.Textbox(label='System prompt', value=self.options.system_prompt, interactive=True, lines=4, elem_id='prompt_enhance_system')
|
||||
with gr.Accordion('Output', open=True, elem_id='prompt_enhance_system_prompt'):
|
||||
with gr.Row():
|
||||
prompt_output = gr.Textbox(label='Enhanced prompt', value='', interactive=True, lines=4)
|
||||
with gr.Row():
|
||||
clear_btn = gr.Button(value='Clear', elem_id='prompt_enhance_clear', variant='secondary')
|
||||
clear_btn.click(fn=lambda: '', inputs=[], outputs=[prompt_output])
|
||||
copy_btn = gr.Button(value='Set prompt', elem_id='prompt_enhance_copy', variant='secondary')
|
||||
copy_btn.click(fn=lambda x: x, inputs=[prompt_output], outputs=[self.prompt])
|
||||
apply_btn.click(fn=self.apply, inputs=[self.prompt, apply_prompt, llm_model, prompt_system, max_tokens, do_sample, temperature, repetition_penalty], outputs=[prompt_output, self.prompt])
|
||||
return [apply_auto, llm_model, prompt_system, max_tokens, do_sample, temperature, repetition_penalty]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user