import google.generativeai as genai
from typing import Union, IO
import gradio as gr
import numpy as np
import PIL.Image
import os

api_key = os.getenv("GEMINI_API_KEY")
genai.configure(api_key = api_key)

def ImageStory(image: Union[np.ndarray, str, IO]) -> str:
    """
    Writes a short story about the uploaded image.

    Args:
        image: The image being uploaded. It can be a NumPy array, file path, or file-like object.

    Returns:
        A short story about the uploaded image.
    """  
        
    # load model
    model = genai.GenerativeModel("gemini-1.5-flash")

    # check image file and convert to a PIL image
    if isinstance(image, np.ndarray):
        img = Image.fromarray(image)
    else:
        img = Image.open(image)

    response = model.generate_content(["write a short story about the image", img])

    return response.text


app = gr.Interface(ImageStory,
                   inputs = gr.Image(label = "Image"),
                   outputs = gr.Text(label = "Story"),
                   examples = ["rubiks cube.jpg", "giraffe.jpg", "street.jpg"],
                   title = "Image To Story",
                   theme = "patrickosornio/my_theme1")



if __name__ == "__main__":
    app.launch(mcp_server = True)