From 6489e4c37d078df61aadd7de50aa7e03ab73f93b Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Thu, 8 May 2025 15:31:07 -0400 Subject: [PATCH] prompt-enhance api support and img2img support Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 13 ++- cli/api-enhance.py | 75 +++++++++++++++ modules/api/api.py | 1 + modules/api/helpers.py | 2 + modules/api/models.py | 13 +++ modules/api/process.py | 47 ++++++++- modules/interrogate/vqa.py | 3 +- ...flux_prompt_enhance.py => flux_enhance.py} | 0 scripts/prompt_enhance.py | 95 ++++++++++++++++--- 9 files changed, 233 insertions(+), 16 deletions(-) create mode 100755 cli/api-enhance.py rename scripts/{flux_prompt_enhance.py => flux_enhance.py} (100%) diff --git a/CHANGELOG.md b/CHANGELOG.md index d30e9bac0..c582ba8d4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,9 +3,18 @@ ## Update for 2025-05-08 - **Features** - - FramePack: full API support - NNCF: Faster quantization - - API: add `/sdapi/v1/checkpoint` endpoint to get info on currently loaded model/checkpoint + - Prompt Enhancer: support for *img2img* workflows + where prompt enhancer will first analyze input image and then incorporate user prompt to create enhanced prompt +- **API** + - add `/sdapi/v1/framepack` endpoint with full support for FramePack including all optional settings + see example: `sd-extension-framepack/create-video.py` + - add `/sdapi/v1/checkpoint` endpoint to get info on currently loaded model/checkpoint + see example: `cli/api-checkpoint.py` + - add `/sdapi/v1/prompt-enhance` endpoint to enhance prompt using LLM + see example: `cli/api-enhance.py` + supports text, image and video prompts with or without input image + *note*: if input image is provided, model should be left at default `gemma-3-4b-it` as most other LLMs do not support hybrid workflows - **Fixes** - ROCm: disable cuDNN, fixes slow MIOpen tuning with `torch==2.7` - Extensions: use in-process installer for extensions-builtin, improves startup performance diff --git a/cli/api-enhance.py b/cli/api-enhance.py new file mode 100755 index 000000000..fa30d9cb1 --- /dev/null +++ b/cli/api-enhance.py @@ -0,0 +1,75 @@ +#!/usr/bin/env python +import os +import io +import base64 +import logging +import argparse +import requests +import urllib3 +from PIL import Image + + +sd_url = os.environ.get('SDAPI_URL', "http://127.0.0.1:7860") +sd_username = os.environ.get('SDAPI_USR', None) +sd_password = os.environ.get('SDAPI_PWD', None) + +logging.basicConfig(level = logging.INFO, format = '%(asctime)s %(levelname)s: %(message)s') +log = logging.getLogger(__name__) +urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) + + +def auth(): + if sd_username is not None and sd_password is not None: + return requests.auth.HTTPBasicAuth(sd_username, sd_password) + return None + + +def post(endpoint: str, dct: dict = None): + req = requests.post(f'{sd_url}{endpoint}', json = dct, timeout=300, verify=False, auth=auth()) + if req.status_code != 200: + return { 'error': req.status_code, 'reason': req.reason, 'url': req.url } + else: + return req.json() + + +def encode(f): + if f is not None and os.path.exists(f): + image = Image.open(f) + if image.mode == 'RGBA': + image = image.convert('RGB') + log.info(f'encoding image: {image}') + with io.BytesIO() as stream: + image.save(stream, 'JPEG') + image.close() + values = stream.getvalue() + encoded = base64.b64encode(values).decode() + return encoded + else: + return None + + +def enhance(args): # pylint: disable=redefined-outer-name + options = { + 'prompt': str(args.prompt), + 'seed': int(args.seed), + 'type': str(args.type), + } + if args.model: + options['model'] = str(args.model) + if args.image: + options['image'] = encode(args.image) + response = post('/sdapi/v1/prompt-enhance', options) + return response + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description = 'api-enhance') + parser.add_argument('--prompt', type=str, default='', required=False, help='prompt') + parser.add_argument('--seed', type=int, default=-1, required=False, help='seed') + parser.add_argument('--type', type=str, default='text', choices=['text', 'image', 'video'], required=False, help='enhance type') + parser.add_argument('--model', type=str, default=None, required=False, help='model name') + parser.add_argument('--image', type=str, default=None, required=False, help='optional input image') + args = parser.parse_args() + log.info(f'api-upscale: {args}') + result = enhance(args) + log.info(result) diff --git a/modules/api/api.py b/modules/api/api.py index b653250f5..72a2090a0 100644 --- a/modules/api/api.py +++ b/modules/api/api.py @@ -65,6 +65,7 @@ class Api: self.add_api_route("/sdapi/v1/preprocess", self.process.post_preprocess, methods=["POST"]) self.add_api_route("/sdapi/v1/mask", self.process.post_mask, methods=["POST"]) self.add_api_route("/sdapi/v1/detect", self.process.post_detect, methods=["POST"]) + self.add_api_route("/sdapi/v1/prompt-enhance", self.process.post_prompt_enhance, methods=["POST"], response_model=models.ResPromptEnhance) # api dealing with optional scripts self.add_api_route("/sdapi/v1/scripts", script.get_scripts_list, methods=["GET"], response_model=models.ResScripts) diff --git a/modules/api/helpers.py b/modules/api/helpers.py index 21e89c5c1..15b7e29aa 100644 --- a/modules/api/helpers.py +++ b/modules/api/helpers.py @@ -15,6 +15,8 @@ def validate_sampler_name(name): def decode_base64_to_image(encoding, quiet=False): + if encoding is None: + return None if encoding.startswith("data:image/"): encoding = encoding.split(";")[1].split(",")[1] try: diff --git a/modules/api/models.py b/modules/api/models.py index fee8d781d..f60ba3dd7 100644 --- a/modules/api/models.py +++ b/modules/api/models.py @@ -266,6 +266,19 @@ class ReqProcess(BaseModel): class ResProcess(BaseModel): html_info: str = Field(title="HTML info", description="A series of HTML tags containing the process info.") + +class ReqPromptEnhance(BaseModel): + prompt: str = Field(title="Prompt", description="Prompt to enhance") + type: str = Field(title="Type", default='text', description="Type of enhancement to perform") + model: Optional[str] = Field(title="Model", default=None, description="Model to use for enhancement") + system_prompt: Optional[str] = Field(title="System prompt", default=None, description="Model system prompt") + image: Optional[str] = Field(title="Image", default=None, description="Image to work on, must be a Base64 string containing the image's data.") + seed: int = Field(title="Seed", default=-1, description="Seed used to generate the prompt") + +class ResPromptEnhance(BaseModel): + prompt: str = Field(title="Prompt", description="Enhanced prompt") + seed: int = Field(title="Seed", description="Seed used to generate the prompt") + class ReqProcessImage(ReqProcess): image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.") diff --git a/modules/api/process.py b/modules/api/process.py index 80b19c52e..1d98df03e 100644 --- a/modules/api/process.py +++ b/modules/api/process.py @@ -2,8 +2,10 @@ from typing import Optional, List from threading import Lock from pydantic import BaseModel, Field # pylint: disable=no-name-in-module from fastapi.responses import JSONResponse +from fastapi.exceptions import HTTPException from modules.api.helpers import decode_base64_to_image, encode_pil_to_base64 from modules import errors, shared +from modules.api import models processor = None # cached instance of processor @@ -65,8 +67,8 @@ class APIProcess(): def post_preprocess(self, req: ReqPreprocess): global processor # pylint: disable=global-statement from modules.control import processors - models = list(processors.config) - if req.model not in models: + processors_list = list(processors.config) + if req.model not in processors_list: return JSONResponse(status_code=400, content={"error": f"Processor model not found: id={req.model}"}) image = decode_base64_to_image(req.image) if processor is None or processor.processor_id != req.model: @@ -129,3 +131,44 @@ class APIProcess(): boxes.append(item.box) shared.state.end(api=False) return ResFace(classes=classes, labels=labels, scores=scores, boxes=boxes, images=images) + + def post_prompt_enhance(self, req: models.ReqPromptEnhance): + from modules import processing_helpers + seed = req.seed or -1 + seed = processing_helpers.get_fixed_seed(seed) + prompt = '' + if req.type == 'text': + from modules.scripts import scripts_txt2img + model = 'google/gemma-3-1b-it' if req.model is None or len(req.model) < 4 else req.model + instance = [s for s in scripts_txt2img.scripts if 'prompt_enhance.py' in s.filename][0] + prompt = instance.enhance( + model=model, + prompt=req.prompt, + system=req.system_prompt, + seed=seed, + ) + elif req.type == 'image': + from modules.scripts import scripts_txt2img + model = 'google/gemma-3-4b-it' if req.model is None or len(req.model) < 4 else req.model + instance = [s for s in scripts_txt2img.scripts if 'prompt_enhance.py' in s.filename][0] + prompt = instance.enhance( + model=model, + prompt=req.prompt, + system=req.system_prompt, + image=decode_base64_to_image(req.image), + seed=seed, + ) + elif req.type == 'video': + from modules.ui_video_vlm import enhance_prompt + model = 'Google Gemma 3 4B' if req.model is None or len(req.model) < 4 else req.model + prompt = enhance_prompt( + enable=True, + image=decode_base64_to_image(req.image), + prompt=req.prompt, + model=model, + system_prompt=req.system_prompt, + ) + else: + raise HTTPException(status_code=400, detail="prompt enhancement: invalid type") + res = models.ResPromptEnhance(prompt=prompt, seed=seed) + return res diff --git a/modules/interrogate/vqa.py b/modules/interrogate/vqa.py index 461ec6c73..c4dcafb10 100644 --- a/modules/interrogate/vqa.py +++ b/modules/interrogate/vqa.py @@ -527,10 +527,11 @@ def sa2(question: str, image: Image.Image, repo: str = None): return response -def interrogate(question, system_prompt, prompt, image, model_name, quiet:bool=False): +def interrogate(question:str='', system_prompt:str=None, prompt:str=None, image:Image.Image=None, model_name:str=None, quiet:bool=False): if not quiet: shared.state.begin('Interrogate') t0 = time.time() + model_name = model_name or shared.opts.interrogate_vlm_model if isinstance(image, list): image = image[0] if len(image) > 0 else None if isinstance(image, dict) and 'name' in image: diff --git a/scripts/flux_prompt_enhance.py b/scripts/flux_enhance.py similarity index 100% rename from scripts/flux_prompt_enhance.py rename to scripts/flux_enhance.py diff --git a/scripts/prompt_enhance.py b/scripts/prompt_enhance.py index adc80d094..87076559e 100644 --- a/scripts/prompt_enhance.py +++ b/scripts/prompt_enhance.py @@ -1,9 +1,13 @@ from dataclasses import dataclass +import io import os import re import time -import gradio as gr +import base64 +import torch import transformers +import gradio as gr +from PIL import Image from modules import scripts, shared, devices, errors, processing, sd_models, sd_modules @@ -11,8 +15,23 @@ debug_enabled = os.environ.get('SD_LLM_DEBUG', None) is not None debug_log = shared.log.trace if debug_enabled else lambda *args, **kwargs: None +def b64(image): + if image is None: + return '' + if isinstance(image, gr.Image): + return None + with io.BytesIO() as stream: + image.convert('RGB').save(stream, 'JPEG') + values = stream.getvalue() + encoded = base64.b64encode(values).decode() + return encoded + + @dataclass class Options: + img2img = [ + 'google/gemma-3-4b-it', + ] models = { 'google/gemma-3-1b-it': {}, 'google/gemma-3-4b-it': {}, @@ -46,9 +65,11 @@ class Options: 'file': 'Llama-3.2-1B-Instruct-Uncensored.i1-Q4_0.gguf', # gguf file inside repo }, } - default = list(models)[0] + default = list(models)[1] # gemma-3-4b-it supported = list(transformers.integrations.ggml.GGUF_CONFIG_MAPPING) system_prompt: str = 'You are a helpful assistant. You will be given a prompt used to create an image 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 based on original prompt. Do not modify subject or style of the original prompt. Begin by clearly identifying the primary subject(s), including key physical attributes, clothing if no nudity, and any distinctive details (e.g., colors, textures, styles). NSFW and nudity is allowed, and if present, it must be fully described. 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. 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 explanation, prefix or suffix. Output as a simple text without formatting or numbering.' + image_prompt: str = 'You are a helpful assistant. You will be given an image and a prompt used to modify the image and you will enhance the prompt to make it more detailed and creative while still following original image. Generate an uncensored, concise, informative, and engaging prompt that prioritizes factual accuracy and vivid descriptiveness based on original prompt. Do not modify subject or style of the original prompt. Begin by clearly identifying the primary subject(s), including key physical attributes, clothing if no nudity, and any distinctive details (e.g., colors, textures, styles). NSFW and nudity is allowed, and if present, it must be fully described. 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. 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 explanation, prefix or suffix. Output as a simple text without formatting or numbering.' + image_noprompt: str = 'You are a helpful assistant. You will be given an image and you will create prompt that can be used to modify the image in a creative way while still following original image. Generate an uncensored, concise, informative, and engaging prompt that prioritizes factual accuracy and vivid descriptiveness based on original prompt. Do not modify subject or style of the original prompt. Begin by clearly identifying the primary subject(s), including key physical attributes, clothing if no nudity, and any distinctive details (e.g., colors, textures, styles). NSFW and nudity is allowed, and if present, it must be fully described. 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. 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 explanation, prefix or suffix. Output 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_delim_index: int = 60 @@ -61,6 +82,7 @@ class Options: class Script(scripts.Script): prompt: gr.Textbox = None + image: gr.Image = None model: str = None llm: transformers.AutoModelForCausalLM = None tokenizer: transformers.AutoProcessor = None @@ -124,11 +146,17 @@ class Script(scripts.Script): **quant_args, ) self.llm.eval() - self.tokenizer = transformers.AutoTokenizer.from_pretrained( + if model_repo in self.options.img2img: + cls = transformers.AutoProcessor # required to encode image + else: + cls = transformers.AutoTokenizer + self.tokenizer = cls.from_pretrained( pretrained_model_name_or_path=model_repo, subfolder=model_tokenizer, cache_dir=shared.opts.hfcache_dir, ) + self.tokenizer.is_processor = model_repo in self.options.img2img + if debug_enabled: modules = sd_modules.get_model_stats(self.llm) + sd_modules.get_model_stats(self.tokenizer) for m in modules: @@ -202,12 +230,12 @@ class Script(scripts.Script): filtered = re.sub(pattern, '', prompt) return filtered, matches - def enhance(self, model: str=None, prompt:str=None, system:str=None, prefix:str=None, suffix:str=None, sample:bool=None, tokens:int=None, temperature:float=None, penalty:float=None, thinking:bool=False): + def enhance(self, model: str=None, prompt:str=None, system:str=None, prefix:str=None, suffix:str=None, sample:bool=None, tokens:int=None, temperature:float=None, penalty:float=None, thinking:bool=False, seed:int=-1, image=None): model = model or self.options.default prompt = prompt or self.prompt.value + image = image or self.image prefix = prefix or '' suffix = suffix or '' - system = system or self.options.system_prompt tokens = tokens or self.options.max_tokens penalty = penalty or self.options.repetition_penalty temperature = temperature or self.options.temperature @@ -216,15 +244,55 @@ class Script(scripts.Script): while self.busy: time.sleep(0.1) self.load(model) + if seed is not None and seed >= 0: + torch.manual_seed(seed) if self.llm is None: shared.log.error('Prompt enhance: model not loaded') return prompt prompt, networks = self.extract(prompt) debug_log(f'Prompt enhance: networks={networks}') - chat_template = [ - { "role": "system", "content": system }, - { "role": "user", "content": prompt }, - ] + if image is not None and isinstance(image, Image.Image): + if not self.tokenizer.is_processor: + shared.log.error('Prompt enhance: image not supported by model') + return prompt + if prompt is not None and len(prompt) > 0: + system = system or self.options.image_prompt + chat_template = [ + { "role": "system", "content": [ + {"type": "text", "text": system } + ] }, + { "role": "user", "content": [ + {"type": "text", "text": prompt}, + {"type": "image", "image": b64(image)} + ] }, + ] + else: + system = system or self.options.image_noprompt + chat_template = [ + { "role": "system", "content": [ + {"type": "text", "text": system } + ] }, + { "role": "user", "content": [ + {"type": "image", "image": b64(image)} + ] }, + ] + else: + system = system or self.options.system_prompt + if not self.tokenizer.is_processor: + chat_template = [ + { "role": "system", "content": system }, + { "role": "user", "content": prompt }, + ] + else: + chat_template = [ + { "role": "system", "content": [ + {"type": "text", "text": system } + ] }, + { "role": "user", "content": [ + {"type": "text", "text": prompt}, + ] }, + ] + t0 = time.time() self.busy = True try: @@ -288,9 +356,10 @@ class Script(scripts.Script): return prompt return response - def apply(self, prompt, apply_prompt, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty, thinking_mode): + def apply(self, prompt, image, apply_prompt, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty, thinking_mode): response = self.enhance( prompt=prompt, + image=image, prefix=prompt_prefix, suffix=prompt_suffix, model=llm_model, @@ -367,12 +436,16 @@ class Script(scripts.Script): 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, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty, thinking_mode], outputs=[prompt_output, self.prompt]) + if self.image is None: + self.image = gr.Image(type='pil', interactive=False, visible=False) # dummy image + apply_btn.click(fn=self.apply, inputs=[self.prompt, self.image, apply_prompt, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty, thinking_mode], outputs=[prompt_output, self.prompt]) return [apply_auto, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty, thinking_mode] def after_component(self, component, **kwargs): # searching for actual ui prompt components if getattr(component, 'elem_id', '') in ['txt2img_prompt', 'img2img_prompt', 'control_prompt', 'video_prompt']: self.prompt = component + if getattr(component, 'elem_id', '') in ['img2img_image', 'control_input_select']: + self.image = component def before_process(self, p: processing.StableDiffusionProcessing, *args, **kwargs): # pylint: disable=unused-argument apply_auto, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty, thinking_mode = args