mirror of
https://github.com/GoogleCloudPlatform/vertex-ai-samples.git
synced 2026-09-26 14:42:04 +00:00
Add Pic2Word serving dockerfile and handler (#2134)
* Added Pic2Word notebook to Vertex AI Model Garden. * Added jismailyan to pic2word notebook codeowners * Pic2Word update. * Formatted notebook using lint script. * Add Pic2Word serving dockerfile and handler. * Add jismailyan to CODEOWNERS for model OSS pic2word * Fix filename typo --------- Co-authored-by: Andrew Ferlitsch <aferlitsch@google.com>
This commit is contained in:
jismailyan-google
co-authored by
Andrew Ferlitsch
Andrew Ferlitsch
parent
ac4bf93914
commit
90da7214c7
@@ -14,3 +14,4 @@
|
||||
/vertex_vision_model_garden/model_oss/diffusers @weigary
|
||||
/vertex_vision_model_garden/model_oss/keras @dstnluong-google
|
||||
/vertex_vision_model_garden/model_oss/transformers @dstnluong-google
|
||||
/vertex_vision_model_garden/model_oss/pic2word @jismailyan-google
|
||||
|
||||
+115
@@ -0,0 +1,115 @@
|
||||
FROM pytorch/torchserve:0.7.1-gpu
|
||||
|
||||
USER root
|
||||
|
||||
ENV infer_port=7080
|
||||
ENV mng_port=7081
|
||||
ENV model_name="pic2word"
|
||||
ENV PATH="/home/model-server/:${PATH}"
|
||||
|
||||
# Copy license.
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
wget
|
||||
|
||||
RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
|
||||
|
||||
# Install dependencies.
|
||||
ENV PIP_ROOT_USER_ACTION=ignore
|
||||
RUN python3 -m pip install --upgrade pip
|
||||
RUN pip install google-cloud-storage==2.7.0
|
||||
RUN pip install open_clip_torch==2.20.0
|
||||
RUN pip install numpy==1.22.0
|
||||
RUN pip install scikit-image==0.21.0
|
||||
RUN pip install scikit-learn==1.0.2
|
||||
RUN pip install torch==2.0.0
|
||||
RUN pip install torchvision==0.15.2
|
||||
RUN pip install tensorboard==2.13.0
|
||||
RUN pip install ase==3.21.1
|
||||
RUN pip install braceexpand==0.1.7
|
||||
RUN pip install cached-property==1.5.2
|
||||
RUN pip install configparser==5.0.2
|
||||
RUN pip install cycler==0.10.0
|
||||
RUN pip install decorator==4.4.2
|
||||
RUN pip install docker-pycreds==0.4.0
|
||||
RUN pip install gitdb==4.0.7
|
||||
RUN pip install gitpython==3.1.30
|
||||
RUN pip install googledrivedownloader==0.4
|
||||
RUN pip install h5py==3.1.0
|
||||
RUN pip install isodate==0.6.0
|
||||
RUN pip install jinja2==3.0.1
|
||||
RUN pip install kiwisolver==1.3.1
|
||||
RUN pip install littleutils==0.2.2
|
||||
RUN pip install llvmlite==0.36.0
|
||||
RUN pip install markupsafe==2.0.1
|
||||
RUN pip install matplotlib==3.3.4
|
||||
RUN pip install networkx==2.5.1
|
||||
RUN pip install numba==0.53.1
|
||||
RUN pip install ogb==1.3.1
|
||||
RUN pip install outdated==0.2.1
|
||||
RUN pip install pathtools==0.1.2
|
||||
RUN pip install promise==2.3
|
||||
RUN pip install psutil==5.8.0
|
||||
RUN pip install pyarrow==4.0.0
|
||||
RUN pip install pyparsing==2.4.7
|
||||
RUN pip install python-louvain==0.15
|
||||
RUN pip install pyyaml==5.4.1
|
||||
RUN pip install rdflib==5.0.0
|
||||
RUN pip install sentry-sdk==1.14.0
|
||||
RUN pip install shortuuid==1.0.1
|
||||
RUN pip install sklearn==0.0
|
||||
RUN pip install smmap==4.0.0
|
||||
RUN pip install subprocess32==3.5.4
|
||||
RUN pip install torch-geometric==1.7.0
|
||||
RUN pip install wandb==0.10.30
|
||||
RUN pip install wilds==1.1.0
|
||||
RUN pip install ftfy==6.1.1
|
||||
RUN pip install regex==2023.6.3
|
||||
RUN pip install webdataset==0.2.48
|
||||
RUN pip install requests==2.31.0
|
||||
RUN pip install hydra-core==1.3.2
|
||||
RUN pip install omegaconf==2.3.0
|
||||
RUN pip install fairseq==0.10.0
|
||||
RUN pip install bitarray==2.7.6
|
||||
|
||||
# Get 'composed_image_retrieval' repository from github.
|
||||
RUN git clone https://github.com/google-research/composed_image_retrieval
|
||||
# Set workdir to composed_image_retrieval.
|
||||
WORKDIR ./composed_image_retrieval
|
||||
# Using git reset command to pin it down to a specific version.
|
||||
RUN git reset --hard 8c053297c2fae9cd17ddcded48445a4f47208dbd
|
||||
|
||||
# Fix issue introduced by installing composed_image_retrieval
|
||||
# https://github.com/huggingface/transformers/issues/8638#issuecomment-790772391
|
||||
RUN pip uninstall dataclasses -y
|
||||
|
||||
# Copy model artifacts.
|
||||
COPY model_oss/pic2word/handler.py /home/model-server/handler.py
|
||||
|
||||
# Create torchserve configuration file.
|
||||
RUN echo \
|
||||
"default_response_timeout=1800\n" \
|
||||
"service_envelope=json\n" \
|
||||
"inference_address=http://0.0.0.0:${infer_port}\n" \
|
||||
"management_address=http://0.0.0.0:${mng_port}" >> /home/model-server/config.properties
|
||||
|
||||
# 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"]
|
||||
@@ -0,0 +1,167 @@
|
||||
"""Custom handler for Pic2Word."""
|
||||
|
||||
from argparse import Namespace # pylint: disable=g-importing-member
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
from absl import logging
|
||||
from data import CustomFolder
|
||||
from eval_utils import visualize_results
|
||||
from model.clip import load
|
||||
from model.model import convert_weights
|
||||
from model.model import IM2TEXT
|
||||
from params import get_project_root
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from ts.torch_handler.base_handler import BaseHandler
|
||||
|
||||
from google3.cloud.ml.applications.vision.model_garden.model_oss.util import fileutils
|
||||
|
||||
# The COCO dataset is stored in a publicly accessible bucket.
|
||||
_COCO_STORAGE_DIR = "gs://pic2word-bucket/data/coco/"
|
||||
_COCO_LOCAL_DIR = "/home/model-server/composed_image_retrieval/data/coco/"
|
||||
_COCO_VAL2017_PATH = "coco/val2017"
|
||||
_COCO_DATASET_NAME = "coco"
|
||||
_MODEL_NAME = "ViT-L/14"
|
||||
_LOCAL_QUERY_PATH = "./query/"
|
||||
_IMAGE_OUTPUT_LOCAL_DIR = "demo_out/images"
|
||||
_OUTPUT_LOCAL_DIR = "/demo_out/"
|
||||
_DATA_DIR = "data"
|
||||
_CHECKPOINT_DIR = "checkpoint/pic2word_model.pt"
|
||||
_REQUEST_PROMPTS = "prompts"
|
||||
_REQUEST_OUTPUT_STORAGE_DIR = "output_storage_dir"
|
||||
_REQUEST_IMAGE_PATH = "image_path"
|
||||
_REQUEST_IMAGE_FILE_NAME = "image_file_name"
|
||||
_RESPONSE_MSG = "Successfully retrieved images."
|
||||
|
||||
|
||||
class ModelHandler(BaseHandler):
|
||||
"""A custom model handler implementation."""
|
||||
|
||||
def __init__(self):
|
||||
self.initialized = False
|
||||
self.gpu = 0
|
||||
self.model = None
|
||||
self.dataloader = None
|
||||
self.prompt = None
|
||||
self.output_storage_dir = None
|
||||
|
||||
def initialize(self, context: Any):
|
||||
"""Initialize."""
|
||||
logging.info("Initializing pic2word.")
|
||||
|
||||
# Download COCO dataset. The model looks for this folder specifically
|
||||
# during image retrieval to generate a response for each request.
|
||||
# This is a publicly accessible bucket.
|
||||
fileutils.download_gcs_dir_to_local(
|
||||
_COCO_STORAGE_DIR,
|
||||
_COCO_LOCAL_DIR,
|
||||
)
|
||||
|
||||
# Load the model.
|
||||
|
||||
self.initialized = True
|
||||
|
||||
torch.cuda.set_device(self.gpu)
|
||||
model, _, preprocess_val = load(_MODEL_NAME, jit=False)
|
||||
|
||||
img2text = IM2TEXT(
|
||||
embed_dim=model.embed_dim,
|
||||
output_dim=model.token_embedding.weight.shape[1],
|
||||
)
|
||||
|
||||
model.cuda(self.gpu)
|
||||
img2text.cuda(self.gpu)
|
||||
|
||||
convert_weights(model)
|
||||
convert_weights(img2text)
|
||||
|
||||
self.model = model
|
||||
self.img2text = img2text
|
||||
|
||||
# Load the dataset
|
||||
logging.info("Loading dataset.")
|
||||
|
||||
root_project = os.path.join(get_project_root(), _DATA_DIR)
|
||||
dataset = CustomFolder(
|
||||
os.path.join(root_project, _COCO_VAL2017_PATH), transform=preprocess_val
|
||||
)
|
||||
|
||||
# Initialize the dataloader. This is used to create the pickle file from
|
||||
# the dataset.
|
||||
dataloader = DataLoader(
|
||||
dataset,
|
||||
batch_size=64,
|
||||
shuffle=False,
|
||||
num_workers=1,
|
||||
pin_memory=True,
|
||||
drop_last=False,
|
||||
)
|
||||
|
||||
self.dataloader = dataloader
|
||||
|
||||
logging.info("Finished initializing Pic2Word server.")
|
||||
|
||||
def preprocess(self, data: Any) -> str:
|
||||
"""Preprocess input data."""
|
||||
logging.info("Preprocessing Pic2Word inference request.")
|
||||
query = data[0]
|
||||
|
||||
self.output_storage_dir = query[_REQUEST_OUTPUT_STORAGE_DIR]
|
||||
prompts = query[_REQUEST_PROMPTS]
|
||||
prompts = prompts.split(",")
|
||||
self.prompt = prompts
|
||||
|
||||
image_path = query[_REQUEST_IMAGE_PATH]
|
||||
# The query image is only supported via GCS bucket upload.
|
||||
fileutils.download_gcs_dir_to_local(image_path, _LOCAL_QUERY_PATH)
|
||||
image_file_name = query[_REQUEST_IMAGE_FILE_NAME]
|
||||
|
||||
query_file = f"./query/{image_file_name}"
|
||||
|
||||
logging.info("Setting model args.")
|
||||
|
||||
args = {
|
||||
"openai-pretrained": True,
|
||||
"resume": _CHECKPOINT_DIR,
|
||||
"retrieval_data": _COCO_DATASET_NAME,
|
||||
"query_file": query_file,
|
||||
"demo_out": _OUTPUT_LOCAL_DIR,
|
||||
"prompts": prompts,
|
||||
"distributed": False,
|
||||
"dp": False,
|
||||
"gpu": 0,
|
||||
"model": _MODEL_NAME,
|
||||
"world_size": 1,
|
||||
}
|
||||
model_input = Namespace(**args)
|
||||
|
||||
logging.info("Finished preprocessing Pic2Word inference request.")
|
||||
return model_input
|
||||
|
||||
def inference(self, model_input: Any):
|
||||
"""Runs inference."""
|
||||
logging.info("Running model-inference.")
|
||||
visualize_results(
|
||||
model=self.model,
|
||||
img2text=self.img2text,
|
||||
args=model_input,
|
||||
prompt=self.prompt,
|
||||
dataloader=self.dataloader,
|
||||
)
|
||||
|
||||
def postprocess(self):
|
||||
"""Upload the output images to the bucket."""
|
||||
logging.info("Running request postprocess.")
|
||||
fileutils.upload_local_dir_to_gcs(
|
||||
_IMAGE_OUTPUT_LOCAL_DIR, self.output_storage_dir
|
||||
)
|
||||
|
||||
def handle(self, data: Any, context: Any) -> str: # pylint: disable=unused-argument
|
||||
"""Runs preprocess, inference, and post-processing."""
|
||||
logging.info("Received Pic2Word inference request")
|
||||
model_input = self.preprocess(data)
|
||||
self.inference(model_input)
|
||||
self.postprocess()
|
||||
logging.info("Done handling input.")
|
||||
return _RESPONSE_MSG
|
||||
Reference in New Issue
Block a user