xiaohei/model_garden
1
0
Fork 0
mirror of https://github.com/GoogleCloudPlatform/vertex-ai-samples.git synced 2026-09-26 14:42:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

Open-source ImageBind serving Docker image. (#2401)

Browse Source 
* Open-source ImageBind serving Docker image.

* Update CODEOWNERS to cover ImageBind.
This commit is contained in:
Kathy Yu
2023-10-20 18:19:35 +00:00
committed by GitHub
parent d8a6659fd2
commit ee46a01c83
4 changed files with 373 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
community-content/CODEOWNERS
+1
View File
@@ -22,4 +22,5 @@
/vertex_vision_model_garden/model_oss/lm-evaluation-harness @kathyyu-google
/vertex_vision_model_garden/model_oss/tfvision @dstnluong-google
/vertex_vision_model_garden/model_oss/fvlm @minwoo33park
/vertex_vision_model_garden/model_oss/imagebind @kathyyu-google
community-content/vertex_vision_model_garden/model_oss/imagebind/config.properties
+21
View File
@@ -0,0 +1,21 @@
number_of_netty_threads=32
job_queue_size=1000
model_store=/home/model-server/model-store
workflow_store=/home/model-server/wf-store
default_response_timeout=1800
service_envelope=json
inference_address=http://0.0.0.0:7080
management_address=http://0.0.0.0:7081
metrics_address=http://0.0.0.0:7082
models={\
"imagebind_serving": {\
"1.0": {\
"defaultVersion": true,\
"marName": "imagebind_serving.mar",\
"minWorkers": 1,\
"maxWorkers": 1,\
"batchSize": 1\
}\
}\
}
community-content/vertex_vision_model_garden/model_oss/imagebind/dockerfile/serve.Dockerfile
+74
View File
@@ -0,0 +1,74 @@
# Dockerfile for the serving docker for ImageBind.
#
# To build:
# docker build -f model_oss/imagebind/dockerfile/serve.Dockerfile . -t ${YOUR_IMAGE_TAG}
#
# To push to gcr:
# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
FROM pytorch/torchserve:0.7.0-gpu
USER root
ENV infer_port=7080
ENV mng_port=7081
ENV model_name="imagebind_serving"
ENV PATH="/home/model-server/:${PATH}"
# Install tools.
RUN apt-get update && apt-get -y upgrade && apt-get install -y --no-install-recommends \
curl \
wget \
vim \
git \
libgeos-dev
# Install libraries.
ENV PIP_ROOT_USER_ACTION=ignore
RUN python3 -m pip install --upgrade pip
RUN pip install absl-py==1.4.0
RUN pip install google-cloud-storage==2.7.0
# Install ImageBind and dependencies.
RUN git clone https://github.com/facebookresearch/ImageBind.git
WORKDIR ImageBind
# Pin the commit at 07/14/2023.
RUN git reset --hard 95d27c7fd5a8362f3527e176c3a80ae5a4d880c0
# Modify tokenizer file path from ImageBind repo to work with the server.
RUN sed -i '25d' imagebind/data.py
RUN sed -i '25 i\BPE_PATH = "/home/model-server/ImageBind/bpe/bpe_simple_vocab_16e6.txt.gz"' imagebind/data.py
RUN pip install .
WORKDIR /home/model-server
# Copy license.
RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
# Copy model artifacts.
COPY model_oss/imagebind/handler.py /home/model-server/handler.py
COPY model_oss/imagebind/config.properties /home/model-server/config.properties
COPY model_oss/util/ /home/model-server/util/
ENV PYTHONPATH /home/model-server/
# Expose ports.
EXPOSE ${infer_port}
EXPOSE ${mng_port}
# Archive model artifacts and dependencies.
# Do not set --model-file and --serialized-file because model and checkpoint
# will be dynamically loaded in handler.py.
RUN torch-model-archiver \
--model-name=${model_name} \
--version=1.0 \
--handler=/home/model-server/handler.py \
--runtime=python3 \
--export-path=/home/model-server/model-store \
--archive-format=default \
--force
# Run Torchserve HTTP serve to respond to prediction requests.
CMD ["torchserve", "--start", \
"--ts-config", "/home/model-server/config.properties", \
"--models", "${model_name}=${model_name}.mar", \
"--model-store", "/home/model-server/model-store"]
community-content/vertex_vision_model_garden/model_oss/imagebind/handler.py
+277
View File
@@ -0,0 +1,277 @@
"""Custom handler for the ImageBind model."""
import logging
import os
from typing import Any, Dict, List
from imagebind import data as data_util
from imagebind.models import imagebind_model
from imagebind.models.imagebind_model import ModalityType
from PIL import Image
import torch
from torchvision import transforms
from ts.torch_handler import base_handler
from util import constants
from util import fileutils
_VIDEO_KEY_TO_AVOID_CONFLICT_WITH_IMAGE = "video"
class ImageBindHandler(base_handler.BaseHandler):
"""Custom handler for the ImageBind model.
Attributes:
map_location: Mapping storage location.
device: Device on which to run inference.
manifest: TorchServe manifest.
task: Task for which to run the ImageBind model.
model: ImageBind model instance.
"""
def initialize(self, context: Any) -> None:
"""Initializes the ImageBind model handler.
Args:
context: TorchServe context, which contains system information and the
manifest.
Raises:
ValueError: A task that is unsupported by the handler.
"""
properties = context.system_properties
self.map_location = (
"cuda"
if torch.cuda.is_available() and properties.get("gpu_id") is not None
else "cpu"
)
self.device = torch.device(
self.map_location + ":" + str(properties.get("gpu_id"))
if torch.cuda.is_available() and properties.get("gpu_id") is not None
else self.map_location
)
self.manifest = context.manifest
self.task = os.environ.get("TASK", constants.FEATURE_EMBEDDING_GENERATION)
if self.task not in [
constants.FEATURE_EMBEDDING_GENERATION,
constants.ZERO_SHOT_CLASSIFICATION,
]:
raise ValueError(f"Invalid task: {self.task}.")
logging.info(
"Handler initializing ImageBind pretrained model for task %s.",
self.task,
)
self.model = imagebind_model.imagebind_huge(pretrained=True)
self.model.eval()
self.model.to(self.device)
logging.info("Initialized ImageBind pretrained model.")
self.initialized = True
def preprocess(self, data: Any) -> List[Dict[str, Any]]:
"""Preprocesses input data, including text, image, audio and video data.
Args:
data: Input data.
Returns:
A list of processed data samples, with each sample being a dictionary of
modality (key): input (value) pairs.
"""
logging.info("Preprocessing: %d instances received.", len(data))
preprocessed_sample_list = []
for item in data:
preprocessed_sample = {}
if ModalityType.TEXT in item:
preprocessed_sample[ModalityType.TEXT] = (
data_util.load_and_transform_text(
item[ModalityType.TEXT], self.device
)
)
for image_modality in [
ModalityType.VISION,
ModalityType.DEPTH,
ModalityType.THERMAL,
]:
if image_modality in item:
image_paths = item[image_modality]
local_image_paths = fileutils.download_gcs_file_list_to_local(
image_paths, constants.LOCAL_DATA_DIR
)
is_depth_or_thermal = image_modality in [
ModalityType.DEPTH,
ModalityType.THERMAL,
]
preprocessed_sample[image_modality] = (
self._load_and_transform_image_data(
local_image_paths,
self.device,
is_depth_or_thermal=is_depth_or_thermal,
)
)
if ModalityType.AUDIO in item:
audio_paths = item[ModalityType.AUDIO]
local_audio_paths = fileutils.download_gcs_file_list_to_local(
audio_paths, constants.LOCAL_DATA_DIR
)
preprocessed_sample[ModalityType.AUDIO] = (
data_util.load_and_transform_audio_data(
local_audio_paths, self.device
)
)
if _VIDEO_KEY_TO_AVOID_CONFLICT_WITH_IMAGE in item:
video_paths = item[_VIDEO_KEY_TO_AVOID_CONFLICT_WITH_IMAGE]
local_video_paths = fileutils.download_gcs_file_list_to_local(
video_paths, constants.LOCAL_DATA_DIR
)
preprocessed_sample[_VIDEO_KEY_TO_AVOID_CONFLICT_WITH_IMAGE] = (
data_util.load_and_transform_video_data(
local_video_paths, self.device
)
)
if ModalityType.IMU in item:
# Input data in the IMU modality are expected in shape [B, 6, 2000].
preprocessed_sample[ModalityType.IMU] = torch.tensor(
item[ModalityType.IMU], dtype=torch.float32, device=self.device
)
if preprocessed_sample:
preprocessed_sample_list.append(preprocessed_sample)
return preprocessed_sample_list
def _load_and_transform_image_data(
self,
image_paths: List[str],
device: torch.device,
is_depth_or_thermal: bool = False,
) -> torch.Tensor:
"""Loads and transforms 3-channel images, depth images and thermal images.
Args:
image_paths: A list of image paths.
device: Device onto which to load images.
is_depth_or_thermal: Whether the images are depth or thermal images.
Returns:
A list of processed tensors corresponding to the input images.
Raises:
ValueError: The input image_paths is None.
"""
if image_paths is None:
raise ValueError("image_paths must not be None.")
image_outputs = []
for image_path in image_paths:
transforms_list = [
transforms.Resize(
224, interpolation=transforms.InterpolationMode.BICUBIC
),
transforms.CenterCrop(224),
transforms.ToTensor(),
]
if not is_depth_or_thermal:
transforms_list.append(
transforms.Normalize(
mean=(0.48145466, 0.4578275, 0.40821073),
std=(0.26862954, 0.26130258, 0.27577711),
)
)
data_transform = transforms.Compose(transforms_list)
with open(image_path, "rb") as fopen:
if is_depth_or_thermal:
image = Image.open(fopen).convert("L")
else:
image = Image.open(fopen).convert("RGB")
image = data_transform(image).to(device)
image_outputs.append(image)
return torch.stack(image_outputs, dim=0)
def inference(
self, data: List[Dict[str, Any]], *args, **kwargs
) -> List[Dict[str, Any]]:
"""Runs inference using the ImageBind model.
Args:
data: A list of processed data samples, with each sample being a
dictionary of modality (key): input (value) pairs.
*args: Additional inference args.
**kwargs: Additional inference kwargs.
Returns:
A list of model outputs, with each output being a dictionary of
modality (key): embedding (value) pairs.
"""
output_list = []
with torch.no_grad():
for inputs in data:
if _VIDEO_KEY_TO_AVOID_CONFLICT_WITH_IMAGE in inputs:
# Allows inference on both image and video data, which both fall under
# ModalityType.VISION.
video_inputs = {
ModalityType.VISION: inputs[
_VIDEO_KEY_TO_AVOID_CONFLICT_WITH_IMAGE
]
}
video_embeddings = self.model(video_inputs)
video_embeddings[_VIDEO_KEY_TO_AVOID_CONFLICT_WITH_IMAGE] = (
video_embeddings[ModalityType.VISION]
)
del video_embeddings[ModalityType.VISION]
del inputs[_VIDEO_KEY_TO_AVOID_CONFLICT_WITH_IMAGE]
else:
video_embeddings = {}
embeddings = self.model(inputs)
embeddings.update(video_embeddings)
output_list.append(embeddings)
return output_list
def postprocess(self, output_list: List[Dict[str, Any]]) -> List[Any]:
"""Postprocesses model outputs for the task of interest.
For feature embedding generation, returns the embeddings for each modality
for each input.
For zero-shot classification, generates classification probabilities
between the inputs of a pair of modalities for all possible pairings.
Args:
output_list: A list of model outputs, with each output being a dictionary
of modality (key): embedding (value) pairs.
Returns:
A list of postprocessed model outputs for the task of interest, with each
output corresponding to an input.
Raises:
ValueError: Fewer than two modalities are provided for zero-shot
classification, or the task is not supported.
"""
preds = []
if self.task == constants.FEATURE_EMBEDDING_GENERATION:
for item in output_list:
preds.append({k: v.tolist() for k, v in item.items()})
elif self.task == constants.ZERO_SHOT_CLASSIFICATION:
for item in output_list:
modalities = list(item.keys())
if len(modalities) < 2:
raise ValueError(
"Two or more modalities are needed for task"
f" {constants.ZERO_SHOT_CLASSIFICATION}."
)
pairwise_probs = {}
for m1 in modalities:
for m2 in modalities:
if m1 == m2:
continue
probs = torch.softmax(item[m1] @ item[m2].T, dim=-1)
pairwise_probs[
f"Classify each input in {m1} (row) against inputs in"
f" {m2} (column)"
] = probs.tolist()
preds.append(pairwise_probs)
else:
raise ValueError(f"Task {self.task} is not supported by the handler.")
return preds