Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-20 09:38:23 +02:00
Code Issues Packages Projects Releases Wiki Activity

prompt-enhance optional auto-apply

Browse Source
This commit is contained in:
Vladimir Mandic
2024-09-10 08:35:12 -04:00
parent 7a6d93134d
commit ef12df3e98
2 changed files with 28 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/prompt_enhance.py
+19 -8
View File
@@ -1,9 +1,10 @@
# repo: https://huggingface.co/gokaygokay/Flux-Prompt-Enhance/discussions/2
# repo: https://huggingface.co/gokaygokay/Flux-Prompt-Enhance
import time
import random
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import gradio as gr
from modules import shared, scripts, devices
from modules import shared, scripts, devices, processing
repo_id = "gokaygokay/Flux-Prompt-Enhance"
@@ -16,6 +17,7 @@ class Script(scripts.Script):
model: AutoModelForSeq2SeqLM = None
prefix: str = "enhance prompt: "
button: gr.Button = None
auto_apply: gr.Checkbox = None
max_length: gr.Slider = None
temperature: gr.Slider = None
repetition_penalty: gr.Slider = None
@@ -35,7 +37,7 @@ class Script(scripts.Script):
shared.log.info(f'Prompt enhance: model="{repo_id}"')
self.model = AutoModelForSeq2SeqLM.from_pretrained('gokaygokay/Flux-Prompt-Enhance', cache_dir=shared.opts.diffusers_dir).to(device=devices.cpu, dtype=devices.dtype)
def enhance(self, prompt, temperature, repetition_penalty, max_length):
def enhance(self, prompt, auto_apply: bool = False, temperature: float = 0.7, repetition_penalty: float = 1.2, max_length: int = 128):
self.load()
t0 = time.time()
input_text = self.prefix + prompt
@@ -57,28 +59,37 @@ class Script(scripts.Script):
prompts = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
prompts = [[p] for p in prompts]
t1 = time.time()
shared.log.info(f'Prompt enhance: temperature={temperature} repetition={repetition_penalty} length={max_length} sequences={num_return_sequences} time={t1-t0:.2f}s')
shared.log.info(f'Prompt enhance: temperature={temperature} repetition={repetition_penalty} length={max_length} sequences={num_return_sequences} apply={auto_apply} time={t1-t0:.2f}s')
return prompts
def select(self, cell: gr.SelectData, _table):
prompt = cell.value
prompt = cell.value if hasattr(cell, 'value') else cell
shared.log.info(f'Prompt enhance: prompt="{prompt}"')
return prompt
def ui(self, _is_img2img):
with gr.Row():
self.button = gr.Button(value='Enhance prompt')
self.auto_apply = gr.Checkbox(label='Auto apply', default=False)
with gr.Row():
self.max_length = gr.Slider(label='Length', minimum=64, maximum=512, step=1, value=128)
self.temperature = gr.Slider(label='Temperature', minimum=0.1, maximum=2.0, step=0.05, value=0.7)
self.repetition_penalty = gr.Slider(label='Penalty', minimum=0.1, maximum=2.0, step=0.05, value=1.2)
with gr.Row():
self.table = gr.DataFrame(self.prompts, label='', show_label=False, interactive=False, wrap=True, datatype="str", col_count=1, row_count=num_return_sequences, headers=['Prompts'])
self.table = gr.DataFrame(self.prompts, label='', show_label=False, interactive=False, wrap=True, datatype="str", col_count=1, max_rows=num_return_sequences, headers=['Prompts'])
if self.prompt is not None:
self.button.click(fn=self.enhance, inputs=[self.prompt, self.temperature, self.repetition_penalty, self.max_length], outputs=[self.table])
self.button.click(fn=self.enhance, inputs=[self.prompt, self.auto_apply, self.temperature, self.repetition_penalty, self.max_length], outputs=[self.table])
self.table.select(fn=self.select, inputs=[self.table], outputs=[self.prompt])
return []
return [self.auto_apply, self.temperature, self.repetition_penalty, self.max_length]
def run(self, p: processing.StableDiffusionProcessing, auto_apply, temperature, repetition_penalty, max_length): # pylint: disable=arguments-differ
if self.auto_apply:
p.prompt = shared.prompt_styles.apply_styles_to_prompt(p.prompt, p.styles)
shared.log.debug(f'Prompt enhance: source="{p.prompt}"')
prompts = self.enhance(p.prompt, auto_apply, temperature, repetition_penalty, max_length)
p.prompt = random.choice(prompts)[0]
shared.log.debug(f'Prompt enhance: prompt="{p.prompt}"')
def after_component(self, component, **kwargs): # searching for actual ui prompt components
if getattr(component, 'elem_id', '') in ['txt2img_prompt', 'img2img_prompt', 'control_prompt']: