Spaces:
Running
Running
| 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 | |
| ) |