401 Deadline exception when using Imagen Python

[eridiumeng]

Example code:

vertexai.init(project=PROJECT_ID, location=LOCATION)
model: ImageGenerationModel = ImageGenerationModel.from_pretrained(model_name=MODEL_NAME)

Common.log_info(message='Handling image generation request')
Common.log_https_request(req=req)
data: dict = req.data
prompt: str = data.get(PROMPT_KEY, '')
user_id: str = req.auth.uid
if not user_id:
    Common.log_info(message='Invalid user id, not generating image')
    return
if not prompt:
    Common.log_info(message='Invalid prompt, not generating image')
    return
number_of_images: int = data.get(NUMBER_OF_IMAGES_KEY, 1)
guidance_scale: int = data.get(GUIDANCE_SCALE_KEY, MEDIUM_GUIDANCE_SCALE)
output_image_path: str = (GENERATED_IMAGE_PREFIX + user_id +
                          '/' +
                          Common.get_current_datetime() + '/' +
                          Common.replace_all_whitespace_with_underscore(
                              input_string=prompt) + JPEG_SUFFIX)
output_gcs_uri: str = BUCKET_URI_PREFIX + output_image_path
Common.log_info(message=f'Generating image at: {output_gcs_uri}')
model.generate_images(
    prompt=prompt,
    number_of_images=number_of_images,
    aspect_ratio=Common.convert_to_aspect_ratio(input_string=data.get(ASPECT_RATIO_KEY, '')),
    output_gcs_uri=BUCKET_URI_PREFIX + output_image_path,
    language=AUTO_LANGUAGE_DETECTION,
    guidance_scale=guidance_scale,
    add_watermark=False,
    safety_filter_level=STRICT_SAFETY_FILTER,
    person_generation=ALLOW_ADULT_PERSON_GENERATION,
)
Common.log_info(message=f'Generated image at: {output_image_path}')
return {
    GENERATED_IMAGE_PATH_KEY: output_image_path
}

Log output

{
insertId: "668e115800009112eb018410"
labels: {
execution_id: "RCyLEQaZkqtT"
goog-managed-by: "cloudfunctions"
instanceId: "0087244a805cf6d58e414f53f12334887c32a9193fd30e864537965aed11aef51420d9a28124761e498af43bce228ec8099599bc142b0e59e58e9ea6dac43d80"
}
logName: "projects/claptrap-project/logs/run.googleapis.com%2Fstderr"
receiveTimestamp: "2024-07-10T04:43:04.043022212Z"
resource: {
labels: {
configuration_name: "generate-images"
location: "us-central1"
project_id: "claptrap-project"
revision_name: "generate-images-00004-mil"
service_name: "generate-images"
}
type: "cloud_run_revision"
}
severity: "ERROR"
spanId: "10602671325904409459"
textPayload: "An exception occurred: 401 Image generation failed with the following error: Deadline"
timestamp: "2024-07-10T04:43:04.037138Z"
}

[johanpicard]

+1 same issue even with a single image. Generation works fine through the Studio though for the same Imagegeneration@006 model, even for 4 images.

[glasnt]

Based on the log output, you're running this on either Cloud Run or Cloud Functions. The default resources on these services will be less than a dedicated environment like (presumably from johanpicard's comment) Vertex AI Studio

You may need to increase either the CPU or memory in order to ensure the computation completes before the deadline.

[johanpicard]

eridumeng is however, as I am, calling the Vertex AI API which is timing out here, with no ties to the client's configuration which is just waiting idly ?

After further troubleshooting, I realized that the issue was coming from the "output_gcs_uri" parameter which is supposed to output the generated images to GCS - very handy on paper.
Very fast without, hangs and then "Deadline" when included however.

I fixed it by not using it, generating the images, then grabbing the GeneratedImage elements and wrote them myself to GCS using the following code :

` images = model.generate_images(
prompt=prompt,
number_of_images=config.IMAGE_COUNT,
language="en",
aspect_ratio="1:1",
safety_filter_level="block_some",
# output_gcs_uri=output_folder !!! Does not work !!!
)

storage_client = storage.Client()
bucket = storage_client.bucket(config.BUCKET_NAME)

if images:
    images_urls = []
    clean_concept = concept.replace(" ", "_")  # Replace spaces with underscores

    # Process each generated image
    for i, generated_image in enumerate(images):
        # Create a blob in the concept-specific folder

        blob_name = f"{clean_concept}/{clean_concept}_{i}.png" 
        blob = bucket.blob(blob_name)

        # Upload the generated image content
        blob.upload_from_string(generated_image._image_bytes)

        # Store the GCS URL
        image_url = f"gs://{bucket_name}/{blob_name}"
        image_urls.append(image_url)

    return image_urls
else:
    return []`