From ffef3c7b1dc5f3678a03abfa98fd54547de23796 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Thu, 4 Dec 2025 10:44:35 -0800 Subject: [PATCH] Image: Make the LLM Variations prompt configurable --- extensions/openai/images.py | 1 + extensions/openai/typing.py | 1 + modules/shared.py | 1 + modules/ui.py | 2 ++ modules/ui_image_generation.py | 36 ++++++++++++++++++++++++++++------ 5 files changed, 35 insertions(+), 6 deletions(-) diff --git a/extensions/openai/images.py b/extensions/openai/images.py index 1ecb1e63..0bb91a1e 100644 --- a/extensions/openai/images.py +++ b/extensions/openai/images.py @@ -36,6 +36,7 @@ def generations(request): 'image_batch_count': request.batch_count, 'image_cfg_scale': request.cfg_scale, 'image_llm_variations': request.llm_variations, + 'image_llm_variations_prompt': request.llm_variations_prompt or shared.settings.get('image_llm_variations_prompt', ''), }) # Exhaust generator, keep final result diff --git a/extensions/openai/typing.py b/extensions/openai/typing.py index dfdb9a7e..31a5dc6d 100644 --- a/extensions/openai/typing.py +++ b/extensions/openai/typing.py @@ -276,6 +276,7 @@ class ImageGenerationRequestParams(BaseModel): n: int = Field(default=1, ge=1, description="Alias for batch_size (OpenAI compatibility)") batch_count: int = Field(default=1, ge=1, description="Sequential batch count") llm_variations: bool = False + llm_variations_prompt: str | None = None # OpenAI compatibility (unused) model: str | None = None diff --git a/modules/shared.py b/modules/shared.py index 1ecc0d28..7fe9dbcf 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -320,6 +320,7 @@ settings = { 'image_batch_size': 1, 'image_batch_count': 1, 'image_llm_variations': False, + 'image_llm_variations_prompt': 'Please create a creative variation of the image generation prompt above. Keep the same general subject and style, but vary the details, composition, lighting, or mood. Respond with only the new prompt, nothing else.', 'image_model_menu': 'None', 'image_dtype': 'bfloat16', 'image_attn_backend': 'sdpa', diff --git a/modules/ui.py b/modules/ui.py index d95f7bc6..919a5740 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -294,6 +294,7 @@ def list_interface_input_elements(): 'image_batch_size', 'image_batch_count', 'image_llm_variations', + 'image_llm_variations_prompt', 'image_model_menu', 'image_dtype', 'image_attn_backend', @@ -549,6 +550,7 @@ def setup_auto_save(): 'image_batch_size', 'image_batch_count', 'image_llm_variations', + 'image_llm_variations_prompt', 'image_model_menu', 'image_dtype', 'image_attn_backend', diff --git a/modules/ui_image_generation.py b/modules/ui_image_generation.py index 6ac0bc24..6eeb3d51 100644 --- a/modules/ui_image_generation.py +++ b/modules/ui_image_generation.py @@ -406,14 +406,25 @@ def create_ui(): info="Z-Image Turbo: 0.0 | Qwen: 4.0" ) shared.gradio['image_seed'] = gr.Number(label="Seed", value=shared.settings['image_seed'], precision=0, info="-1 = Random") + with gr.Column(): shared.gradio['image_batch_size'] = gr.Slider(1, 32, value=shared.settings['image_batch_size'], step=1, label="Batch Size (VRAM Heavy)", info="Generates N images at once.") shared.gradio['image_batch_count'] = gr.Slider(1, 128, value=shared.settings['image_batch_count'], step=1, label="Sequential Count (Loop)", info="Repeats the generation N times.") - shared.gradio['image_llm_variations'] = gr.Checkbox( - value=shared.settings['image_llm_variations'], - label='LLM Prompt Variations', - info='Use the loaded LLM to generate creative prompt variations for each sequential batch.' - ) + + gr.Markdown("### LLM Variations") + shared.gradio['image_llm_variations'] = gr.Checkbox( + value=shared.settings['image_llm_variations'], + label='Activate', + info='Use the loaded LLM to generate creative prompt variations for each sequential batch.' + ) + shared.gradio['image_llm_variations_prompt'] = gr.Textbox( + value=shared.settings['image_llm_variations_prompt'], + label='Variation Prompt', + lines=3, + placeholder='Instructions for generating prompt variations...', + visible=shared.settings['image_llm_variations'], + info='The instruction given to the LLM for generating variations.' + ) with gr.Column(scale=6, min_width=500): with gr.Column(elem_classes=["viewport-container"]): @@ -668,6 +679,14 @@ def create_event_handlers(): show_progress=False ) + # LLM Variations visibility toggle + shared.gradio['image_llm_variations'].change( + lambda x: gr.update(visible=x), + gradio('image_llm_variations'), + gradio('image_llm_variations_prompt'), + show_progress=False + ) + def generate_prompt_variation(state): """Generate a creative variation of the image prompt using the LLM.""" @@ -682,7 +701,12 @@ def generate_prompt_variation(state): logger.warning("No LLM loaded for prompt variation. Using original prompt.") return prompt - augmented_message = f"{prompt}\n\n=====\n\nPlease create a creative variation of the image generation prompt above. Keep the same general subject and style, but vary the details, composition, lighting, or mood. Respond with only the new prompt, nothing else." + # Get the custom variation prompt or use default + variation_instruction = state.get('image_llm_variations_prompt', '') + if not variation_instruction: + variation_instruction = 'Please create a creative variation of the image generation prompt above. Keep the same general subject and style, but vary the details, composition, lighting, or mood. Respond with only the new prompt, nothing else.' + + augmented_message = f"{prompt}\n\n=====\n\n{variation_instruction}" # Use minimal state for generation var_state = state.copy()