app.py

import gradio as gr
import requests
import json
import os
import time
from PIL import Image
import io
from datetime import datetime

# Get API key from Hugging Face secrets
def get_api_key():
    return os.environ.get("PIXAZO_API_KEY")

# Call Pixazo API
def call_pixazo_api(prompt, num_steps=4, seed=None, height=512, width=512, 
                   style_preset=None, guidance_scale=None):
    
    api_key = get_api_key()
    
    if not api_key:
        return None, {
            "error": "API key not configured",
            "instructions": "Please set PIXAZO_API_KEY in Hugging Face Space secrets",
            "timestamp": datetime.now().isoformat()
        }
    
    url = "https://gateway.pixazo.ai/flux-1-schnell/v1/getData"
    
    headers = {
        'Content-Type': 'application/json',
        'Cache-Control': 'no-cache',
        'Ocp-Apim-Subscription-Key': api_key,
    }
    
    # Prepare request body
    body = {
        "prompt": prompt,
        "num_steps": num_steps,
        "height": height,
        "width": width,
    }
    
    if seed and seed > 0:
        body["seed"] = seed
    if style_preset and style_preset != "none":
        body["style_preset"] = style_preset
    if guidance_scale:
        body["guidance_scale"] = guidance_scale
    
    try:
        print(f"🌐 Sending request to Pixazo API...")
        
        response = requests.post(
            url, 
            headers=headers, 
            json=body,
            timeout=120
        )
        
        print(f"✅ Response status: {response.status_code}")
        
        if response.status_code == 200:
            result = response.json()
            
            if "output" in result:
                image_url = result["output"]
                
                # Download the image from the URL
                try:
                    print(f"⬇️ Downloading image from URL...")
                    img_response = requests.get(image_url, timeout=30, headers={
                        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
                    })
                    
                    if img_response.status_code == 200:
                        image = Image.open(io.BytesIO(img_response.content))
                        
                        # Convert to RGB if necessary
                        if image.mode in ('RGBA', 'LA', 'P'):
                            image = image.convert('RGB')
                        
                        return image, {
                            "success": True,
                            "status_code": 200,
                            "image_url_preview": image_url[:80] + "..." if len(image_url) > 80 else image_url,
                            "image_size": f"{image.size[0]}x{image.size[1]}",
                            "timestamp": datetime.now().isoformat()
                        }
                    else:
                        return None, {
                            "error": f"Failed to download image: HTTP {img_response.status_code}",
                            "image_url": image_url[:80] + "..." if len(image_url) > 80 else image_url,
                            "download_status": img_response.status_code,
                            "timestamp": datetime.now().isoformat()
                        }
                        
                except Exception as e:
                    return None, {
                        "error": f"Error downloading image: {str(e)}",
                        "image_url": image_url[:80] + "..." if len(image_url) > 80 else image_url,
                        "timestamp": datetime.now().isoformat()
                    }
            else:
                return None, {
                    "error": "Response missing 'output' field",
                    "response_keys": list(result.keys()),
                    "timestamp": datetime.now().isoformat()
                }
        else:
            try:
                error_data = response.json()
                error_msg = error_data.get("error", str(error_data))
            except:
                error_msg = response.text[:200] if response.text else "No error message"
            
            return None, {
                "error": f"API returned status {response.status_code}: {error_msg}",
                "status_code": response.status_code,
                "timestamp": datetime.now().isoformat()
            }
            
    except requests.exceptions.Timeout:
        return None, {
            "error": "Request timed out after 120 seconds",
            "suggestion": "Try reducing image size or number of steps",
            "timestamp": datetime.now().isoformat()
        }
    except requests.exceptions.RequestException as e:
        return None, {
            "error": f"Network error: {str(e)}",
            "timestamp": datetime.now().isoformat()
        }
    except Exception as e:
        return None, {
            "error": f"Unexpected error: {str(e)}",
            "timestamp": datetime.now().isoformat()
        }

# Batch generate images
def generate_images(prompt, num_steps, seed, height, width, style_preset, 
                   num_images, guidance_scale):
    
    images = []
    all_results = []
    
    for i in range(num_images):
        try:
            current_seed = seed + i if seed > 0 else 0
            
            image, result = call_pixazo_api(
                prompt=prompt,
                num_steps=num_steps,
                seed=current_seed,
                height=height,
                width=width,
                style_preset=style_preset,
                guidance_scale=guidance_scale
            )
            
            if image:
                images.append(image)
            
            all_results.append(result)
            
            status_text = f"Generated {i+1}/{num_images} images"
            if "error" in result:
                status_text = f"Image {i+1} error: {result['error'][:50]}..."
            
            yield images if images else None, status_text, json.dumps(result, indent=2)
            
            if i < num_images - 1:
                time.sleep(1)
                
        except Exception as e:
            error_result = {"error": str(e), "image_index": i, "timestamp": datetime.now().isoformat()}
            all_results.append(error_result)
            yield images if images else None, f"Exception on image {i+1}: {str(e)[:50]}", json.dumps(error_result, indent=2)
            break
    
    # Final summary
    success_count = len(images)
    
    if success_count == 0:
        final_status = f"❌ No images generated ({success_count}/{num_images} successful)"
        gallery_val = None
    elif success_count < num_images:
        final_status = f"⚠️ Partial success: {success_count}/{num_images} images"
        gallery_val = images
    else:
        final_status = f"🎉 Success! All {success_count} images generated"
        gallery_val = images
    
    summary = {
        "summary": {
            "total_attempts": num_images,
            "successful": success_count,
            "failed": num_images - success_count,
            "timestamp": datetime.now().isoformat()
        },
        "parameters_used": {
            "prompt_length": len(prompt),
            "steps": num_steps,
            "resolution": f"{width}x{height}",
            "style": style_preset
        }
    }
    
    yield gallery_val, final_status, json.dumps(summary, indent=2)

# Custom CSS - Keep the beauty alive! 🎨
custom_css = """
.gradio-container {
    max-width: 1200px !important;
    margin: 0 auto !important;
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, sans-serif;
}

/* Beautiful Gradient Header */
.header-gradient {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    color: white !important;
    padding: 40px 30px !important;
    border-radius: 16px !important;
    margin-bottom: 30px !important;
    text-align: center !important;
    box-shadow: 0 20px 60px rgba(102, 126, 234, 0.3) !important;
}

.header-title {
    font-size: 2.8em !important;
    font-weight: 700 !important;
    margin-bottom: 10px !important;
}

.header-subtitle {
    font-size: 1.2em !important;
    opacity: 0.9 !important;
    max-width: 600px !important;
    margin: 0 auto !important;
}

/* API Status */
.api-status {
    padding: 16px 20px !important;
    border-radius: 12px !important;
    margin-bottom: 25px !important;
    font-weight: 600 !important;
    font-size: 1.1em !important;
    border-left: 6px solid !important;
}

.api-status-success {
    background: linear-gradient(135deg, #d1fae5 0%, #a7f3d0 100%) !important;
    color: #065f46 !important;
    border-left-color: #10b981 !important;
}

.api-status-error {
    background: linear-gradient(135deg, #fee2e2 0%, #fecaca 100%) !important;
    color: #7f1d1d !important;
    border-left-color: #ef4444 !important;
}

/* Input Sections */
.input-section {
    background: white !important;
    padding: 30px !important;
    border-radius: 16px !important;
    border: 1px solid #e5e7eb !important;
    box-shadow: 0 4px 20px rgba(0,0,0,0.05) !important;
    margin-bottom: 25px !important;
}

.section-title {
    color: #333 !important;
    font-size: 1.4em !important;
    font-weight: 600 !important;
    margin-bottom: 20px !important;
    display: flex !important;
    align-items: center !important;
    gap: 10px !important;
}

/* Example Cards */
.example-card {
    background: #f8fafc !important;
    padding: 20px !important;
    border-radius: 12px !important;
    border: 2px solid #e2e8f0 !important;
    margin: 12px 0 !important;
    cursor: pointer !important;
    transition: all 0.3s ease !important;
    text-align: left !important;
    width: 100% !important;
}

.example-card:hover {
    background: #e2e8f0 !important;
    transform: translateY(-4px) !important;
    box-shadow: 0 10px 30px rgba(0,0,0,0.1) !important;
    border-color: #4f46e5 !important;
}

.example-title {
    color: #4f46e5 !important;
    font-weight: 600 !important;
    font-size: 1.1em !important;
    margin-bottom: 8px !important;
}

.example-desc {
    color: #666 !important;
    font-size: 0.95em !important;
    line-height: 1.4 !important;
}

/* Beautiful Generate Button */
.generate-btn {
    background: linear-gradient(135deg, #4f46e5 0%, #7c3aed 100%) !important;
    color: white !important;
    font-weight: 700 !important;
    padding: 18px 40px !important;
    font-size: 1.2em !important;
    border-radius: 12px !important;
    border: none !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 8px 25px rgba(79, 70, 229, 0.4) !important;
}

.generate-btn:hover {
    transform: translateY(-3px) !important;
    box-shadow: 0 15px 35px rgba(79, 70, 229, 0.5) !important;
}

.secondary-btn {
    background: #f1f5f9 !important;
    color: #475569 !important;
    border: 2px solid #e2e8f0 !important;
}

/* Gallery */
.gallery-container {
    background: white !important;
    padding: 25px !important;
    border-radius: 16px !important;
    border: 1px solid #e5e7eb !important;
    box-shadow: 0 4px 20px rgba(0,0,0,0.05) !important;
}

/* Status Box */
.status-box {
    background: #f8fafc !important;
    padding: 20px !important;
    border-radius: 12px !important;
    border-left: 5px solid #4f46e5 !important;
    margin-bottom: 20px !important;
    font-size: 1.1em !important;
}

/* Footer */
.footer {
    text-align: center !important;
    color: #666 !important;
    font-size: 0.9em !important;
    margin-top: 40px !important;
    padding-top: 20px !important;
    border-top: 1px solid #e5e7eb !important;
}

/* Make sure dark mode doesn't break our beauty */
.dark .input-section,
.dark .gallery-container {
    background: #1f2937 !important;
    border-color: #374151 !important;
}

.dark .example-card {
    background: #374151 !important;
    border-color: #4b5563 !important;
}

.dark .status-box {
    background: #374151 !important;
}
"""

# Create the interface
with gr.Blocks(title="Pixazo Image Generator", theme=gr.themes.Soft()) as demo:
    
    # Beautiful Header
    gr.HTML("""
    <div class="header-gradient">
        <div class="header-title">
            🎨 Pixazo Image Generator
        </div>
        <div class="header-subtitle">
            Generate stunning AI images with FLUX-1 Schnell model
        </div>
    </div>
    """)
    
    # API Status
    api_key = get_api_key()
    if api_key:
        gr.HTML(f"""
        <div class="api-status api-status-success">
            ✅ API Status: Connected • Key length: {len(api_key)} characters
        </div>
        """)
    else:
        gr.HTML("""
        <div class="api-status api-status-error">
            ❌ API Status: Not Configured
            <br><small>Go to Space Settings → Repository secrets → Add: <code>PIXAZO_API_KEY</code> = your API key</small>
        </div>
        """)
    
    with gr.Row():
        # Left Column - Inputs
        with gr.Column(scale=1):
            # Prompt Section
            with gr.Group(elem_classes="input-section"):
                gr.Markdown("### ✨ Image Description")
                prompt = gr.Textbox(
                    label="",
                    placeholder="Describe the image you want to generate...",
                    lines=4,
                    show_label=False
                )
            
            # Parameters Section
            with gr.Group(elem_classes="input-section"):
                gr.Markdown("### ⚙️ Generation Parameters")
                
                with gr.Row():
                    style_preset = gr.Dropdown(
                        label="Style Preset",
                        choices=[
                            "none", "cyberpunk", "fantasy", "anime", "photographic", 
                            "digital-art", "comic", "3d-model", "pixel-art", 
                            "isometric", "watercolor", "oil-painting", "sketch"
                        ],
                        value="none"
                    )
                
                with gr.Row():
                    guidance_scale = gr.Slider(
                        label="Guidance Scale",
                        minimum=1.0,
                        maximum=20.0,
                        value=7.5,
                        step=0.5
                    )
                
                with gr.Row():
                    num_steps = gr.Slider(
                        label="Steps",
                        minimum=1,
                        maximum=50,
                        value=4,
                        step=1
                    )
                    
                    seed = gr.Number(
                        label="Seed",
                        value=42,
                        minimum=0,
                        maximum=999999
                    )
                
                with gr.Row():
                    width = gr.Slider(
                        label="Width",
                        minimum=256,
                        maximum=1024,
                        value=512,
                        step=64
                    )
                    
                    height = gr.Slider(
                        label="Height",
                        minimum=256,
                        maximum=1024,
                        value=512,
                        step=64
                    )
                
                num_images = gr.Slider(
                    label="Number of Images",
                    minimum=1,
                    maximum=8,
                    value=1,
                    step=1
                )
            
            # Action Buttons
            with gr.Row():
                generate_btn = gr.Button(
                    "✨ Generate Images", 
                    variant="primary",
                    elem_classes="generate-btn",
                    scale=2
                )
                clear_btn = gr.Button("🗑️ Clear", variant="secondary", scale=1)
            
            # Examples Section
            with gr.Group(elem_classes="input-section"):
                gr.Markdown("### 💡 Example Prompts")
                
                examples = [
                    ("Cyberpunk City", "A futuristic cyberpunk city at night with neon signs and flying cars", "cyberpunk"),
                    ("Fantasy Dragon", "A majestic dragon on a mountain peak overlooking a magical kingdom", "fantasy"),
                    ("Anime Character", "A cute anime girl with magical powers in a cherry blossom garden", "anime"),
                    ("Space Explorer", "An astronaut on Mars with detailed reflection in visor", "photographic"),
                    ("Isometric Scene", "Cozy isometric bedroom with plants and warm lighting", "isometric")
                ]
                
                for title, desc, style in examples:
                    # Create example buttons
                    btn = gr.Button(
                        f"📌 {title}",
                        variant="secondary",
                        size="sm",
                        elem_classes="example-card"
                    )
                    
                    def use_example(p=desc, s=style):
                        return p, s
                    
                    btn.click(
                        use_example,
                        outputs=[prompt, style_preset]
                    )
        
        # Right Column - Outputs
        with gr.Column(scale=2):
            # Status Display
            status = gr.Textbox(
                label="Status",
                value="🎯 Ready to generate images. Enter a prompt above!",
                interactive=False,
                elem_classes="status-box"
            )
            
            # Gallery
            with gr.Group(elem_classes="gallery-container"):
                gallery = gr.Gallery(
                    label="🎨 Generated Images",
                    columns=3,
                    height="auto",
                    show_label=True
                )
            
            # Info Tabs
            with gr.Tabs():
                with gr.TabItem("📊 Response Details"):
                    json_output = gr.JSON(
                        label="API Response",
                        container=True
                    )
                
                with gr.TabItem("ℹ️ About"):
                    gr.Markdown(f"""
                    ### About This App
                    
                    **Version:** 1.0.0
                    **Gradio:** {gr.__version__}
                    **API Format:** `{"output": "IMAGE_URL"}`
                    
                    **Features:**
                    - 🖼️ Generate multiple images at once
                    - 🎭 Various style presets
                    - ⚡ Adjustable quality settings
                    - 🔒 Secure API key handling
                    - 📱 Responsive design
                    - 🔄 Real-time progress updates
                    
                    **Credits:**
                    - Powered by Pixazo AI
                    - Built with Gradio
                    - Hosted on Hugging Face Spaces
                    """)
    
    # Footer
    gr.HTML("""
    <div class="footer">
        <p>Built with ❤️ using Gradio • Powered by Pixazo FLUX-1 Schnell</p>
        <p>Model: flux-1-schnell • Response Format: <code>{"output": "IMAGE_URL"}</code></p>
    </div>
    """)
    
    # Event Handlers
    def on_generate(prompt, num_steps, seed, height, width, style_preset, num_images, guidance_scale):
        if not prompt.strip():
            yield None, "❌ Please enter a prompt first", {"error": "No prompt provided"}
            return
        
        for gallery_val, status_val, json_val in generate_images(
            prompt, num_steps, seed, height, width, style_preset, num_images, guidance_scale
        ):
            yield gallery_val, status_val, json_val
    
    def on_clear():
        return None, "🎯 Ready to generate images. Enter a prompt above!", None
    
    # Connect buttons
    generate_btn.click(
        fn=on_generate,
        inputs=[prompt, num_steps, seed, height, width, style_preset, num_images, guidance_scale],
        outputs=[gallery, status, json_output]
    )
    
    clear_btn.click(
        fn=on_clear,
        outputs=[gallery, status, json_output]
    )

# Launch with beauty preserved! 🎨
if __name__ == "__main__":
    print("=" * 60)
    print("🚀 Launching Pixazo Image Generator")
    print("🎨 Beautiful UI Version")
    print(f"📦 Gradio Version: {gr.__version__}")
    print(f"🔑 API Key: {'✅ Configured' if get_api_key() else '❌ Not Configured'}")
    print("=" * 60)
    
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        css=custom_css
    )