notebooks/community/model_garden/model_garden_pytorch_pic2word.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ad41232f-4ac9-4607-8ee6-9b1d98d0f5c3"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "30e2cb63"
      },
      "source": [
        "# Vertex AI Model Garden - Pic2Word\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_pic2word.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_pic2word.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_pic2word.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5d5af88c-a3c8-46fa-b90a-1b6737a91534"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to use the [Pic2Word](https://github.com/google-research/composed_image_retrieval) model in Vertex AI Model Garden.\n",
        "\n",
        "### Objective\n",
        "\n",
        "Following the notebook you will conduct experiments using the pre-built docker image on Vertex AI.\n",
        "\n",
        "- Deploy pretrained Pic2Word models in Google Cloud Vertex AI\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
        "- Vertex AI Model Registry\n",
        "- Vertex AI Online Prediction\n",
        "\n",
        "### Dataset\n",
        "\n",
        "We use the COCO validation set (5,000 images) for evaluation.\n",
        "\n",
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "- Vertex AI\n",
        "- Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing?_ga=2.46650789.-341051769.1686949237) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing?_ga=2.46650789.-341051769.1686949237), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/?_ga=2.247379078.-341051769.1686949237) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e75fe500-5dc1-4ab3-8829-3eeada5be7f7"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "90a615fe"
      },
      "source": [
        "### Colab only\n",
        "Run the following commands for Colab and skip this section if you are using Workbench or Colab Enterprise."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3d4e50e1"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth as google_auth\n",
        "\n",
        "    google_auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e8b0ef65"
      },
      "source": [
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
        "\n",
        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
        "\n",
        "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs.\n",
        "\n",
        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with **Vertex AI User** and **Storage Object Admin** roles for deploying the model to a Vertex AI endpoint.[link text](https://)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LkoD6vdxbOeL"
      },
      "source": [
        "### Import the necessary packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oS17k46BbVIQ"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import sys\n",
        "\n",
        "from google.cloud import aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dwuw1ZsrbSGw"
      },
      "source": [
        "Set the following variables for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the specified region (`REGION`). Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\")."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_nSN9otVbhqE"
      },
      "outputs": [],
      "source": [
        "# Cloud project ID.\n",
        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Region for launching jobs.\n",
        "# TPU deployment is only supported in us-west1.\n",
        "REGION = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Cloud Storage bucket for storing experiment outputs.\n",
        "# Start with gs:// prefix, e.g. gs://foo_bucket.\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"gemma\")\n",
        "\n",
        "# Initialize Vertex AI API.\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
        "\n",
        "# The service account looks like:\n",
        "# '@.iam.gserviceaccount.com'\n",
        "# Please visit https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console\n",
        "# and create a service account with `Vertex AI User` and `Storage Object Admin` roles.\n",
        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Provision GCS bucket permissions to the SERVICE_ACCOUNT.\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "! gcloud config set project $PROJECT_ID\n",
        "! gcloud services enable language.googleapis.com"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "315598c1-3ff4-4bc8-9b82-7e2fe55d4678"
      },
      "source": [
        "### Define common variables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6d26ddb2-04ad-4973-8965-9b25eb24993f"
      },
      "outputs": [],
      "source": [
        "# Prediction constants.\n",
        "PREDICTION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pic2word-serve:latest\"\n",
        "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_T4\"\n",
        "PREDICTION_MACHINE_TYPE = \"n1-standard-8\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0daec6e5-7642-458f-bd93-60a15431b0a1"
      },
      "source": [
        "## Deploy model for online prediction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b286ce6b-cb00-441b-8249-93497de17974"
      },
      "source": [
        "This section uploads the model to Vertex Model Registry and deploys it on an Endpoint resource. This will take around 15 minutes to finish."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "01f686b2-9227-496b-9c4d-a63becb59fb8"
      },
      "source": [
        "### Upload and deploy model to Vertex AI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8885d5be-5007-4778-9562-976102fefe73"
      },
      "outputs": [],
      "source": [
        "# Upload model.\n",
        "serving_env = {\"MODEL_ID\": \"googleai-pic2word\", \"DEPLOY_SOURCE\": \"notebook\"}\n",
        "\n",
        "model = aiplatform.Model.upload(\n",
        "    display_name=\"pic2word-model\",\n",
        "    serving_container_image_uri=PREDICTION_DOCKER_URI,\n",
        "    serving_container_ports=[7080],\n",
        "    serving_container_predict_route=\"/predictions/pic2word\",\n",
        "    serving_container_health_route=\"/ping\",\n",
        "    serving_container_environment_variables=serving_env,\n",
        ")\n",
        "\n",
        "# Create an endpoint.\n",
        "endpoint = aiplatform.Endpoint.create(display_name=\"pic2word-endpoint\")\n",
        "\n",
        "# Deploy model to endpoint.\n",
        "model.deploy(\n",
        "    endpoint=endpoint,\n",
        "    machine_type=PREDICTION_MACHINE_TYPE,\n",
        "    accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
        "    accelerator_count=1,\n",
        "    traffic_percentage=100,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f0bd36bb-f988-479a-be40-18c54980ca1f"
      },
      "source": [
        "You can manage your uploaded models in the [Model Registry](https://pantheon.corp.google.com/vertex-ai/models) and your endpoints in the [Endpoints](https://pantheon.corp.google.com/vertex-ai/endpoints)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "95edbcdd-8246-444a-aef2-a9cadbc03000"
      },
      "source": [
        "## Send a prediction request to the endpoint"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3a83d6bb"
      },
      "source": [
        "Set the query text and upload your desired query image to the specified INPUT_BUCKET. Use an asterisk in the query text to specify which token gets replaced by the image token. The retrieved images will be uploaded to the specified output bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f7c6fc99-42cd-4544-86df-b14f49189bfd"
      },
      "outputs": [],
      "source": [
        "# Text prompt.\n",
        "PROMPT = \"a bunch of *\"\n",
        "\n",
        "# File name of your query image. Example: \"cat.png\".\n",
        "IMAGE_FILE_NAME = \"cat.png\"\n",
        "\n",
        "# gs:// path to the directory where the query image is stored.\n",
        "INPUT_BUCKET = os.path.join(BUCKET_URI, \"input\")\n",
        "\n",
        "# gs:// path to the directory where the result images will be uploaded to.\n",
        "OUTPUT_BUCKET = os.path.join(BUCKET_URI, \"output\")\n",
        "\n",
        "instances = [\n",
        "    {\n",
        "        \"prompts\": PROMPT,\n",
        "        \"image_path\": INPUT_BUCKET,\n",
        "        \"image_file_name\": IMAGE_FILE_NAME,\n",
        "        \"output_storage_dir\": OUTPUT_BUCKET,\n",
        "    },\n",
        "]\n",
        "\n",
        "response = endpoint.predict(instances=instances)\n",
        "print(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8e71ee23-967c-4871-802f-ccb39958386e"
      },
      "source": [
        "## Cleaning Up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can delete the Google Cloud project you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d9be2a18-db11-4108-9622-647b33be2594"
      },
      "outputs": [],
      "source": [
        "# Delete endpoint resource.\n",
        "endpoint.delete(force=True)\n",
        "\n",
        "# Delete model resource.\n",
        "model.delete()\n",
        "\n",
        "# Delete Cloud Storage objects that were created.\n",
        "delete_bucket = False\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_pytorch_pic2word.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}