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_with_team(user_input):
    """
    Enhanced prompt generation using team collaboration approach:
    Director -> Engineer (with web search) -> Designer
    """
    if not FIREWORKS_API_KEY:
        return user_input
    
    try:
        url = "https://api.fireworks.ai/inference/v1/chat/completions"
        
        # Team collaboration system prompt
        system_prompt = """You are a product design team consisting of three roles working in sequence:

## TEAM COLLABORATION PROCESS:

### 1. DIRECTOR (Creative Vision Leader)
The Director analyzes the user's input and establishes:
- Core product category and target market
- Key functional requirements
- Design philosophy and brand positioning
- Initial design direction and constraints
- Questions for the engineering team

### 2. ENGINEER (Technical Feasibility Expert)
The Engineer responds to the Director's vision by:
- Researching current technologies and materials
- Identifying manufacturing possibilities
- Suggesting innovative technical features
- Evaluating feasibility of proposed concepts
- Providing data on similar successful products
- Recommending technical specifications

### 3. DESIGNER (Visual Design Specialist)
The Designer synthesizes all inputs to create the final prompt:
- Translating technical specs into visual language
- Adding aesthetic details and finishing touches
- Specifying materials, textures, and colors
- Defining the rendering style
- Creating the complete visual description

## OUTPUT FORMAT:
Structure the team discussion as follows:

**DIRECTOR'S VISION:**
[Director's analysis and initial direction]

**ENGINEER'S TECHNICAL ASSESSMENT:**
[Research findings and technical recommendations]

**DIRECTOR'S REFINEMENT:**
[Approval/modification of engineering inputs]

**DESIGNER'S FINAL PROMPT:**
[Complete product design prompt with all visual details]

Focus exclusively on PHYSICAL PRODUCT DESIGN - no services, apps, or abstract concepts.
Examples: furniture, electronics, vehicles, appliances, tools, wearables, etc."""
        
        # Create the collaborative dialogue
        user_prompt = f"""Create a product design for: {user_input}

Execute the full team collaboration process:
1. Director establishes the vision
2. Engineer researches and provides technical input (simulate web research findings)
3. Director refines based on engineering feedback
4. Designer creates the final detailed visual prompt

Remember: Focus ONLY on physical product design."""
        
        payload = {
            "model": "accounts/fireworks/models/qwen3-235b-a22b-instruct-2507",
            "max_tokens": 800,
            "top_p": 0.9,
            "top_k": 40,
            "presence_penalty": 0.3,
            "frequency_penalty": 0.2,
            "temperature": 0.8,
            "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=15)
        
        if response.status_code == 200:
            result = response.json()
            enhanced = result.get('choices', [{}])[0].get('message', {}).get('content', '')
            if enhanced:
                # Extract only the designer's final prompt from the response
                if "DESIGNER'S FINAL PROMPT:" in enhanced:
                    parts = enhanced.split("DESIGNER'S FINAL PROMPT:")
                    if len(parts) > 1:
                        return parts[1].strip()
                return enhanced
        
        return user_input
        
    except Exception as e:
        print(f"Team collaboration 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
    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
    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_product_design(prompt, enhance_prompt_flag, image1, image2=None):
    """
    Process product design with team collaboration approach
    """
    if not os.getenv('REPLICATE_API_TOKEN'):
        return None, prompt, "Please set REPLICATE_API_TOKEN", ""
    
    try:
        # Store team discussion for display
        team_discussion = ""
        
        # Enhance prompt with team collaboration if requested
        final_prompt = prompt
        if enhance_prompt_flag:
            # This will contain the full team discussion
            full_response = enhance_prompt_with_team(prompt)
            final_prompt = full_response
            
            # Store the discussion for display
            team_discussion = f"🎬 **Director → Engineer → Designer Collaboration**\n\n{full_response}"
        
        # Prepare input for model
        model_input = {
            "prompt": final_prompt
        }
        
        # Add style 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 from concept!"
        
        # Run model
        output = replicate.run(
            "google/nano-banana",
            input=model_input
        )
        
        if output is None:
            return None, final_prompt, "No output received", team_discussion
        
        # 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, team_discussion
        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, team_discussion
        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, team_discussion
        
        return None, final_prompt, "Could not process output", team_discussion
        
    except Exception as e:
        error_msg = str(e)
        if "image_input" in error_msg.lower():
            return None, prompt, "Note: Model may require at least one image.", ""
        return None, prompt, f"Error: {error_msg[:100]}", ""

# Enhanced CSS for Product Design focus
css = """
.gradio-container {
    background: linear-gradient(135deg, #1e3c72 0%, #2a5298 100%);
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
    min-height: 100vh;
}
.header-container {
    background: linear-gradient(135deg, #434343 0%, #000000 100%);
    padding: 2.5rem;
    border-radius: 24px;
    margin-bottom: 2.5rem;
    box-shadow: 0 20px 60px rgba(0, 0, 0, 0.4);
}
.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.5);
}
.subtitle {
    color: #e0e0e0;
    text-align: center;
    font-size: 1.2rem;
    margin-top: 0.5rem;
    font-weight: 300;
}
.mode-indicator {
    background: rgba(255, 255, 255, 0.1);
    backdrop-filter: blur(10px);
    border-radius: 12px;
    padding: 0.8rem 1.5rem;
    margin-top: 1rem;
    text-align: center;
    font-weight: 600;
    color: #ffd700;
    border: 1px solid rgba(255, 215, 0, 0.3);
}
.main-content {
    background: rgba(255, 255, 255, 0.98);
    backdrop-filter: blur(20px);
    border-radius: 24px;
    padding: 2.5rem;
    box-shadow: 0 10px 40px rgba(0, 0, 0, 0.2);
}
.gr-button-primary {
    background: linear-gradient(135deg, #434343 0%, #000000 100%) !important;
    border: 2px solid #ffd700 !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(255, 215, 0, 0.3) !important;
    background: linear-gradient(135deg, #ffd700 0%, #ffed4e 100%) !important;
    color: #000 !important;
}
.gr-input, .gr-textarea {
    background: #ffffff !important;
    border: 2px solid #d0d0d0 !important;
    border-radius: 12px !important;
    color: #2d3436 !important;
    font-size: 1rem !important;
    padding: 0.8rem 1rem !important;
}
.gr-input:focus, .gr-textarea:focus {
    border-color: #434343 !important;
    box-shadow: 0 0 0 4px rgba(67, 67, 67, 0.1) !important;
}
.team-discussion {
    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
    border-radius: 12px;
    padding: 1.5rem;
    margin-top: 1rem;
    border-left: 4px solid #434343;
    font-family: 'Monaco', 'Courier New', monospace;
    font-size: 0.9rem;
    max-height: 400px;
    overflow-y: auto;
}
.info-box {
    background: linear-gradient(135deg, #ffd700 0%, #ffed4e 100%);
    border-radius: 12px;
    padding: 1rem;
    margin-bottom: 1rem;
    border-left: 4px solid #000;
    color: #000;
    font-weight: 500;
}
.enhanced-prompt-box {
    background: #f0f0f0;
    border-radius: 12px;
    padding: 1rem;
    margin-top: 1rem;
    border-left: 4px solid #ffd700;
}
footer {
    display: none !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">🏭 PRODUCT DESIGN STUDIO</h1>
            <p class="subtitle">AI-Powered Industrial & Product Design System</p>
            <div class="mode-indicator">
                👥 Team Collaboration: Director → Engineer → Designer
            </div>
        """)

    with gr.Column(elem_classes="main-content"):
        # Info box
        gr.HTML("""
            <div class="info-box">
                <strong>🎯 Professional Product Design Process:</strong><br>
                • <b>Director:</b> Establishes design vision and requirements<br>
                • <b>Engineer:</b> Researches technology and feasibility<br>
                • <b>Designer:</b> Creates detailed visual specifications<br>
                • Focus: Physical products only (electronics, furniture, vehicles, tools, etc.)
            </div>
        """)
        
        with gr.Row(equal_height=True):
            # Left Column - Inputs
            with gr.Column(scale=1):
                prompt = gr.Textbox(
                    label="Product Concept",
                    placeholder="Examples: 'ergonomic office chair', 'smart home speaker', 'electric bicycle', 'modular kitchen appliance'...",
                    lines=3,
                    value="professional mechanical keyboard",
                    elem_classes="prompt-input"
                )
                
                enhance_prompt_checkbox = gr.Checkbox(
                    label="👥 Enable Team Collaboration (Director + Engineer + Designer)",
                    value=True,
                    info="Activates multi-role design process with technical research"
                )
                
                with gr.Row(equal_height=True):
                    image1 = gr.Image(
                        label="Style/Material Reference 1",
                        type="pil",
                        height=200,
                        elem_classes="image-container"
                    )
                    image2 = gr.Image(
                        label="Style/Material Reference 2",
                        type="pil",
                        height=200,
                        elem_classes="image-container"
                    )
                
                generate_btn = gr.Button(
                    "Generate Product Design 🏭",
                    variant="primary",
                    size="lg"
                )
            
            # Right Column - Output
            with gr.Column(scale=1):
                output_image = gr.Image(
                    label="Product Design Visualization",
                    type="pil",
                    height=420,
                    elem_classes="image-container"
                )
                
                enhanced_prompt_display = gr.Textbox(
                    label="Final Design Specifications",
                    interactive=False,
                    lines=4,
                    elem_classes="enhanced-prompt-box",
                    visible=True,
                    value="Design specifications will appear here..."
                )
                
                status = gr.Textbox(
                    label="Status",
                    interactive=False,
                    lines=1,
                    elem_classes="status-text",
                    value="Ready to design products..."
                )
        
        # Team Discussion Display
        with gr.Row():
            team_discussion_display = gr.Markdown(
                label="Team Collaboration Process",
                value="",
                elem_classes="team-discussion"
            )
    
    # 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 Design Examples
    gr.Examples(
        examples=[
            ["ergonomic gaming mouse", True, None, None],
            ["portable coffee maker", True, None, None],
            ["modular storage system", True, None, None],
            ["smart bicycle helmet", True, None, None],
            ["industrial power tool", True, None, None],
            ["luxury watch design", True, None, None],
            ["compact home gym equipment", True, None, None],
            ["sustainable packaging solution", True, None, None],
            ["professional camera lens", True, None, None],
            ["electric vehicle charging station", True, None, None],
        ],
        inputs=[prompt, enhance_prompt_checkbox, image1, image2],
        label="Product Design Examples (Team will collaborate on these concepts)"
    )

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