Image Generation Agent (Streaming)

1from io import BytesIO
2
3from kern.agent import Agent, RunOutput  # noqa
4from kern.models.google import Gemini
5from PIL import Image
6
7# No system message should be provided
8agent = Agent(
9    model=Gemini(
10        id="gemini-2.0-flash-exp-image-generation",
11        response_modalities=["Text", "Image"],
12    )
13)
14
15# Print the response in the terminal
16response = agent.run("Make me an image of a cat in a tree.", stream=True)
17
18for chunk in response:
19    if hasattr(chunk, "images") and chunk.images:  # type: ignore
20        images = chunk.images  # type: ignore
21        if images and isinstance(images, list):
22            for image_response in images:
23                image_bytes = image_response.content
24                if image_bytes:
25                    image = Image.open(BytesIO(image_bytes))
26                    image.show()
27                    # Save the image to a file
28                    # image.save("generated_image.png")