[Google Study Jam X IT’s Study] Build Real World AI Applications with Gemini and Imagen (2)

Build an AI Image Generator app using Imagen on Vertex AI

Code

텍스트 프롬프트를 입력으로 받아 이미지를 생성한다.

import argparse

import vertexai
from vertexai.preview.vision_models import ImageGenerationModel

def generate_image(
    project_id: str, location: str, output_file: str, prompt: str
) -> vertexai.preview.vision_models.ImageGenerationResponse:
    """Generate an image using a text prompt.
    Args:
      project_id: Google Cloud project ID, used to initialize Vertex AI.
      location: Google Cloud region, used to initialize Vertex AI.
      output_file: Local path to the output image file.
      prompt: The text prompt describing what you want to see."""

    vertexai.init(project=project_id, location=location)

	# Load Model
    model = ImageGenerationModel.from_pretrained("imagen-3.0-generate-002")

	# Generate images
    images = model.generate_images(
        prompt=prompt,
        # Optional parameters
        number_of_images=1,
        seed=1,
        add_watermark=False,
    )

    images[0].save(location=output_file)

    return images

generate_image(
    project_id='<project-id>',
    location='<REGION>',
    output_file='image.jpeg',
    prompt='<Text Prompt>',
    )