def query_imagen()

in projects/imagen-object-changer/imagen_object_changer.py [0:0]
49 lines of code
3 McCabe index (conditional complexity)

def query_imagen(prompt, input_img, mask_img, output_json, token, project_id):
    """Queries GenAI Imagen API for mask-based image editing.

    Uses Imagen to replace parts of the original image. The image mask
    restricts the image generation work area.

    Args:
      prompt: the Imagen mask-based editing GenAI prompt
      input_img: the original source image
      mask_img: the editing mask image
      output_json: output file for writing Imagen response
      token: Gcloud access token within the GCP project
      project_id: GCP project ID

    Returns:
      Boolean, whether Imagen query was successful or not

    Raises:
      None
    """
    # Base64 encode the image and mask files
    img = base64.b64encode(input_img)
    mask_img = base64.b64encode(mask_img)

    # Create the JSON request body
    data = {
        "instances": [
            {
                "prompt": prompt,
                "image": {"bytesBase64Encoded": img.decode("utf-8")},
                "mask": {
                    "image": {
                        "bytesBase64Encoded": mask_img.decode("utf-8"),
                    }
                },
            }
        ],
        "parameters": {"sampleCount": 4, "sampleImageSize": "1024"},
    }

    # Make the API request
    headers = {
        "Authorization": "Bearer {}".format(token),
        "Content-Type": "application/json",
        "User-Agent": "Mozilla/5.0",
        "Accept-Encoding": "identity",
        "Accept": "*/*",
    }

    url = (
        "https://us-central1-aiplatform.googleapis.com/v1/projects/"
        + project_id
        + "/locations/us-central1/publishers/google/models/"
        + "imagegeneration@002:predict"
    )

    print("Querying Imagen...")
    response = requests.post(
        url,
        headers=headers,
        data=json.dumps(
            data, sort_keys=False, indent=2, separators=(",", ": ")
        ),
        verify=True,
        timeout=None,
    )

    imagen_success = True
    if not response:
        print("No or empty response from Imagen. Exiting")
        imagen_success = False
    else:
        with open(output_json, mode="w", encoding="utf-8") as f:
            f.write(response.text)
        print("Wrote Imagen response to: {}".format(output_json))
    return imagen_success