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("""
🏭 PRODUCT DESIGN STUDIO
AI-Powered Industrial & Product Design System with Web Research
👥 Director → 🔍 Engineer (with Brave Search) → 🎨 Designer
""")
with gr.Column(elem_classes="main-content"):
gr.HTML("""
🎯 Professional Product Design Process:
• Director: Establishes vision and requirements
• Engineer: Searches current trends & technologies via Brave API
• Designer: Creates detailed visual specifications
• Focus: Physical products only (electronics, furniture, vehicles, appliances)
""")
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
)