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

# Set up Gemini API key
GEMINI_API_KEY = "your-api-key-here"
genai.configure(api_key=GEMINI_API_KEY)

# Optimized function for fast & accurate image generation
def generate_image(text, width=768, height=768):
    try:
        model = genai.GenerativeModel("gemini-pro-vision")
        
        # Fast inference configuration
        response = model.generate_content(
            text,
            generation_config={
                "response_mime_type": "image/png",
                "width": width,
                "height": height
            }
        )

        # Decode base64 image
        img_data = base64.b64decode(response.text)
        image = Image.open(io.BytesIO(img_data))

        return image
    except Exception as e:
        return f"Error: {str(e)}"

# Create Gradio UI with advanced options
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(placeholder="Describe the image (e.g., Astronaut in a jungle)"),
        gr.Slider(minimum=256, maximum=1024, step=128, value=768, label="Width"),
        gr.Slider(minimum=256, maximum=1024, step=128, value=768, label="Height"),
    ],
    outputs="image",
    title="Fast Gemini AI Text-to-Image Generator",
    description="Enter a prompt and generate high-quality images with fast response times."
)

iface.launch()