import os
import gradio as gr
from google import genai
from google.genai import types

# Set your Google API key
os.environ["GOOGLE_API_KEY"] = "AIzaSyDL5Rilo7ptJpUOZdY6wy8PJYUcVcnDADs"

# Initialize the Google Generative AI client
client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])

MODEL_NAME = "gemini-2.5-flash-image-preview"

def remix_images(image1, image2, prompt):
    try:
        # Run request to Google Generative AI
        response = client.models.generate_images(
            model=MODEL_NAME,
            prompt=prompt,
            images=[image1, image2]
        )

        # Extract the first generated image
        if response.generated_images and len(response.generated_images) > 0:
            img = response.generated_images[0]
            return img.image  # returns PIL.Image
        else:
            return None
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🖼️ Google Gemini 2.5 Flash Image Preview (via Gradio)")

    with gr.Row():
        img1 = gr.Image(type="filepath", label="Upload Image 1")
        img2 = gr.Image(type="filepath", label="Upload Image 2")

    prompt = gr.Textbox(label="Prompt", placeholder="Describe how to remix the images...")

    btn = gr.Button("Generate")
    output = gr.Image(label="Generated Image")

    btn.click(fn=remix_images, inputs=[img1, img2, prompt], outputs=output)

if __name__ == "__main__":
    demo.launch()