app.py

import gradio as gr
import google.generativeai as genai
from PIL import Image
import io
import base64

# Mapping styles to instructions
STYLE_INSTRUCTIONS = {
    "General": (
        "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."
    ),
    "Realistic": (
        "Describe this design with a focus on photorealistic details, including textures, lighting, and depth, "
        "ensuring accuracy in representation. Avoid illustrations or cartoon-like interpretations."
    ),
    "Kawaii/Cartoon": (
        "Describe this design with an emphasis on cute, rounded shapes, playful expressions, and vibrant, fun aesthetics "
        "typical of the kawaii/cartoon style. Focus on light-hearted and visually charming elements."
    ),
    "Vector": (
        "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. "
        "Specify that the design should be created in a clean, professional vector and illustration style, "
        "emphasizing sharp lines, geometric shapes, and smooth gradients where applicable. "
        "The entire design must be in black and white, using only shades of gray for depth and contrast. "
        "Avoid any use of color unless it is essential to the design's purpose. "
        "Avoid cartoonish or overly stylized elements unless they align with the design's purpose. "
        "Use clear and direct language to convey the design's theme, humor, or purpose, if applicable. "
        "Ensure the description is compact, well-structured, and optimized for creating a graphic. "
        "Include only essential details and avoid unnecessary elaboration."
    ),
    "Silhouette": (
        "Describe this design using high-contrast black and white silhouettes, focusing on bold shapes and clear outlines "
        "to convey the subject effectively. Eliminate interior details unless crucial to recognition."
    ),
}

# Function to check API key validity
def check_api(api_key):
    try:
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel("gemini-1.5-pro")
        response = model.generate_content("Hello, can you respond?")
        return "✅ API Key is valid!" if response.text else "⚠️ API responded but empty."
    except Exception as e:
        return f"❌ Error: {str(e)}"

# Function to generate a prompt using Gemini
def generate_prompt(image, api_key, style):
    if not api_key:
        return "❌ Error: Please enter your Google Gemini API key."

    try:
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel("gemini-1.5-pro")

        # Convert image to base64
        image_bytes = io.BytesIO()
        image.save(image_bytes, format=image.format if image.format else "PNG")
        image_base64 = base64.b64encode(image_bytes.getvalue()).decode('utf-8')

        instruction = STYLE_INSTRUCTIONS.get(style, STYLE_INSTRUCTIONS["General"])

        response = model.generate_content([
            instruction,
            {"mime_type": "image/png", "data": image_base64}
        ])

        return response.text if response.text else "⚠️ No response generated."
    except Exception as e:
        return f"❌ Error generating prompt: {str(e)}"

# Gradio Interface
def main():
    with gr.Blocks() as app:
        gr.Markdown("## 🎨 Flux Prompt Generator with Gemini 1.5 Pro")

        api_key = gr.Textbox(label="🔑 Google Gemini API Key", type="password", placeholder="Paste your Gemini API key here")
        check_button = gr.Button("✅ Check API Key")
        check_output = gr.Textbox(label="API Key Status")

        with gr.Row():
            with gr.Column():
                image_input = gr.Image(label="📷 Upload Design Image", type="pil", width=300)
                style_dropdown = gr.Dropdown(label="🎨 Style", choices=list(STYLE_INSTRUCTIONS.keys()), value="General")

        generate_button = gr.Button("🧠 Generate Prompt")
        prompt_output = gr.Textbox(label="📝 Generated Flux Prompt", lines=6)
        copy_button = gr.Button("📋 Copy to Clipboard")

        def copy_prompt(prompt):
            return gr.update(value=prompt, interactive=True)

        check_button.click(check_api, inputs=api_key, outputs=check_output)
        generate_button.click(generate_prompt, inputs=[image_input, api_key, style_dropdown], outputs=prompt_output)
        copy_button.click(copy_prompt, inputs=prompt_output, outputs=prompt_output)

    return app

# Launch the app
app = main()
app.launch()