app.py

"""
AI Image Generator - Gradio Web App

A web application for generating images using Stable Diffusion
with a user-friendly Gradio interface.
"""

import json
import logging
import os
from typing import List, Optional, Tuple

import gradio as gr
import torch
from PIL import Image

from utils.generation import ImageGenerator, save_image_with_metadata


# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Global image generator instance
image_generator: Optional[ImageGenerator] = None


def load_examples() -> dict:
    """Load example prompts from JSON file."""
    try:
        with open('examples.json', 'r', encoding='utf-8') as f:
            return json.load(f)
    except FileNotFoundError:
        logger.warning("examples.json not found, using default examples")
        return {
            "examples": []
        }
    except Exception as e:
        logger.error(f"Failed to load examples: {str(e)}")
        return {
            "examples": [],
        }


def initialize_generator() -> ImageGenerator:
    """
    Initialize the image generator with error handling.
    
    Returns:
        ImageGenerator: The initialized generator instance.
    
    Raises:
        RuntimeError: If initialization fails.
    """
    try:
        logger.info("Initializing Stable Diffusion model...")
        generator = ImageGenerator()
        logger.info("Model initialized successfully")
        return generator
    except Exception as e:
        logger.error(f"Failed to initialize model: {str(e)}")
        raise RuntimeError(f"Model initialization failed: {str(e)}")


def generate_images_interface(
    prompt: str,
    num_inference_steps: int,
    # guidance_scale: float,  # Commented out - using default value 0.0
    seed: Optional[int],
    width: int,
    height: int,
    num_images: int
) -> Tuple[List[str], str]:
    """
    Generate images based on user input parameters.
    
    Args:
        prompt: Text prompt for image generation.
        num_inference_steps: Number of denoising steps.
        # guidance_scale: How closely to follow the prompt.  # Commented out
        seed: Random seed for reproducibility.
        width: Image width.
        height: Image height.
        num_images: Number of images to generate.
    
    Returns:
        Tuple[List[str], str]: (list of image paths, status message)
    """
    global image_generator
    
    try:
        # Initialize generator if not already done
        if image_generator is None:
            image_generator = initialize_generator()
        
        # Validate inputs
        if not prompt or not prompt.strip():
            return [], "❌ Please enter a prompt"
        
        logger.info(f"Generating {num_images} image(s) with prompt: '{prompt[:50]}...'")
        
        # Generate images
        images = image_generator.generate_images(
            prompt=prompt,
            num_images=num_images,
            num_inference_steps=num_inference_steps,
            guidance_scale=0.0,  # Using default value since guidance_scale is removed from UI
            width=width,
            height=height,
            seed=seed
        )
        
        # Save images and collect paths
        saved_paths = []
        for i, image in enumerate(images):
            metadata = {
                "num_inference_steps": num_inference_steps,
                "guidance_scale": 0.0,  # Using default value
                "seed": seed,
                "width": width,
                "height": height,
                "image_number": i + 1
            }
            
            filepath = save_image_with_metadata(
                image=image,
                prompt=prompt,
                metadata=metadata
            )
            saved_paths.append(filepath)
        
        status_msg = f"✅ Successfully generated {len(images)} image(s)!"
        logger.info(f"Generation completed: {status_msg}")
        
        return saved_paths, status_msg
        
    except Exception as e:
        error_msg = f"❌ Generation failed: {str(e)}"
        logger.error(f"Generation error: {str(e)}")
        return [], error_msg


def create_interface() -> gr.Blocks:
    """
    Create the Gradio interface for the AI Image Generator.
    
    Returns:
        gr.Blocks: The configured Gradio interface.
    """
    
    # Custom CSS for better styling
    css = """
    .gradio-container {
        max-width: 100% !important;
        margin: 0 !important;
        padding: 0 !important;
    }
    .main-header {
        text-align: center;
        margin-bottom: 2rem;
    }
    .param-section {
        background: #f8f9fa;
        padding: 1rem;
        border-radius: 8px;
        margin-bottom: 1rem;
    }
    .examples-section {
        background: #f0f8ff;
        padding: 1.5rem;
        border-radius: 12px;
        margin-bottom: 2rem;
        border: 2px solid #e1f5fe;
    }
    .gradio-row {
        gap: 2rem !important;
        max-width: 100% !important;
    }
    .gradio-column {
        min-width: 0 !important;
        flex: 1 !important;
    }
    .gradio-blocks {
        max-width: 100% !important;
    }
    """
    
    with gr.Blocks(css=css, title="AI Image Generator") as interface:
        
        # Header
        gr.HTML("""
        <div class="main-header">
            <h1>🎨 AI Image Generator — powered by Stable Diffusion</h1>
            <p style="font-size: 1.1em; color: #666;">Enter a prompt to generate an image in seconds</p>
        </div>
        """)
        
        # Load examples data first
        examples_data = load_examples()
        
        # Create examples for different categories
        all_examples = []

        # Add examples from categories
        if "examples" in examples_data:
            for category in examples_data["examples"]:
                if "prompts" in category:
                    for prompt_data in category["prompts"][:2]:  # Take first 2 from each category
                        if isinstance(prompt_data, dict):
                            all_examples.append([
                                prompt_data["text"],
                                prompt_data["num_inference_steps"],
                                # prompt_data["guidance_scale"],  # Commented out - using default value
                                prompt_data["seed"],
                                prompt_data["width"],
                                prompt_data["height"],
                                1  # num_images default
                            ])
                        else:
                            all_examples.append([prompt_data, 2, None, 1024, 1024, 1])  # Updated to match new parameters
        
        # Limit to 8 examples for better UI
        all_examples = all_examples[:8]
        
        with gr.Row():
            # Left column - Input parameters
            with gr.Column(scale=2):
                gr.Markdown("### 📝 Input Parameters")
                
                # Main prompt input
                prompt_input = gr.Textbox(
                    label="Prompt",
                    placeholder="A beautiful sunset over mountains, digital art",
                    lines=3,
                    max_lines=5
                )
                
                # Generation parameters
                with gr.Group():
                    gr.Markdown("#### ⚙️ Generation Settings")
                    
                    num_inference_steps = gr.Slider(
                        minimum=1,
                        maximum=4,
                        value=2,
                        step=1,
                        label="Inference Steps",
                        info="More steps = better quality, slower generation"
                    )
                    
                    # guidance_scale = gr.Slider(
                    #     minimum=1.0,
                    #     maximum=20.0,
                    #     value=7.5,
                    #     step=0.1,
                    #     label="Guidance Scale",
                    #     info="How closely to follow the prompt"
                    # )
                    
                    seed = gr.Number(
                        label="Seed (optional)",
                        value=None,
                        precision=0,
                        info="Leave empty for random generation"
                    )
                
                # Image parameters
                with gr.Group():
                    gr.Markdown("#### 🖼️ Image Settings")
                    
                    with gr.Row():
                        width = gr.Number(
                            label="Width",
                            value=1024,
                            minimum=512,
                            maximum=1024,
                            step=64,
                            precision=0
                        )
                        
                        height = gr.Number(
                            label="Height", 
                            value=1024,
                            minimum=512,
                            maximum=1024,
                            step=64,
                            precision=0
                        )
                    
                    num_images = gr.Slider(
                        minimum=1,
                        maximum=4,
                        value=1,
                        step=1,
                        label="Number of Images",
                        info="Generate multiple variations"
                    )
                
                # Generate button
                generate_btn = gr.Button(
                    "🎨 Generate Images",
                    variant="primary",
                    size="lg"
                )
            
            # Right column - Output
            with gr.Column(scale=2):
                gr.Markdown("### 🖼️ Generated Images")
                
                # Status message
                status_output = gr.Textbox(
                    label="Status",
                    interactive=False,
                    value="Ready to generate images!"
                )
                
                # Image gallery
                gallery = gr.Gallery(
                    label="Generated Images",
                    show_label=True,
                    elem_id="gallery",
                    columns=2,
                    rows=2,
                    height="auto",
                    object_fit="contain"
                )
        
        # Footer
        gr.HTML("""
        <div style="text-align: center; margin-top: 2rem; padding: 1rem; border-top: 1px solid #eee;">
            <p>
                Powered by <a href="https://huggingface.co/docs/diffusers/using-diffusers/sdxl_turbo" target="_blank">Stable Diffusion XL Turbo</a> | 
                Built with <a href="https://gradio.app" target="_blank">Gradio</a>
            </p>
        </div>
        """)
        
        # Examples section - now that all components are defined
        with gr.Group():
            gr.Markdown("### 💡 Example Prompts - Click to Auto-Fill Form")
            gr.Examples(
                examples=all_examples,
                inputs=[
                    prompt_input,
                    num_inference_steps,
                    # guidance_scale,  # Commented out - removed from UI
                    seed,
                    width,
                    height,
                    num_images
                ],
                label="",
                elem_id="examples-section"
            )
        
        # Event handlers
        generate_btn.click(
            fn=generate_images_interface,
            inputs=[
                prompt_input,
                num_inference_steps,
                # guidance_scale,  # Commented out - removed from UI
                seed,
                width,
                height,
                num_images
            ],
            outputs=[gallery, status_output]
        )
        
    
    return interface


def main():
    """Main function to launch the Gradio interface."""
    try:
        logger.info("Starting AI Image Generator...")
        
        # Create and launch interface
        interface = create_interface()
        
        # Launch with appropriate settings
        interface.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=False,
            show_error=True,
            quiet=False
        )
        
    except Exception as e:
        logger.error(f"Failed to start application: {str(e)}")
        raise


if __name__ == "__main__":
    main()