app.py

import gradio as gr
import requests
from PIL import Image
from io import BytesIO
import os
from datetime import datetime

# Get token from environment variable or use hardcoded token
HF_API_TOKEN = os.getenv("HF_API_KEY", "")

# Image generation models
MODELS = {
    "Stable Diffusion XL": "stabilityai/stable-diffusion-xl-base-1.0",
    "Stable Diffusion 2.1": "stabilityai/stable-diffusion-2-1",
    "Realistic Vision": "SG161222/Realistic_Vision_V2.0",
    "Dreamlike Photoreal": "dreamlike-art/dreamlike-photoreal-2.0",
}

# Create output directory
OUTPUT_DIR = "generated_images"
os.makedirs(OUTPUT_DIR, exist_ok=True)

headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}

def generate_image(prompt, model_name, negative_prompt="", progress=gr.Progress()):
    """Generate image using Hugging Face API"""
    if not prompt:
        return None, "⚠️ Please enter a prompt!", None
    
    if not HF_API_TOKEN:
        return None, "❌ Error: HF_TOKEN not found. Please set it in environment.", None
    
    progress(0.1, desc="🎨 Initializing image generation...")
    
    model_url = f"https://api-inference.huggingface.co/models/{MODELS[model_name]}"
    
    # Prepare payload
    payload = {
        "inputs": prompt,
    }
    
    if negative_prompt:
        payload["parameters"] = {
            "negative_prompt": negative_prompt
        }
    
    try:
        progress(0.3, desc="🖌️ Generating your image...")
        
        response = requests.post(model_url, headers=headers, json=payload, timeout=60)
        
        if response.status_code == 200:
            progress(0.8, desc="✨ Finalizing image...")
            
            image = Image.open(BytesIO(response.content))
            
            # Save image
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"generated_{timestamp}.png"
            filepath = os.path.join(OUTPUT_DIR, filename)
            image.save(filepath)
            
            progress(1.0, desc="✅ Image generated successfully!")
            
            success_msg = f"""✅ **Image Generated Successfully!**

📝 Prompt: {prompt[:100]}{'...' if len(prompt) > 100 else ''}
🤖 Model: {model_name}
💾 Saved as: {filename}
📁 Location: {OUTPUT_DIR}/

Ready to download!"""
            
            return image, success_msg, filepath
        
        elif response.status_code == 503:
            return None, "⏳ Model is loading... Please wait 20 seconds and try again.", None
        else:
            return None, f"❌ Error {response.status_code}: {response.text}", None
            
    except Exception as e:
        return None, f"❌ Error: {str(e)}", None

def generate_random_prompt():
    """Generate a random creative prompt"""
    import random
    
    subjects = [
        "a majestic dragon",
        "a futuristic city",
        "a magical forest",
        "an astronaut",
        "a cyberpunk samurai",
        "a floating island",
        "a steampunk airship",
        "a phoenix rising",
        "a crystal palace",
        "a mystical portal"
    ]
    
    settings = [
        "at sunset with golden hour lighting",
        "in a neon-lit cyberpunk environment",
        "under the northern lights",
        "in a misty morning atmosphere",
        "with dramatic cinematic lighting",
        "in a fantasy landscape",
        "at night with moonlight",
        "in an ethereal dreamscape",
        "with volumetric fog",
        "in a surreal dimension"
    ]
    
    styles = [
        "ultra detailed, 8K resolution",
        "digital art masterpiece",
        "photorealistic, highly detailed",
        "concept art, trending on artstation",
        "cinematic composition",
        "hyper realistic",
        "fantasy art style",
        "vibrant colors, high contrast",
        "professional photography",
        "artistic illustration"
    ]
    
    subject = random.choice(subjects)
    setting = random.choice(settings)
    style = random.choice(styles)
    
    return f"{subject} {setting}, {style}"

# Custom CSS for modern, responsive design
custom_css = """
/* Global Styles */
.gradio-container {
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important;
    max-width: 1400px !important;
    margin: 0 auto !important;
}

/* Header Styling */
.header-container {
    text-align: center;
    padding: 2rem 1rem;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    border-radius: 16px;
    margin-bottom: 2rem;
    box-shadow: 0 10px 40px rgba(102, 126, 234, 0.3);
}

.header-container h1 {
    color: white !important;
    font-size: 3rem !important;
    font-weight: 800 !important;
    margin-bottom: 0.5rem !important;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.2);
}

.header-container p {
    color: rgba(255, 255, 255, 0.95) !important;
    font-size: 1.2rem !important;
    font-weight: 400 !important;
}

/* Input Section */
.input-section {
    background: white;
    padding: 2rem;
    border-radius: 16px;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.08);
    margin-bottom: 1rem;
}

/* Buttons */
.generate-btn {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    color: white !important;
    font-size: 1.1rem !important;
    font-weight: 600 !important;
    padding: 1rem 2rem !important;
    border-radius: 12px !important;
    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.4) !important;
    transition: all 0.3s ease !important;
}

.generate-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(102, 126, 234, 0.6) !important;
}

.random-btn {
    background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%) !important;
    border: none !important;
    color: white !important;
    font-weight: 600 !important;
    border-radius: 12px !important;
}

.clear-btn {
    background: linear-gradient(135deg, #4facfe 0%, #00f2fe 100%) !important;
    border: none !important;
    color: white !important;
    font-weight: 600 !important;
    border-radius: 12px !important;
}

/* Image Gallery */
.image-gallery {
    background: white;
    padding: 2rem;
    border-radius: 16px;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.08);
}

/* Info Cards */
.info-card {
    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
    padding: 1.5rem;
    border-radius: 12px;
    margin: 1rem 0;
    border-left: 4px solid #667eea;
}

/* Status Messages */
.status-success {
    color: #10b981 !important;
    font-weight: 600 !important;
}

.status-error {
    color: #ef4444 !important;
    font-weight: 600 !important;
}

/* Responsive Design */
@media (max-width: 768px) {
    .header-container h1 {
        font-size: 2rem !important;
    }
    
    .header-container p {
        font-size: 1rem !important;
    }
    
    .input-section, .image-gallery {
        padding: 1rem;
    }
    
    .generate-btn {
        font-size: 1rem !important;
        padding: 0.8rem 1.5rem !important;
    }
}

/* Textbox Styling */
.gr-textbox {
    border-radius: 12px !important;
    border: 2px solid #e5e7eb !important;
    transition: all 0.3s ease !important;
}

.gr-textbox:focus {
    border-color: #667eea !important;
    box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
}

/* Dropdown Styling */
.gr-dropdown {
    border-radius: 12px !important;
    border: 2px solid #e5e7eb !important;
}

/* Image Container */
.gr-image {
    border-radius: 12px !important;
    overflow: hidden !important;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.1) !important;
}

/* Examples Grid */
.gr-examples {
    gap: 1rem !important;
}

.gr-example {
    border-radius: 8px !important;
    transition: all 0.3s ease !important;
}

.gr-example:hover {
    transform: scale(1.02) !important;
    box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1) !important;
}

/* Footer */
.footer {
    text-align: center;
    padding: 2rem;
    color: #6b7280;
    font-size: 0.9rem;
}

/* Animations */
@keyframes fadeIn {
    from { opacity: 0; transform: translateY(20px); }
    to { opacity: 1; transform: translateY(0); }
}

.input-section, .image-gallery {
    animation: fadeIn 0.6s ease-out;
}
"""

# Create Gradio Interface
with gr.Blocks(css=custom_css, theme=gr.themes.Soft(), title="AI Image Generator") as demo:
    
    # Header
    with gr.Row(elem_classes="header-container"):
        gr.HTML("""
            <div>
                <h1>🎨 AI Image Generator</h1>
                <p>Transform your imagination into stunning visuals with AI</p>
            </div>
        """)
    
    with gr.Row():
        # Left Column - Input Section
        with gr.Column(scale=1, elem_classes="input-section"):
            gr.Markdown("## ✨ Create Your Image")
            
            prompt_input = gr.Textbox(
                label="🖊️ Describe Your Image",
                placeholder="e.g., A futuristic city at sunset, ultra detailed, cinematic lighting, 8K...",
                lines=4,
                max_lines=6,
                info="Be creative and detailed for best results!"
            )
            
            with gr.Row():
                random_btn = gr.Button("🎲 Random Prompt", size="sm", elem_classes="random-btn")
                clear_btn = gr.Button("🗑️ Clear", size="sm", elem_classes="clear-btn")
            
            model_dropdown = gr.Dropdown(
                choices=list(MODELS.keys()),
                value="Stable Diffusion XL",
                label="🤖 AI Model",
                info="Choose your preferred AI model"
            )
            
            negative_prompt_input = gr.Textbox(
                label="🚫 Negative Prompt (Optional)",
                placeholder="e.g., blurry, low quality, distorted...",
                lines=2,
                info="Describe what you DON'T want in the image"
            )
            
            with gr.Accordion("💡 Pro Tips & Examples", open=False):
                gr.Markdown("""
                ### 🎯 Tips for Better Results:
                - **Be specific**: Include details about style, lighting, mood
                - **Use quality tags**: "8K", "ultra detailed", "professional"
                - **Describe composition**: "centered", "close-up", "wide angle"
                - **Add artistic style**: "digital art", "oil painting", "photorealistic"
                
                ### 🌟 Example Prompts:
                """)
                
                gr.Examples(
                    examples=[
                        ["A majestic dragon flying over a medieval castle at sunset, fantasy art, ultra detailed, epic composition, 8K"],
                        ["Cyberpunk city street with neon lights, rain-soaked, futuristic, blade runner style, cinematic"],
                        ["A serene Japanese garden with cherry blossoms, spring morning, soft lighting, peaceful atmosphere"],
                        ["Portrait of a steampunk explorer, brass goggles, leather jacket, Victorian era, professional photography"],
                        ["Underwater coral reef scene, tropical fish, sun rays, vibrant colors, nature photography, crystal clear"],
                        ["Magical library with floating books, ethereal lighting, fantasy interior, highly detailed, concept art"],
                        ["Space station orbiting Earth, sci-fi, realistic, stars in background, NASA style, 4K"],
                        ["Ancient temple ruins in jungle, overgrown with vines, golden hour, adventure, cinematic lighting"]
                    ],
                    inputs=prompt_input,
                    label="Click to use these examples"
                )
            
            generate_btn = gr.Button("✨ Generate Image", variant="primary", size="lg", elem_classes="generate-btn")
            
            status_text = gr.Markdown("Ready to create amazing images! 🚀")
        
        # Right Column - Output Section
        with gr.Column(scale=1, elem_classes="image-gallery"):
            gr.Markdown("## 🖼️ Your Generated Image")
            
            image_output = gr.Image(
                label="Generated Artwork",
                type="pil",
                height=500,
                show_label=False
            )
            
            file_output = gr.File(
                label="📥 Download Image",
                type="filepath",
                visible=True
            )
            
            with gr.Accordion("📊 Generation Info", open=False):
                gr.Markdown("""
                Your images are automatically saved in the `generated_images` folder.
                
                **Supported Features:**
                - Multiple AI models
                - High-resolution output
                - Negative prompts
                - Auto-save functionality
                """)
    
    # Footer
    with gr.Row():
        gr.HTML("""
            <div class="footer">
                <p>
                    🚀 Powered by Stable Diffusion & Hugging Face | 
                    💡 Made with Gradio | 
                    ⭐ Transform your ideas into art
                </p>
            </div>
        """)
    
    # Event Handlers
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt_input, model_dropdown, negative_prompt_input],
        outputs=[image_output, status_text, file_output]
    )
    
    # Also allow Enter key to generate
    prompt_input.submit(
        fn=generate_image,
        inputs=[prompt_input, model_dropdown, negative_prompt_input],
        outputs=[image_output, status_text, file_output]
    )
    
    # Random prompt button
    random_btn.click(
        fn=generate_random_prompt,
        inputs=[],
        outputs=prompt_input
    )
    
    # Clear button
    clear_btn.click(
        fn=lambda: ("", "", None, "Ready to create! ✨", None),
        inputs=[],
        outputs=[prompt_input, negative_prompt_input, image_output, status_text, file_output]
    )

# Launch the app
if __name__ == "__main__":
    print("=" * 60)
    print("🎨 AI Image Generator - Starting...")
    print("=" * 60)
    print("✨ Modern & Responsive UI")
    print("🖼️ Multiple AI Models Available")
    print("📱 Optimized for all devices")
    print("=" * 60)
    print("🌐 Opening in your browser...")
    print("=" * 60)
    demo.launch(
        share=False,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )