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


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.
  """  
  api_key = os.getenv("GEMINI_API_KEY")
    
  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 black female superhero with braided hair flying over a over an active volcano with a fiery red sky",
                             "A black jazz musician playing his saxophone and surrounded by colorful musical notes."])


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