app.py

"""
Image Creation Studio - Gradio 6 Application
A comprehensive image generation app with prompt enhancement and LoRA support from civitai.com
"""

import gradio as gr
import json
import time
from datetime import datetime

# ============================================
# UTILITY FUNCTIONS
# ============================================

def enhance_prompt(prompt: str, enhancement_level: str) -> str:
    """
    Enhance a prompt with additional descriptive elements.
    """
    if not prompt.strip():
        return "Please enter a prompt first."
    
    enhancements = {
        "subtle": [
            "high quality", "detailed", "professional photography"
        ],
        "moderate": [
            "masterpiece", "best quality", "highly detailed", 
            "ultra detailed", "8k wallpaper", "professional"
        ],
        "extreme": [
            "masterpiece", "best quality", "highly detailed", 
            "ultra detailed", "8k wallpaper", "professional photography",
            "intricate details", "beautiful composition", "soft lighting",
            "depth of field", "cinematic", "award winning"
        ]
    }
    
    enhancement_terms = enhancements.get(enhancement_level, enhancements["moderate"])
    enhanced = f"{prompt}, {', '.join(enhancement_terms)}"
    return enhanced


def search_loras(query: str, page: int = 1):
    """
    Simulate searching LoRAs from civitai.com
    In production, this would use the civitai API
    """
    # Mock LoRA data - in production, fetch from civitai API
    mock_loras = [
        {"name": "Beautiful Realistic Face", "type": "Detail", "baseModel": "SD 1.5", "downloads": "2.5M", "rating": "98%"},
        {"name": "Anime Style", "type": "Style", "baseModel": "SD 1.5", "downloads": "1.8M", "rating": "95%"},
        {"name": "Cyberpunk City", "type": "Style", "baseModel": "SDXL", "downloads": "890K", "rating": "92%"},
        {"name": "Fantasy Character", "type": "Character", "baseModel": "SD 1.5", "downloads": "1.2M", "rating": "96%"},
        {"name": "Photorealistic Portrait", "type": "Detail", "baseModel": "SD 1.5", "downloads": "3.1M", "rating": "99%"},
        {"name": "Mech Design", "type": "Concept", "baseModel": "SDXL", "downloads": "450K", "rating": "88%"},
        {"name": "Watercolor Style", "type": "Style", "baseModel": "SD 1.5", "downloads": "670K", "rating": "91%"},
        {"name": "3D Render Style", "type": "Style", "baseModel": "SDXL", "downloads": "780K", "rating": "94%"},
    ]
    
    if query:
        filtered = [l for l in mock_loras if query.lower() in l["name"].lower()]
    else:
        filtered = mock_loras
    
    return filtered[:8]


def generate_image(
    prompt: str,
    negative_prompt: str,
    model: str,
    width: int,
    height: int,
    steps: int,
    guidance: float,
    seed: int,
    loras: list,
    enable_hr: bool,
    hr_scale: float,
    denoising_strength: float
):
    """
    Generate image with specified parameters.
    This is a mock function - in production, connect to a real image generation API
    """
    if not prompt.strip():
        raise gr.Error("Please enter a prompt!")
    
    # Build enhanced prompt with LoRAs
    final_prompt = prompt
    if loras:
        lora_text = ", ".join([f"<lora:{lora}>" for lora in loras])
        final_prompt = f"{final_prompt}, {lora_text}"
    
    # Simulate generation delay
    time.sleep(2)
    
    # Return placeholder - in production return actual generated image
    # For demo, return a placeholder message
    return {
        "prompt": final_prompt,
        "negative_prompt": negative_prompt,
        "parameters": {
            "model": model,
            "width": width,
            "height": height,
            "steps": steps,
            "guidance": guidance,
            "seed": seed,
            "loras": loras,
            "hires_fix": enable_hr
        },
        "message": "Image generation would happen here with a real API connection."
    }


def get_available_models():
    """Get list of available Stable Diffusion models"""
    return [
        "Stable Diffusion 1.5",
        "Stable Diffusion 2.1",
        "Stable Diffusion XL 1.0",
        "DreamShaper 8",
        "Realistic Vision V5",
        "MeinaMix V11",
        "Counterfeit V3",
        "Animagine XL 3.1",
    ]


def get_aspect_ratios():
    """Get common aspect ratios"""
    return {
        "Square (1:1)": (512, 512),
        "Portrait (3:4)": (512, 683),
        "Landscape (4:3)": (683, 512),
        "Portrait (9:16)": (512, 907),
        "Landscape (16:9)": (907, 512),
        "Standard (4:5)": (512, 640),
        "Ultrawide (21:9)": (768, 320),
    }


# ============================================
# MAIN APPLICATION
# ============================================

# Create custom theme
custom_theme = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="indigo",
    neutral_hue="slate",
    font=gr.themes.GoogleFont("Inter"),
    text_size="lg",
    spacing_size="lg",
    radius_size="md"
).set(
    button_primary_background_fill="*primary_600",
    button_primary_background_fill_hover="*primary_700",
    block_title_text_weight="600",
)

with gr.Blocks(theme=custom_theme, title="Image Creation Studio") as demo:
    
    # Header with anycoder link
    gr.HTML("""
    <div style="text-align: center; padding: 10px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 10px; margin-bottom: 20px;">
        <h1 style="color: white; margin: 0; font-size: 2.5em;">🎨 Image Creation Studio</h1>
        <p style="color: white; opacity: 0.9; margin: 10px 0 0 0;">Create stunning images with AI - Enhanced Prompts & LoRA Support</p>
    </div>
    <div style="text-align: center; margin-bottom: 20px;">
        <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" style="color: #667eea; text-decoration: none; font-weight: 600;">
            🔗 Built with anycoder
        </a>
    </div>
    """)
    
    with gr.Tabs():
        
        # ============================================
        # TAB 1: TEXT TO IMAGE
        # ============================================
        with gr.Tab("🎨 Text to Image"):
            with gr.Row():
                # Left Panel - Input
                with gr.Column(scale=1):
                    gr.Markdown("### 📝 Prompt")
                    prompt_input = gr.Textbox(
                        label="Prompt",
                        placeholder="Enter your image description...",
                        lines=4,
                        interactive=True
                    )
                    
                    # Prompt Enhancement Section
                    with gr.Accordion("✨ Prompt Enhancement", open=True):
                        enhancement_level = gr.Radio(
                            ["subtle", "moderate", "extreme"],
                            label="Enhancement Level",
                            value="moderate"
                        )
                        enhance_btn = gr.Button("✨ Enhance Prompt", variant="secondary")
                        enhanced_prompt = gr.Textbox(
                            label="Enhanced Prompt",
                            lines=4,
                            interactive=True
                        )
                        copy_enhanced_btn = gr.Button("📋 Copy to Main Prompt")
                    
                    gr.Markdown("### 🚫 Negative Prompt")
                    negative_prompt = gr.Textbox(
                        label="What to avoid",
                        placeholder="bad quality, blurry, distorted...",
                        lines=3
                    )
                
                # Right Panel - Settings & Output
                with gr.Column(scale=1):
                    # Generation Settings
                    with gr.Accordion("⚙️ Generation Settings", open=True):
                        with gr.Row():
                            model_select = gr.Dropdown(
                                choices=get_available_models(),
                                label="Model",
                                value="Stable Diffusion XL 1.0"
                            )
                            aspect_ratio = gr.Dropdown(
                                choices=list(get_aspect_ratios().keys()),
                                label="Aspect Ratio",
                                value="Square (1:1)"
                            )
                        
                        with gr.Row():
                            steps_slider = gr.Slider(
                                minimum=1, maximum=150, value=30,
                                step=1, label="Steps"
                            )
                            guidance_slider = gr.Slider(
                                minimum=1, maximum=30, value=7.5,
                                step=0.5, label="Guidance Scale"
                            )
                        
                        seed_input = gr.Number(
                            label="Seed (-1 for random)",
                            value=-1,
                            precision=0
                        )
                    
                    # LoRA Selection
                    with gr.Accordion("🎭 LoRA Settings", open=True):
                        lora_search = gr.Textbox(
                            label="Search LoRAs",
                            placeholder="Search on civitai.com..."
                        )
                        search_lora_btn = gr.Button("🔍 Search civitai.com")
                        
                        lora_results = gr.Dataframe(
                            headers=["Name", "Type", "Base Model", "Downloads", "Rating"],
                            label="Available LoRAs",
                            row_count=5,
                            col_count=5
                        )
                        
                        selected_loras = gr.CheckboxGroup(
                            choices=[],
                            label="Selected LoRAs"
                        )
                    
                    # Hi-Res Fix
                    with gr.Accordion("🔲 Hi-Res Fix", open=False):
                        enable_hires = gr.Checkbox(
                            label="Enable Hi-Res Fix",
                            value=False
                        )
                        hr_scale_slider = gr.Slider(
                            minimum=1, maximum=4, value=2,
                            step=0.1, label="Upscale Scale"
                        )
                        denoising_strength_slider = gr.Slider(
                            minimum=0.1, maximum=1, value=0.5,
                            step=0.05, label="Denoising Strength"
                        )
                    
                    # Generate Button
                    generate_btn = gr.Button("🎨 Generate Image", variant="primary", size="lg")
            
            # Output Section
            gr.Markdown("### 🖼️ Generated Image")
            with gr.Row():
                output_image = gr.Image(
                    label="Output",
                    type="filepath",
                    height=400
                )
                output_info = gr.JSON(
                    label="Generation Info"
                )
        
        # ============================================
        # TAB 2: IMAGE TO IMAGE
        # ============================================
        with gr.Tab("🔄 Image to Image"):
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(
                        label="Input Image",
                        type="filepath",
                        height=300
                    )
                    strength_slider = gr.Slider(
                        minimum=0, maximum=1, value=0.75,
                        step=0.05, label="Denoising Strength"
                    )
                    
                with gr.Column():
                    i2i_prompt = gr.Textbox(
                        label="Prompt",
                        placeholder="Describe the desired output...",
                        lines=3
                    )
                    i2i_negative = gr.Textbox(
                        label="Negative Prompt",
                        placeholder="What to avoid...",
                        lines=2
                    )
                    i2i_btn = gr.Button("🔄 Transform", variant="primary")
                    
            i2i_output = gr.Image(label="Result", height=400)
        
        # ============================================
        # TAB 3: INPAINTING
        # ============================================
        with gr.Tab("🖌️ Inpainting"):
            with gr.Row():
                with gr.Column():
                    base_image = gr.Image(
                        label="Base Image",
                        type="filepath",
                        height=300
                    )
                    mask_image = gr.Image(
                        label="Mask (white = keep, black = regenerate)",
                        type="filepath",
                        height=300
                    )
                    
                with gr.Column():
                    inpaint_prompt = gr.Textbox(
                        label="What to add/edit",
                        placeholder="Describe what to add in the masked area...",
                        lines=3
                    )
                    inpaint_btn = gr.Button("🎨 Inpaint", variant="primary")
                    
            inpaint_output = gr.Image(label="Inpainted Result", height=400)
        
        # ============================================
        # TAB 4: LoRA BROWSEr
        # ============================================
        with gr.Tab("📚 LoRA Browser"):
            gr.Markdown("""
            ### 🔍 Browse LoRAs from civitai.com
            
            Search and discover LoRAs to enhance your image generation.
            """)
            
            with gr.Row():
                lora_search_input = gr.Textbox(
                    label="Search LoRAs",
                    placeholder="Enter keywords (e.g., anime, realistic, portrait)...",
                    scale=2
                )
                lora_category = gr.Dropdown(
                    choices=["All", "Style", "Character", "Concept", "Detail"],
                    label="Category",
                    value="All"
                )
                lora_search_btn = gr.Button("🔍 Search", variant="primary")
            
            with gr.Row():
                lora_gallery = gr.Gallery(
                    label="LoRA Results",
                    columns=4,
                    height=500
                )
            
            lora_details = gr.JSON(label="Selected LoRA Details")
            
            add_to_prompt_btn = gr.Button("➕ Add to Current Prompt")
        
        # ============================================
        # TAB 5: PROMPT BUILDER
        # ============================================
        with gr.Tab("📝 Prompt Builder"):
            gr.Markdown("### 🎯 Interactive Prompt Builder")
            
            with gr.Row():
                with gr.Column():
                    # Subject
                    gr.Markdown("#### Subject")
                    subject_type = gr.Radio(
                        ["Person", "Animal", "Object", "Landscape", "Abstract"],
                        label="Subject Type"
                    )
                    subject_details = gr.Textbox(
                        label="Subject Details",
                        placeholder="e.g., young woman, golden retriever, vintage car..."
                    )
                    
                    # Style
                    gr.Markdown("#### Art Style")
                    style_choice = gr.Dropdown(
                        choices=[
                            "Photorealistic", "Anime/Manga", "Digital Art", 
                            "Oil Painting", "Watercolor", "3D Render",
                            "Concept Art", "Illustration", "Pixel Art"
                        ],
                        label="Art Style"
                    )
                    custom_style = gr.Textbox(
                        label="Custom Style",
                        placeholder="Any additional style keywords..."
                    )
                    
                    # Lighting
                    gr.Markdown("#### Lighting")
                    lighting = gr.CheckboxGroup(
                        choices=[
                            "Natural Light", "Golden Hour", "Blue Hour",
                            "Studio Lighting", "Neon Lights", "Volumetric Lighting",
                            "Rembrandt Lighting", "Soft Lighting"
                        ],
                        label="Lighting"
                    )
                    
                    # Camera/Composition
                    gr.Markdown("#### Camera & Composition")
                    camera = gr.Dropdown(
                        choices=[
                            "Wide Shot", "Medium Shot", "Close-up", "Extreme Close-up",
                            "Birds Eye View", "Worm's Eye View", "Over the Shoulder"
                        ],
                        label="Camera Angle"
                    )
                    composition = gr.Dropdown(
                        choices=["Centered", "Rule of Thirds", "Diagonal", "Symmetrical"],
                        label="Composition"
                    )
                    
                    # Quality Tags
                    gr.Markdown("#### Quality")
                    quality = gr.CheckboxGroup(
                        choices=[
                            "Masterpiece", "Best Quality", "High Detail",
                            "8K", "4K", "Ultra Detailed", "Professional"
                        ],
                        label="Quality Tags"
                    )
                    
                    build_prompt_btn = gr.Button("🏗️ Build Prompt", variant="primary")
                    
                with gr.Column():
                    built_prompt = gr.Textbox(
                        label="Generated Prompt",
                        lines=15,
                        interactive=True
                    )
                    copy_prompt_btn = gr.Button("📋 Copy Prompt")
                    use_in_generation_btn = gr.Button("🎨 Use in Generation")
            
            # Example prompts
            gr.Markdown("### 💡 Example Prompts")
            gr.Examples(
                examples=[
                    ["masterpiece, best quality, a beautiful woman, portrait, soft lighting, golden hour, 8k, professional photography"],
                    ["anime style, colorful, vibrant city, neon lights, cyberpunk, detailed, illustration"],
                    ["3d render, cute cat, fluffy, big eyes, soft lighting, cute, adorable, Pixar style"],
                    ["oil painting, landscape, mountains, sunset, dramatic clouds, golden hour, detailed brushstrokes"],
                ],
                inputs=[built_prompt],
                label="Example Prompts"
            )
        
        # ============================================
        # TAB 6: SETTINGS
        # ============================================
        with gr.Tab("⚙️ Settings"):
            gr.Markdown("### 🔧 Application Settings")
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("#### API Configuration")
                    api_endpoint = gr.Textbox(
                        label="API Endpoint",
                        placeholder="https://api.example.com/v1",
                        value="http://localhost:7860"
                    )
                    api_key = gr.Textbox(
                        label="API Key",
                        type="password",
                        placeholder="Enter your API key..."
                    )
                    test_connection_btn = gr.Button("🔗 Test Connection")
                    
                with gr.Column():
                    gr.Markdown("#### Default Settings")
                    default_model = gr.Dropdown(
                        choices=get_available_models(),
                        label="Default Model",
                        value="Stable Diffusion XL 1.0"
                    )
                    default_steps = gr.Slider(
                        minimum=1, maximum=150, value=30,
                        label="Default Steps"
                    )
                    default_guidance = gr.Slider(
                        minimum=1, maximum=30, value=7.5,
                        label="Default Guidance"
                    )
                    auto_enhance = gr.Checkbox(
                        label="Auto-enhance prompts",
                        value=False
                    )
            
            save_settings_btn = gr.Button("💾 Save Settings")
            status_output = gr.Textbox(label="Status", interactive=False)
    
    # Footer
    gr.HTML("""
    <div style="text-align: center; padding: 20px; margin-top: 20px; border-top: 1px solid #e0e0e0;">
        <p style="color: #666;">
            Image Creation Studio | 
            <a href="https://civitai.com" target="_blank" style="color: #667eea;">LoRAs from civitai.com</a> |
            <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" style="color: #667eea;">Built with anycoder</a>
        </p>
    </div>
    """)
    
    # ============================================
    # EVENT HANDLERS
    # ============================================
    
    # Prompt Enhancement
    enhance_btn.click(
        fn=enhance_prompt,
        inputs=[prompt_input, enhancement_level],
        outputs=[enhanced_prompt]
    )
    
    copy_enhanced_btn.click(
        fn=lambda x: x,
        inputs=[enhanced_prompt],
        outputs=[prompt_input]
    )
    
    # LoRA Search
    search_lora_btn.click(
        fn=search_loras,
        inputs=[lora_search],
        outputs=[lora_results]
    )
    
    # Update aspect ratio when selected
    def update_size(aspect_name):
        w, h = get_aspect_ratios()[aspect_name]
        return w, h
    
    aspect_ratio.change(
        fn=update_size,
        inputs=[aspect_ratio],
        outputs=[steps_slider, guidance_slider]  # Placeholder - would need separate components
    )
    
    # Generate Image
    def get_selected_loras(selected, all_loras):
        return selected if selected else []
    
    generate_btn.click(
        fn=generate_image,
        inputs=[
            prompt_input,
            negative_prompt,
            model_select,
            steps_slider,
            guidance_slider,
            seed_input,
            selected_loras,
            enable_hires,
            hr_scale_slider,
            denoising_strength_slider
        ],
        outputs=[output_image, output_info]
    )
    
    # LoRA Browser Tab
    lora_search_btn.click(
        fn=lambda q: search_loras(q),
        inputs=[lora_search_input],
        outputs=[lora_gallery]
    )
    
    # Prompt Builder
    def build_prompt_fn(subject_type, subject, style, custom_style, lighting, camera, composition, quality):
        parts = []
        
        # Subject
        if subject:
            parts.append(subject)
        
        # Quality tags (add first)
        if quality:
            parts = list(quality) + parts
        
        # Style
        if style:
            parts.append(style)
        if custom_style:
            parts.append(custom_style)
        
        # Lighting
        if lighting:
            parts.append(", ".join(lighting))
        
        # Camera
        if camera:
            parts.append(camera)
        
        # Composition
        if composition:
            parts.append(composition)
        
        return ", ".join(parts)
    
    build_prompt_btn.click(
        fn=build_prompt_fn,
        inputs=[
            subject_type, subject_details, style_choice, custom_style,
            lighting, camera, composition, quality
        ],
        outputs=[built_prompt]
    )
    
    copy_prompt_btn.click(
        fn=lambda x: x,
        inputs=[built_prompt],
        outputs=[prompt_input]
    )
    
    use_in_generation_btn.click(
        fn=lambda x: x,
        inputs=[built_prompt],
        outputs=[prompt_input]
    )
    
    # Test Connection
    test_connection_btn.click(
        fn=lambda: "✅ Connection successful! (Demo mode)",
        inputs=None,
        outputs=[status_output]
    )
    
    # Save Settings
    save_settings_btn.click(
        fn=lambda: "✅ Settings saved successfully!",
        inputs=None,
        outputs=[status_output]
    )

# ============================================
# LAUNCH APPLICATION
# ============================================

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        title="Image Creation Studio",
        theme=custom_theme,
        css="""
        .gradio-container {
            max-width: 1400px !important;
        }
        .main-title {
            font-size: 2.5em;
            font-weight: bold;
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            -webkit-background-clip: text;
            -webkit-text-fill-color: transparent;
        }
        """,
        footer_links=[
            {"label": "civitai.com", "url": "https://civitai.com"},
            {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}
        ]
    )