import gradio as gr
from google import genai
from PIL import Image
from io import BytesIO
import base64

# Configure the Gemini API client with the hardcoded API key
GOOGLE_API_KEY = "AIzaSyDL5Rilo7ptJpUOZdY6wy8PJYUcVcnDADs"
client = genai.Client(api_key=GOOGLE_API_KEY)

def process_image(image, prompt):
    try:
        # Convert Gradio image (PIL Image) to base64
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        img_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")

        # Prepare the content for the Gemini API
        contents = [
            {
                "parts": [
                    {"text": prompt},
                    {
                        "inline_data": {
                            "mime_type": "image/png",
                            "data": img_base64
                        }
                    }
                ]
            }
        ]

        # Call the Gemini API
        response = client.models.generate_content(
            model="gemini-2.5-flash-image-preview",
            contents=contents
        )

        # Process the response
        for part in response.candidates[0].content.parts:
            if part.inline_data is not None:
                # Decode the generated image
                img_data = base64.b64decode(part.inline_data.data)
                img = Image.open(BytesIO(img_data))
                return img
            elif part.text is not None:
                return f"Text response: {part.text}"
        
        return "No image or text returned by the model."
    
    except Exception as e:
        return f"Error: {str(e)}"

# Create the Gradio interface
title = "Gemini 2.5 Flash Image Editor (Nano Banana)"
description = "Upload an image and provide a text prompt to generate or edit images using Google's Gemini 2.5 Flash Image Preview model."
interface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Textbox(label="Prompt", placeholder="e.g., 'Add a nano-banana to the image in a fancy restaurant setting'")
    ],
    outputs=gr.Image(type="pil", label="Generated Image"),
    title=title,
    description=description,
    examples=[
        [None, "Create a minimalist composition with a single red maple leaf in the bottom-right, soft lighting, square image"],
        [None, "Generate an image of a cat eating a nano-banana in a fancy restaurant"]
    ]
)

# Launch the app
if __name__ == "__main__":
    interface.launch()