Spaces:

faizee07
/

AI-UI-Designer-with-Consistent-Theming

Build error

App Files Files Community

faizee07 commited on Nov 2, 2025

Commit

ae01dc3

verified ·

1 Parent(s): e645415

Update app.py

Browse files

Files changed (1) hide show

app.py +197 -428

app.py CHANGED Viewed

@@ -5,465 +5,234 @@ import re
 from typing import Dict, List, Tuple
 import os
 class UIDesignAgent:
     def __init__(self, api_key: str):
         self.api_key = api_key
         self.base_url = "https://openrouter.ai/api/v1/chat/completions"
         self.design_system = {}
-    def extract_design_system(self, prompt: str) -> Dict:
-        """Extract and maintain consistent design system across pages"""
-        system_prompt = """You are an expert UI/UX designer. Create a consistent design system with:
-        - Primary color: #2563eb
-        - Secondary color: #7c3aed
-        - Accent color: #06b6d4
-        - Background: #ffffff
-        - Text: #1e293b
-        - Font: Inter, system-ui, sans-serif
-        Use these EXACT colors for ALL pages. Return the design system as JSON."""
-        headers = {
-            "Authorization": f"Bearer {self.api_key}",
-            "Content-Type": "application/json",
-            "HTTP-Referer": "https://huggingface.co",
-            "X-Title": "AI UI Designer"
-        }
-        data = {
-            "model": "minimax/minimax-m2:free",
-            "messages": [
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": f"Create a design system for: {prompt}"}
-            ],
-            "temperature": 0.7,
-            "max_tokens": 10000
         }
         try:
-            response = requests.post(self.base_url, headers=headers, json=data)
-            if response.status_code == 200:
-                result = response.json()
-                if 'choices' in result and len(result['choices']) > 0:
-                    design_system_text = result['choices'][0]['message']['content']
-                    # Try to extract JSON
-                    json_match = re.search(r'\{.*\}', design_system_text, re.DOTALL)
-                    if json_match:
-                        try:
-                            self.design_system = json.loads(json_match.group())
-                        except:
-                            self.design_system = self.get_default_design_system()
-                    else:
-                        self.design_system = self.get_default_design_system()
-                else:
-                    self.design_system = self.get_default_design_system()
-            else:
-                print(f"API Error: {response.status_code}")
-                self.design_system = self.get_default_design_system()
             return self.design_system
         except Exception as e:
-            print(f"Error: {e}")
             return self.get_default_design_system()
-    def get_default_design_system(self) -> Dict:
-        """Return default design system"""
-        return {
-            "colors": {
-                "primary": "#2563eb",
-                "secondary": "#7c3aed",
-                "accent": "#06b6d4",
-                "background": "#ffffff",
-                "surface": "#f8fafc",
-                "text": "#1e293b",
-                "textLight": "#64748b",
-                "border": "#e2e8f0"
-            },
-            "typography": {
-                "fontFamily": "'Inter', system-ui, -apple-system, sans-serif",
-                "h1": "2.5rem",
-                "h2": "2rem",
-                "h3": "1.5rem",
-                "body": "1rem",
-                "small": "0.875rem"
-            },
-            "spacing": {
-                "xs": "0.25rem",
-                "sm": "0.5rem",
-                "md": "1rem",
-                "lg": "1.5rem",
-                "xl": "2rem",
-                "xxl": "3rem"
-            },
-            "borderRadius": {
-                "sm": "0.25rem",
-                "md": "0.5rem",
-                "lg": "0.75rem",
-                "xl": "1rem"
-            }
-        }
-    def generate_page_html(self, page_description: str, page_number: int) -> str:
-        """Generate HTML for a specific page using consistent design system"""
-        colors = self.design_system.get('colors', self.get_default_design_system()['colors'])
-        typography = self.design_system.get('typography', self.get_default_design_system()['typography'])
-        system_prompt = f"""You are an expert UI developer. Create a complete HTML page.
-        MANDATORY DESIGN RULES:
-        - Primary color: {colors['primary']}
-        - Secondary color: {colors['secondary']}
-        - Background: {colors['background']}
-        - Text color: {colors['text']}
-        - Font family: {typography['fontFamily']}
-        Use ONLY these colors. DO NOT use any other colors.
-        Create a modern, responsive design with these exact colors.
-        Include all CSS inline in the HTML.
         """
-        headers = {
-            "Authorization": f"Bearer {self.api_key}",
-            "Content-Type": "application/json",
-            "HTTP-Referer": "https://huggingface.co",
-            "X-Title": "AI UI Designer"
-        }
-        data = {
-            "model": "minimax/minimax-m2:free",
-            "messages": [
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": f"Create page {page_number}: {page_description}. Use the EXACT colors provided. Return complete HTML with inline CSS."}
-            ],
-            "temperature": 0.7,
-            "max_tokens": 2000
-        }
         try:
-            response = requests.post(self.base_url, headers=headers, json=data)
-            if response.status_code == 200:
-                result = response.json()
-                if 'choices' in result and len(result['choices']) > 0:
-                    html_content = result['choices'][0]['message']['content']
-                    # Extract HTML from code blocks
-                    html_match = re.search(r'```html?\n?(.*?)\n?```', html_content, re.DOTALL)
-                    if html_match:
-                        html_content = html_match.group(1)
-                    # Inject consistent CSS
-                    html_content = self.inject_design_system_css(html_content)
-                    return html_content
-                else:
-                    return self.create_error_page(f"No response for page {page_number}")
-            else:
-                return self.create_error_page(f"API Error: {response.status_code}")
         except Exception as e:
-            return self.create_error_page(f"Error: {str(e)}")
-    def inject_design_system_css(self, html: str) -> str:
-        """Inject consistent CSS variables into HTML"""
-        colors = self.design_system.get('colors', self.get_default_design_system()['colors'])
-        typography = self.design_system.get('typography', self.get_default_design_system()['typography'])
-        css_injection = f"""
-        <style>
-            :root {{
-                --primary: {colors['primary']};
-                --secondary: {colors['secondary']};
-                --accent: {colors['accent']};
-                --background: {colors['background']};
-                --surface: {colors['surface']};
-                --text: {colors['text']};
-                --text-light: {colors['textLight']};
-                --border: {colors['border']};
-                --font-family: {typography['fontFamily']};
-            }}
-            * {{
-                margin: 0;
-                padding: 0;
-                box-sizing: border-box;
-            }}
-            body {{
-                font-family: var(--font-family);
-                background-color: var(--background);
-                color: var(--text);
-                line-height: 1.6;
-            }}
-            h1, h2, h3, h4, h5, h6 {{
-                color: var(--text);
-                margin-bottom: 1rem;
-            }}
-            a {{
-                color: var(--primary);
-                text-decoration: none;
-            }}
-            a:hover {{
-                text-decoration: underline;
-            }}
-            button, .btn {{
-                background-color: var(--primary);
-                color: white;
-                padding: 0.75rem 1.5rem;
-                border: none;
-                border-radius: 0.5rem;
-                cursor: pointer;
-                font-family: var(--font-family);
-                transition: all 0.3s ease;
-            }}
-            button:hover, .btn:hover {{
-                transform: translateY(-2px);
-                box-shadow: 0 4px 12px rgba(37, 99, 235, 0.3);
-            }}
-        </style>
-        """
-        # Add viewport meta tag if not present
-        viewport_meta = '<meta name="viewport" content="width=device-width, initial-scale=1.0">'
-        if '<head>' in html:
-            html = html.replace('<head>', f'<head>\n{viewport_meta}\n{css_injection}')
-        elif '<html>' in html:
-            html = html.replace('<html>', f'<html>\n<head>\n{viewport_meta}\n{css_injection}\n</head>')
-        else:
-            html = f"""<!DOCTYPE html>
-<html>
-<head>
-{viewport_meta}
-{css_injection}
-</head>
-<body>
-{html}
-</body>
-</html>"""
-        return html
-    def create_error_page(self, error_message: str) -> str:
-        """Create an error page with consistent styling"""
-        return f"""
-        <!DOCTYPE html>
-        <html>
-        <head>
-            <style>
-                body {{
-                    font-family: system-ui, -apple-system, sans-serif;
-                    display: flex;
-                    align-items: center;
-                    justify-content: center;
-                    height: 100vh;
-                    margin: 0;
-                    background: #f8fafc;
-                }}
-                .error {{
-                    background: white;
-                    padding: 2rem;
-                    border-radius: 0.5rem;
-                    box-shadow: 0 1px 3px rgba(0,0,0,0.1);
-                    max-width: 500px;
-                }}
-                h2 {{
-                    color: #ef4444;
-                    margin-bottom: 1rem;
-                }}
-                p {{
-                    color: #64748b;
-                }}
-            </style>
-        </head>
-        <body>
-            <div class="error">
-                <h2>Error Generating Page</h2>
-                <p>{error_message}</p>
-            </div>
-        </body>
-        </html>
-        """
-    def extract_pages_from_prompt(self, prompt: str) -> List[str]:
-        """Extract individual page descriptions from prompt"""
-        # Simple extraction based on common patterns
-        pages = []
-        # Check for numbered pages
-        numbered_pattern = r'\d+\)([^0-9]+?)(?=\d+\)|$)'
-        matches = re.findall(numbered_pattern, prompt)
-        if matches:
-            pages = [match.strip() for match in matches]
-        else:
-            # Check for comma-separated pages
-            if ',' in prompt:
-                pages = [p.strip() for p in prompt.split(',')]
-            else:
-                # Treat as single page
-                pages = [prompt]
-        # Ensure we have at least 3 pages
-        while len(pages) < 3:
-            pages.append(f"Additional page {len(pages) + 1}")
-        return pages[:3]  # Limit to 3 pages
     def generate_multi_page_design(self, prompt: str) -> Tuple[List[str], Dict]:
-        """Generate multiple pages with consistent design"""
-        # Extract design system first
         self.design_system = self.extract_design_system(prompt)
-        # Extract pages from prompt
-        pages = self.extract_pages_from_prompt(prompt)
-        # Generate HTML for each page
-        generated_pages = []
-        for i, page_desc in enumerate(pages[:3], 1):  # Limit to 3 pages
-            html = self.generate_page_html(page_desc, i)
-            generated_pages.append(html)
-        # Ensure we have exactly 3 pages
-        while len(generated_pages) < 3:
-            generated_pages.append(self.create_error_page("Page not specified"))
         return generated_pages, self.design_system
-def create_ui_design(prompt: str, api_key: str):
-    """Main function to create UI designs"""
-    if not api_key or api_key.strip() == "":
-        error_html = """
-        <div style='padding: 20px; color: #ef4444; background: #fee2e2; border-radius: 8px;'>
-            <h3>API Key Required</h3>
-            <p>Please provide your OpenRouter API key to generate designs.</p>
-            <p>Get your key from <a href='https://openrouter.ai' target='_blank'>OpenRouter.ai</a></p>
-        </div>
-        """
-        return error_html, error_html, error_html, "{}"
-    if not prompt or prompt.strip() == "":
-        error_html = """
-        <div style='padding: 20px; color: #f59e0b; background: #fef3c7; border-radius: 8px;'>
-            <h3>Prompt Required</h3>
-            <p>Please describe the website you want to create.</p>
-        </div>
-        """
-        return error_html, error_html, error_html, "{}"
-    try:
-        agent = UIDesignAgent(api_key)
-        pages, design_system = agent.generate_multi_page_design(prompt)
-        # Format design system for display
-        design_system_display = json.dumps(design_system, indent=2)
-        return pages[0], pages[1], pages[2], design_system_display
-    except Exception as e:
-        error_html = f"""
-        <div style='padding: 20px; color: #ef4444; background: #fee2e2; border-radius: 8px;'>
-            <h3>Generation Error</h3>
-            <p>{str(e)}</p>
-        </div>
-        """
-        return error_html, error_html, error_html, "{}"
-# Create Gradio interface
 def create_interface():
-    with gr.Blocks(title="AI UI Designer - Consistent Multi-Page", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("""
-        # 🎨 AI UI Designer - Consistent Multi-Page Generator
-        ### Generate multiple website pages with consistent design system
-        This tool ensures **consistent colors and styling** across all generated pages.
-        """)
         with gr.Row():
             with gr.Column(scale=1):
-                api_key_input = gr.Textbox(
-                    label="OpenRouter API Key",
-                    placeholder="sk-or-...",
-                    type="password",
-                    info="Get your key from openrouter.ai"
-                )
-                prompt_input = gr.Textbox(
-                    label="Website Description",
-                    placeholder="Example: Create a tech startup website with: 1) Landing page with hero section and features, 2) About page with team section, 3) Contact page with form",
-                    lines=4
-                )
-                with gr.Row():
-                    generate_btn = gr.Button("🚀 Generate Design", variant="primary", scale=2)
-                    clear_btn = gr.Button("🗑️ Clear", scale=1)
-                gr.Markdown("### 📋 Design System")
-                design_system_output = gr.Code(
-                    language="json",
-                    lines=10
-                )
             with gr.Column(scale=2):
                 with gr.Tabs():
-                    with gr.TabItem("📄 Page 1"):
-                        page1_output = gr.HTML()
-                    with gr.TabItem("📄 Page 2"):
-                        page2_output = gr.HTML()
-                    with gr.TabItem("📄 Page 3"):
-                        page3_output = gr.HTML()
-        gr.Markdown("""
-        ---
-        ### 💡 Tips for Best Results:
-        - Clearly describe each page you want
-        - Specify key components (hero, navbar, footer, etc.)
-        - Mention any specific features you need
-        ### 🎯 Features:
-        - ✅ Consistent color scheme across all pages
-        - ✅ Unified typography and spacing
-        - ✅ Modern, responsive design
-        - ✅ Professional UI components
-        """)
-        # Set up event handlers
-        generate_btn.click(
-            fn=create_ui_design,
-            inputs=[prompt_input, api_key_input],
-            outputs=[page1_output, page2_output, page3_output, design_system_output]
-        )
-        clear_btn.click(
-            fn=lambda: ("", "", "", "", "{}"),
-            outputs=[prompt_input, page1_output, page2_output, page3_output, design_system_output]
-        )
-        # Add examples - FIXED STRING LITERALS
-        example_prompts = [
-            "Create a SaaS website: 1) Landing page with hero and pricing cards, 2) Features page with grid layout, 3) Contact page with form",
-            "Design a portfolio: 1) Homepage with projects showcase, 2) About page with skills, 3) Contact page",
-            "Build an e-commerce site: 1) Homepage with products, 2) Product detail page, 3) Cart page"
-        ]
         gr.Examples(
-            examples=example_prompts,
             inputs=prompt_input
         )
     return demo
 if __name__ == "__main__":
     demo = create_interface()
-    demo.launch(share=True)  # Add share=True for public link

 from typing import Dict, List, Tuple
 import os
+# Let's use a model that's good at this task and very cheap/fast
+SELECTED_MODEL = "minimax/minimax-m2:free"
+# <<< NEW FUNCTION TO CREATE THE VISUAL STYLE GUIDE >>>
+def visualize_design_system(design_system: Dict) -> str:
+    """
+    Takes a design system dictionary and returns an HTML string
+    that visually represents it as a style guide.
+    """
+    if not design_system or not isinstance(design_system, dict):
+        return "<div class='error-box'>No valid design system to display.</div>"
+    # Safely get data with defaults
+    colors = design_system.get('colors', {})
+    typography = design_system.get('typography', {})
+    spacing = design_system.get('spacing', {})
+    border_radius = design_system.get('borderRadius', {})
+    # ---- HTML & CSS for the visualizer itself ----
+    html = """
+    <!DOCTYPE html>
+    <html lang="en">
+    <head>
+        <meta charset="UTF-8">
+        <meta name="viewport" content="width=device-width, initial-scale=1.0">
+        <title>Design System</title>
+        <style>
+            body {
+                font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Helvetica, Arial, sans-serif;
+                background-color: #f9fafb;
+                color: #374151;
+                padding: 1.5rem;
+                line-height: 1.6;
+            }
+            .section { margin-bottom: 2.5rem; }
+            h2 {
+                font-size: 1.5rem; font-weight: 600; color: #111827;
+                border-bottom: 1px solid #e5e7eb; padding-bottom: 0.5rem; margin-bottom: 1.5rem;
+            }
+            .swatch-container { display: flex; flex-wrap: wrap; gap: 1.5rem; }
+            .swatch {
+                width: 120px; height: 120px; border-radius: 0.75rem;
+                display: flex; flex-direction: column; justify-content: flex-end;
+                padding: 0.75rem; font-size: 0.875rem;
+                box-shadow: 0 4px 6px -1px rgba(0,0,0,0.1), 0 2px 4px -2px rgba(0,0,0,0.1);
+            }
+            .swatch .name { font-weight: 600; }
+            .swatch .hex { opacity: 0.8; }
+            .type-container .sample { padding: 1rem; background-color: #fff; border: 1px solid #e5e7eb; border-radius: 0.5rem; margin-bottom: 1rem; }
+            .type-container .label { font-size: 0.9rem; color: #6b7280; margin-bottom: 0.25rem; }
+            .component-container { display: flex; flex-direction: column; gap: 2rem; background-color: #fff; padding: 2rem; border-radius: 0.75rem; border: 1px solid #e5e7eb;}
+            .error-box { padding: 1rem; background: #fee2e2; color: #b91c1c; border-radius: 0.5rem; }
+        </style>
+    </head>
+    <body>
+    """
+    # ---- 1. Color Palette Section ----
+    html += "<div class='section'><h2>Color Palette</h2><div class='swatch-container'>"
+    for name, hex_val in colors.items():
+        # Determine if text should be light or dark based on background
+        try:
+            r, g, b = int(hex_val[1:3], 16), int(hex_val[3:5], 16), int(hex_val[5:7], 16)
+            luminance = (0.299 * r + 0.587 * g + 0.114 * b) / 255
+            text_color = '#ffffff' if luminance < 0.5 else '#000000'
+        except:
+            text_color = '#000000'
+        html += f"""
+        <div class="swatch" style="background-color: {hex_val}; color: {text_color};">
+            <span class="name">{name.capitalize()}</span>
+            <span class="hex">{hex_val}</span>
+        </div>
+        """
+    html += "</div></div>"
+    # ---- 2. Typography Section ----
+    font_family = typography.get('fontFamily', 'sans-serif')
+    html += f"<div class='section type-container'><h2>Typography</h2><p style='margin-bottom: 1.5rem;'>Font Family: <code>{font_family}</code></p>"
+    type_samples = {'h1': 'Heading 1', 'h2': 'Heading 2', 'body': 'This is a paragraph of body text. It is the default font size for long-form content.', 'small': 'Small utility text'}
+    for name, sample_text in type_samples.items():
+        font_size = typography.get(name, '1rem')
+        html += f"""
+        <div class="sample">
+            <div class="label">{name.capitalize()} ({font_size})</div>
+            <div style="font-family: {font_family}; font-size: {font_size};">
+                {sample_text}
+            </div>
+        </div>
+        """
+    html += "</div>"
+    # ---- 3. Components Section ----
+    primary_color = colors.get('primary', '#3b82f6')
+    surface_color = colors.get('surface', '#f9fafb')
+    text_color = colors.get('text', '#111827')
+    border_color = colors.get('border', '#e5e7eb')
+    radius = border_radius.get('lg', '0.75rem')
+    html += f"""
+    <div class='section'>
+        <h2>Components Preview</h2>
+        <div class="component-container" style="background-color: {surface_color};">
+            <div>
+                <h3 style="font-family:{font_family}; font-size:{typography.get('h3', '1.25rem')}; color:{text_color}; margin-bottom: 1rem;">Example Card</h3>
+                <div style="background-color: #fff; border: 1px solid {border_color}; border-radius: {radius}; padding: 1.5rem; box-shadow: 0 1px 3px 0 rgba(0,0,0,0.1);">
+                    <p style="font-family:{font_family}; font-size:{typography.get('body', '1rem')}; color:{text_color}; margin-bottom: 1.5rem;">This is a sample card using the design system's surface, border, and radius values.</p>
+                    <button style="font-family:{font_family}; background-color: {primary_color}; color: white; border: none; padding: 0.75rem 1.5rem; border-radius: {border_radius.get('md', '0.5rem')}; cursor: pointer; font-size: 1rem;">Primary Button</button>
+                </div>
+            </div>
+        </div>
+    </div>
+    """
+    html += "</body></html>"
+    return html
 class UIDesignAgent:
+    # ... (The class code remains mostly the same, just ensuring it returns a proper dictionary) ...
     def __init__(self, api_key: str):
         self.api_key = api_key
         self.base_url = "https://openrouter.ai/api/v1/chat/completions"
+        self.model = SELECTED_MODEL
         self.design_system = {}
+    def get_default_design_system(self) -> Dict:
+        return {
+            "colors": { "primary": "#2563eb", "secondary": "#7c3aed", "accent": "#06b6d4", "background": "#ffffff", "surface": "#f8fafc", "text": "#1e293b", "textLight": "#64748b", "border": "#e2e8f0" },
+            "typography": { "fontFamily": "'Inter', system-ui, sans-serif", "h1": "2.5rem", "h2": "2rem", "h3": "1.5rem", "body": "1rem", "small": "0.875rem" },
+            "spacing": { "md": "1rem", "lg": "1.5rem", "xl": "2rem" },
+            "borderRadius": { "sm": "0.25rem", "md": "0.5rem", "lg": "0.75rem" }
         }
+    def extract_design_system(self, prompt: str) -> Dict:
+        system_prompt = """You are an expert UI/UX designer. Create a consistent design system based on the user's prompt.
+        The design system must include: 'colors' (primary, secondary, background, surface, text, border),
+        'typography' (fontFamily, h1, h2, body, small), and 'borderRadius' (sm, md, lg).
+        Return ONLY a single, valid JSON object and nothing else."""
+        headers = { "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json" }
+        data = { "model": self.model, "messages": [{"role": "system", "content": system_prompt}, {"role": "user", "content": f"Create a design system for: {prompt}"}], "response_format": {"type": "json_object"}}
         try:
+            response = requests.post(self.base_url, headers=headers, json=data, timeout=30)
+            result = response.json()
+            content = result['choices'][0]['message']['content']
+            self.design_system = json.loads(content)
             return self.design_system
         except Exception as e:
+            print(f"Error extracting design system: {e}. Using default.")
             return self.get_default_design_system()
+    def generate_page_html(self, page_description: str) -> str:
+        # Use the stored design system
+        design_system_str = json.dumps(self.design_system)
+        system_prompt = f"""You are an expert UI developer. Create a complete, single HTML file.
+        Use this exact JSON design system for all styling: {design_system_str}
+        RULES:
+        1. Create a modern, responsive page.
+        2. Embed all CSS in a `<style>` tag in the `<head>`.
+        3. Your entire response must be ONLY raw HTML code, starting with `<!DOCTYPE html>`.
         """
+        headers = { "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json" }
+        data = { "model": self.model, "messages": [{"role": "system", "content": system_prompt}, {"role": "user", "content": f"Design page: {page_description}."}], "max_tokens": 3500 }
         try:
+            response = requests.post(self.base_url, headers=headers, json=data, timeout=60)
+            return response.json()['choices'][0]['message']['content']
         except Exception as e:
+            return f"<html><body>Error generating page: {e}</body></html>"
     def generate_multi_page_design(self, prompt: str) -> Tuple[List[str], Dict]:
         self.design_system = self.extract_design_system(prompt)
+        pages_desc = re.findall(r'\d+\)([^0-9]+?)(?=\d+\)|$)', prompt) or [prompt]
+        while len(pages_desc) < 3: pages_desc.append(f"Page {len(pages_desc)+1} placeholder")
+        generated_pages = [self.generate_page_html(desc) for desc in pages_desc[:3]]
         return generated_pages, self.design_system
+def create_ui_design(prompt: str, api_key: str, progress=gr.Progress(track_tqdm=True)):
+    if not api_key: return ("Missing API Key",) * 5
+    agent = UIDesignAgent(api_key)
+    progress(0.2, desc="Generating Design System...")
+    pages, design_system = agent.generate_multi_page_design(prompt)
+    progress(0.8, desc="Creating Visualizer...")
+    # <<< CALL THE NEW VISUALIZER FUNCTION >>>
+    visualizer_html = visualize_design_system(design_system)
+    design_system_json_string = json.dumps(design_system, indent=2)
+    progress(1.0, desc="Done!")
+    return pages[0], pages[1], pages[2], design_system_json_string, visualizer_html
 def create_interface():
+    with gr.Blocks(title="AI UI Designer", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("# 🎨 AI UI Designer - Consistent Multi-Page Generator")
         with gr.Row():
             with gr.Column(scale=1):
+                api_key_input = gr.Textbox(label="OpenRouter API Key", placeholder="sk-or-...", type="password")
+                prompt_input = gr.Textbox(label="Website Description", placeholder="e.g., A 3-page site for a coffee shop: 1) landing page, 2) menu, 3) contact us", lines=4)
+                generate_btn = gr.Button("🚀 Generate Design", variant="primary")
+                # <<< UPDATED UI WITH TABS FOR JSON and VISUALIZER >>>
+                with gr.Tabs():
+                    with gr.TabItem("Visualizer"):
+                        design_system_visualizer = gr.HTML(label="Design System Visualizer")
+                    with gr.TabItem("JSON"):
+                        design_system_output = gr.Code(label="Design System JSON", language="json", lines=15)
             with gr.Column(scale=2):
                 with gr.Tabs():
+                    with gr.TabItem("📄 Page 1"): page1_output = gr.HTML()
+                    with gr.TabItem("📄 Page 2"): page2_output = gr.HTML()
+                    with gr.TabItem("📄 Page 3"): page3_output = gr.HTML()
         gr.Examples(
+            examples=["A dark-mode site for a DJ: 1) homepage with event list, 2) music gallery, 3) booking form", "A bright, clean site for a bakery: 1) landing page with featured items, 2) full menu page, 3) online order form"],
             inputs=prompt_input
         )
+        # <<< UPDATE THE OUTPUTS TO INCLUDE THE NEW VISUALIZER COMPONENT >>>
+        generate_btn.click(
+            fn=create_ui_design,
+            inputs=[prompt_input, api_key_input],
+            outputs=[page1_output, page2_output, page3_output, design_system_output, design_system_visualizer]
+        )
     return demo
 if __name__ == "__main__":
     demo = create_interface()
+    demo.launch()