401 Deadline exception when using Imagen Python #12020

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
}

{
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"
}