Formatting and refactoring

PiperOrigin-RevId: 767152796
This commit is contained in:
Vertex MG Team
2025-06-04 08:18:32 -07:00
committed by Copybara-Service
parent a6e69a4561
commit 7d4fb0ff8d
@@ -149,6 +149,9 @@
"\n",
"models, endpoints = {}, {}\n",
"\n",
"# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
"use_dedicated_endpoint = True # @param {type:\"boolean\"}\n",
"\n",
"# @markdown Click \"Show Code\" to see more details."
]
},
@@ -185,9 +188,19 @@
"ACCELERATOR_COUNT = 1\n",
"\n",
"\n",
"def deploy_model(model_id, task, accelerator_type, machine_type, accelerator_count):\n",
"def deploy_model(\n",
" model_id: str,\n",
" task: str,\n",
" machine_type: str = \"n1-standard-8\",\n",
" accelerator_type: str = \"NVIDIA_TESLA_T4\",\n",
" accelerator_count: int = 1,\n",
" use_dedicated_endpoint: bool = False,\n",
"):\n",
" model_name = \"owl-vit\"\n",
" endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
" endpoint = aiplatform.Endpoint.create(\n",
" display_name=f\"{model_name}-endpoint\",\n",
" dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
" )\n",
" serving_env = {\n",
" \"MODEL_ID\": model_id,\n",
" \"TASK\": task,\n",
@@ -225,14 +238,19 @@
")\n",
"\n",
"\n",
"models[\"model\"], endpoints[\"endpoint\"] = deploy_model(\n",
"LABEL = \"owl-vit\"\n",
"models[LABEL], endpoints[LABEL] = deploy_model(\n",
" model_id=\"google/owlvit-base-patch32\",\n",
" task=\"zero-shot-object-detection\",\n",
" accelerator_type=ACCELERATOR_TYPE,\n",
" machine_type=MACHINE_TYPE,\n",
" accelerator_type=ACCELERATOR_TYPE,\n",
" accelerator_count=ACCELERATOR_COUNT,\n",
" use_dedicated_endpoint=use_dedicated_endpoint,\n",
")\n",
"\n",
"endpoint = endpoints[LABEL]\n",
"model = models[LABEL]\n",
"\n",
"# @markdown Click \"Show Code\" to see more details."
]
},
@@ -266,7 +284,9 @@
"instances = [\n",
" {\"image\": common_util.image_to_base64(image), \"text\": text},\n",
"]\n",
"preds = endpoints[\"endpoint\"].predict(instances=instances).predictions\n",
"preds = endpoint.predict(\n",
" instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
").predictions\n",
"\n",
"\n",
"def draw_image_with_boxes(image, boxes):\n",