# 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
"""This sample demonstrates how to use Gemini VertexAI to describe and draw images.
Key features demonstrated in this sample:
| 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()])` |
"""
import base64
import os
from io import BytesIO
from PIL import Image
from genkit.ai import Genkit
from genkit.blocks.model import GenerateResponseWrapper
from genkit.plugins.google_genai import VertexAI
if 'GCLOUD_PROJECT' not in os.environ:
os.environ['GCLOUD_PROJECT'] = input('Please enter your GCLOUD_PROJECT: ')
ai = Genkit(plugins=[VertexAI()])
@ai.flow()
async def draw_image_with_imagen() -> GenerateResponseWrapper:
"""Draw an image.
Returns:
The image.
"""
config = {
'number_of_images': 1,
'language': 'en',
'seed': 20,
'add_watermark': False,
}
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:
print('No image generated.')
return
media_url = result.message.content[0].root.media.url if result.message.content[0].root.media else ''
if not media_url:
print('No media URL found in response.')
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:
print('Invalid media URL.')
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())
