Innovideo / ui.py
Bossmarc747's picture
oj
32cd713
"""
Gradio UI components and layout for the image generation application.
This module defines the user interface using Gradio components,
including input controls, output displays, and event handlers.
"""
import gradio as gr
import time
from typing import Callable, Dict, Any, List, Tuple
from config import (
APP_TITLE,
APP_DESCRIPTION,
EXAMPLE_PROMPTS,
CSS,
MAX_IMAGE_SIZE,
MAX_SEED,
DEFAULT_WIDTH,
DEFAULT_HEIGHT,
DEFAULT_GUIDANCE_SCALE,
DEFAULT_INFERENCE_STEPS
)
from utils import save_image, format_generation_info, GenerationHistory
class ImageGenUI:
"""Manages the Gradio UI for the image generation application."""
def __init__(self, generate_func: Callable):
"""
Initialize the UI with the image generation function.
Args:
generate_func: Function to call for image generation
"""
self.generate_func = generate_func
self.history = GenerationHistory(max_history=10)
self.demo = None
def build_ui(self) -> gr.Blocks:
"""
Build and configure the Gradio UI.
Returns:
Configured Gradio Blocks interface
"""
with gr.Blocks(css=CSS) as demo:
gr.Markdown(f"# {APP_TITLE}")
gr.Markdown(APP_DESCRIPTION)
with gr.Row():
with gr.Column(scale=3):
# Input controls
with gr.Group():
prompt = gr.Text(
label="Prompt",
placeholder="Describe the image you want to generate",
lines=2
)
negative_prompt = gr.Text(
label="Negative Prompt",
placeholder="Describe what you want to avoid in the image",
lines=2
)
with gr.Row():
generate_btn = gr.Button("Generate Image", variant="primary")
clear_btn = gr.Button("Clear")
# Advanced settings
with gr.Accordion("Advanced Settings", open=False):
with gr.Row():
with gr.Column():
seed = gr.Slider(
label="Seed",
minimum=0,
maximum=MAX_SEED,
step=1,
value=0
)
randomize_seed = gr.Checkbox(
label="Randomize seed",
value=True
)
with gr.Column():
width = gr.Slider(
label="Width",
minimum=256,
maximum=MAX_IMAGE_SIZE,
step=32,
value=DEFAULT_WIDTH
)
height = gr.Slider(
label="Height",
minimum=256,
maximum=MAX_IMAGE_SIZE,
step=32,
value=DEFAULT_HEIGHT
)
with gr.Row():
guidance_scale = gr.Slider(
label="Guidance Scale",
minimum=0.0,
maximum=10.0,
step=0.1,
value=DEFAULT_GUIDANCE_SCALE,
info="How closely to follow the prompt (higher = more faithful)"
)
num_inference_steps = gr.Slider(
label="Inference Steps",
minimum=1,
maximum=50,
step=1,
value=DEFAULT_INFERENCE_STEPS,
info="More steps = higher quality but slower generation"
)
with gr.Column(scale=4):
# Output display
with gr.Group():
result_image = gr.Image(
label="Generated Image",
elem_classes=["output-image"]
)
image_info = gr.Markdown(label="Image Details")
with gr.Row():
save_btn = gr.Button("Save Image")
save_status = gr.Markdown("")
# Example prompts
gr.Examples(
examples=EXAMPLE_PROMPTS,
inputs=prompt,
label="Example Prompts"
)
# Generation history
with gr.Accordion("Generation History", open=False):
history_gallery = gr.Gallery(
label="Previous Generations",
show_label=True,
elem_id="history-gallery",
columns=5,
height="auto"
)
refresh_history_btn = gr.Button("Refresh History")
# Footer
gr.Markdown(
"Made with ❤️ using Gradio and Hugging Face Diffusers",
elem_classes=["footer"]
)
# Event handlers
def generate_image(
prompt_text,
negative_prompt_text,
seed_val,
randomize,
width_val,
height_val,
guidance,
steps,
progress=gr.Progress(track_tqdm=True)
):
"""Handle image generation and update UI."""
# Generate the image
image, used_seed = self.generate_func(
prompt_text,
negative_prompt_text,
seed_val,
randomize,
width_val,
height_val,
guidance,
steps,
progress_callback=progress.tqdm
)
# Update info text
info = format_generation_info(
prompt_text,
negative_prompt_text,
used_seed,
width_val,
height_val,
guidance,
steps
)
# Add to history
if image is not None:
self.history.add(
image,
prompt_text,
negative_prompt_text,
used_seed,
width_val,
height_val,
guidance,
steps
)
return image, info, used_seed
def save_current_image(image, prompt_text):
"""Save the current image and return status."""
if image is None:
return "No image to save"
try:
filepath = save_image(image, prompt_text)
return f"Image saved to {filepath}"
except Exception as e:
return f"Error saving image: {str(e)}"
def update_history():
"""Update the history gallery."""
entries = self.history.get_latest(10)
if not entries:
return []
# Format for gallery
images = [entry["image"] for entry in entries]
labels = [f"{entry['prompt'][:30]}..." for entry in entries]
return gr.Gallery.update(value=images, label=labels)
def clear_inputs():
"""Clear all input fields."""
return [
gr.Text.update(value=""), # prompt
gr.Text.update(value=""), # negative_prompt
gr.Slider.update(value=0), # seed
gr.Checkbox.update(value=True), # randomize_seed
gr.Markdown.update(value="") # image_info
]
# Connect event handlers
generate_btn.click(
fn=generate_image,
inputs=[
prompt,
negative_prompt,
seed,
randomize_seed,
width,
height,
guidance_scale,
num_inference_steps
],
outputs=[result_image, image_info, seed]
)
prompt.submit(
fn=generate_image,
inputs=[
prompt,
negative_prompt,
seed,
randomize_seed,
width,
height,
guidance_scale,
num_inference_steps
],
outputs=[result_image, image_info, seed]
)
save_btn.click(
fn=save_current_image,
inputs=[result_image, prompt],
outputs=[save_status]
)
refresh_history_btn.click(
fn=update_history,
inputs=[],
outputs=[history_gallery]
)
clear_btn.click(
fn=clear_inputs,
inputs=[],
outputs=[prompt, negative_prompt, seed, randomize_seed, image_info]
)
self.demo = demo
return demo
def launch(self, **kwargs):
"""Launch the Gradio interface with the specified parameters."""
if self.demo is None:
self.build_ui()
self.demo.launch(**kwargs)