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

api_key = os.getenv("GEMINI_API_KEY")

def generate_image(prompt):  

  """
  Transforms text into an image.

  Args:
      prompt: A natural language instruction that is used to generate an image.

  Returns:
        An image.
  """  
    
  client = genai.Client(api_key=api_key)

  prompt = tuple(prompt)

  response = client.models.generate_content(
    model="gemini-2.0-flash-preview-image-generation",
    contents=prompt,
    config=types
    .GenerateContentConfig(
      response_modalities=['TEXT', 'IMAGE']
    )
  )

  for part in response.candidates[0].content.parts:
    if part.text is not None:
      pass
    elif part.inline_data is not None:
      image = Image.open(BytesIO((part.inline_data.data)))
      # image.save('gemini-native-image.png')
      # display(image)
  return image 

# build gradio interface
app = gr.Interface(fn = generate_image, 
                   inputs = gr.Text(label="Prompt",placeholder="Type your prompt here. . ."),
                   outputs = gr.Image(label="Generated Image"),
                   title="Gemini Image Generator",
                   examples=["A 3D rendering of a little black girl wearing a colorful dress and smiling broadly at the camera",
                             "Disney and Pixar-style playful bunny skipping about in a garden full of carrots and lettuce",
                             "A strong black female superhero with braided hair flying in the bright blue sky",
                             "A jazz musician playing his saxophone and surrounded by colorful musical notes."])


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