app.py

import gradio as gr
import replicate
import os
from PIL import Image
import requests
from io import BytesIO
import time
import tempfile
import base64
import json

# Set up API keys from environment variables
os.environ['REPLICATE_API_TOKEN'] = os.getenv('REPLICATE_API_TOKEN')
FIREWORKS_API_KEY = os.getenv('FIREWORKS_API_KEY', '')

def enhance_prompt(user_input):
    """
    Enhance user prompt using Fireworks AI to create detailed creative design prompts
    """
    if not FIREWORKS_API_KEY:
        # If no API key, return the original prompt
        return user_input
    
    try:
        url = "https://api.fireworks.ai/inference/v1/chat/completions"
        
        system_prompt = """You are a creative design prompt engineer. Transform user input into detailed, creative design prompts.
        Format your response as a structured design concept following this pattern:
        - Start with a clear main concept description
        - Include specific design elements with bullet points
        - Mention materials, textures, or finishes
        - Specify artistic style or rendering technique
        - Add atmospheric or contextual details
        
        Example format:
        "A sleek industrial design concept for [object]:
        - Curved metallic body with minimal bezel
        - Touchscreen panel for settings
        - Modern matte black finish
        - Hand-drawn concept sketch style"
        
        Be creative, detailed, and specific. Focus on visual design elements."""
        
        payload = {
            "model": "accounts/fireworks/models/qwen3-235b-a22b-instruct-2507",
            "max_tokens": 300,
            "top_p": 1,
            "top_k": 40,
            "presence_penalty": 0,
            "frequency_penalty": 0,
            "temperature": 0.7,
            "messages": [
                {
                    "role": "system",
                    "content": system_prompt
                },
                {
                    "role": "user",
                    "content": f"Create a detailed creative design prompt for: {user_input}"
                }
            ]
        }
        
        headers = {
            "Accept": "application/json",
            "Content-Type": "application/json",
            "Authorization": f"Bearer {FIREWORKS_API_KEY}"
        }
        
        response = requests.post(url, headers=headers, data=json.dumps(payload), timeout=10)
        
        if response.status_code == 200:
            result = response.json()
            enhanced = result.get('choices', [{}])[0].get('message', {}).get('content', '')
            if enhanced:
                return enhanced
        
        return user_input
        
    except Exception as e:
        print(f"Prompt enhancement error: {str(e)}")
        return user_input

def upload_image_to_hosting(image):
    """
    Upload image to multiple hosting services with fallback
    """
    # Method 1: Try imgbb.com (most reliable)
    try:
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        buffered.seek(0)
        img_base64 = base64.b64encode(buffered.getvalue()).decode()
        
        response = requests.post(
            "https://api.imgbb.com/1/upload",
            data={
                'key': '6d207e02198a847aa98d0a2a901485a5',
                'image': img_base64,
            }
        )
        
        if response.status_code == 200:
            data = response.json()
            if data.get('success'):
                return data['data']['url']
    except:
        pass
    
    # Method 2: Try 0x0.st (simple and reliable)
    try:
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        buffered.seek(0)
        
        files = {'file': ('image.png', buffered, 'image/png')}
        response = requests.post("https://0x0.st", files=files)
        
        if response.status_code == 200:
            return response.text.strip()
    except:
        pass
    
    # Method 3: Fallback to base64
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    buffered.seek(0)
    img_base64 = base64.b64encode(buffered.getvalue()).decode()
    return f"data:image/png;base64,{img_base64}"

def process_images(prompt, enhance_prompt_flag, image1, image2=None):
    """
    Process with Creative Design - works with or without images
    """
    if not os.getenv('REPLICATE_API_TOKEN'):
        return None, prompt, "Please set REPLICATE_API_TOKEN"
    
    try:
        # Enhance prompt if requested
        final_prompt = prompt
        if enhance_prompt_flag:
            final_prompt = enhance_prompt(prompt)
        
        # Prepare input for model
        model_input = {
            "prompt": final_prompt
        }
        
        # Only add image_input if images are provided
        if image1 or image2:
            image_urls = []
            
            if image1:
                url1 = upload_image_to_hosting(image1)
                image_urls.append(url1)
            
            if image2:
                url2 = upload_image_to_hosting(image2)
                image_urls.append(url2)
            
            model_input["image_input"] = image_urls
            status_msg = "✨ Creative design generated with style transfer!"
        else:
            # No images - text-only generation
            status_msg = "✨ Creative design generated from text!"
        
        # Run model
        output = replicate.run(
            "google/nano-banana",
            input=model_input
        )
        
        if output is None:
            return None, final_prompt, "No output received"
        
        # Get the generated image
        try:
            if hasattr(output, 'read'):
                img_data = output.read()
                img = Image.open(BytesIO(img_data))
                return img, final_prompt, status_msg
        except:
            pass
        
        try:
            if hasattr(output, 'url'):
                output_url = output.url()
                response = requests.get(output_url, timeout=30)
                if response.status_code == 200:
                    img = Image.open(BytesIO(response.content))
                    return img, final_prompt, status_msg
        except:
            pass
        
        output_url = None
        if isinstance(output, str):
            output_url = output
        elif isinstance(output, list) and len(output) > 0:
            output_url = output[0]
        
        if output_url:
            response = requests.get(output_url, timeout=30)
            if response.status_code == 200:
                img = Image.open(BytesIO(response.content))
                return img, final_prompt, status_msg
        
        return None, final_prompt, "Could not process output"
        
    except Exception as e:
        error_msg = str(e)
        if "image_input" in error_msg.lower():
            return None, prompt, "Note: This model may require at least one image. Try uploading an image or check the error: " + error_msg[:100]
        return None, prompt, f"Error: {error_msg[:100]}"

# Enhanced CSS with Creative Design theme
css = """
.gradio-container {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
    min-height: 100vh;
}
.header-container {
    background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
    padding: 2.5rem;
    border-radius: 24px;
    margin-bottom: 2.5rem;
    box-shadow: 0 20px 60px rgba(240, 147, 251, 0.35);
}
.logo-text {
    font-size: 3.5rem;
    font-weight: 900;
    color: white;
    text-align: center;
    margin: 0;
    letter-spacing: -2px;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.2);
}
.subtitle {
    color: white;
    text-align: center;
    font-size: 1.1rem;
    margin-top: 0.5rem;
    opacity: 0.95;
}
.mode-indicator {
    background: rgba(255, 255, 255, 0.25);
    backdrop-filter: blur(10px);
    border-radius: 12px;
    padding: 0.6rem 1.2rem;
    margin-top: 1rem;
    text-align: center;
    font-weight: 600;
    color: white;
}
.main-content {
    background: rgba(255, 255, 255, 0.97);
    backdrop-filter: blur(20px);
    border-radius: 24px;
    padding: 2.5rem;
    box-shadow: 0 10px 40px rgba(0, 0, 0, 0.15);
}
.gr-button-primary {
    background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%) !important;
    border: none !important;
    color: white !important;
    font-weight: 700 !important;
    font-size: 1.1rem !important;
    padding: 1.2rem 2rem !important;
    border-radius: 14px !important;
    transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
    text-transform: uppercase;
    letter-spacing: 1px;
    width: 100%;
    margin-top: 1rem !important;
}
.gr-button-primary:hover {
    transform: translateY(-3px) !important;
    box-shadow: 0 15px 40px rgba(245, 87, 108, 0.4) !important;
}
.gr-input, .gr-textarea {
    background: #ffffff !important;
    border: 2px solid #e1e8ed !important;
    border-radius: 14px !important;
    color: #2d3436 !important;
    font-size: 1rem !important;
    padding: 0.8rem 1rem !important;
}
.gr-input:focus, .gr-textarea:focus {
    border-color: #f093fb !important;
    box-shadow: 0 0 0 4px rgba(240, 147, 251, 0.15) !important;
}
.gr-form {
    background: transparent !important;
    border: none !important;
}
.gr-panel {
    background: #ffffff !important;
    border: 2px solid #e1e8ed !important;
    border-radius: 16px !important;
    padding: 1.5rem !important;
}
.gr-box {
    border-radius: 14px !important;
    border-color: #e1e8ed !important;
}
label {
    color: #636e72 !important;
    font-weight: 600 !important;
    font-size: 0.85rem !important;
    text-transform: uppercase;
    letter-spacing: 0.5px;
    margin-bottom: 0.5rem !important;
}
.status-text {
    font-family: 'SF Mono', 'Monaco', monospace;
    color: #00b894;
    font-size: 0.9rem;
}
.image-container {
    border-radius: 14px !important;
    overflow: hidden;
    border: 2px solid #e1e8ed !important;
    background: #fafbfc !important;
}
footer {
    display: none !important;
}
.image-upload {
    min-height: 200px !important;
    max-height: 200px !important;
}
.output-image {
    min-height: 420px !important;
    max-height: 420px !important;
}
.gr-row {
    gap: 1rem !important;
}
.gr-column {
    gap: 1rem !important;
}
.info-box {
    background: linear-gradient(135deg, #ffeaa7 0%, #fdcb6e 100%);
    border-radius: 12px;
    padding: 1rem;
    margin-bottom: 1rem;
    border-left: 4px solid #f5576c;
    color: #2d3436;
}
.enhanced-prompt-box {
    background: linear-gradient(135deg, #a8edea 0%, #fed6e3 100%);
    border-radius: 12px;
    padding: 1rem;
    margin-top: 1rem;
    border-left: 4px solid #f093fb;
}
.gr-checkbox {
    background: white !important;
}
"""

with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:
    with gr.Column(elem_classes="header-container"):
        gr.HTML("""
            <h1 class="logo-text">🎨 Creative Design</h1>
            <p class="subtitle">AI-Powered Creative Design Generation & Style Transfer</p>
            <div class="mode-indicator">
                💡 Intelligent prompt enhancement for stunning creative designs!
            </div>
            <div style="display: flex; justify-content: center; align-items: center; gap: 10px; margin-top: 20px;">
                <a href="https://huggingface.co/spaces/ginigen/Nano-Banana-PRO" target="_blank">
                    <img src="https://img.shields.io/static/v1?label=CREATIVE%20DESIGN&message=PRO&color=%23f093fb&labelColor=%23764ba2&logo=HUGGINGFACE&logoColor=white&style=for-the-badge" alt="badge">
                </a>    
                <a href="https://huggingface.co/spaces/openfree/Nano-Banana-Upscale" target="_blank">
                    <img src="https://img.shields.io/static/v1?label=CREATIVE%20DESIGN&message=UPSCALE&color=%23f093fb&labelColor=%23764ba2&logo=GOOGLE&logoColor=white&style=for-the-badge" alt="Creative Design Upscale">
                </a>                    
                <a href="https://huggingface.co/spaces/aiqtech/Nano-Banana-API" target="_blank">
                    <img src="https://img.shields.io/static/v1?label=CREATIVE%20DESIGN&message=API&color=%23f093fb&labelColor=%23764ba2&logo=GOOGLE&logoColor=white&style=for-the-badge" alt="Creative Design API">
                </a>
                <a href="https://huggingface.co/spaces/ginigen/Nano-Banana-Video" target="_blank">
                    <img src="https://img.shields.io/static/v1?label=CREATIVE%20DESIGN&message=VIDEO&color=%23f093fb&labelColor=%23764ba2&logo=GOOGLE&logoColor=white&style=for-the-badge" alt="Creative Design VIDEO">
                </a>                    
                <a href="https://discord.gg/openfreeai" target="_blank">
                    <img src="https://img.shields.io/static/v1?label=Discord&message=Openfree%20AI&color=%23f093fb&labelColor=%23764ba2&logo=discord&logoColor=white&style=for-the-badge" alt="Discord Openfree AI">
                </a>
            </div>
        """)

    with gr.Column(elem_classes="main-content"):
        # Info box
        gr.HTML("""
            <div class="info-box">
                <strong>🚀 How to use Creative Design:</strong><br>
                • <b>Simple Input:</b> Enter a basic idea or concept<br>
                • <b>AI Enhancement:</b> Enable prompt enhancement for detailed creative designs<br>
                • <b>Style Transfer:</b> Add images to apply specific visual styles<br>
                • The AI transforms simple ideas into professional design concepts!
            </div>
        """)
        
        with gr.Row(equal_height=True):
            # Left Column - Inputs
            with gr.Column(scale=1):
                prompt = gr.Textbox(
                    label="Design Concept / Idea",
                    placeholder="e.g., 'modern coffee machine', 'futuristic chair', 'minimalist lamp'...",
                    lines=3,
                    value="sleek smartphone design",
                    elem_classes="prompt-input"
                )
                
                enhance_prompt_checkbox = gr.Checkbox(
                    label="🎨 Enhance Prompt with AI (Create detailed design prompt)",
                    value=True,
                    info="Transform simple ideas into professional design prompts"
                )
                
                with gr.Row(equal_height=True):
                    image1 = gr.Image(
                        label="Style Reference 1 (Optional)",
                        type="pil",
                        height=200,
                        elem_classes="image-container image-upload"
                    )
                    image2 = gr.Image(
                        label="Style Reference 2 (Optional)",
                        type="pil",
                        height=200,
                        elem_classes="image-container image-upload"
                    )
                
                generate_btn = gr.Button(
                    "Create Design ✨",
                    variant="primary",
                    size="lg"
                )
            
            # Right Column - Output
            with gr.Column(scale=1):
                output_image = gr.Image(
                    label="Generated Design",
                    type="pil",
                    height=420,
                    elem_classes="image-container output-image"
                )
                
                enhanced_prompt_display = gr.Textbox(
                    label="Enhanced Creative Prompt",
                    interactive=False,
                    lines=4,
                    elem_classes="enhanced-prompt-box",
                    visible=True,
                    value="Your enhanced prompt will appear here..."
                )
                
                status = gr.Textbox(
                    label="Status",
                    interactive=False,
                    lines=1,
                    elem_classes="status-text",
                    value="Ready to create amazing designs..."
                )
    
    # Event handler
    generate_btn.click(
        fn=process_images,
        inputs=[prompt, enhance_prompt_checkbox, image1, image2],
        outputs=[output_image, enhanced_prompt_display, status]
    )
    
    # Creative Design Examples
    gr.Examples(
        examples=[
            ["minimalist desk lamp", True, None, None],
            ["futuristic motorcycle", True, None, None],
            ["organic architecture building", True, None, None],
            ["smart home assistant device", True, None, None],
            ["eco-friendly water bottle", True, None, None],
            ["modular furniture system", True, None, None],
            ["wearable health monitor", True, None, None],
            ["urban electric scooter", True, None, None],
        ],
        inputs=[prompt, enhance_prompt_checkbox, image1, image2],
        label="Creative Design Examples (AI will enhance these into detailed prompts!)"
    )

# Launch
if __name__ == "__main__":
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860
    )