""" 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)