app.py

import os
import gradio as gr
import google.generativeai as genai
from PIL import Image, ImageEnhance, ImageFilter
import io
import base64
import json
import time
try:
    import pyperclip
except ImportError:
    os.system('pip install pyperclip')
    import pyperclip

# --- Security First Configuration ---
GEMINI_KEY = os.environ.get("GEMINI_KEY", "")  # Load from environment variable

# --- Template Optimization System ---
BASE_TEMPLATE = (
    "Describe this design as a single, cohesive, and concise prompt suitable for Flux 1.1 Pro. "
    "Focus on key elements such as text, symbols, layout, and overall style."
)

STYLE_INSTRUCTIONS = {
    "General": BASE_TEMPLATE,
    "Realistic": (
        f"{BASE_TEMPLATE} Focus on photorealistic details: textures, lighting, depth. "
        "Avoid illustrations/cartoon styles. Use technical photography terms."
    ),
    "Kawaii/Cartoon": (
        f"{BASE_TEMPLATE} Emphasize cute, rounded shapes, playful expressions, "
        "vibrant colors. Use anime/kawaii terminology."
    ),
    "Vector": (
        f"{BASE_TEMPLATE} Specify clean vector style with sharp lines, geometric shapes. "
        "Black/white only with gray gradients. Use design technical terms."
    ),
    "Silhouette": (
        f"{BASE_TEMPLATE} Use high-contrast black/white silhouettes. "
        "Focus on bold shapes and outlines. Eliminate interior details."
    ),
}

# --- Flux Configuration Engine ---
FLUX_SPECS = {
    "aspect_ratios": ["1:1", "16:9", "4:3", "9:16", "Custom"],
    "output_formats": ["SVG", "PNG-300dpi", "PDF", "EPS"],
    "color_modes": ["B&W", "CMYK", "Spot Colors", "RGB"],
    "complexity_levels": ["Minimal", "Medium", "High", "Ultra"]
}

# --- Quality Control Systems ---
class QualityValidator:
    VALIDATION_PROMPT = """Analyze this Flux prompt on:
1. Technical specificity (1-5)
2. Style adherence (1-5)
3. Flux compatibility (1-5)
Respond ONLY as JSON: {"total": x/15, "issues": [list]}"""

    @classmethod
    def validate(cls, prompt, model):
        try:
            response = model.generate_content([cls.VALIDATION_PROMPT, prompt])
            return json.loads(response.text)
        except:
            return {"total": 0, "issues": ["Validation failed"]}

# --- Image Processing Pipeline ---
def preprocess_image(img):
    """Enhance image quality before analysis"""
    img = img.convert("RGB")
    img = ImageEnhance.Contrast(img).enhance(1.2)
    img = img.filter(ImageFilter.SHARPEN)
    return img

# --- Core Application Logic ---
def generate_prompt(image, api_key, style, creativity, neg_prompt, flux_specs):
    genai.configure(api_key=api_key or GEMINI_KEY)
    model = genai.GenerativeModel("gemini-1.5-pro", generation_config={
        "temperature": creativity,
        "top_p": 0.95
    })

    # Pre-process image
    img = preprocess_image(image)
    img_bytes = io.BytesIO()
    img.save(img_bytes, format="PNG")
    img_b64 = base64.b64encode(img_bytes.getvalue()).decode()

    # Build instruction
    instruction = f"{STYLE_INSTRUCTIONS[style]}\nAVOID: {neg_prompt}\nFLUX SPECS: {flux_specs}"

    # Generate initial prompt
    response = model.generate_content([instruction, {"mime_type": "image/png", "data": img_b64}])
    raw_prompt = response.text

    # Quality validation
    validation = QualityValidator.validate(raw_prompt, model)
    if validation["total"] < 10:  # Regenerate if low quality
        response = model.generate_content([
            f"Improve this prompt addressing: {validation['issues']}\n\n{raw_prompt}"
        ])
        raw_prompt = response.text

    return raw_prompt, validation

# --- UI Components ---
def create_advanced_params():
    with gr.Accordion("⚙️ Advanced Parameters", open=False):
        with gr.Row():
            creativity = gr.Slider(0.0, 1.0, 0.7, label="Creativity")
            neg_prompt = gr.Textbox(label="🚫 Avoid in Prompt", placeholder="e.g., no text, avoid gradients")
        with gr.Row():
            aspect = gr.Dropdown(FLUX_SPECS["aspect_ratios"], value="1:1", label="Aspect Ratio")
            fmt = gr.Dropdown(FLUX_SPECS["output_formats"], value="SVG", label="Output Format")
            color = gr.Dropdown(FLUX_SPECS["color_modes"], value="B&W", label="Color Mode")
            complexity = gr.Dropdown(FLUX_SPECS["complexity_levels"], value="Medium", label="Complexity")
    return [creativity, neg_prompt, aspect, fmt, color, complexity]

# --- Gradio Interface ---
def build_interface():
    with gr.Blocks(title="Flux Pro Generator") as app:
        # Security First
        api_key = gr.Textbox(label="🔑 Gemini API Key", value=GEMINI_KEY, type="password", 
                           info="Set GEMINI_KEY environment variable for production")

        # Main Workflow
        with gr.Row():
            with gr.Column(scale=1):
                img_input = gr.Image(label="🖼️ Upload Design", type="pil", sources=["upload", "clipboard"])
                style = gr.Dropdown(list(STYLE_INSTRUCTIONS), value="General", label="🎨 Target Style")
                adv_params = create_advanced_params()
                
            with gr.Column(scale=2):
                prompt_output = gr.Textbox(label="📝 Optimized Prompt", lines=8, interactive=False)
                with gr.Row():
                    gen_btn = gr.Button("✨ Generate", variant="primary")
                    copy_btn = gr.Button("📋 Copy")
                quality_report = gr.JSON(label="🔍 Quality Report")

        # Enterprise Features
        token_counter = gr.Textbox(label="🧮 Token Usage", interactive=False)
        history = gr.State([])

        # Event Handling
        gen_btn.click(
            generate_prompt,
            inputs=[img_input, api_key, style] + adv_params,
            outputs=[prompt_output, quality_report]
        )
        
        copy_btn.click(
            lambda x: (pyperclip.copy(x), x),
            inputs=prompt_output,
            outputs=prompt_output
        )

    return app

# --- Production Deployment ---
if __name__ == "__main__":
    app = build_interface()
    app.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", 7860)),
        share=False,
        auth=(
            os.getenv("APP_USER", "admin"), 
            os.getenv("APP_PWD", "admin")
        ) if os.getenv("ENV") == "prod" else None
    )