app.py

import gradio as gr
import replicate
import os
from PIL import Image
import requests
from io import BytesIO
import time
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', '')
BRAVE_API_KEY = os.getenv('BRAVE_API_KEY', '')  # Add Brave API key

def search_design_trends(product_type):
    """
    Use Brave Search API to find current design trends and technologies
    """
    if not BRAVE_API_KEY:
        return "No search results available (API key not set)"
    
    try:
        # Construct search query for product design
        query = f"{product_type} design trends 2024 2025 materials technology innovation"
        
        url = "https://api.search.brave.com/res/v1/web/search"
        headers = {
            "Accept": "application/json",
            "X-Subscription-Token": BRAVE_API_KEY
        }
        params = {
            "q": query,
            "count": 5
        }
        
        response = requests.get(url, headers=headers, params=params, timeout=10)
        
        if response.status_code == 200:
            data = response.json()
            results = []
            
            if 'web' in data and 'results' in data['web']:
                for item in data['web']['results'][:3]:
                    results.append(f"- {item.get('title', '')}: {item.get('description', '')[:150]}")
            
            if results:
                return "\n".join(results)
            else:
                return "Limited search results available"
        else:
            return "Search unavailable at the moment"
            
    except Exception as e:
        print(f"Brave search error: {str(e)}")
        return "Search results unavailable"

def enhance_prompt_with_team(user_input):
    """
    Enhanced prompt generation using team collaboration with real web search
    """
    if not FIREWORKS_API_KEY:
        # Fallback to basic enhancement
        return create_basic_design_prompt(user_input)
    
    try:
        # First, get real search results for the engineer
        search_results = search_design_trends(user_input)
        
        url = "https://api.fireworks.ai/inference/v1/chat/completions"
        
        # Enhanced system prompt with search integration
        system_prompt = """You are a product design team with three specialized roles:

## TEAM ROLES:

### 1. DIRECTOR (Creative Vision Leader)
- Analyzes user requirements and market positioning
- Sets design philosophy and constraints
- Makes final decisions on features

### 2. ENGINEER (Technical Expert with Research Data)
- Uses provided search results to inform decisions
- Evaluates manufacturing feasibility
- Recommends materials and technologies
- Suggests innovative features based on trends

### 3. DESIGNER (Visual Specialist)
- Creates detailed visual specifications
- Specifies exact materials, finishes, colors
- Defines rendering style and atmosphere

## OUTPUT FORMAT:

**DIRECTOR'S VISION:**
[Brief vision statement - 2-3 sentences]

**ENGINEER'S TECHNICAL ASSESSMENT:**
[Based on search data, provide 3-4 technical recommendations]

**DESIGNER'S FINAL PROMPT:**
[Detailed visual description for image generation - focus on visual elements]

Keep responses concise but specific. Focus on PHYSICAL PRODUCTS only."""
        
        # Create collaborative prompt with search results
        user_prompt = f"""Product concept: {user_input}

ENGINEER'S RESEARCH DATA:
{search_results}

Based on this input and research, execute the team collaboration:
1. Director establishes vision (brief)
2. Engineer uses the research data above to recommend features
3. Designer creates a detailed visual prompt

Focus on creating a manufacturable, innovative product design."""
        
        payload = {
            "model": "accounts/fireworks/models/qwen3-235b-a22b-instruct-2507",  # Original model maintained
            "max_tokens": 500,
            "top_p": 0.9,
            "temperature": 0.7,
            "messages": [
                {
                    "role": "system",
                    "content": system_prompt
                },
                {
                    "role": "user",
                    "content": user_prompt
                }
            ]
        }
        
        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=30)
        
        if response.status_code == 200:
            result = response.json()
            enhanced = result.get('choices', [{}])[0].get('message', {}).get('content', '')
            if enhanced:
                # Extract designer's final prompt
                if "DESIGNER'S FINAL PROMPT:" in enhanced:
                    parts = enhanced.split("DESIGNER'S FINAL PROMPT:")
                    if len(parts) > 1:
                        # Return both the final prompt and the full discussion
                        return parts[1].strip(), enhanced
                return enhanced, enhanced
        
        # Fallback if API fails
        return create_basic_design_prompt(user_input), "Team collaboration unavailable - using basic prompt"
        
    except requests.exceptions.Timeout:
        print("Fireworks API timeout - using fallback")
        return create_basic_design_prompt(user_input), "API timeout - using fallback prompt generation"
    except Exception as e:
        print(f"Team collaboration error: {str(e)}")
        return create_basic_design_prompt(user_input), f"Error: {str(e)[:100]}"

def create_basic_design_prompt(user_input):
    """
    Fallback function for basic design prompt generation
    """
    return f"""A professional product design concept for {user_input}:
- Modern, innovative industrial design
- Premium materials and finishes
- Clean lines with functional aesthetics
- Photorealistic rendering with studio lighting
- Detailed surface textures and materials
- Professional product photography style
- High-end manufacturing quality
- Shown from 3/4 perspective angle"""

def upload_image_to_hosting(image):
    """
    Upload image to hosting service
    """
    try:
        # Try imgbb first
        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,
            },
            timeout=10
        )
        
        if response.status_code == 200:
            data = response.json()
            if data.get('success'):
                return data['data']['url']
    except:
        pass
    
    # 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_product_design(prompt, enhance_prompt_flag, image1, image2=None):
    """
    Process product design with team collaboration
    """
    if not os.getenv('REPLICATE_API_TOKEN'):
        return None, prompt, "⚠️ Please set REPLICATE_API_TOKEN", ""
    
    try:
        team_discussion = ""
        final_prompt = prompt
        
        # Apply team collaboration if enabled
        if enhance_prompt_flag:
            result = enhance_prompt_with_team(prompt)
            if isinstance(result, tuple):
                final_prompt, team_discussion = result
            else:
                final_prompt = result
                team_discussion = "Team collaboration completed"
            
            # Format team discussion for display
            if "DIRECTOR'S VISION:" in team_discussion:
                team_discussion = f"### 🎬 Team Collaboration Process\n\n{team_discussion}"
        
        # Prepare model input
        model_input = {
            "prompt": final_prompt
        }
        
        # Add reference images if 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 = "✅ Product design generated with style references!"
        else:
            status_msg = "✅ Product design generated successfully!"
        
        # Generate image with Replicate
        try:
            output = replicate.run(
                "black-forest-labs/flux-schnell",  # Alternative model that works better
                input={
                    "prompt": final_prompt,
                    "num_outputs": 1,
                    "aspect_ratio": "1:1",
                    "output_format": "webp",
                    "output_quality": 90
                }
            )
        except:
            # Fallback to original model
            output = replicate.run(
                "stability-ai/sdxl:39ed52f2a78e934b3ba6e2a89f5b1c712de7dfea535525255b1aa35c5565e08b",
                input={
                    "prompt": final_prompt,
                    "negative_prompt": "low quality, blurry, distorted",
                    "width": 1024,
                    "height": 1024,
                    "num_outputs": 1
                }
            )
        
        if output is None:
            return None, final_prompt, "❌ No output received from model", team_discussion
        
        # Process output
        output_url = None
        if isinstance(output, str):
            output_url = output
        elif isinstance(output, list) and len(output) > 0:
            output_url = output[0]
        elif hasattr(output, 'url'):
            output_url = output.url
        
        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, team_discussion
        
        return None, final_prompt, "❌ Could not process output", team_discussion
        
    except Exception as e:
        error_msg = str(e)
        return None, prompt, f"❌ Error: {error_msg[:200]}", ""

# Professional Product Design CSS
css = """
.gradio-container {
    background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
    min-height: 100vh;
}
.header-container {
    background: linear-gradient(135deg, #2d2d2d 0%, #0f0f0f 100%);
    padding: 2rem;
    border-radius: 20px;
    margin-bottom: 2rem;
    box-shadow: 0 10px 40px rgba(0, 0, 0, 0.5);
    border: 1px solid rgba(255, 215, 0, 0.2);
}
.logo-text {
    font-size: 2.8rem;
    font-weight: 800;
    color: #ffd700;
    text-align: center;
    margin: 0;
    letter-spacing: -1px;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.8);
}
.subtitle {
    color: #b0b0b0;
    text-align: center;
    font-size: 1rem;
    margin-top: 0.5rem;
    font-weight: 400;
    letter-spacing: 0.5px;
}
.mode-indicator {
    background: rgba(255, 215, 0, 0.1);
    border: 1px solid rgba(255, 215, 0, 0.3);
    border-radius: 10px;
    padding: 0.6rem 1.2rem;
    margin: 1rem auto;
    text-align: center;
    font-weight: 500;
    color: #ffd700;
    max-width: 500px;
    font-size: 0.9rem;
}
.main-content {
    background: rgba(255, 255, 255, 0.98);
    border-radius: 20px;
    padding: 2rem;
    box-shadow: 0 5px 20px rgba(0, 0, 0, 0.3);
}
.gr-button-primary {
    background: linear-gradient(135deg, #2d2d2d 0%, #0f0f0f 100%) !important;
    border: 1px solid #ffd700 !important;
    color: #ffd700 !important;
    font-weight: 600 !important;
    font-size: 1rem !important;
    padding: 1rem 2rem !important;
    border-radius: 10px !important;
    transition: all 0.3s ease !important;
    text-transform: uppercase;
    letter-spacing: 0.5px;
    width: 100%;
    margin-top: 1rem !important;
}
.gr-button-primary:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 25px rgba(255, 215, 0, 0.25) !important;
    background: linear-gradient(135deg, #ffd700 0%, #ffed4e 100%) !important;
    color: #000 !important;
}
.info-box {
    background: linear-gradient(135deg, #fff3cd 0%, #ffeaa7 100%);
    border-radius: 10px;
    padding: 1rem;
    margin-bottom: 1.5rem;
    border-left: 4px solid #ffd700;
    color: #333;
    font-size: 0.9rem;
}
.team-discussion {
    background: #f8f9fa;
    border-radius: 10px;
    padding: 1.2rem;
    margin-top: 1rem;
    border: 1px solid #dee2e6;
    font-family: 'Courier New', monospace;
    font-size: 0.85rem;
    max-height: 300px;
    overflow-y: auto;
    white-space: pre-wrap;
}
.enhanced-prompt-box {
    background: #f0f0f0;
    border-radius: 10px;
    padding: 1rem;
    margin-top: 0.5rem;
    border-left: 3px solid #ffd700;
    font-size: 0.9rem;
}
.gr-input, .gr-textarea {
    background: #ffffff !important;
    border: 1px solid #d0d0d0 !important;
    border-radius: 8px !important;
    font-size: 0.95rem !important;
    padding: 0.7rem !important;
}
.gr-input:focus, .gr-textarea:focus {
    border-color: #ffd700 !important;
    box-shadow: 0 0 0 3px rgba(255, 215, 0, 0.1) !important;
}
.image-container {
    border-radius: 10px !important;
    overflow: hidden;
    border: 1px solid #e0e0e0 !important;
}
footer {
    display: none !important;
}
"""

# Create Gradio interface
with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:
    with gr.Column(elem_classes="header-container"):
        gr.HTML("""
            <h1 class="logo-text">🏭 PRODUCT DESIGN STUDIO</h1>
            <p class="subtitle">AI-Powered Industrial & Product Design System with Web Research</p>
            <div class="mode-indicator">
                👥 Director → 🔍 Engineer (with Brave Search) → 🎨 Designer
            </div>
        """)

    with gr.Column(elem_classes="main-content"):
        gr.HTML("""
            <div class="info-box">
                <strong>🎯 Professional Product Design Process:</strong><br>
                • <b>Director:</b> Establishes vision and requirements<br>
                • <b>Engineer:</b> Searches current trends & technologies via Brave API<br>
                • <b>Designer:</b> Creates detailed visual specifications<br>
                • <b>Focus:</b> Physical products only (electronics, furniture, vehicles, appliances)
            </div>
        """)
        
        with gr.Row(equal_height=True):
            with gr.Column(scale=1):
                prompt = gr.Textbox(
                    label="📝 Product Concept",
                    placeholder="Enter product type: 'ergonomic keyboard', 'smart watch', 'office chair', 'coffee maker'...",
                    lines=2,
                    value="wireless earbuds",
                    elem_classes="prompt-input"
                )
                
                enhance_prompt_checkbox = gr.Checkbox(
                    label="🚀 Enable Team Collaboration with Web Research",
                    value=True,
                    info="Activates Director + Engineer (Brave Search) + Designer process"
                )
                
                with gr.Row():
                    image1 = gr.Image(
                        label="Style Reference 1 (Optional)",
                        type="pil",
                        height=180
                    )
                    image2 = gr.Image(
                        label="Style Reference 2 (Optional)",
                        type="pil",
                        height=180
                    )
                
                generate_btn = gr.Button(
                    "🎨 Generate Product Design",
                    variant="primary",
                    size="lg"
                )
            
            with gr.Column(scale=1):
                output_image = gr.Image(
                    label="Generated Product Design",
                    type="pil",
                    height=400,
                    elem_classes="image-container"
                )
                
                enhanced_prompt_display = gr.Textbox(
                    label="📋 Final Design Specifications",
                    interactive=False,
                    lines=3,
                    elem_classes="enhanced-prompt-box"
                )
                
                status = gr.Textbox(
                    label="Status",
                    interactive=False,
                    elem_classes="status-text",
                    value="Ready to design..."
                )
        
        with gr.Row():
            team_discussion_display = gr.Markdown(
                value="",
                elem_classes="team-discussion"
            )
    
    # Connect event handler
    generate_btn.click(
        fn=process_product_design,
        inputs=[prompt, enhance_prompt_checkbox, image1, image2],
        outputs=[output_image, enhanced_prompt_display, status, team_discussion_display]
    )
    
    # Product examples
    gr.Examples(
        examples=[
            ["wireless earbuds", True, None, None],
            ["ergonomic office chair", True, None, None],
            ["smart home speaker", True, None, None],
            ["electric toothbrush", True, None, None],
            ["portable coffee maker", True, None, None],
            ["gaming keyboard", True, None, None],
            ["fitness tracker", True, None, None],
            ["desk lamp", True, None, None],
        ],
        inputs=[prompt, enhance_prompt_checkbox, image1, image2],
        label="💡 Product Design Examples"
    )

# Launch settings
if __name__ == "__main__":
    # Check for required API keys
    if not os.getenv('REPLICATE_API_TOKEN'):
        print("⚠️ Warning: REPLICATE_API_TOKEN not set")
    if not os.getenv('FIREWORKS_API_KEY'):
        print("⚠️ Warning: FIREWORKS_API_KEY not set - will use fallback prompts")
    if not os.getenv('BRAVE_API_KEY'):
        print("⚠️ Warning: BRAVE_API_KEY not set - web search unavailable")
    
    demo.launch(
        share=False,  # Set to False for Hugging Face Spaces
        server_name="0.0.0.0",
        server_port=7860,
        ssr_mode=False  # Disable SSR to avoid warning
    )