main.py

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

# Create Gemini Client
client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))


def generate_image(prompt: str, style: str):
    try:
        full_prompt = (
            f"{prompt}, style: {style}, ultra high quality, "
            "sharp focus, cinematic lighting, professional photography"
        )

        result = client.models.generate_images(
            model="imagen-3.0-generate-002",
            prompt=full_prompt,
            config=genai.types.GenerateImagesConfig(
                number_of_images=1,
                aspect_ratio="1:1"
            )
        )

        image_bytes = result.generated_images[0].image.image_bytes
        return Image.open(BytesIO(image_bytes))

    except Exception as e:
        raise RuntimeError(f"Image generation failed: {e}")