AI Gateway: Use Google AI SDK for using Vertex AI's Imagen model

[tomek-labuk]

How to for generating images using Vertex AI's Imagen model through Kong AI Gateway. Set response_modalities to ["IMAGE"] and extract the image data from the response

Python script:

from google import genai

client = genai.Client(
    vertexai=True,
    project="your-project-id",
    location="us-central1",
    http_options={
        "api_version": "v1",
        "base_url": "http://localhost:8000/gemini"
    }
)

response = client.models.generate_content(
    model="imagen-3.0-generate-001",
    contents="A serene landscape with mountains and a lake at sunset",
    config={
        "response_modalities": ["IMAGE"]
    }
)

# Extract and save image
for part in response.candidates[0].content.parts:
    if part.inline_data:
        with open("generated_image.png", "wb") as f:
            f.write(part.inline_data.data)
        print("Image saved!")