diff --git a/scripts/prompt_enhance.py b/scripts/prompt_enhance.py new file mode 100644 index 000000000..881594a1b --- /dev/null +++ b/scripts/prompt_enhance.py @@ -0,0 +1,85 @@ +# repo: https://huggingface.co/gokaygokay/Flux-Prompt-Enhance/discussions/2 + +import time +from transformers import AutoTokenizer, AutoModelForSeq2SeqLM +import gradio as gr +from modules import shared, scripts, devices + + +repo_id = "gokaygokay/Flux-Prompt-Enhance" +num_return_sequences = 5 + + +class Script(scripts.Script): + prompts = [['']] + tokenizer: AutoTokenizer = None + model: AutoModelForSeq2SeqLM = None + prefix: str = "enhance prompt: " + button: gr.Button = None + max_length: gr.Slider = None + temperature: gr.Slider = None + repetition_penalty: gr.Slider = None + table: gr.DataFrame = None + prompt: gr.Textbox = None + + def title(self): + return 'Prompt enhance' + + def show(self, is_img2img): + return shared.native + + def load(self): + if self.tokenizer is None: + self.tokenizer = AutoTokenizer.from_pretrained('gokaygokay/Flux-Prompt-Enhance', cache_dir=shared.opts.diffusers_dir) + if self.model is None: + 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): + self.load() + t0 = time.time() + input_text = self.prefix + prompt + input_ids = self.tokenizer(input_text, return_tensors="pt").input_ids.to(devices.device) + self.model = self.model.to(devices.device) + kwargs = { + 'max_length': max_length, + 'num_return_sequences': num_return_sequences, + 'do_sample': True, + 'temperature': temperature, + 'repetition_penalty': repetition_penalty + } + try: + outputs = self.model.generate(input_ids, **kwargs) + except Exception as e: + shared.log.error(f'Prompt enhance: error="{e}"') + return [['']] + self.model = self.model.to(devices.cpu) + 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') + return prompts + + def select(self, cell: gr.SelectData, _table): + prompt = cell.value + 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') + 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']) + + 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.table.select(fn=self.select, inputs=[self.table], outputs=[self.prompt]) + return [] + + def after_component(self, component, **kwargs): # searching for actual ui prompt components + if getattr(component, 'elem_id', '') in ['txt2img_prompt', 'img2img_prompt', 'control_prompt']: + self.prompt = component diff --git a/scripts/prompt_matrix.py b/scripts/prompt_matrix.py index 39c0fe993..5babf20e7 100644 --- a/scripts/prompt_matrix.py +++ b/scripts/prompt_matrix.py @@ -39,7 +39,7 @@ def draw_xy_grid(xs, ys, x_label, y_label, cell): class Script(scripts.Script): def title(self): - return "Prompt Matrix" + return "Prompt matrix" def ui(self, is_img2img): with gr.Row(): diff --git a/scripts/prompts_from_file.py b/scripts/prompts_from_file.py index 3e9745f6b..465a44974 100644 --- a/scripts/prompts_from_file.py +++ b/scripts/prompts_from_file.py @@ -92,7 +92,7 @@ def load_prompt_file(file): class Script(scripts.Script): def title(self): - return "Prompts from File" + return "Prompts from file" def ui(self, is_img2img): with gr.Row():