A fix in the prediction section

PiperOrigin-RevId: 701847844

Author: vertex-mg-bot

notebooks/community/model_garden/model_garden_mammut.ipynb

@@ -107,13 +107,14 @@
         "# @markdown ### Prerequisites\n",
         "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
         "\n",
-        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. 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\"). If not set, a unique GCS bucket will be created instead.\n",
+        "# @markdown 2. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. 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\"). If not set, a unique GCS bucket will be created instead.\n",
         "\n",
         "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
         "! pip install -q gradio==4.21.0\n",
         "\n",
         "import importlib\n",
         "import os\n",
+        "import uuid\n",
         "from datetime import datetime\n",
         "from typing import Tuple\n",
         "\n",
@@ -132,21 +133,31 @@
         "# Get the default region for launching jobs.\n",
         "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "\n",
+        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
+        "# @markdown | ----------- | ----------- | ----------- |\n",
+        "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
+        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "\n",
+        "\n",
         "# Cloud Storage bucket for storing the experiment artifacts.\n",
         "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
         "# prefer using your own GCS bucket, change the value yourself below.\n",
         "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
         "BUCKET_URI = \"gs://\"  # @param {type: \"string\"}\n",
+        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
         "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
         "\n",
         "# Create a unique GCS bucket for this notebook, if not specified by the user.\n",
         "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
         "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
-        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
-        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
         "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
         "else:\n",
-        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
         "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
         "    bucket_region = shell_output[0].strip().lower()\n",
         "    if bucket_region != REGION:\n",
@@ -186,23 +197,6 @@
         "models, endpoints = {}, {}\n",
         "\n",
         "\n",
-        "def resize_image(image: Image.Image, new_width: int = 512) -> Image.Image:\n",
-        "    width, height = image.size\n",
-        "    new_height = int(height * new_width / width)\n",
-        "    new_image = image.resize((new_width, new_height))\n",
-        "    return new_image\n",
-        "\n",
-        "\n",
-        "def load_image(image_url):\n",
-        "    if image_url.startswith(\"gs://\"):\n",
-        "        local_image_path = \"./images/test_image.jpg\"\n",
-        "        common_util.download_gcs_file_to_local(image_url, local_image_path)\n",
-        "        image = common_util.load_img(local_image_path)\n",
-        "    else:\n",
-        "        image = common_util.download_image(image_url)\n",
-        "    return image\n",
-        "\n",
-        "\n",
         "def deploy_mammut(\n",
         "    task: str, machine_type: str, accelerator_type: str, accelerator_count: int\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
@@ -268,7 +262,7 @@
         "):\n",
         "    \"\"\"Generates predictions based on the input image and text using an Endpoint.\"\"\"\n",
         "    # Resize and convert image to base64 string.\n",
-        "    resized_image = resize_image(image, new_width)\n",
+        "    resized_image = common_util.resize_image(image, new_width=512)\n",
         "    instances = [\n",
         "        {\n",
         "            \"image_bytes\": {\"b64\": common_util.image_to_base64(resized_image)},\n",
@@ -354,6 +348,17 @@
         "# @markdown This can be either a Cloud Storage path (gs://\\<image-path\\>) or a public url (http://\\<image-path\\>)\n",
         "image_url = \"https://images.pexels.com/photos/4012966/pexels-photo-4012966.jpeg\"  # @param {type:\"string\"}\n",
         "\n",
+        "\n",
+        "def load_image(image_url):\n",
+        "    if image_url.startswith(\"gs://\"):\n",
+        "        local_image_path = \"./images/test_image.jpg\"\n",
+        "        common_util.download_gcs_file_to_local(image_url, local_image_path)\n",
+        "        image = common_util.load_img(local_image_path)\n",
+        "    else:\n",
+        "        image = common_util.download_image(image_url)\n",
+        "    return image\n",
+        "\n",
+        "\n",
         "image = load_image(image_url)\n",
         "display(image)\n",
         "\n",
@@ -608,6 +613,7 @@
         "# @markdown This can be either a Cloud Storage path (gs://\\<image-path\\>) or a public url (http://\\<image-path\\>)\n",
         "image_url = \"https://images.pexels.com/photos/20427316/pexels-photo-20427316/free-photo-of-a-moped-parked-in-front-of-a-blue-door.jpeg?auto=compress&cs=tinysrgb&w=630&h=375&dpr=2\"  # @param {type:\"string\"}\n",
         "\n",
+        "\n",
         "image = load_image(image_url)\n",
         "display(image)\n",
         "\n",
@@ -683,7 +689,7 @@
         "text_embeddings = []\n",
         "image_embeddings = []\n",
         "for image in images:\n",
-        "    prediction = predict(retrieval_endpoint, image, text)\n",
+        "    prediction = predict(endpoints[\"retrieval\"], image, text)\n",
         "    image_embeddings.append(np.array(prediction[\"normalized_image_embedding\"]))\n",
         "    text_embeddings.append(np.array(prediction[\"normalized_text_embedding\"]))\n",
         "\n",
@@ -730,7 +736,7 @@
         "# Delete Cloud Storage objects that were created.\n",
         "delete_bucket = False  # @param {type:\"boolean\"}\n",
         "if delete_bucket:\n",
-        "    ! gsutil -m rm -r $BUCKET_URI"
+        "    ! gsutil -m rm -r $BUCKET_NAME"
       ]
     }
   ],