Image: Make the LLM Variations prompt configurable

This commit is contained in:
oobabooga 2025-12-04 10:44:35 -08:00
parent 5763947c37
commit ffef3c7b1d
5 changed files with 35 additions and 6 deletions

View file

@ -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

View file

@ -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

View file

@ -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',

View file

@ -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',

View file

@ -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()