Use Imagen models on Vertex AI for image generation and captioning

vertex_ai_image_generation.py

# pip install google-vertex-haystack

import io
import PIL.Image as Image

from google_vertex_haystack.generators.image_generator import VertexAIImageGenerator
from google_vertex_haystack.generators.captioner import VertexAIImageCaptioner
from haystack.dataclasses.byte_stream import ByteStream

## generate the image
image_generator = VertexAIImageGenerator(model="imagegeneration", project_id=project_id)
image_generator_result = image_generator.run("""new years eve celebration with red and golden fireworks behind snowy mountains, 4K, has big 2024 text in the center in neon colors""")

## save the image
image = Image.open(io.BytesIO(image_generator_result["images"][0].data))
image.save("output.png")

## caption the image
image_captioner = VertexAIImageCaptioner(model='imagetext', project_id=project_id)
image = ByteStream.from_file_path("output.png")
image_captioner_result = image_captioner.run(image=image)
print(image_captioner_result["captions"])