app.py

# AI Image Creator: Enhanced UI and UX
# Part 1: Core Setup, Model Classes and API Configuration

import gradio as gr
import logging
import sys
import random
import time
import os
from huggingface_hub import InferenceClient
from PIL import Image
import io
import base64

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger("ai_image_creator")

# =============== MODEL CLIENTS SETUP ===============
def setup_client(api_key, provider=None):
    """Initialize and return API client"""
    try:
        if provider:
            client = InferenceClient(provider=provider, api_key=api_key)
            logger.info(f"{provider} client initialized successfully")
        else:
            client = InferenceClient(api_key=api_key)
            logger.info("Hugging Face client initialized successfully")
        return client
    except Exception as e:
        logger.error(f"Error initializing client: {str(e)}")
        return None

# Initialize clients
try:
    # Replace with your actual HF API key
    hf_api_key = os.getenv("HF_API_KEY1")

    hf_client = setup_client(hf_api_key)
    logger.info("Hugging Face client created successfully")

    # Set up Llama client if API key is provided
    llama_api_key = os.getenv("HF_API_KEY2")  # Replace with actual key if available
    try:
        llama_client = setup_client(llama_api_key, "sambanova")
        use_llama = True
        logger.info("Llama client created successfully")
    except Exception as e:
        logger.warning(f"Llama client not available: {str(e)}. Will use fallback enhancement.")
        llama_client = None
        use_llama = False
except Exception as e:
    logger.error(f"Failed to create Hugging Face client: {str(e)}")
    hf_client = None
    llama_client = None
    use_llama = False

# =============== DATA MODELS ===============

# Image Models with friendly names, descriptions and icons
IMAGE_MODELS = {
    "stabilityai/stable-diffusion-xl-base-1.0": {
        "display_name": "SDXL 1.0",
        "description": "Best overall quality, slower generation",
        "icon": "⭐",
        "speed": "slow",
        "quality": "excellent"
    },
    "runwayml/stable-diffusion-v1-5": {
        "display_name": "SD 1.5",
        "description": "Good for general purpose, faster generation",
        "icon": "🚀",
        "speed": "fast",
        "quality": "good"
    },
    "stabilityai/stable-diffusion-2-1": {
        "display_name": "SD 2.1",
        "description": "Improved details, balanced speed and quality",
        "icon": "✨",
        "speed": "medium",
        "quality": "very good"
    },
    "prompthero/openjourney": {
        "display_name": "OpenJourney",
        "description": "Midjourney-like stylized results",
        "icon": "🎨",
        "speed": "medium",
        "quality": "stylized"
    },
    "dreamlike-art/dreamlike-diffusion-1.0": {
        "display_name": "Dreamlike",
        "description": "Artistic style with dreamy aesthetics",
        "icon": "💫",
        "speed": "medium",
        "quality": "artistic"
    }
}

# Creation types with icons and detailed descriptions
CREATION_TYPES = {
    "Realistic Photo": {
        "description": "Create a photorealistic image with natural details and lighting",
        "icon": "📷",
        "prompt_hint": "Try to include details about lighting, time of day, and environment"
    },
    "Digital Art": {
        "description": "Create colorful digital artwork with clean lines and vibrant colors",
        "icon": "🖌️",
        "prompt_hint": "Consider specifying color palette and mood for better results"
    },
    "Fantasy Illustration": {
        "description": "Create magical and fantastical scenes with otherworldly elements",
        "icon": "🧙",
        "prompt_hint": "Describe magical elements, creatures, and environments in detail"
    },
    "Concept Art": {
        "description": "Create professional concept art for characters, environments or objects",
        "icon": "🎮",
        "prompt_hint": "Include details about perspective, purpose, and design influences"
    },
    "Anime/Manga": {
        "description": "Create Japanese anime or manga style illustration",
        "icon": "🍙",
        "prompt_hint": "Specify anime aesthetics like shading style and character features"
    },
    "Oil Painting": {
        "description": "Create an image with oil painting textures and artistic brushstrokes",
        "icon": "🖼️",
        "prompt_hint": "Consider describing texture, brushwork style, and canvas feel"
    },
    "Watercolor": {
        "description": "Create a soft watercolor illustration with subtle color blending",
        "icon": "💧",
        "prompt_hint": "Mention color blending, paper texture, and watercolor-specific effects"
    },
    "Sketch": {
        "description": "Create a detailed sketch or drawing with line art focus",
        "icon": "✏️",
        "prompt_hint": "Describe line weight, hatching style, and sketch medium (pencil, charcoal, etc.)"
    },
    "3D Rendering": {
        "description": "Create an image that looks like a 3D rendered scene with realistic lighting",
        "icon": "💻",
        "prompt_hint": "Include details about lighting setup, materials, and camera perspective"
    },
    "Pixel Art": {
        "description": "Create retro-style pixel art with limited color palette",
        "icon": "👾",
        "prompt_hint": "Specify resolution, color limitations, and pixel art style (e.g., 16-bit, 8-bit)"
    }
}

# Art styles with icons and detailed descriptions
ART_STYLES = {
    "Photorealistic": {
        "description": "Detailed realistic style that resembles a photograph with accurate lighting and textures",
        "icon": "📸",
        "examples": "Works by Chuck Close, Richard Estes, or modern 3D renderings"
    },
    "Impressionist": {
        "description": "Soft brushstrokes that capture light and atmosphere over precise details, like Monet",
        "icon": "🌈",
        "examples": "Claude Monet, Pierre-Auguste Renoir, Camille Pissarro"
    },
    "Surrealist": {
        "description": "Dreamlike quality with impossible or irrational scenes, like Salvador Dali",
        "icon": "🌀",
        "examples": "Salvador Dali, René Magritte, Frida Kahlo"
    },
    "Pop Art": {
        "description": "Bold colors, sharp lines and popular culture references, like Andy Warhol",
        "icon": "🎭",
        "examples": "Andy Warhol, Roy Lichtenstein, Keith Haring"
    },
    "Minimalist": {
        "description": "Simplified forms, limited color palette, and clean composition with minimal elements",
        "icon": "⬜",
        "examples": "Piet Mondrian, Kazimir Malevich, Agnes Martin"
    },
    "Abstract": {
        "description": "Non-representational style using shapes, colors, and forms to express ideas",
        "icon": "🔶",
        "examples": "Wassily Kandinsky, Jackson Pollock, Mark Rothko"
    },
    "Cubist": {
        "description": "Geometric shapes and multiple perspectives shown simultaneously, like Picasso",
        "icon": "📐",
        "examples": "Pablo Picasso, Georges Braque, Juan Gris"
    },
    "Art Nouveau": {
        "description": "Ornate, flowing lines inspired by natural forms with decorative elegance",
        "icon": "🌿",
        "examples": "Alphonse Mucha, Gustav Klimt, Antoni Gaudí"
    },
    "Gothic": {
        "description": "Dark, medieval-inspired aesthetic with dramatic lighting and architectural elements",
        "icon": "🏰",
        "examples": "Zdzisław Beksiński, H.R. Giger, medieval architecture"
    },
    "Cyberpunk": {
        "description": "Futuristic dystopian style with neon colors, technology, and urban decay",
        "icon": "🤖",
        "examples": "Blade Runner, Ghost in the Shell, Akira"
    },
    "Steampunk": {
        "description": "Victorian-era aesthetic combined with steam-powered technology and brass elements",
        "icon": "⚙️",
        "examples": "Works by James Ng, Keith Thompson, retrofuturistic Jules Verne adaptations"
    },
    "Retro/Vintage": {
        "description": "Nostalgic style reminiscent of past decades with period-appropriate elements",
        "icon": "📺",
        "examples": "1950s advertisements, vintage travel posters, pulp magazine covers"
    },
    "Art Deco": {
        "description": "Geometric patterns, bold colors, and luxurious materials in a symmetrical style",
        "icon": "🏢",
        "examples": "Works from the 1920s-30s, Chrysler Building, Tamara de Lempicka paintings"
    },
    "Baroque": {
        "description": "Dramatic, ornate style with rich details, contrast, and dynamic composition",
        "icon": "👑",
        "examples": "Caravaggio, Rembrandt, Peter Paul Rubens"
    },
    "Ukiyo-e": {
        "description": "Traditional Japanese woodblock print style with flat areas of color and strong outlines",
        "icon": "🌊",
        "examples": "Hokusai's Great Wave, Hiroshige's landscapes, traditional Japanese prints"
    },
    "Comic Book": {
        "description": "Bold outlines, bright colors, and action-oriented composition like classic comics",
        "icon": "💥",
        "examples": "Jack Kirby, Steve Ditko, modern Marvel/DC art styles"
    },
    "Psychedelic": {
        "description": "Vibrant, swirling colors with abstract patterns inspired by 1960s art",
        "icon": "🌈",
        "examples": "1960s concert posters, Peter Max, Alex Grey"
    },
    "Vaporwave": {
        "description": "Glitch aesthetics with pastel colors, 80s/90s nostalgia and digital elements",
        "icon": "📼",
        "examples": "Retrowave aesthetics, 80s digital graphics, glitch art"
    },
    "Studio Ghibli": {
        "description": "Whimsical, detailed animation style inspired by Japanese animated films",
        "icon": "🐉",
        "examples": "Spirited Away, My Neighbor Totoro, Howl's Moving Castle"
    },
    "Hyperrealism": {
        "description": "Extremely detailed realism that exceeds photograph-like precision",
        "icon": "🔍",
        "examples": "Works by Roberto Bernardi, Denis Peterson, Gottfried Helnwein"
    }
}

# Moods with icons and descriptions
MOODS = {
    "Happy": {
        "description": "Bright, cheerful atmosphere with warm colors",
        "icon": "😊",
        "color_palette": "Warm and vibrant colors: yellows, bright oranges, light blues"
    },
    "Sad": {
        "description": "Melancholic atmosphere with muted colors",
        "icon": "😢",
        "color_palette": "Muted blues, grays, desaturated colors, cool tones"
    },
    "Mysterious": {
        "description": "Enigmatic atmosphere with shadows and hidden elements",
        "icon": "🔮",
        "color_palette": "Deep purples, dark blues, hints of teal, selective lighting"
    },
    "Peaceful": {
        "description": "Serene, calm atmosphere with gentle lighting",
        "icon": "🕊️",
        "color_palette": "Soft blues, gentle greens, pastel colors, balanced light"
    },
    "Tense": {
        "description": "Suspenseful atmosphere with dramatic lighting",
        "icon": "😰",
        "color_palette": "High contrast, dark shadows, selective reds, strong highlights"
    },
    "Whimsical": {
        "description": "Playful, imaginative atmosphere with fanciful elements",
        "icon": "🦄",
        "color_palette": "Pastels, candy colors, unexpected color combinations"
    },
    "Dark": {
        "description": "Gloomy atmosphere with deep shadows and low lighting",
        "icon": "🌑",
        "color_palette": "Dark blues, blacks, deep greens, minimal highlights"
    },
    "Energetic": {
        "description": "Dynamic, vibrant atmosphere with strong colors and movement",
        "icon": "⚡",
        "color_palette": "Saturated primary colors, bold contrasts, vibrant hues"
    },
    "Romantic": {
        "description": "Soft, dreamy atmosphere with warm, gentle lighting",
        "icon": "❤️",
        "color_palette": "Soft pinks, gentle reds, golden highlights, warm tones"
    },
    "Epic": {
        "description": "Grand, impressive atmosphere with dramatic scale and lighting",
        "icon": "🏔️",
        "color_palette": "Bold colors, dramatic contrast, atmospheric lighting, expansive scale"
    }
}

# Example prompts with rich metadata
EXAMPLE_PROMPTS = [
    {
        "text": "A serene lake at sunset with mountains in the background and a small wooden boat floating nearby",
        "thumbnail_desc": "Peaceful lake scene at sunset",
        "creation_type": "Realistic Photo",
        "art_style": "Photorealistic",
        "mood": "Peaceful",
        "tags": ["nature", "landscape", "water", "sunset"]
    },
    {
        "text": "A futuristic cityscape with flying cars, neon lights, and tall skyscrapers under a night sky with two moons",
        "thumbnail_desc": "Futuristic city with flying cars",
        "creation_type": "Concept Art",
        "art_style": "Cyberpunk",
        "mood": "Mysterious",
        "tags": ["scifi", "future", "urban", "night"]
    },
    {
        "text": "A close-up portrait of an elderly craftsman with weathered hands working on an intricate wooden carving in his workshop",
        "thumbnail_desc": "Elderly craftsman working with wood",
        "creation_type": "Oil Painting",
        "art_style": "Hyperrealism",
        "mood": "Peaceful",
        "tags": ["portrait", "craftsmanship", "elderly", "detail"]
    },
    {
        "text": "A magical forest with glowing mushrooms, fairy lights, and a small cottage with smoke coming from the chimney",
        "thumbnail_desc": "Magical forest with glowing elements",
        "creation_type": "Fantasy Illustration",
        "art_style": "Studio Ghibli",
        "mood": "Whimsical",
        "tags": ["fantasy", "forest", "magic", "cottage"]
    },
    {
        "text": "A cybernetic samurai with glowing blue circuits standing in a rainy Tokyo street at night",
        "thumbnail_desc": "Cybernetic samurai in rainy Tokyo",
        "creation_type": "Digital Art",
        "art_style": "Cyberpunk",
        "mood": "Dark",
        "tags": ["character", "cyberpunk", "samurai", "rain"]
    },
    {
        "text": "A cute cat with dragon wings and tiny horns sleeping on a pile of gold coins",
        "thumbnail_desc": "Cat with dragon features on gold",
        "creation_type": "Fantasy Illustration",
        "art_style": "Comic Book",
        "mood": "Whimsical",
        "tags": ["creature", "fantasy", "cute", "treasure"]
    },
    {
        "text": "An ancient temple covered in vines and moss, partially sunken into a crystal-clear cenote in the jungle",
        "thumbnail_desc": "Ancient temple in jungle cenote",
        "creation_type": "Concept Art",
        "art_style": "Photorealistic",
        "mood": "Mysterious",
        "tags": ["architecture", "ruins", "jungle", "water"]
    },
    {
        "text": "A cozy coffee shop interior with rain falling outside the windows, soft lighting, and a few people reading books",
        "thumbnail_desc": "Cozy rainy day in coffee shop",
        "creation_type": "Digital Art",
        "art_style": "Impressionist",
        "mood": "Peaceful",
        "tags": ["interior", "rainy", "cozy", "urban"]
    }
]

# CSS for enhanced UI styling - Will be included in part 2

# =============== HELPER FUNCTIONS ===============

# Helper function to format dropdown choices with icons
def format_dropdown_choices(options_dict):
    return [f"{options_dict[key].get('icon', '•')} {key}" for key in options_dict.keys()]

# Helper function to extract the key from formatted choice
def extract_key(formatted_choice):
    # Skip the icon and space at the beginning
    parts = formatted_choice.split(' ', 1)
    if len(parts) > 1:
        return parts[1]
    return formatted_choice

# Function to load example prompt
def load_example(example_index):
    if example_index < 0 or example_index >= len(EXAMPLE_PROMPTS):
        return "", "", "", ""

    example = EXAMPLE_PROMPTS[example_index]
    creation = f"{CREATION_TYPES[example['creation_type']]['icon']} {example['creation_type']}"
    art = f"{ART_STYLES[example['art_style']]['icon']} {example['art_style']}"
    mood = f"{MOODS[example['mood']]['icon']} {example['mood']}"

    return example["text"], creation, art, mood

# Get model key from formatted display name
def get_model_key(formatted_name):
    # Extract display name without the icon
    if ' ' in formatted_name:
        display_name = formatted_name.split(' ', 1)[1]
        # Find the corresponding key
        for key, info in IMAGE_MODELS.items():
            if info['display_name'] == display_name:
                return key
    return list(IMAGE_MODELS.keys())[0]  # Default to first model if not found

# AI Image Creator: Enhanced UI and UX
# Part 2: Enhanced UI Components and Styling

# CSS for styling the interface - Comprehensive styling for modern UI
css = """
/* Main theme colors with CSS variables for better customization */
:root {
    --primary-color: #4F46E5;
    --primary-hover: #4338CA;
    --secondary-color: #7C3AED;
    --secondary-hover: #6D28D9;
    --background-color: #F8FAFC;
    --card-color: #FFFFFF;
    --text-color: #1E293B;
    --text-muted: #64748B;
    --accent-color: #3B82F6;
    --success-color: #10B981;
    --success-hover: #059669;
    --warning-color: #F59E0B;
    --error-color: #EF4444;
    --error-hover: #DC2626;
    --border-color: #E2E8F0;
    --border-hover: #CBD5E1;
    --shadow-sm: 0 1px 2px 0 rgba(0, 0, 0, 0.05);
    --shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06);
    --shadow-lg: 0 10px 15px -3px rgba(0, 0, 0, 0.1), 0 4px 6px -2px rgba(0, 0, 0, 0.05);
    --radius-sm: 0.375rem;
    --radius: 0.5rem;
    --radius-lg: 0.75rem;
    --radius-xl: 1rem;
    --font-sans: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
}

/* Global styles and resets */
body, html {
    font-family: var(--font-sans);
    color: var(--text-color);
    background-color: var(--background-color);
    line-height: 1.5;
    margin: 0;
    padding: 0;
}

/* Container with responsive padding */
.container {
    max-width: 1400px;
    margin: 0 auto;
    padding: 1rem;
}

@media (max-width: 640px) {
    .container {
        padding: 0.5rem;
    }
}

/* Card styling with elevation and hover effects */
.gr-panel {
    border-radius: var(--radius) !important;
    border: 1px solid var(--border-color) !important;
    box-shadow: var(--shadow) !important;
    overflow: hidden;
    transition: transform 0.2s, box-shadow 0.2s;
    background-color: var(--card-color) !important;
}

.gr-panel:hover {
    transform: translateY(-2px);
    box-shadow: var(--shadow-lg) !important;
}

/* Button styling with gradient and hover states */
button.primary {
    background: linear-gradient(135deg, var(--primary-color), var(--secondary-color)) !important;
    color: white !important;
    border: none !important;
    border-radius: var(--radius) !important;
    font-weight: 600 !important;
    letter-spacing: 0.025em !important;
    padding: 0.75rem 1.5rem !important;
    transition: all 0.3s ease !important;
    box-shadow: var(--shadow-sm) !important;
    outline: none !important;
    text-transform: none !important;
}

button.primary:hover {
    transform: translateY(-1px);
    box-shadow: var(--shadow) !important;
    opacity: 0.9;
}

button.primary:active {
    transform: translateY(0);
    opacity: 0.8;
}

button.primary[disabled], button.primary[disabled]:hover {
    opacity: 0.5;
    cursor: not-allowed;
    transform: none;
}

/* Secondary button styling */
button.secondary {
    background-color: transparent !important;
    color: var(--primary-color) !important;
    border: 1px solid var(--primary-color) !important;
    border-radius: var(--radius) !important;
    font-weight: 500 !important;
    padding: 0.625rem 1.25rem !important;
    transition: all 0.2s ease !important;
    text-transform: none !important;
}

button.secondary:hover {
    background-color: rgba(79, 70, 229, 0.05) !important;
    border-color: var(--primary-hover) !important;
}

/* Style for the example buttons */
.example-button {
    font-size: 0.875rem !important;
    padding: 0.5rem 0.75rem !important;
    background-color: transparent !important;
    border: 1px solid var(--border-color) !important;
    border-radius: var(--radius) !important;
    transition: all 0.2s !important;
    text-align: left !important;
    justify-content: flex-start !important;
    height: auto !important;
    text-overflow: ellipsis;
    overflow: hidden;
    white-space: nowrap;
    width: 100%;
    color: var(--text-color) !important;
}

.example-button:hover {
    background-color: rgba(79, 70, 229, 0.05) !important;
    border-color: var(--primary-color) !important;
    transform: translateY(-1px);
}

/* Form controls styling */
.gr-input, .gr-textarea, .gr-dropdown {
    border-radius: var(--radius) !important;
    border: 1px solid var(--border-color) !important;
    transition: border-color 0.2s, box-shadow 0.2s !important;
    font-family: var(--font-sans) !important;
    color: var(--text-color) !important;
}

.gr-input:focus, .gr-textarea:focus, .gr-dropdown:focus-within {
    border-color: var(--primary-color) !important;
    box-shadow: 0 0 0 3px rgba(79, 70, 229, 0.2) !important;
    outline: none !important;
}

.gr-form {
    gap: 1rem !important;
}

.gr-input-label, .gr-dropdown-label, .gr-textarea-label, .gr-checkbox-label, .gr-radio-label {
    font-size: 0.875rem !important;
    font-weight: 500 !important;
    color: var(--text-color) !important;
    margin-bottom: 0.25rem !important;
}

/* Input placeholder styling */
.gr-input::placeholder, .gr-textarea::placeholder {
    color: var(--text-muted) !important;
    opacity: 0.7;
}

/* Input and textarea styling */
textarea, input[type="text"] {
    border-radius: var(--radius) !important;
    border: 1px solid var(--border-color) !important;
    padding: 0.75rem 1rem !important;
    transition: border-color 0.2s, box-shadow 0.2s !important;
    font-family: var(--font-sans) !important;
}

textarea:focus, input[type="text"]:focus {
    border-color: var(--primary-color) !important;
    box-shadow: 0 0 0 3px rgba(79, 70, 229, 0.2) !important;
    outline: none !important;
}

/* Dropdown styling */
.gr-dropdown {
    border-radius: var(--radius) !important;
    border: 1px solid var(--border-color) !important;
    background-color: var(--card-color) !important;
}

.gr-dropdown > div {
    border-radius: var(--radius) !important;
    min-height: 38px !important;
}

.gr-dropdown > div > span {
    font-size: 0.9375rem !important;
}

/* Dropdown menu styling */
.gr-dropdown ul {
    background-color: var(--card-color) !important;
    border: 1px solid var(--border-color) !important;
    border-radius: var(--radius) !important;
    box-shadow: var(--shadow) !important;
}

.gr-dropdown ul li {
    padding: 0.5rem 0.75rem !important;
}

.gr-dropdown ul li:hover {
    background-color: rgba(79, 70, 229, 0.05) !important;
}

/* Custom header with gradient background */
.app-header {
    text-align: center;
    padding: 1.75rem 1rem;
    margin-bottom: 2rem;
    background: linear-gradient(135deg, rgba(79, 70, 229, 0.08), rgba(124, 58, 237, 0.08));
    border-radius: var(--radius-lg);
    border-bottom: 3px solid var(--primary-color);
    position: relative;
    overflow: hidden;
}

.app-header::before {
    content: '';
    position: absolute;
    top: -50px;
    left: -50px;
    right: -50px;
    height: 100px;
    background: linear-gradient(135deg, rgba(79, 70, 229, 0.2), rgba(124, 58, 237, 0.2));
    transform: rotate(-5deg);
    z-index: 0;
}

.app-header h1 {
    font-size: 2.5rem !important;
    font-weight: 800 !important;
    background: linear-gradient(135deg, var(--primary-color), var(--secondary-color));
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin-bottom: 0.5rem !important;
    position: relative;
    z-index: 1;
}

.app-header p {
    font-size: 1.25rem !important;
    color: var(--text-color);
    opacity: 0.8;
    max-width: 40rem;
    margin: 0 auto;
    position: relative;
    z-index: 1;
}

/* Responsive header */
@media (max-width: 640px) {
    .app-header h1 {
        font-size: 2rem !important;
    }

    .app-header p {
        font-size: 1rem !important;
    }
}

/* Examples gallery with grid layout */
.example-gallery {
    display: grid;
    grid-template-columns: repeat(auto-fill, minmax(140px, 1fr));
    gap: 1rem;
    margin: 1rem 0;
}

.example-item {
    border-radius: var(--radius);
    overflow: hidden;
    cursor: pointer;
    border: 2px solid transparent;
    transition: all 0.2s;
    display: flex;
    flex-direction: column;
    background-color: var(--card-color);
}

.example-item:hover {
    transform: translateY(-2px);
    border-color: var(--accent-color);
    box-shadow: var(--shadow);
}

.example-item:active {
    transform: translateY(0);
}

.example-item-image {
    width: 100%;
    aspect-ratio: 1;
    background-color: #f0f0f0;
    display: flex;
    align-items: center;
    justify-content: center;
    color: var(--text-muted);
    font-size: 1.5rem;
    border-bottom: 1px solid var(--border-color);
}

.example-item-caption {
    padding: 0.5rem;
    font-size: 0.75rem;
    line-height: 1.25;
    color: var(--text-color);
    overflow: hidden;
    display: -webkit-box;
    -webkit-line-clamp: 2;
    -webkit-box-orient: vertical;
}

/* Loading indicator with animation */
.loading-indicator {
    display: flex;
    align-items: center;
    justify-content: center;
    height: 100%;
    width: 100%;
    position: absolute;
    top: 0;
    left: 0;
    background-color: rgba(255, 255, 255, 0.8);
    z-index: 1000;
    backdrop-filter: blur(2px);
    border-radius: var(--radius);
}

.spinner {
    width: 40px;
    height: 40px;
    border: 4px solid rgba(79, 70, 229, 0.2);
    border-radius: 50%;
    border-top-color: var(--primary-color);
    animation: spin 1s linear infinite;
}

@keyframes spin {
    to {
        transform: rotate(360deg);
    }
}

/* Info cards with subtle styling */
.info-card {
    background-color: var(--card-color);
    border-radius: var(--radius);
    padding: 1rem;
    border: 1px solid var(--border-color);
    margin-bottom: 1rem;
    transition: all 0.2s;
}

.info-card:hover {
    border-color: var(--border-hover);
    box-shadow: var(--shadow-sm);
}

.info-card h3 {
    margin-top: 0;
    margin-bottom: 0.5rem;
    font-size: 1rem;
    font-weight: 600;
    color: var(--primary-color);
    display: flex;
    align-items: center;
    gap: 0.5rem;
}

.info-card p {
    margin: 0;
    font-size: 0.875rem;
    color: var(--text-muted);
    line-height: 1.5;
}

/* Model info panel with branded styling */
.model-info {
    background-color: rgba(79, 70, 229, 0.05);
    border-left: 3px solid var(--primary-color);
    padding: 0.75rem 1rem;
    border-radius: 0 var(--radius) var(--radius) 0;
    margin: 1rem 0;
}

.model-info h3 {
    display: flex;
    align-items: center;
    gap: 0.5rem;
    margin-top: 0;
    margin-bottom: 0.5rem;
    font-size: 1rem;
    font-weight: 600;
    color: var(--primary-color);
}

.model-info p {
    margin: 0 0 0.5rem 0;
    font-size: 0.875rem;
    color: var(--text-color);
}

.model-info .model-id {
    font-size: 0.75rem;
    color: var(--text-muted);
    font-family: monospace;
    background-color: rgba(0, 0, 0, 0.03);
    padding: 0.25rem 0.5rem;
    border-radius: 4px;
    margin-top: 0.5rem;
    word-break: break-all;
}

/* Parameter pills for displaying selections */
.parameter-pill {
    display: inline-flex;
    align-items: center;
    background-color: rgba(79, 70, 229, 0.1);
    color: var(--primary-color);
    border-radius: 16px;
    padding: 0.25rem 0.75rem;
    margin-right: 0.5rem;
    margin-bottom: 0.5rem;
    font-size: 0.75rem;
    font-weight: 500;
    user-select: none;
}

.parameter-pill .icon {
    margin-right: 0.25rem;
}

/* Badge for showing model speed/quality */
.badge {
    display: inline-flex;
    align-items: center;
    justify-content: center;
    border-radius: 9999px;
    padding: 0.125rem 0.5rem;
    font-size: 0.75rem;
    font-weight: 500;
    margin-left: 0.5rem;
}

.badge-success {
    background-color: rgba(16, 185, 129, 0.1);
    color: var(--success-color);
}

.badge-warning {
    background-color: rgba(245, 158, 11, 0.1);
    color: var(--warning-color);
}

.badge-error {
    background-color: rgba(239, 68, 68, 0.1);
    color: var(--error-color);
}

.badge-info {
    background-color: rgba(59, 130, 246, 0.1);
    color: var(--accent-color);
}

/* Enhanced accordion styling */
.gr-accordion {
    border: 1px solid var(--border-color) !important;
    border-radius: var(--radius) !important;
    margin-bottom: 1rem !important;
    overflow: hidden;
}

.gr-accordion > div:first-child {
    padding: 0.75rem 1rem !important;
    background-color: rgba(0, 0, 0, 0.02) !important;
    font-weight: 600 !important;
    color: var(--text-color) !important;
    font-size: 0.9375rem !important;
}

.gr-accordion > div:last-child {
    padding: 1rem !important;
}

/* Tooltip styling */
.tooltip {
    position: relative;
    display: inline-block;
    cursor: help;
}

.tooltip .tooltiptext {
    visibility: hidden;
    width: 200px;
    background-color: var(--text-color);
    color: white;
    text-align: center;
    border-radius: var(--radius-sm);
    padding: 0.5rem 0.75rem;
    position: absolute;
    z-index: 1000;
    bottom: 125%;
    left: 50%;
    transform: translateX(-50%);
    opacity: 0;
    transition: opacity 0.3s;
    font-size: 0.75rem;
    box-shadow: var(--shadow);
    pointer-events: none;
}

.tooltip .tooltiptext::after {
    content: "";
    position: absolute;
    top: 100%;
    left: 50%;
    margin-left: -5px;
    border-width: 5px;
    border-style: solid;
    border-color: var(--text-color) transparent transparent transparent;
}

.tooltip:hover .tooltiptext {
    visibility: visible;
    opacity: 1;
}

/* History item styling */
.history-item {
    display: flex;
    align-items: center;
    padding: 0.75rem;
    border-radius: var(--radius);
    margin-bottom: 0.75rem;
    background-color: var(--card-color);
    border: 1px solid var(--border-color);
    cursor: pointer;
    transition: all 0.2s;
}

.history-item:hover {
    background-color: rgba(79, 70, 229, 0.05);
    transform: translateY(-1px);
}

.history-item-image {
    width: 48px;
    height: 48px;
    border-radius: var(--radius-sm);
    object-fit: cover;
    margin-right: 0.75rem;
    background-color: #f0f0f0;
    display: flex;
    align-items: center;
    justify-content: center;
    color: var(--text-muted);
    font-size: 1.25rem;
}

.history-item-content {
    flex: 1;
    overflow: hidden;
}

.history-item-title {
    margin: 0;
    font-size: 0.875rem;
    font-weight: 500;
    white-space: nowrap;
    overflow: hidden;
    text-overflow: ellipsis;
    color: var(--text-color);
}

.history-item-subtitle {
    margin: 0;
    font-size: 0.75rem;
    color: var(--text-muted);
    margin-top: 0.25rem;
}

/* Tabs styling */
.tabs {
    display: flex;
    border-bottom: 1px solid var(--border-color);
    margin-bottom: 1rem;
}

.tab {
    padding: 0.75rem 1rem;
    cursor: pointer;
    border-bottom: 2px solid transparent;
    font-weight: 500;
    font-size: 0.9375rem;
    color: var(--text-muted);
    transition: all 0.2s;
}

.tab:hover {
    color: var(--primary-color);
}

.tab.active {
    color: var(--primary-color);
    border-bottom-color: var(--primary-color);
}

/* Progress bar */
.progress-container {
    width: 100%;
    height: 8px;
    background-color: rgba(79, 70, 229, 0.1);
    border-radius: 4px;
    overflow: hidden;
    margin: 0.5rem 0;
}

.progress-bar {
    height: 100%;
    background: linear-gradient(to right, var(--primary-color), var(--secondary-color));
    width: 0%;
    transition: width 0.3s ease;
    border-radius: 4px;
}

/* Image output container */
.image-output-container {
    position: relative;
    border-radius: var(--radius);
    overflow: hidden;
    transition: all 0.2s;
    background-color: #f5f5f5;
    min-height: 300px;
}

.image-output-container img {
    width: 100%;
    display: block;
    border-radius: var(--radius);
}

.image-placeholder {
    position: absolute;
    top: 0;
    left: 0;
    width: 100%;
    height: 100%;
    display: flex;
    flex-direction: column;
    align-items: center;
    justify-content: center;
    color: var(--text-muted);
    font-size: 1rem;
    text-align: center;
    padding: 1rem;
}

.image-placeholder .icon {
    font-size: 3rem;
    margin-bottom: 1rem;
    opacity: 0.6;
}

/* Image controls overlay */
.image-controls {
    position: absolute;
    bottom: 0;
    left: 0;
    right: 0;
    padding: 0.75rem;
    background: linear-gradient(to top, rgba(0,0,0,0.7), rgba(0,0,0,0));
    display: flex;
    justify-content: flex-end;
    opacity: 0;
    transition: opacity 0.2s;
}

.image-output-container:hover .image-controls {
    opacity: 1;
}

.image-control-button {
    background-color: rgba(255, 255, 255, 0.9);
    border: none;
    border-radius: 50%;
    width: 36px;
    height: 36px;
    display: flex;
    align-items: center;
    justify-content: center;
    cursor: pointer;
    margin-left: 0.5rem;
    color: var(--text-color);
    transition: all 0.2s;
}

.image-control-button:hover {
    background-color: white;
    transform: translateY(-2px);
    box-shadow: var(--shadow);
}

/* Character counter */
.character-counter {
    text-align: right;
    font-size: 0.75rem;
    color: var(--text-muted);
    margin-top: 0.25rem;
    transition: color 0.2s;
}

.character-counter.warning {
    color: var(--warning-color);
}

.character-counter.error {
    color: var(--error-color);
}

/* Status message */
.status-message {
    padding: 0.75rem 1rem;
    border-radius: var(--radius);
    margin: 1rem 0;
    font-size: 0.875rem;
    display: flex;
    align-items: center;
}

.status-message .icon {
    margin-right: 0.75rem;
    font-size: 1.25rem;
}

.status-success {
    background-color: rgba(16, 185, 129, 0.1);
    color: var(--success-color);
    border-left: 3px solid var(--success-color);
}

.status-error {
    background-color: rgba(239, 68, 68, 0.1);
    color: var(--error-color);
    border-left: 3px solid var(--error-color);
}

.status-warning {
    background-color: rgba(245, 158, 11, 0.1);
    color: var(--warning-color);
    border-left: 3px solid var(--warning-color);
}

.status-info {
    background-color: rgba(59, 130, 246, 0.1);
    color: var(--accent-color);
    border-left: 3px solid var(--accent-color);
}

/* Responsive adjustments */
@media (max-width: 768px) {
    .example-gallery {
        grid-template-columns: repeat(auto-fill, minmax(110px, 1fr));
    }

    .gr-panel {
        padding: 0.75rem !important;
    }

    .gr-accordion > div:first-child {
        padding: 0.625rem 0.75rem !important;
    }

    .gr-accordion > div:last-child {
        padding: 0.75rem !important;
    }

    button.primary {
        padding: 0.625rem 1.25rem !important;
    }
}
"""

# =============== ENHANCED UI COMPONENTS ===============

# Function to create and show application UI
def create_ui():
    with gr.Blocks(title="", css=css) as interface:
        # Custom header with branding
        with gr.Row(elem_classes="app-header"):
            with gr.Column():
                gr.HTML("""
                    <h1>✨Memory Magic Studio</h1>
                    <p>Capture the magic of your imagination and watch it come to life.🪄</p>
                """)

        # Main content area
        with gr.Row(equal_height=False):
            # Left column - Input controls
            with gr.Column(scale=1, min_width=380):
                # Description input with character counter
                with gr.Group(elem_classes="input-group"):
                    description_input = gr.Textbox(
                        label="Describe what you want to see",
                        placeholder="Be detailed and specific about colors, composition, lighting, and subject...",
                        lines=4,
                        max_lines=8,
                        elem_id="description-input"
                    )

                    # Character counter with dynamic updates
                    char_counter = gr.HTML(
                        value="<div class='character-counter'>0 characters</div>",
                        elem_classes="char-counter-container"
                    )

                # Creation settings with enhanced dropdowns
                with gr.Group(elem_classes="settings-group"):
                    gr.HTML("<h3 style='margin-top: 0; font-size: 1rem; margin-bottom: 0.75rem;'>🛠️ Image Settings</h3>")

                    # Creation Type and Art Style in one row
                    with gr.Row():
                        # Creation type dropdown with icons
                        creation_type = gr.Dropdown(
                            choices=format_dropdown_choices(CREATION_TYPES),
                            value=f"{CREATION_TYPES['Digital Art']['icon']} Digital Art",
                            label="Creation Type",
                            elem_classes="enhanced-dropdown"
                        )

                        # Art style dropdown with icons
                        art_style = gr.Dropdown(
                            choices=format_dropdown_choices(ART_STYLES),
                            value=f"{ART_STYLES['Photorealistic']['icon']} Photorealistic",
                            label="Art Style",
                            elem_classes="enhanced-dropdown"
                        )

                    # Mood and Model in one row
                    with gr.Row():
                        # Mood dropdown with icons
                        mood_dropdown = gr.Dropdown(
                            choices=format_dropdown_choices(MOODS),
                            value=f"{MOODS['Peaceful']['icon']} Peaceful",
                            label="Mood",
                            elem_classes="enhanced-dropdown"
                        )

                        # Model selector with display names
                        formatted_models = [f"{info['icon']} {info['display_name']}" for model_key, info in IMAGE_MODELS.items()]
                        model_selector = gr.Dropdown(
                            choices=formatted_models,
                            value=f"{IMAGE_MODELS['stabilityai/stable-diffusion-xl-base-1.0']['icon']} {IMAGE_MODELS['stabilityai/stable-diffusion-xl-base-1.0']['display_name']}",
                            label="Model",
                            elem_classes="enhanced-dropdown"
                        )

                # Examples gallery with clear visual structure
                with gr.Group(elem_classes="examples-group"):
                    gr.HTML("<h3 style='margin-top: 0; font-size: 1rem; margin-bottom: 0.75rem;'>🌟 Try an example:</h3>")

                    # Gallery of examples
                    with gr.Row(elem_classes="example-gallery"):
                        for i, example in enumerate(EXAMPLE_PROMPTS):
                            with gr.Column(elem_classes="example-item"):
                                # Example card with visual element and caption
                                example_card = gr.Button(
                                    example["thumbnail_desc"],
                                    elem_classes="example-button"
                                )

                                # Event handler for example selection
                                example_card.click(
                                    fn=lambda idx=i: load_example(idx),
                                    outputs=[description_input, creation_type, art_style, mood_dropdown]
                                )

                # Information panels for selected options
                with gr.Group(elem_classes="info-panels"):
                    # Creation type info
                    creation_info = gr.HTML(value="", elem_classes="option-info")

                    # Art style info
                    art_info = gr.HTML(value="", elem_classes="option-info")

                    # Model info panel
                    model_info = gr.HTML(value="", elem_classes="option-info")

                # Generate button with clear call to action
                with gr.Group(elem_classes="action-group"):
                    # Generate button with progress indicator
                    generate_button = gr.Button(
                        "✨ Generate Image",
                        variant="primary",
                        elem_classes="primary",
                        elem_id="generate-button"
                    )

                    # Generation status message
                    generation_status = gr.HTML(value="", elem_classes="generation-status")

                # Tips section in collapsible accordion
                with gr.Accordion("📝 Tips for better results", open=True):
                    gr.HTML("""
                    <div class="tips-container">
                        <h3>Tips for better results:</h3>
                        <ul>
                            <li><strong>Be specific and detailed</strong> - Include information about subjects, environment, lighting, colors, perspective, time of day, etc.</li>
                            <li><strong>Use descriptive adjectives</strong> - Words like "vibrant", "moody", "ethereal", "weathered" help set the tone.</li>
                            <li><strong>Experiment with art styles</strong> - Different styles can dramatically change the look and feel.</li>
                            <li><strong>Combine with the right model</strong> - SDXL provides the highest quality but takes longer.</li>
                            <li><strong>Think about composition</strong> - Mention foreground/background elements and their relationships.</li>
                        </ul>
                    </div>
                    """)

            # Right column - Output display
            with gr.Column(scale=1, min_width=480):
                # Image display area with placeholder
                with gr.Group(elem_classes="output-container"):
                    # Output image with interactive elements
                    image_output = gr.Image(
                        label="Generated Image",
                        elem_id="image-output",
                        type="pil",
                        height=512
                    )

                # Image generation details
                with gr.Accordion("Image Details", open=True):
                    parameters_display = gr.HTML(value="")

                # Enhanced prompt display
                with gr.Accordion("Enhanced Prompt", open=False):
                    prompt_output = gr.Textbox(
                        label="AI-Enhanced Prompt Used",
                        lines=5,
                        elem_id="prompt-output",
                        elem_classes="prompt-display"
                    )

                # Technical details for advanced users
                with gr.Accordion("Technical Details", open=False):
                    technical_info = gr.HTML("""
                    <div class="technical-info">
                        <h4>How Image Generation Works</h4>
                        <p>Images are generated using Stable Diffusion, a latent text-to-image diffusion model. Your text prompt is:
                        <ol>
                            <li>Enhanced with AI to add descriptive details and quality terms</li>
                            <li>Processed through a neural network that gradually transforms random noise into an image</li>
                            <li>Refined based on the parameters you select (model, style, mood)</li>
                        </ol>
                        </p>
                        <p>Different models have different strengths:
                        <ul>
                            <li><strong>SDXL 1.0</strong>: Highest quality, best composition and details</li>
                            <li><strong>SD 1.5</strong>: Faster generation, good for general purpose</li>
                            <li><strong>SD 2.1</strong>: Better with human faces, improved consistency</li>
                            <li><strong>OpenJourney</strong>: Midjourney-like stylized artistic results</li>
                            <li><strong>Dreamlike</strong>: Dreamy, ethereal aesthetic with artistic flair</li>
                        </ul>
                        </p>
                    </div>
                    """)

        # Set up event handlers within the Blocks context
        description_input.change(
            fn=update_char_count,
            inputs=description_input,
            outputs=char_counter
        )

        creation_type.change(
            fn=update_creation_info,
            inputs=creation_type,
            outputs=creation_info
        )

        art_style.change(
            fn=update_art_style_info,
            inputs=art_style,
            outputs=art_info
        )

        model_selector.change(
            fn=update_model_info,
            inputs=model_selector,
            outputs=model_info
        )

        generate_button.click(
            fn=generate_with_status,
            inputs=[
                description_input,
                creation_type,
                art_style,
                mood_dropdown,
                model_selector
            ],
            outputs=[
                image_output,
                generation_status,
                prompt_output,
                parameters_display
            ]
        )

        # Load default values on page load
        interface.load(
            fn=lambda: (
                update_creation_info(f"{CREATION_TYPES['Digital Art']['icon']} Digital Art"),
                update_art_style_info(f"{ART_STYLES['Photorealistic']['icon']} Photorealistic"),
                update_model_info(f"{IMAGE_MODELS['stabilityai/stable-diffusion-xl-base-1.0']['icon']} {IMAGE_MODELS['stabilityai/stable-diffusion-xl-base-1.0']['display_name']}")
            ),
            outputs=[creation_info, art_info, model_info]
        )

        return interface, description_input, creation_type, art_style, mood_dropdown, model_selector, generate_button, image_output, generation_status, prompt_output, parameters_display, char_counter, creation_info, art_info, model_info

# Helper function to update character count with color coding
def update_char_count(text):
    count = len(text)
    if count == 0:
        color_class = ""
    elif count < 20:
        color_class = "warning"
    elif count > 300:
        color_class = "warning"
    elif count > 500:
        color_class = "error"
    else:
        color_class = ""

    return f"<div class='character-counter {color_class}'>{count} characters</div>"

# Helper function to update creation type info
def update_creation_info(choice):
    key = extract_key(choice)
    if key in CREATION_TYPES:
        info = CREATION_TYPES[key]
        return f"""<div class="info-card">
            <h3>{info['icon']} {key}</h3>
            <p>{info['description']}</p>
            <p style="margin-top: 0.5rem; font-style: italic; opacity: 0.8;">💡 {info['prompt_hint']}</p>
        </div>"""
    return ""

# Helper function to update art style info
def update_art_style_info(choice):
    key = extract_key(choice)
    if key in ART_STYLES:
        info = ART_STYLES[key]
        return f"""<div class="info-card">
            <h3>{info['icon']} {key}</h3>
            <p>{info['description']}</p>
            <p style="margin-top: 0.5rem; font-style: italic; opacity: 0.8;">🎨 Examples: {info['examples']}</p>
        </div>"""
    return ""

# Helper function to update model info
def update_model_info(formatted_choice):
    # Extract display name without the icon
    if ' ' in formatted_choice:
        display_name = formatted_choice.split(' ', 1)[1]
        # Find the corresponding key and info
        for key, info in IMAGE_MODELS.items():
            if info['display_name'] == display_name:
                # Create speed badge
                speed_badge = ""
                if info.get('speed') == 'fast':
                    speed_badge = '<span class="badge badge-success">Fast</span>'
                elif info.get('speed') == 'medium':
                    speed_badge = '<span class="badge badge-warning">Medium</span>'
                elif info.get('speed') == 'slow':
                    speed_badge = '<span class="badge badge-error">Slower</span>'

                # Create quality badge
                quality_badge = ""
                if info.get('quality') == 'excellent':
                    quality_badge = '<span class="badge badge-success">Excellent Quality</span>'
                elif info.get('quality') == 'very good':
                    quality_badge = '<span class="badge badge-success">Very Good Quality</span>'
                elif info.get('quality') == 'good':
                    quality_badge = '<span class="badge badge-info">Good Quality</span>'
                elif info.get('quality') == 'stylized':
                    quality_badge = '<span class="badge badge-info">Stylized</span>'
                elif info.get('quality') == 'artistic':
                    quality_badge = '<span class="badge badge-info">Artistic</span>'

                return f"""<div class="model-info">
                    <h3>{info['icon']} {info['display_name']} {speed_badge} {quality_badge}</h3>
                    <p>{info['description']}</p>
                    <div class="model-id">{key}</div>
                </div>"""
    return ""

# Helper function to update status message
def update_status(message, is_error=False, is_warning=False, is_info=False):
    if is_error:
        status_class = "status-error"
        icon = "❌"
    elif is_warning:
        status_class = "status-warning"
        icon = "⚠️"
    elif is_info:
        status_class = "status-info"
        icon = "ℹ️"
    else:
        status_class = "status-success"
        icon = "✅"

    return f"""<div class="status-message {status_class}">
        <span class="icon">{icon}</span>
        <span>{message}</span>
    </div>"""

# Helper function to format parameters display as pills
def format_parameters(creation_type_val, art_style_val, mood_val, model_name):
    creation_key = extract_key(creation_type_val)
    art_key = extract_key(art_style_val)
    mood_key = extract_key(mood_val)

    # Get model info
    model_display_name = "Unknown Model"
    model_id = ""
    model_icon = "🤖"
    for key, info in IMAGE_MODELS.items():
        if f"{info['icon']} {info['display_name']}" == model_name:
            model_display_name = info['display_name']
            model_id = key
            model_icon = info['icon']
            break

    html = """<div style="margin-bottom: 1rem;">
        <div style="font-weight: 500; margin-bottom: 0.75rem;">Generated with these parameters:</div>
        <div style="display: flex; flex-wrap: wrap; gap: 0.5rem;">"""

    # Add creation type pill
    if creation_key in CREATION_TYPES:
        html += f"""<div class="parameter-pill">
            <span class="icon">{CREATION_TYPES[creation_key]['icon']}</span> {creation_key}
        </div>"""

    # Add art style pill
    if art_key in ART_STYLES:
        html += f"""<div class="parameter-pill">
            <span class="icon">{ART_STYLES[art_key]['icon']}</span> {art_key}
        </div>"""

    # Add mood pill
    if mood_key in MOODS:
        html += f"""<div class="parameter-pill">
            <span class="icon">{MOODS[mood_key]['icon']}</span> {mood_key}
        </div>"""

    # Add model pill
    html += f"""<div class="parameter-pill">
        <span class="icon">{model_icon}</span> {model_display_name}
    </div>"""

    # Close container
    html += """</div>
        <div style="margin-top: 1rem; font-size: 0.75rem; color: var(--text-muted);">
            Image generated on {timestamp}
        </div>
    </div>""".replace("{timestamp}", time.strftime("%Y-%m-%d at %H:%M:%S"))

    return html

# AI Image Creator: Enhanced UI and UX
# Part 3: Processing Logic, Prompt Enhancement, and Application Flow

# =============== PROMPT ENHANCEMENT LOGIC ===============

# Function to enhance prompt with Llama 4 with improved logical understanding
def enhance_prompt_with_llama(user_input, creation_type, art_style, mood):
    """
    Enhance user input with Llama 4 model to create detailed image generation prompts

    Args:
        user_input (str): User's original description
        creation_type (str): Selected creation type (e.g., "Digital Art")
        art_style (str): Selected art style (e.g., "Photorealistic")
        mood (str): Selected mood (e.g., "Peaceful")

    Returns:
        str: Enhanced prompt optimized for image generation
    """
    try:
        if not use_llama or llama_client is None:
            logger.warning("Llama enhancement not available, using fallback")
            return enhance_prompt_fallback(user_input, creation_type, art_style, mood)

        logger.info(f"Enhancing prompt with Llama 4 for creation type: {creation_type}, art style: {art_style}")

        # Enhanced Llama 4 system prompt with detailed instructions
        system_prompt = """You are a world-class prompt engineer who specializes in creating detailed, effective prompts for text-to-image AI models.

Your task is to transform a user's simple description into a comprehensive, detailed image generation prompt that will create stunning visuals. Consider all the provided elements (description, creation type, art style, mood) and combine them into a cohesive, detailed prompt.

MOST IMPORTANTLY - ADD LOGICAL DETAILS:
- Analyze what the user wants and add logical details that would make the scene realistic or coherent
- If describing something fantastical (e.g., "flying cat"), add logical details about how this could work (e.g., "a cat with majestic feathered wings spread wide")
- Think about environment, lighting, perspective, time of day, weather, and other contextual elements
- Create a vivid, imaginable scene with spatial relationships clearly defined

PROMPT STRUCTURE GUIDELINES:
1. Start with the core subject and its primary characteristics
2. Add environment and setting details
3. Describe lighting, atmosphere, and mood
4. Include specific visual style and artistic technique references
5. Add technical quality terms (8K, detailed, masterful, etc.)

FORMAT YOUR RESPONSE AS A SINGLE PARAGRAPH with no additional comments, explanations, or bullet points. Use natural language without awkward comma separations. Aim for 75-150 words.

AVOID:
- Do not include quotation marks in your response
- Do not preface with "here's a prompt" or similar text
- Do not use placeholders
- Do not add negative prompts
- Do not write in list format or use bullet points

Respond only with the enhanced prompt and nothing else."""

        # Get creation type description
        creation_info = CREATION_TYPES.get(creation_type, {"description": "Create a detailed image", "icon": "🎨"})
        creation_description = creation_info["description"]

        # Get art style description
        style_info = ART_STYLES.get(art_style, {"description": "with detailed and professional quality", "icon": "🖌️"})
        style_description = style_info["description"]

        # Get mood description
        mood_info = MOODS.get(mood, {"description": "atmospheric", "icon": "✨"})
        mood_description = mood_info["description"]

        # Prepare the user prompt for Llama
        user_prompt = f"""Description: {user_input}
Creation Type: {creation_type} - {creation_description}
Art Style: {art_style} - {style_description}
Mood: {mood} - {mood_description}

Please create a comprehensive, detailed image generation prompt that combines all these elements."""

        try:
            # Request enhancement from Llama 4
            completion = llama_client.chat.completions.create(
                model="meta-llama/Llama-4-Scout-17B-16E-Instruct",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                max_tokens=500,
                temperature=0.7,  # Slight creativity while maintaining coherence
            )
            enhanced = completion.choices[0].message.content
            logger.info(f"Llama 4 enhanced prompt: {enhanced[:100]}...")
            return enhanced if enhanced else user_input
        except Exception as e:
            logger.error(f"Error during Llama enhancement: {str(e)}")
            return enhance_prompt_fallback(user_input, creation_type, art_style, mood)
    except Exception as e:
        logger.error(f"Error in Llama enhancement: {str(e)}")
        return enhance_prompt_fallback(user_input, creation_type, art_style, mood)

# Fallback prompt enhancement without Llama
def enhance_prompt_fallback(user_input, creation_type, art_style, mood):
    """
    Enhance user input without requiring Llama API using rule-based enhancement

    Args:
        user_input (str): User's original description
        creation_type (str): Selected creation type (e.g., "Digital Art")
        art_style (str): Selected art style (e.g., "Photorealistic")
        mood (str): Selected mood (e.g., "Peaceful")

    Returns:
        str: Enhanced prompt using predefined rules and templates
    """
    logger.info(f"Using fallback enhancement for: {user_input[:50]}...")

    # Quality terms by creation type
    quality_terms = {
        "Realistic Photo": [
            "photorealistic", "high resolution", "detailed",
            "natural lighting", "sharp focus", "professional photography",
            "crisp details", "realistic textures", "DSLR photo"
        ],
        "Digital Art": [
            "vibrant colors", "clean lines", "digital illustration",
            "polished", "professional digital art", "detailed rendering",
            "digital painting", "colorful", "vector-like precision"
        ],
        "Fantasy Illustration": [
            "magical atmosphere", "fantasy art", "detailed illustration",
            "epic", "otherworldly", "imaginative scene",
            "fantasy environment", "magical lighting", "mythical qualities"
        ],
        "Concept Art": [
            "professional concept art", "detailed design", "conceptual illustration",
            "industry standard", "visual development", "production artwork",
            "concept design", "detailed environment", "character design"
        ],
        "Anime/Manga": [
            "anime style", "manga illustration", "cel shaded",
            "Japanese animation", "2D character art", "anime aesthetic",
            "clean linework", "anime proportions", "stylized features"
        ],
        "Oil Painting": [
            "oil on canvas", "textured brushwork", "rich colors",
            "traditional painting", "artistic brushstrokes", "gallery quality",
            "glazed layers", "impasto technique", "classical painting style"
        ],
        "Watercolor": [
            "watercolor painting", "soft color bleeding", "delicate washes",
            "transparent layers", "loose brushwork", "gentle transitions",
            "watercolor paper texture", "wet-on-wet technique", "fluid color blending"
        ],
        "Sketch": [
            "detailed sketch", "pencil drawing", "line art",
            "hand-drawn", "fine details", "shading techniques",
            "graphite", "charcoal texture", "gestural lines"
        ],
        "3D Rendering": [
            "3D render", "volumetric lighting", "ray tracing",
            "3D modeling", "realistic textures", "computer graphics",
            "physically based rendering", "global illumination", "ambient occlusion"
        ],
        "Pixel Art": [
            "pixel art", "8-bit style", "retro game aesthetic",
            "limited color palette", "pixelated", "nostalgic game art",
            "16-bit look", "pixel perfect", "dithering effects"
        ]
    }

    # Style modifiers for different art styles - more detailed descriptions
    style_modifiers = {
        "Photorealistic": "highly detailed photorealistic style with perfect lighting, natural shadows, and lifelike textures",
        "Impressionist": "impressionist style with visible brushstrokes capturing light and atmosphere over precise details, reminiscent of Claude Monet",
        "Surrealist": "surrealist style with dreamlike and impossible elements, juxtaposed reality, inspired by Salvador Dali",
        "Pop Art": "pop art style with bold colors, sharp lines, halftone patterns and cultural references, like Andy Warhol",
        "Minimalist": "minimalist style with simplified forms, limited color palette, clean composition, and essential elements only",
        "Abstract": "abstract style using non-representational shapes, colors, and forms to express emotion rather than reality",
        "Cubist": "cubist style with geometric forms, multiple perspectives shown simultaneously, fractured surfaces, like Pablo Picasso",
        "Art Nouveau": "art nouveau style with ornate flowing lines inspired by natural forms, decorative elegance, and organic shapes",
        "Gothic": "gothic style with dark atmosphere, dramatic elements, pointed arches, and medieval-inspired architecture",
        "Cyberpunk": "cyberpunk style with neon colors, high tech low life aesthetic, futuristic technology, and urban decay",
        "Steampunk": "steampunk style with Victorian aesthetics, brass machinery, steam-powered technology, and retrofuturistic design",
        "Retro/Vintage": "retro style with nostalgic elements from past decades, aged texture, and period-appropriate colors and design",
        "Art Deco": "art deco style with geometric patterns, bold colors, symmetry, luxurious materials, and streamlined forms",
        "Baroque": "baroque style with dramatic lighting, rich details, contrast, dynamic composition, and ornate decorations",
        "Ukiyo-e": "ukiyo-e style Japanese woodblock print aesthetic with flat areas of color, strong outlines, and traditional subjects",
        "Comic Book": "comic book style with bold outlines, vibrant colors, dynamic action poses, and expressive characters",
        "Psychedelic": "psychedelic style with vibrant swirling colors, abstract patterns, distorted perspective, and 1960s-inspired visuals",
        "Vaporwave": "vaporwave aesthetic with glitch art, pastel colors, 80s/90s nostalgia, ancient statues, and digital elements",
        "Studio Ghibli": "Studio Ghibli anime style with whimsical detailed environments, soft colors, and charming character design",
        "Hyperrealism": "hyperrealistic style with extreme detail beyond photography, perfect textures, and meticulous precision"
    }

    # Mood modifiers for different moods - enhanced descriptions
    mood_modifiers = {
        "Happy": "bright cheerful atmosphere with warm golden lighting, vibrant colors, and uplifting elements",
        "Sad": "melancholic atmosphere with muted colors, soft shadows, and somber emotional tone",
        "Mysterious": "enigmatic atmosphere with shadows, fog, hidden elements, and dramatic lighting contrasts",
        "Peaceful": "serene calm atmosphere with gentle lighting, soft colors, and tranquil composition",
        "Tense": "suspenseful atmosphere with dramatic lighting, stark contrasts, and unsettling composition",
        "Whimsical": "playful whimsical atmosphere with imaginative elements, saturated colors, and fantastical details",
        "Dark": "dark gloomy atmosphere with deep shadows, limited lighting, and ominous elements",
        "Energetic": "dynamic vibrant atmosphere with strong colors, motion effects, and active composition",
        "Romantic": "soft romantic atmosphere with dreamy lighting, gentle colors, and intimate ambiance",
        "Epic": "grand epic atmosphere with dramatic scale, sweeping vistas, and majestic lighting"
    }

    # Get terms for the specific creation type, or use generic terms
    type_terms = quality_terms.get(creation_type, [
        "high quality", "detailed", "professional", "masterful", "high resolution", "sharp details"
    ])

    # Common quality terms enhanced with trending and technical terms
    common_terms = [
        "8K resolution", "highly detailed", "professional", "masterpiece",
        "trending on artstation", "award winning", "stunning", "intricate details",
        "perfect composition", "cinematic lighting"
    ]

    # Get style modifier
    style_modifier = style_modifiers.get(art_style, "detailed professional style")

    # Get mood modifier
    mood_modifier = mood_modifiers.get(mood, "atmospheric")

    # Basic prompt structure - core subject and style elements
    prompt_parts = [
        user_input,
        style_modifier,
        mood_modifier
    ]

    # Add randomly selected quality terms for variety
    selected_type_terms = random.sample(type_terms, min(3, len(type_terms)))
    selected_common_terms = random.sample(common_terms, min(3, len(common_terms)))

    # Combine terms
    quality_description = ", ".join(selected_type_terms + selected_common_terms)

    # Final enhanced prompt
    enhanced_prompt = f"{', '.join(prompt_parts)}, {quality_description}"

    logger.info(f"Fallback enhanced prompt: {enhanced_prompt[:100]}...")
    return enhanced_prompt

# =============== IMAGE GENERATION FUNCTIONS ===============

# Generate image function with loading state handling and retry mechanism
def generate_image(description, creation_type, art_style, mood, model_name, retries=1):
    """
    Generate image based on user inputs by enhancing prompt and calling image model API

    Args:
        description (str): User's original description
        creation_type (str): Selected creation type
        art_style (str): Selected art style
        mood (str): Selected mood
        model_name (str): Model identifier
        retries (int): Number of retries if generation fails

    Returns:
        tuple: (image, status_message, enhanced_prompt)
    """
    try:
        # Validate input
        if not description.strip():
            return None, "Please enter a description for your image", ""

        logger.info(f"Generating image with model: {model_name}")

        # Enhance prompt with Llama or fallback
        enhanced_prompt = enhance_prompt_with_llama(description, creation_type, art_style, mood)

        # Validate client availability
        if hf_client is None:
            logger.error("Hugging Face client not available")
            return None, "Error: Unable to connect to image generation service. Please try again later.", enhanced_prompt

        # Add negative prompt to avoid common issues
        negative_prompt = "low quality, blurry, distorted, deformed, disfigured, bad anatomy, watermark, signature, text, poorly drawn, amateur, ugly"

        try:
            # Generate image with progress tracking
            logger.info(f"Sending request to model {model_name} with prompt: {enhanced_prompt[:100]}...")

            # Log start time for performance tracking
            start_time = time.time()

            # Generate the image
            image = hf_client.text_to_image(
                prompt=enhanced_prompt,
                model=model_name,
                negative_prompt=negative_prompt
            )

            # Calculate generation time
            generation_time = time.time() - start_time
            logger.info(f"Image generated successfully in {generation_time:.2f} seconds")

            # Success message with generation details
            if use_llama:
                enhancement_method = "Llama 4 AI"
            else:
                enhancement_method = "rule-based enhancement"

            success_message = f"Image created successfully in {generation_time:.1f}s using {model_name.split('/')[-1]} model and {enhancement_method}"

            return image, success_message, enhanced_prompt

        except Exception as e:
            error_message = str(e)
            logger.error(f"Error during image generation: {error_message}")

            # Retry logic for transient errors
            if retries > 0:
                logger.info(f"Retrying image generation, {retries} attempts remaining")
                time.sleep(1)  # Small delay before retry
                return generate_image(description, creation_type, art_style, mood, model_name, retries - 1)

            # Format user-friendly error message
            if "429" in error_message:
                friendly_error = "Server is currently busy. Please try again in a few moments."
            elif "401" in error_message or "403" in error_message:
                friendly_error = "Authentication error with the image service. Please check API settings."
            elif "timeout" in error_message.lower():
                friendly_error = "Request timed out. The server might be under heavy load."
            else:
                friendly_error = f"Error generating image: {error_message}"

            return None, friendly_error, enhanced_prompt
    except Exception as e:
        logger.error(f"Unexpected error in generate_image: {str(e)}")
        return None, f"Unexpected error: {str(e)}", ""

# Wrapper function for generate_image with status updates
def generate_with_status(description, creation_type_val, art_style_val, mood_val, model_name):
    """
    Wrapper for generate_image that handles UI status updates and parameter formatting

    Args:
        description (str): User's original description
        creation_type_val (str): Formatted creation type with icon
        art_style_val (str): Formatted art style with icon
        mood_val (str): Formatted mood with icon
        model_name (str): Formatted model name with icon

    Returns:
        tuple: (image, status_html, enhanced_prompt, parameters_html)
    """
    # Check if description is empty
    if not description.strip():
        return None, update_status("Please enter a description", is_error=True), "", ""

    # Extract keys from formatted values
    creation_key = extract_key(creation_type_val)
    art_key = extract_key(art_style_val)
    mood_key = extract_key(mood_val)

    # Get model key from formatted name
    model_key = None
    for key, info in IMAGE_MODELS.items():
        if f"{info['icon']} {info['display_name']}" == model_name:
            model_key = key
            break

    if not model_key:
        return None, update_status("Invalid model selection", is_error=True), "", ""

    try:
        # Generate the image
        image, message, enhanced_prompt = generate_image(
            description, creation_key, art_key, mood_key, model_key
        )

        if image is None:
            return None, update_status(message, is_error=True), "", ""

        # Format parameters display
        params_html = format_parameters(creation_type_val, art_style_val, mood_val, model_name)

        # Success message
        success_message = update_status(message)
        return image, success_message, enhanced_prompt, params_html

    except Exception as e:
        error_message = str(e)
        logger.error(f"Error in generate_with_status: {error_message}")
        return None, update_status(f"Error: {error_message}", is_error=True), "", ""

# =============== MAIN APPLICATION FLOW ===============

def main():
    """
    Main application entry point - creates UI and sets up event handlers
    """
    # Create the UI components - event handlers are now defined inside create_ui
    interface, *_ = create_ui()

    # Launch the interface with appropriate parameters for Gradio version
    # Check Gradio version to decide on the correct parameters to use
    try:
        # Use simple parameters that work across versions
        interface.launch(
            share=False,  # Set to True to create a public link
            debug=False   # Set to True for development
        )
    except Exception as e:
        logger.error(f"Error launching Gradio interface: {str(e)}")
        # Fallback to the most basic launch parameters
        interface.launch()

# =============== APP EXECUTION ===============

if __name__ == "__main__":
    # Start the application
    main()