# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0
"""Vertex AI Imagen sample - Generate images with Google's Imagen model.
This sample demonstrates how to use Vertex AI's Imagen model for
high-quality image generation through Genkit.
Key Concepts (ELI5)::
┌─────────────────────┬────────────────────────────────────────────────────┐
│ Concept │ ELI5 Explanation │
├─────────────────────┼────────────────────────────────────────────────────┤
│ Imagen │ Google's image generation AI. Describe what you │
│ │ want, it draws it for you. │
├─────────────────────┼────────────────────────────────────────────────────┤
│ Text-to-Image │ Type words, get a picture. "A cat on a roof" │
│ │ → image of a cat on a roof. │
├─────────────────────┼────────────────────────────────────────────────────┤
│ Base64 │ A way to encode images as text. Allows embedding │
│ │ images in JSON responses. │
├─────────────────────┼────────────────────────────────────────────────────┤
│ Media URL │ A link to the generated image. Download or │
│ │ display directly in your app. │
└─────────────────────┴────────────────────────────────────────────────────┘
Key Features
============
| Feature Description | Example Function / Code Snippet |
|-----------------------------------------|-------------------------------------|
| Text-to-Image Generation (Imagen) | `draw_image_with_imagen` |
| Image Response Handling | `media_url`, `base64_data` parsing |
| VertexAI Plugin Usage | `ai = Genkit(plugins=[VertexAI()])` |
See README.md for testing instructions.
"""
import base64
import os
from io import BytesIO
from PIL import Image
from rich.traceback import install as install_rich_traceback
from genkit.ai import Genkit
from genkit.blocks.model import GenerateResponseWrapper
from genkit.plugins.google_genai import VertexAI
install_rich_traceback(show_locals=True, width=120, extra_lines=3)
# Check for GCLOUD_PROJECT or GOOGLE_CLOUD_PROJECT
# If GOOGLE_CLOUD_PROJECT is set but GCLOUD_PROJECT isn't, use it
if 'GCLOUD_PROJECT' not in os.environ:
if 'GOOGLE_CLOUD_PROJECT' in os.environ:
os.environ['GCLOUD_PROJECT'] = os.environ['GOOGLE_CLOUD_PROJECT']
else:
os.environ['GCLOUD_PROJECT'] = input('Please enter your GCLOUD_PROJECT_ID: ')
ai = Genkit(plugins=[VertexAI()])
@ai.flow()
async def draw_image_with_imagen() -> GenerateResponseWrapper:
"""Draw an image using Imagen model.
Returns:
The image.
"""
config = {
'number_of_images': 1,
'language': 'en',
'seed': 20,
'add_watermark': False,
}
# pyrefly: ignore[no-matching-overload] - config dict is compatible with dict[str, object]
return await ai.generate(
prompt='Draw a cat in a hat',
model='vertexai/imagegeneration@006',
# optional config; check README for available fields
config=config,
)
async def main() -> None:
"""Main function."""
# Imagen draws an image by description. The model used is available only in
# VertexAI API.
result = await draw_image_with_imagen()
if not result.message or not result.message.content:
return
# pyrefly: ignore[missing-attribute] - MediaModel has url attribute
media = result.message.content[0].root.media
media_url = media.url if media and hasattr(media, 'url') else ''
if not media_url:
return
# Extract base64 data after the comma in "data:image/png;base64,..."
base64_data = media_url.split(',', 1)[1] if ',' in media_url else ''
if not base64_data:
return
decoded_image = BytesIO(base64.b64decode(base64_data))
image = Image.open(decoded_image)
image.show('Image generated by Gemini')
if __name__ == '__main__':
ai.run_main(main())
