Spaces:

aqibali06
/

website_code_generator

Sleeping

File size: 20,457 Bytes

import gradio as gr
import google.generativeai as genai
import os
import tempfile
import zipfile
from PIL import Image
import io
import base64
import re

# Configure Gemini API
def configure_gemini(api_key):
    genai.configure(api_key=api_key)
    return genai.GenerativeModel('gemini-1.5-flash')

# Enhanced prompt for better code generation
def create_prompt(input_type, description=None, multiple_variations=False):
    base_prompt = """
You are an expert frontend developer specializing in modern, accessible, and responsive web design.
Generate clean, semantic HTML5 code with embedded CSS and JavaScript based on the provided design.

CRITICAL REQUIREMENTS:
- Create a complete HTML5 document with proper DOCTYPE, meta tags, and structure
- Use Google Fonts (Poppins, Inter, or Roboto) for modern typography
- Implement CSS custom properties (variables) for consistent theming
- Use semantic HTML5 elements (header, nav, main, section, article, aside, footer)
- Ensure full responsiveness with mobile-first approach using CSS Grid and Flexbox
- Add proper ARIA attributes and accessibility features
- Include smooth animations and hover effects using CSS transitions
- Use modern CSS features (clamp(), min(), max(), custom properties)
- Implement proper color contrast ratios (WCAG AA compliant)
- Add focus states for keyboard navigation
- Use rem/em units for scalability
- Include proper meta viewport and other essential meta tags

CSS STRUCTURE REQUIREMENTS:
- Reset/normalize styles
- CSS custom properties for colors, spacing, typography
- Responsive breakpoints using CSS Grid/Flexbox
- Smooth transitions and hover effects
- Proper z-index management
- Modern button styles with states (hover, focus, active, disabled)

JAVASCRIPT REQUIREMENTS:
- Use modern ES6+ syntax
- Add smooth scrolling behavior
- Implement proper form validation if forms are present
- Add loading states and user feedback
- Handle errors gracefully
- Use event delegation where appropriate

HTML STRUCTURE TEMPLATE:
```html
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <meta name="description" content="Generated UI">
    <title>Generated UI</title>
    <link href="https://fonts.googleapis.com/css2?family=Poppins:wght@300;400;500;600;700&display=swap" rel="stylesheet">
    <style>
        /* CSS goes here with proper structure */
    </style>
</head>
<body>
    <!-- Semantic HTML content -->
    <script>
        /* Modern JavaScript */
    </script>
</body>
</html>
```
"""

    variation_prompt = ""
    if multiple_variations:
        variation_prompt = """
IMPORTANT: Generate 3 different design variations of the same concept:
- Variation 1: Modern/Minimalist style
- Variation 2: Bold/Colorful style  
- Variation 3: Professional/Corporate style

Separate each variation with the exact comment: <!-- VARIATION SPLIT -->
Each variation should be a complete HTML document.
"""

    if input_type == "text":
        return base_prompt + variation_prompt + f"""
Design Description: {description}

Generate the complete HTML code that implements this design following all requirements above.
Make the design visually stunning with proper spacing, typography, colors, and interactions.
Return only the HTML code without any markdown formatting, explanations, or code fences.
"""
    else:
        return base_prompt + variation_prompt + """
Based on the uploaded design image, recreate the UI as accurately as possible while following all modern web standards.
Pay attention to:
- Exact layout and positioning of elements
- Color schemes and visual hierarchy
- Typography styles and font weights
- Interactive elements and their states
- Spacing, margins, and overall proportions
- Modern styling techniques

Generate complete HTML code that matches the design while being fully responsive and accessible.
Return only the HTML code without any markdown formatting, explanations, or code fences.
"""

def generate_ui_code(api_key, input_type, text_description, image_file, generate_variations):
    """Generate UI code based on text description or image"""
    try:
        # Configure Gemini
        if not api_key:
            return "❌ Please provide your Gemini API key", "", "", "", []

        model = configure_gemini(api_key)

        # Prepare the prompt
        prompt = create_prompt(input_type, text_description, generate_variations)

        # Generate based on input type
        if input_type == "Text Description":
            if not text_description:
                return "❌ Please provide a text description", "", "", "", []

            response = model.generate_content(prompt)
            generated_code = response.text

        else:  # Image input
            if not image_file:
                return "❌ Please upload an image", "", "", "", []

            # Process the image
            image = Image.open(image_file)

            # Generate content with image
            response = model.generate_content([prompt, image])
            generated_code = response.text

        # Clean up the generated code
        cleaned_code = clean_generated_code(generated_code)
        
        # Handle variations
        variations = []
        main_code = cleaned_code
        
        if generate_variations and "<!-- VARIATION SPLIT -->" in cleaned_code:
            variations = cleaned_code.split("<!-- VARIATION SPLIT -->")
            variations = [clean_generated_code(var.strip()) for var in variations if var.strip()]
            main_code = variations[0] if variations else cleaned_code
            variation_options = [f"Variation {i+1}" for i in range(len(variations))]
        else:
            variation_options = []

        # Create preview (truncated version for display)
        preview = generate_preview(main_code)

        return "✅ Code generated successfully!", main_code, preview, main_code, variation_options

    except Exception as e:
        return f"❌ Error: {str(e)}", "", "", "", []

def clean_generated_code(code):
    """Clean and format the generated code"""
    if not code:
        return ""
    
    # Remove markdown code blocks
    code = re.sub(r'```html\n?', '', code)
    code = re.sub(r'```\n?', '', code)
    code = re.sub(r'^```.*\n', '', code, flags=re.MULTILINE)
    
    # Remove any leading/trailing whitespace
    code = code.strip()
    
    # If it's not a complete HTML document, wrap it
    if not code.startswith("<!DOCTYPE html>") and not code.startswith("<html"):
        code = f"""<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <meta name="description" content="Generated UI">
    <title>Generated UI</title>
    <link href="https://fonts.googleapis.com/css2?family=Poppins:wght@300;400;500;600;700&display=swap" rel="stylesheet">
    <style>
        * {{
            margin: 0;
            padding: 0;
            box-sizing: border-box;
        }}
        
        :root {{
            --primary-color: #3b82f6;
            --secondary-color: #64748b;
            --accent-color: #f59e0b;
            --background-color: #ffffff;
            --text-color: #1f2937;
            --border-color: #e5e7eb;
            --shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
            --border-radius: 8px;
            --transition: all 0.3s ease;
            --font-family: 'Poppins', sans-serif;
        }}
        
        body {{
            font-family: var(--font-family);
            line-height: 1.6;
            color: var(--text-color);
            background-color: var(--background-color);
        }}
        
        .container {{
            max-width: 1200px;
            margin: 0 auto;
            padding: 0 1rem;
        }}
        
        @media (max-width: 768px) {{
            .container {{
                padding: 0 0.5rem;
            }}
        }}
    </style>
</head>
<body>
{code}
</body>
</html>"""
    
    # Ensure proper Google Fonts link if not present
    if "fonts.googleapis.com" not in code:
        code = code.replace(
            '<head>',
            '<head>\n    <link href="https://fonts.googleapis.com/css2?family=Poppins:wght@300;400;500;600;700&display=swap" rel="stylesheet">'
        )
    
    return code

def generate_preview(code):
    """Generate a preview of the code (first 1000 characters)"""
    if len(code) > 1000:
        return code[:1000] + "..."
    return code

def download_code(code):
    """Create a downloadable file with the generated code"""
    if not code:
        return None

    # Create a temporary file
    with tempfile.NamedTemporaryFile(mode='w', suffix='.html', delete=False) as f:
        f.write(code)
        return f.name

def toggle_input_visibility(choice):
    """Toggle visibility of input components based on choice"""
    if choice == "Text Description":
        return gr.update(visible=True), gr.update(visible=False)
    else:
        return gr.update(visible=False), gr.update(visible=True)

def update_variation_display(choice, variations_data):
    """Update the displayed variation based on selection"""
    if not variations_data or not choice:
        return "", ""
    
    try:
        # Extract variation number from choice (e.g., "Variation 1" -> 0)
        var_index = int(choice.split()[-1]) - 1
        if 0 <= var_index < len(variations_data):
            selected_code = variations_data[var_index]
            preview = generate_preview(selected_code)
            return selected_code, preview
    except:
        pass
    
    return "", ""

def update_live_preview(code):
    """Update the live preview with the generated code"""
    if not code or not code.strip():
        return "<div style='padding: 20px; text-align: center; color: #666;'>No code to preview</div>"
    
    # Return the code directly for iframe rendering
    return code

# Create the Gradio interface
def create_interface():
    with gr.Blocks(
        title="Enhanced UI Code Generator",
        theme=gr.themes.Soft(),
        css="""
        .container { max-width: 1400px; margin: 0 auto; }
        .header { text-align: center; margin-bottom: 30px; padding: 20px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; border-radius: 15px; }
        .input-section { background: #f8f9fa; padding: 25px; border-radius: 15px; margin-bottom: 20px; box-shadow: 0 4px 6px rgba(0,0,0,0.1); }
        .output-section { background: #fff; padding: 25px; border-radius: 15px; border: 1px solid #e9ecef; box-shadow: 0 4px 6px rgba(0,0,0,0.1); }
        .status-success { color: #28a745; font-weight: bold; }
        .status-error { color: #dc3545; font-weight: bold; }
        .gradio-button { background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border: none; }
        .gradio-button:hover { transform: translateY(-2px); box-shadow: 0 8px 15px rgba(0,0,0,0.2); }
        .preview-container { border: 1px solid #e5e7eb; border-radius: 8px; overflow: hidden; }
        .variation-selector { margin-bottom: 15px; }
        """
    ) as interface:

        # Store variations data
        variations_data = gr.State([])

        gr.Markdown("""
        # 🎨 Enhanced UI Code Generator
        Generate modern, accessible, and responsive frontend code from design descriptions or images using AI
        
        **✨ New Features:**
        - 🌐 **Live Preview** - See your generated website in real-time
        - 🎭 **Multiple Variations** - Generate different design styles
        - 🚀 **Modern Standards** - HTML5, CSS Grid, Accessibility, Google Fonts
        - 📱 **Mobile-First** - Fully responsive design
        - ♿ **Accessible** - WCAG compliant with ARIA attributes
        """, elem_classes=["header"])

        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("## ⚙️ Configuration")
                api_key = gr.Textbox(
                    label="Gemini API Key",
                    placeholder="Enter your Google Gemini API key",
                    type="password",
                    info="Get your API key from Google AI Studio"
                )

                gr.Markdown("## 📝 Input Method")
                input_type = gr.Radio(
                    choices=["Text Description", "Image Upload"],
                    value="Text Description",
                    label="Choose input method"
                )

                # Text input (visible by default)
                text_description = gr.Textbox(
                    label="Design Description",
                    placeholder="Describe your UI design in detail (e.g., 'A modern SaaS landing page with hero section, features grid, pricing cards, and footer')",
                    lines=4,
                    visible=True
                )

                # Image input (hidden by default)
                image_file = gr.File(
                    label="Upload Design Image",
                    file_types=["image"],
                    visible=False
                )

                # Variations option
                generate_variations = gr.Checkbox(
                    label="🎭 Generate Multiple Design Variations",
                    value=False,
                    info="Generate 3 different style variations (Modern, Bold, Professional)"
                )

                # Toggle visibility based on input type
                input_type.change(
                    fn=toggle_input_visibility,
                    inputs=[input_type],
                    outputs=[text_description, image_file]
                )

                generate_btn = gr.Button("🚀 Generate Code", variant="primary", size="lg")

            with gr.Column(scale=2):
                gr.Markdown("## 📋 Generated Code & Preview")

                status = gr.Markdown("Ready to generate code...")

                # Variation selector (hidden by default)
                variation_selector = gr.Dropdown(
                    label="🎭 Select Variation",
                    choices=[],
                    visible=False,
                    elem_classes=["variation-selector"]
                )

                with gr.Tabs():
                    with gr.TabItem("🌐 Live Preview"):
                        live_preview = gr.HTML(
                            value="<div style='padding: 40px; text-align: center; color: #666; background: #f8f9fa; border-radius: 10px;'>🎨 Your generated website will appear here</div>",
                            elem_classes=["preview-container"]
                        )

                    with gr.TabItem("📄 Full Code"):
                        generated_code = gr.Code(
                            label="Generated HTML/CSS/JS",
                            language="html",
                            lines=25,
                            show_label=True
                        )

                        download_btn = gr.DownloadButton(
                            label="💾 Download HTML File",
                            variant="secondary",
                            size="lg"
                        )

                    with gr.TabItem("👁️ Code Preview"):
                        code_preview = gr.Code(
                            label="Code Preview (First 1000 characters)",
                            language="html",
                            lines=20,
                            show_label=True
                        )

                        gr.Markdown("""
                        **Note:** This shows the first 1000 characters of the generated code.
                        Use the 'Full Code' tab to see the complete implementation.
                        """)

        # Enhanced Examples section
        gr.Markdown("## 💡 Professional Example Descriptions")
        gr.Examples(
            examples=[
                ["A modern SaaS landing page with animated hero section, feature cards, pricing tiers, testimonials, and newsletter signup"],
                ["A responsive e-commerce product page with image gallery, product details, reviews, and shopping cart functionality"],
                ["A professional portfolio website with animated header, project showcase grid, skills section, and contact form"],
                ["A modern dashboard interface with sidebar navigation, data visualization cards, charts, and user profile section"],
                ["A restaurant website with hero banner, menu sections, reservation form, location map, and social media integration"],
                ["A fitness app landing page with workout tracking features, progress charts, membership plans, and mobile app download"],
                ["A real estate property listing page with photo carousel, property details, virtual tour button, and contact agent form"],
                ["A modern blog homepage with featured articles, category filters, search functionality, and newsletter subscription"]
            ],
            inputs=[text_description],
            label="Click on any example to use it as your prompt"
        )

        # Event handlers
        def handle_generation(api_key, input_type, text_description, image_file, generate_variations):
            """Handle the generation process and update all outputs"""
            status_msg, code, preview, live_code, var_options = generate_ui_code(
                api_key, input_type, text_description, image_file, generate_variations
            )
            
            # Update variation selector visibility and options
            if var_options:
                variations_data_list = live_code.split("<!-- VARIATION SPLIT -->") if "<!-- VARIATION SPLIT -->" in live_code else [live_code]
                variations_data_list = [clean_generated_code(var.strip()) for var in variations_data_list if var.strip()]
                return (
                    status_msg, 
                    code, 
                    preview, 
                    update_live_preview(code),
                    variations_data_list,
                    gr.update(choices=var_options, visible=True, value=var_options[0] if var_options else None)
                )
            else:
                return (
                    status_msg, 
                    code, 
                    preview, 
                    update_live_preview(code),
                    [code] if code else [],
                    gr.update(choices=[], visible=False, value=None)
                )

        generate_btn.click(
            fn=handle_generation,
            inputs=[api_key, input_type, text_description, image_file, generate_variations],
            outputs=[status, generated_code, code_preview, live_preview, variations_data, variation_selector]
        )

        # Handle variation selection
        variation_selector.change(
            fn=lambda choice, variations: update_variation_display(choice, variations),
            inputs=[variation_selector, variations_data],
            outputs=[generated_code, code_preview]
        )

        # Update live preview when code changes
        generated_code.change(
            fn=update_live_preview,
            inputs=[generated_code],
            outputs=[live_preview]
        )

        # Handle download
        generated_code.change(
            fn=download_code,
            inputs=[generated_code],
            outputs=[download_btn]
        )

    return interface

# Launch the application
if __name__ == "__main__":
    # Install required packages (run this in Colab)
    print("🚀 Enhanced UI Code Generator")
    print("=" * 50)
    print("Installing required packages...")
    print("Run this command in your Colab notebook:")
    print("!pip install gradio google-generativeai pillow")
    print()
    print("Features:")
    print("✅ Live website preview in iframe")
    print("✅ Multiple design variations")
    print("✅ Modern HTML5 + CSS Grid + Accessibility")
    print("✅ Google Fonts integration")
    print("✅ Mobile-first responsive design")
    print("✅ Enhanced prompt engineering")
    print("✅ Improved code cleaning")
    print()

    # Create and launch interface
    interface = create_interface()

    # Launch with public sharing for Colab
    interface.launch(
        share=True,  # Creates a public link for Colab
        server_name="0.0.0.0",
        server_port=7860,
        debug=True
    )