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