{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 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": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Hugging Face DLCs: Serving PaliGemma using Pytorch Inference on Vertex AI with Custom Handler\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/open-models/serving/vertex_ai_pytorch_inference_paligemma_with_custom_handler.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fopen-models%2Fserving%2Fvertex_ai_pytorch_inference_paligemma_with_custom_handler.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/open-models/serving/vertex_ai_pytorch_inference_paligemma_with_custom_handler.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/serving/vertex_ai_pytorch_inference_paligemma_with_custom_handler.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/serving/vertex_ai_pytorch_inference_paligemma_with_custom_handler.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/serving/vertex_ai_pytorch_inference_paligemma_with_custom_handler.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/serving/vertex_ai_pytorch_inference_paligemma_with_custom_handler.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/serving/vertex_ai_pytorch_inference_paligemma_with_custom_handler.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/serving/vertex_ai_pytorch_inference_paligemma_with_custom_handler.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "84f0f73a0f76"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "| Author(s) |  [Ivan Nardini](https://github.com/inardini), [Alvaro Bartolome](https://github.com/alvarobartt) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "> [**PaliGemma**](https://ai.google.dev/gemma) is a lightweight open vision-language model (VLM) inspired by PaLI-3, and based on open components like the SigLIP vision model and the Gemma language model.\n",
        "\n",
        "> [**Hugging Face DLCs**](https://github.com/huggingface/Google-Cloud-Containers) are pre-built and optimized Deep Learning Containers (DLCs) maintained by Hugging Face and Google Cloud teams to simplify environment configuration for your ML workloads.\n",
        "\n",
        "> [**Google Vertex AI**](https://cloud.google.com/vertex-ai) is a Machine Learning (ML) platform that lets you train and deploy ML models and AI applications, and customize large language models (LLMs) for use in your AI-powered applications.\n",
        "\n",
        "This notebook showcases how to deploy Google PaliGemma from the Hugging Face Hub on Vertex AI using the Hugging Face Deep Learning Container (DLC) for Pytorch Inference in combination with a [custom handler](https://huggingface.co/docs/inference-endpoints/en/guides/custom_handler#create-custom-inference-handler).\n",
        "\n",
        "By the end of this notebook, you will learn how to:\n",
        "\n",
        "- Create a custom handler and test it\n",
        "- Register any LLM from the Hugging Face Hub on Vertex AI\n",
        "- Deploy an LLM on Vertex AI\n",
        "- Send online predictions on Vertex AI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Vertex AI SDK and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --user --quiet 'torch' 'torchvision' 'torchaudio'\n",
        "%pip install --upgrade --user --quiet 'transformers' 'accelerate>=0.26.0'\n",
        "%pip install --upgrade --user --quiet 'google-cloud-aiplatform[prediction]' 'crcmod' 'etils'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "source": [
        "### Restart runtime\n",
        "\n",
        "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which restarts the current kernel.\n",
        "\n",
        "The restart might take a minute or longer. After it's restarted, continue to the next step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SbmM4z7FOBpM"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. In Colab or Colab Enterprise, you might see an error message that says \"Your session crashed for an unknown reason.\" This is expected. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e0vxVy5J3g41"
      },
      "source": [
        "### Set Hugging Face `HF_HOME` variable"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P5UCCYrV3pEW"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "from etils.epath import Path\n",
        "\n",
        "ROOT_PATH = Path(\".\")\n",
        "TUTORIAL_PATH = ROOT_PATH / \"deploy_paligemma_with_custom_handler_tutorial\"\n",
        "\n",
        "os.environ[\"HF_HOME\"] = str(TUTORIAL_PATH)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f9af3e57f89a"
      },
      "source": [
        "### Authenticate your Hugging Face account\n",
        "\n",
        "As [`google/paligemma-3b-mix-448`](https://huggingface.co/google/paligemma-3b-mix-448) is a gated model, you are required to review and agree to Google usage license on the Hugging Face Hub for any of the models from the [PaliGemma release collection](https://huggingface.co/collections/google/paligemma-release-6643a9ffbf57de2ae0448dda), and the access request will be processed inmediately.\n",
        "\n",
        "Once that's done, you need to generate a new user access token with read-only access so that the weights can be downloaded from the Hub in the Hugging Face DLC for TGI.\n",
        "\n",
        "> Note that the user access token can only be generated via [the Hugging Face Hub UI](https://huggingface.co/settings/tokens/new), where you can either select read-only access to your account, or follow the recommendations and generate a fine-grained token with read-only access to [`google/paligemma-3b-mix-448`](https://huggingface.co/google/paligemma-3b-mix-448)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4c31c7272804"
      },
      "source": [
        "Then you can install the `huggingface_hub` that comes with a CLI that will be used for the authentication with the token generated in advance. So that then the token can be safely retrieved via `huggingface_hub.get_token`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8d836e0210fe"
      },
      "outputs": [],
      "source": [
        "from huggingface_hub import interpreter_login\n",
        "\n",
        "interpreter_login()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c71a4314c250"
      },
      "source": [
        "Read more about [Hugging Face Security](https://huggingface.co/docs/hub/en/security), specifically about [Hugging Face User Access Tokens](https://huggingface.co/docs/hub/en/security-tokens)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dy5VO78rzX8c"
      },
      "source": [
        "### Requirements"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R8Zm9y0hxU5O"
      },
      "source": [
        "#### Set Project ID and Location\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable these APIs](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,artifactregistry.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "# Use the environment variable if the user doesn't provide Project ID.\n",
        "import os\n",
        "\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "PROJECT_NUMBER = !gcloud projects describe {PROJECT_ID} --format=\"get(projectNumber)\"[0]\n",
        "PROJECT_NUMBER = PROJECT_NUMBER[0]\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "#### Set and create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as models."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_NAME = \"[your-bucket-name]\"  # @param {type: \"string\", placeholder: \"[your-bucket-name]\", isTemplate: true}\n",
        "\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a-Yz-SSagqUW"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "set_service_account"
      },
      "source": [
        "#### Set Service Account and permissions\n",
        "\n",
        "You will need to have the following IAM roles set:\n",
        "\n",
        "- Vertex AI User (roles/aiplatform.user)\n",
        "- Artifact Registry Reader (roles/artifactregistry.reader)\n",
        "- Storage Object Admin (roles/storage.objectAdmin)\n",
        "\n",
        "For more information about granting roles, see [Manage access](https://cloud.google.com/iam/docs/granting-changing-revoking-access).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VX9tpdtuQI5L"
      },
      "source": [
        "> If you run following commands using Vertex AI Workbench, run directly in the terminal.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ssUJJqXJJHgC"
      },
      "outputs": [],
      "source": [
        "SERVICE_ACCOUNT = f\"{PROJECT_NUMBER}-compute@developer.gserviceaccount.com\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wqOHg5aid6HP"
      },
      "outputs": [],
      "source": [
        "for role in ['aiplatform.user', 'storage.objectAdmin', 'artifactregistry.reader']:\n",
        "\n",
        "    ! gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
        "      --member=serviceAccount:{SERVICE_ACCOUNT} \\\n",
        "      --role=roles/{role} --condition=None"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EjebjdiNxe_D"
      },
      "source": [
        "### Initiate Vertex AI SDK\n",
        "\n",
        "Initiate Vertex AI client session."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GReulMr0xjZs"
      },
      "outputs": [],
      "source": [
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8XvMcYzm7xmm"
      },
      "source": [
        "### Import libraries\n",
        "\n",
        "Import relevant libraries."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3fXW5p7o7zlo"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "import gc\n",
        "from io import BytesIO\n",
        "import json\n",
        "\n",
        "from PIL import Image\n",
        "from etils import epath\n",
        "from google.cloud import aiplatform\n",
        "from google.cloud.aiplatform import Endpoint, Model\n",
        "from google.cloud.aiplatform.prediction import LocalModel\n",
        "import requests\n",
        "import torch\n",
        "from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E4rjOknk5yvq"
      },
      "source": [
        "### Helpers\n",
        "\n",
        "Define some helpers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MOoaB5xf51BD"
      },
      "outputs": [],
      "source": [
        "def hf_model_path(hf_home_path: str) -> str:\n",
        "    \"\"\"\n",
        "    Copy model files from the Hugging Face cache.\n",
        "    \"\"\"\n",
        "\n",
        "    # Convert and expand the input path\n",
        "    base_path = Path(hf_home_path).expanduser()\n",
        "\n",
        "    if not base_path.exists():\n",
        "        raise FileNotFoundError(f\"Cache path does not exist: {base_path}\")\n",
        "\n",
        "    # Find the deepest snapshot directory\n",
        "    snapshot_path = None\n",
        "    for model_dir in base_path.iterdir():\n",
        "        if not model_dir.is_dir():\n",
        "            continue\n",
        "\n",
        "        # Look for 'snapshots' directory\n",
        "        snapshots_dir = model_dir / \"snapshots\"\n",
        "        if not snapshots_dir.exists():\n",
        "            continue\n",
        "\n",
        "        # Get the first snapshot folder (usually there's only one)\n",
        "        for snapshot in snapshots_dir.iterdir():\n",
        "            if snapshot.is_dir():\n",
        "                snapshot_path = snapshot\n",
        "                break\n",
        "\n",
        "        if snapshot_path:\n",
        "            break\n",
        "    if not snapshot_path:\n",
        "        raise FileNotFoundError(\"No snapshot directory found in the cache\")\n",
        "\n",
        "    return snapshot_path\n",
        "\n",
        "\n",
        "def get_cuda_device_names():\n",
        "    \"\"\"A function to get the list of NVIDIA GPUs\"\"\"\n",
        "    if not torch.cuda.is_available():\n",
        "        return None\n",
        "\n",
        "    return [str(i) for i in range(torch.cuda.device_count())]\n",
        "\n",
        "\n",
        "def empty_gpu_ram():\n",
        "    gc.collect()\n",
        "    torch.cuda.empty_cache()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EdvJRUWRNGHE"
      },
      "source": [
        "## Prepare a custom handler to serve PaliGemma on Vertex AI\n",
        "\n",
        "Custom Handlers are custom classes in Python that define the pre-processing, inference, and post-processing steps required to run the inference using HuggingFace Pytorch Prediction container on Vertex AI.\n",
        "\n",
        "Think of Custom Handlers as personalized instructions for Hugging Face models. They define how to prepare the input data, run the model, and handle the results. In this sense, Custom Handlers add flexibility. They let you customize how data is prepared and processed, add extra steps, and even build in custom measurements or logging. This means you can tailor the process to your exact needs when the standard setup isn't sufficient.\n",
        "\n",
        "These instructions are stored in a file named `handler.py`. If you need additional dependecies, you can list it in a `requirements.txt` file. The PyTorch container automatically finds and uses these files if they're present.\n",
        "\n",
        "Have a look at [🤗 Serve Anything with Inference Endpoints + Custom Handlers](https://huggingface.co/blog/alvarobartt/serve-anything-inference-endpoints-custom-code) to learn more."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CgPwfG07iMCS"
      },
      "source": [
        "### Test the `handler` locally\n",
        "\n",
        "Before to build the handler module, you can test its coding locally."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OVLBAtraJkQ_"
      },
      "source": [
        "#### Load Paligemma from HuggingFace Hub\n",
        "\n",
        "Load a pre-trained model called \"paligemma\" for text generation. You first sets up the text preprocessor, which handles tasks like tokenization (breaking text into words or sub-words). Then you load the actual model itself, optimized for lower memory usage and automatic device placement (likely GPU if available)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QmHijFltikJZ"
      },
      "outputs": [],
      "source": [
        "paligemma_processor = PaliGemmaProcessor.from_pretrained(\"google/paligemma-3b-mix-448\")\n",
        "\n",
        "paligemma_model = PaliGemmaForConditionalGeneration.from_pretrained(\n",
        "    \"google/paligemma-3b-mix-448\",\n",
        "    low_cpu_mem_usage=True,\n",
        "    device_map=\"auto\",\n",
        ").eval()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zOWCRhkpJqlR"
      },
      "source": [
        "#### Collect image for generating predictions\n",
        "\n",
        "Download an image from a URL and opens it using the Pillow library. The code first fetches the image data from the given URL, then checks if the download was successful (raising an error if not). Finally, it opens the downloaded image data directly from memory using BytesIO, making it ready for further processing with Pillow."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "65FfEb0YjITa"
      },
      "outputs": [],
      "source": [
        "image_url = \"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true\"\n",
        "image_response = requests.get(image_url)\n",
        "image_response.raise_for_status()\n",
        "image = Image.open(BytesIO(image_response.content))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j5Qir_dtJ0-X"
      },
      "source": [
        "#### Prepare the prediction request\n",
        "\n",
        "Prepare the image captioning request by creating a dictionary with image data and captioning instructions.\n",
        "\n",
        "The image, downloaded from a previous response (`image_response`), is encoded into a base64 string. `max_new_tokens`limits the caption to 100 words and `do_sample` tells the model to generate the most likely caption rather than a more creative, potentially less accurate one.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6g3Mt_FVyaME"
      },
      "outputs": [],
      "source": [
        "prediction_request = {\n",
        "    \"instances\": [\n",
        "        {\n",
        "            \"prompt\": \"caption it\",\n",
        "            \"image_base64\": base64.b64encode(image_response.content).decode(\"utf-8\"),\n",
        "            \"generation_kwargs\": {\"max_new_tokens\": 100, \"do_sample\": False},\n",
        "        }\n",
        "    ]\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DH-WU2IAJ4yY"
      },
      "source": [
        "#### Generate prediction\n",
        "\n",
        "Process a list of prediction requests, each containing an image and a text prompt.\n",
        "\n",
        "For each request, you decode the base64-encoded image, prepare the image and prompt for a model (paligemma_model) using a processor (paligemma_processor), and then generate a text response using the model. Finally, you decode the generated response and adds it to a list of predictions.\n",
        "\n",
        "Error handling ensures that each request includes the necessary keys and that the image can be loaded correctly. Custom generation parameters can be provided via generation_kwargs for each instance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UyKf2v3vlOtY"
      },
      "outputs": [],
      "source": [
        "predictions = []\n",
        "\n",
        "for instance in prediction_request[\"instances\"]:\n",
        "    if any(key not in instance for key in {\"prompt\", \"image_base64\"}):\n",
        "        raise ValueError(\n",
        "            \"The request body for each instance should contain both the `prompt` and the `image_base64` key with a valid image.\"\n",
        "        )\n",
        "\n",
        "    try:\n",
        "        image_bytes = base64.b64decode(instance[\"image_base64\"])\n",
        "        image_file = BytesIO(image_bytes)\n",
        "        image = Image.open(image_file)\n",
        "    except Exception as e:\n",
        "        raise ValueError(f\"The provided image cannot be loaded (with exception {e}).\")\n",
        "    inputs = paligemma_processor(\n",
        "        text=instance[\"prompt\"], images=image, return_tensors=\"pt\"\n",
        "    ).to(paligemma_model.device)\n",
        "\n",
        "    input_len = inputs[\"input_ids\"].shape[-1]\n",
        "    with torch.inference_mode():\n",
        "        generation_kwargs = instance.get(\n",
        "            \"generation_kwargs\", {\"max_new_tokens\": 100, \"do_sample\": False}\n",
        "        )\n",
        "        generation = paligemma_model.generate(**inputs, **generation_kwargs)\n",
        "        generation = generation[0][input_len:]\n",
        "        response = paligemma_processor.decode(generation, skip_special_tokens=True)\n",
        "        predictions.append(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zzmvzmxgaiFc"
      },
      "source": [
        "Get the prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XlML8-o8lV9X"
      },
      "outputs": [],
      "source": [
        "print(predictions)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yvh-vsUblb-W"
      },
      "source": [
        "###  Define the `handler.py` module\n",
        "\n",
        "After testing the handler code, you assemble the code in a Python module which defines the custom inference handler for PaliGemma.\n",
        "\n",
        "You write this handler code to a file named handler.py within the model directory."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Qa5_u-Sn6OZ4"
      },
      "outputs": [],
      "source": [
        "model_path = hf_model_path(TUTORIAL_PATH / \"hub\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WpBfVVbtoA_m"
      },
      "outputs": [],
      "source": [
        "handler_module = \"\"\"\n",
        "from typing import Any, Dict, List\n",
        "import torch\n",
        "from transformers import PaliGemmaProcessor, PaliGemmaForConditionalGeneration\n",
        "import base64\n",
        "from io import BytesIO\n",
        "from PIL import Image\n",
        "import logging\n",
        "import sys\n",
        "\n",
        "# Configure logging to output to stdout\n",
        "logging.basicConfig(\n",
        "    level=logging.INFO,\n",
        "    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',\n",
        "    handlers=[logging.StreamHandler(sys.stdout)]\n",
        ")\n",
        "logger = logging.getLogger('huggingface_inference_toolkit')\n",
        "\n",
        "class EndpointHandler:\n",
        "    def __init__(\n",
        "        self,\n",
        "        model_dir: str = '/opt/huggingface/model',\n",
        "        **kwargs: Any,\n",
        "    ) -> None:\n",
        "        self.processor = PaliGemmaProcessor.from_pretrained(model_dir)\n",
        "        self.model = PaliGemmaForConditionalGeneration.from_pretrained(\n",
        "            model_dir,\n",
        "            low_cpu_mem_usage=True,\n",
        "            device_map=\"auto\",\n",
        "            revision=\"bfloat16\",\n",
        "            torch_dtype=torch.bfloat16,\n",
        "        ).eval()\n",
        "\n",
        "    def __call__(self, data: Dict[str, Any]) -> Dict[str, List[Any]]:\n",
        "        logger.info(\"Processing new request\")\n",
        "        predictions = []\n",
        "\n",
        "        for instance in data['instances']:\n",
        "            logger.info(f\"Processing instance: {instance.get('prompt', '')[:100]}...\")\n",
        "\n",
        "            if any(key not in instance for key in {\"prompt\", \"image_base64\"}):\n",
        "                error_msg = \"Missing prompt or image_base64 in request body\"\n",
        "                logger.info(error_msg)\n",
        "                raise ValueError(error_msg)\n",
        "\n",
        "            try:\n",
        "                image_bytes = BytesIO(base64.b64decode(instance['image_base64']))\n",
        "                image = Image.open(image_bytes)\n",
        "                logger.info(\"Image loaded successfully\")\n",
        "            except Exception as e:\n",
        "                error_msg = f\"Failed to load image: {str(e)}\"\n",
        "                logger.info(error_msg)\n",
        "                raise ValueError(error_msg)\n",
        "\n",
        "            inputs = self.processor(\n",
        "                text=instance[\"prompt\"], images=image, return_tensors=\"pt\"\n",
        "            ).to(self.model.device)\n",
        "            input_len = inputs[\"input_ids\"].shape[-1]\n",
        "            logger.info(f\"Input processed, length: {input_len}\")\n",
        "\n",
        "            with torch.inference_mode():\n",
        "                generation_kwargs = data.get(\n",
        "                    \"generation_kwargs\", {\"max_new_tokens\": 100, \"do_sample\": False}\n",
        "                )\n",
        "                logger.info(f\"Generation kwargs: {generation_kwargs}\")\n",
        "\n",
        "                generation = self.model.generate(**inputs, **generation_kwargs)\n",
        "                generation = generation[0][input_len:]\n",
        "                response = self.processor.decode(generation, skip_special_tokens=True)\n",
        "                logger.info(f\"Generated response: {response[:100]}...\")\n",
        "                predictions.append(response)\n",
        "\n",
        "        logger.info(f\"Successfully processed {len(predictions)} instances\")\n",
        "        return {\"predictions\": predictions}\n",
        "\"\"\"\n",
        "\n",
        "with open(model_path / \"handler.py\", \"w\") as hf:\n",
        "    hf.write(handler_module)\n",
        "hf.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3zOgAT4DpO8Z"
      },
      "source": [
        "### Copy model with custom handler on Cloud Bucket\n",
        "\n",
        "Efficiently upload the model directory to Google Cloud Storage using `gsutil`.\n",
        "\n",
        "Note that `-m` enables multi-threaded uploads for faster transfer, especially for large directories. `-o GSUtil:parallel_composite_upload_threshold=150M` optimizes large file uploads by splitting them into smaller parts for parallel transfer, significantly speeding up the process for files larger than 150MB."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "daQnz0Kv6Khn"
      },
      "outputs": [],
      "source": [
        "model_uri = epath.Path(BUCKET_URI) / \"model\"\n",
        "model_uri.mkdir(parents=True, exist_ok=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ukQWCwz1pXzn"
      },
      "outputs": [],
      "source": [
        "! gsutil -o GSUtil:parallel_composite_upload_threshold=150M -m cp -r {str(model_path)}/* {str(model_uri)}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xgcPjmMfXwAg"
      },
      "source": [
        "### (Optional) Testing the handler offline using serving container with Vertex AI LocalModel\n",
        "\n",
        "For debugging purpose, Vertex AI provides `LocalModel` class, accessible through the Vertex AI SDK for Python. This class allows you to build and deploy your model locally, simulating the Vertex AI environment. Using LocalModel involves creating a Docker image that encapsulates your custom predictor code and the associated handler.\n",
        "\n",
        "> **Important**: Running the LocalModel class requires a local Docker installation. This allows the model to be encapsulated within a container for consistent execution across different environments.\n",
        "\n",
        "> If you haven't already installed Docker Engine, please refer to the official installation guide: [Install Docker Engine](https://docs.docker.com/engine/install/). This documentation provides detailed instructions for various operating systems and will guide you through the installation process. Ensure Docker is running correctly before proceeding with the LocalModel examples.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Im5wedPQuMoY"
      },
      "source": [
        "#### Create a LocalModel instance\n",
        "\n",
        "Set up a local model by specifying the container image to use (a Hugging Face Transformers model optimized for PyTorch and GPUs) and the port it will listen on (5000).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VS9Z75d9uW3x"
      },
      "outputs": [],
      "source": [
        "local_paligemma_model = LocalModel(\n",
        "    serving_container_image_uri=\"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-pytorch-inference-cu121.2-2.transformers.4-44.ubuntu2204.py311\",\n",
        "    serving_container_ports=[5000],\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3bxcGRl8uas-"
      },
      "source": [
        "#### Create a LocalEndpoint instance\n",
        "\n",
        "Deploy the model to a local endpoint for serving. You load the model from the specified location (model_uri) and start serving the model locally.\n",
        "\n",
        "The `gpu_device_ids` sets leverages available GPUs if present.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hfbRbRHcuee3"
      },
      "outputs": [],
      "source": [
        "local_paligemma_endpoint = local_paligemma_model.deploy_to_local_endpoint(\n",
        "    artifact_uri=str(model_uri), gpu_device_ids=get_cuda_device_names()\n",
        ")\n",
        "\n",
        "local_paligemma_endpoint.serve()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XmnfGmT8sQ-k"
      },
      "source": [
        "#### Monitoring Your Containerized Deployment\n",
        "\n",
        "To keep track of your container's deployment progress and identify any potential issues, you can use the following Docker commands within your terminal:\n",
        "\n",
        "1. **List all containers:** `docker container ls -a` displays a list of all running and stopped containers. Locate the container associated with your local endpoint and copy its ID.  This ID is essential for the next step.\n",
        "\n",
        "2. **Stream container logs:** `docker logs --follow <CONTAINER_ID>`  provides a real-time stream of your container's logs. Replace `<CONTAINER_ID>` with the ID you copied in the previous step. Monitoring these logs allows you to observe the deployment process, identify any errors or warnings, and understand the container's overall health."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7EcVi9Y_ujFj"
      },
      "source": [
        "#### Generate predictions\n",
        "\n",
        "Send a prediction request to a local Vertex AI endpoint.\n",
        "\n",
        "You convert the request data (prediction_request) into a JSON string, send it to the endpoint, and then print the predictions from the JSON response.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Iso7dA7BzQiy"
      },
      "outputs": [],
      "source": [
        "vertex_prediction_request = json.dumps(prediction_request)\n",
        "vertex_prediction_response = local_paligemma_endpoint.predict(\n",
        "    request=vertex_prediction_request, headers={\"Content-Type\": \"application/json\"}\n",
        ")\n",
        "print(vertex_prediction_response.json()[\"predictions\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ecK96GglLTEb"
      },
      "source": [
        "## Register Google Paligemma on Vertex AI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "28d24c673821"
      },
      "source": [
        "To serve PaliGemma with Pytorch Inference on Vertex AI, you start importing the model on Vertex AI Model Registry, a central repository where you can manage the lifecycle of your ML models on Vertex AI.\n",
        "\n",
        "Before going into the code to upload or import a model on Vertex AI, let's quickly review the arguments provided to the `aiplatform.Model.upload` method:\n",
        "\n",
        "* **`display_name`** is the name that will be shown in the Vertex AI Model Registry.\n",
        "\n",
        "* **`serving_container_image_uri`** is the location of the Hugging Face DLC for Pytorch Inference that will be used for serving the model.\n",
        "\n",
        "* (optional) **`serving_container_ports`** is the port where the Vertex AI endpoint |will be exposed, by default 8080.\n",
        "\n",
        "For more information on the supported `aiplatform.Model.upload` arguments, check [its Python reference](https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform.Model#google_cloud_aiplatform_Model_upload)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ef487ce082f3"
      },
      "outputs": [],
      "source": [
        "model = Model.upload(\n",
        "    display_name=\"google--paligemma-3b-mix-448\",\n",
        "    artifact_uri=str(model_uri),\n",
        "    serving_container_image_uri=\"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-pytorch-inference-cu121.2-2.transformers.4-44.ubuntu2204.py311\",\n",
        "    serving_container_ports=[8080],\n",
        ")\n",
        "model.wait()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c427c0a87016"
      },
      "source": [
        "## Deploy Google PaliGemma on Vertex AI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ef4cd160dc09"
      },
      "source": [
        "After the model is registered on Vertex AI, you can deploy the model to an endpoint.\n",
        "\n",
        "You need to first deploy a model to an endpoint before that model can be used to serve online predictions. Deploying a model associates physical resources with the model so it can serve online predictions with low latency.\n",
        "\n",
        "Before going into the code to deploy a model to an endpoint, let's quickly review the arguments provided to the `aiplatform.Model.deploy` method:\n",
        "\n",
        "- **`endpoint`** is the endpoint to deploy the model to, which is optional, and by default will be set to the model display name with the `_endpoint` suffix.\n",
        "- **`machine_type`**, **`accelerator_type`** and **`accelerator_count`** are arguments that define which instance to use, and additionally, the accelerator to use and the number of accelerators, respectively. The `machine_type` and the `accelerator_type` are tied together, so you will need to select an instance that supports the accelerator that you are using and vice-versa. More information about the different instances at [Compute Engine Documentation - GPU machine types](https://cloud.google.com/compute/docs/gpus), and about the `accelerator_type` naming at [Vertex AI Documentation - MachineSpec](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/MachineSpec).\n",
        "\n",
        "For more information on the supported `aiplatform.Model.deploy` arguments, you can check [its Python reference](https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform.Model#google_cloud_aiplatform_Model_deploy)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e777427015db"
      },
      "outputs": [],
      "source": [
        "deployed_model = model.deploy(\n",
        "    endpoint=Endpoint.create(display_name=\"google--paligemma-7b-it-endpoint\"),\n",
        "    machine_type=\"g2-standard-4\",\n",
        "    accelerator_type=\"NVIDIA_L4\",\n",
        "    accelerator_count=1,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5f2380d3ea60"
      },
      "source": [
        "> Note that the model deployment on Vertex AI can take around 15 to 25 minutes; most of the time being the allocation / reservation of the resources, setting up the network and security, and such."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9b3fd1898241"
      },
      "source": [
        "## Online predictions on Vertex AI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2aa9cee03bd0"
      },
      "source": [
        "Once the model is deployed on Vertex AI, you can run the online predictions using the `aiplatform.Endpoint.predict` method, which will send the requests to the running endpoint in the `/predict` route specified within the container following Vertex AI I/O payload formatting."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "79e393e97ddb"
      },
      "source": [
        "### Via Python"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1a404fd19a08"
      },
      "source": [
        "#### Within the same session\n",
        "\n",
        "To run the online prediction via the Vertex AI SDK, you can simply use the `predict` method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4ab63079a69a"
      },
      "outputs": [],
      "source": [
        "output = deployed_model.predict(instances=prediction_request[\"instances\"])\n",
        "output"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b28efb473c7e"
      },
      "source": [
        "#### From a different session\n",
        "\n",
        "To run the online prediction from a different session, you can run the following snippet."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e003ab03f30f"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "from io import BytesIO\n",
        "import os\n",
        "\n",
        "from PIL import Image\n",
        "from google.cloud import aiplatform\n",
        "import requests\n",
        "\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\", isTemplate: true}\n",
        "if PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=LOCATION)\n",
        "\n",
        "ENDPOINT_DISPLAY_NAME = \"google--paligemma-7b-it-endpoint\"\n",
        "\n",
        "# Iterates over all the Vertex AI Endpoints within the current project and keeps the first match (if any), otherwise set to None\n",
        "ENDPOINT_ID = next(\n",
        "    (\n",
        "        endpoint.name\n",
        "        for endpoint in aiplatform.Endpoint.list()\n",
        "        if endpoint.display_name == ENDPOINT_DISPLAY_NAME\n",
        "    ),\n",
        "    None,\n",
        ")\n",
        "assert ENDPOINT_ID, (\n",
        "    \"`ENDPOINT_ID` is not set, please make sure that the `endpoint_display_name` is correct at \"\n",
        "    f\"https://console.cloud.google.com/vertex-ai/online-prediction/endpoints?project={os.getenv('PROJECT_ID')}\"\n",
        ")\n",
        "\n",
        "# Initiate the endpoint\n",
        "endpoint = aiplatform.Endpoint(\n",
        "    f\"projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}\"\n",
        ")\n",
        "\n",
        "# Get the image\n",
        "IMAGE_URL = \"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true\"  # @param {type:\"string\"}\n",
        "image_response = requests.get(IMAGE_URL)\n",
        "image_response.raise_for_status()\n",
        "image = Image.open(BytesIO(image_response.content))\n",
        "\n",
        "# Generate the prediction\n",
        "output = endpoint.predict(\n",
        "    instances=[\n",
        "        {\n",
        "            \"prompt\": \"caption it\",\n",
        "            \"image_base64\": base64.b64encode(image_response.content).decode(\"utf-8\"),\n",
        "            \"generation_kwargs\": {\"max_new_tokens\": 100, \"do_sample\": False},\n",
        "        }\n",
        "    ],\n",
        ")\n",
        "output"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7cc0b19ab641"
      },
      "source": [
        "### Via gcloud"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8fed7007c31c"
      },
      "source": [
        "You can also send the requests using the `gcloud` CLI via the `gcloud ai endpoints` command.\n",
        "\n",
        "> Note that, before proceeding, you should either replace the values or set the following environment variables in advance from the Python variables set in the example, as follows:\n",
        ">\n",
        "> ```python\n",
        "> import os\n",
        "> os.environ[\"PROJECT_ID\"] = PROJECT_ID\n",
        "> os.environ[\"LOCATION\"] = LOCATION\n",
        "> os.environ[\"ENDPOINT_NAME\"] = \"google--paligemma-7b-it-endpoint\"\n",
        "> ```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7fb1547d409f"
      },
      "outputs": [],
      "source": [
        "%%bash\n",
        "# Get Endpoint ID\n",
        "ENDPOINT_ID=$(gcloud ai endpoints list \\\n",
        "  --project=$PROJECT_ID \\\n",
        "  --region=$LOCATION \\\n",
        "  --filter=\"display_name=$ENDPOINT_NAME\" \\\n",
        "  --format=\"value(name)\" \\\n",
        "  | cut -d'/' -f6)\n",
        "\n",
        "# Get the image in base64\n",
        "IMAGE_URL=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true\"  # @param {type:\"string\"}\n",
        "BASE64_IMAGE=$(curl -s \"$IMAGE_URL\" | base64 -w 0)\n",
        "\n",
        "# Generate the prediction\n",
        "echo '{\n",
        "  \"instances\": [\n",
        "    {\n",
        "        \"prompt\": \"caption it\",\n",
        "        \"image_base64\": \"${BASE64_IMAGE}\",\n",
        "        \"generation_kwargs\": {\"max_new_tokens\": 100, \"do_sample\": false}\n",
        "    }\n",
        "  ]\n",
        "}' | gcloud ai endpoints predict $ENDPOINT_ID \\\n",
        "  --project=$PROJECT_ID \\\n",
        "  --region=$LOCATION \\\n",
        "  --json-request=\"-\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7e588a65759e"
      },
      "source": [
        "### Via cURL"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cf3de731bcf4"
      },
      "source": [
        "Alternatively, you can also send the requests via `cURL`.\n",
        "\n",
        "> Note that, before proceeding, you should either replace the values or set the following environment variables in advance from the Python variables set in the example, as follows:\n",
        ">\n",
        "> ```python\n",
        "> import os\n",
        "> os.environ[\"PROJECT_ID\"] = PROJECT_ID\n",
        "> os.environ[\"LOCATION\"] = LOCATION\n",
        "> os.environ[\"ENDPOINT_NAME\"] = \"google--paligemma-7b-it-endpoint\"\n",
        "> ```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a71da4e1152b"
      },
      "outputs": [],
      "source": [
        "%%bash\n",
        "# Get Endpoint ID\n",
        "ENDPOINT_ID=$(gcloud ai endpoints list \\\n",
        "  --project=$PROJECT_ID \\\n",
        "  --region=$LOCATION \\\n",
        "  --filter=\"display_name=$ENDPOINT_NAME\" \\\n",
        "  --format=\"value(name)\" \\\n",
        "  | cut -d'/' -f6)\n",
        "\n",
        "# Get the image in base64\n",
        "IMAGE_URL=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true\" # @param {type:\"string\"}\n",
        "BASE64_IMAGE=$(curl -s \"$IMAGE_URL\" | base64 -w 0)\n",
        "\n",
        "# Generate the prediction\n",
        "curl -X POST \\\n",
        "    -H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        "    -H \"Content-Type: application/json\" \\\n",
        "    \"https://${LOCATION}-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/${LOCATION}/endpoints/${ENDPOINT_ID}:predict\" \\\n",
        "    -d @- <<EOF\n",
        "{\n",
        "    \"instances\": [\n",
        "        {\n",
        "            \"prompt\": \"caption it\",\n",
        "            \"image_base64\": \"${BASE64_IMAGE}\",\n",
        "            \"generation_kwargs\": {\"max_new_tokens\": 100, \"do_sample\": false}\n",
        "        }\n",
        "    ]\n",
        "}\n",
        "EOF"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f6f17f9aff65"
      },
      "source": [
        "## Cleaning up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KFOXWLLTvMvJ"
      },
      "outputs": [],
      "source": [
        "delete_endpoint = False\n",
        "delete_model = False\n",
        "delete_bucket = False\n",
        "\n",
        "if delete_endpoint:\n",
        "    deployed_model.undeploy_all()\n",
        "    deployed_model.delete()\n",
        "\n",
        "if delete_endpoint:\n",
        "    delete_model.delete()\n",
        "\n",
        "if delete_bucket:\n",
        "    ! gsutil rm -r {BUCKET_URI}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "vertex_ai_pytorch_inference_paligemma_with_custom_handler.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
