feat: add llava 1.5 docker and notebook. (#2636)

Co-authored-by: Pooya Moradi <pooyam@google.com>

community-content/CODEOWNERS

@@ -23,6 +23,7 @@
 /vertex_model_garden/model_oss/tfvision @dstnluong-google
 /vertex_model_garden/model_oss/fvlm @minwoo33park
 /vertex_model_garden/model_oss/imagebind @kathyyu-google
+/vertex_model_garden/model_oss/llava @py4
 /vertex_model_garden/benchmarking_reports @lavraicse
 
 

community-content/vertex_model_garden/model_oss/llava/dockerfile/serve.Dockerfile

@@ -0,0 +1,71 @@
+FROM pytorch/torchserve:0.9.0-gpu
+
+USER root
+
+# Install tools.
+RUN apt-get update -y --allow-releaseinfo-change && apt-get -y upgrade && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim \
+        git
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+ENV INFER_PORT=7080
+ENV MNG_PORT=7081
+ENV MODEL_NAME="llava_serving"
+ENV PATH="/home/model-server/:${PATH}"
+ENV PATH="/usr/local/cuda-12.1/bin:${PATH}"
+ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH
+ENV NVIDIA_VISIBLE_DEVICES=all
+
+# Get 'LLaVA' repository from github.
+RUN git clone https://github.com/haotian-liu/LLaVA /home/model-server/LLaVA
+WORKDIR /home/model-server/LLaVA
+# Using git reset command to pin it down to a specific version.
+RUN git reset --hard 7775b12d6b20cd69089be7a18ea02615a59621cd
+
+# Install the package.
+RUN python3 -m pip install --upgrade pip
+RUN pip install google-cloud-storage==2.13.0
+RUN pip install absl-py==2.0.0
+RUN pip install -e .
+
+# Copy model artifacts.
+COPY model_oss/llava/handler.py /home/model-server/handler.py
+COPY model_oss/llava/model_handler_setup.py /home/model-server/model_handler_setup.py
+COPY model_oss/util/ /home/model-server/util/
+ENV PYTHONPATH /home/model-server
+WORKDIR /home/model-server
+
+# 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}\n"  \
+    "default_workers_per_model=DEFAULT_WORKERS_PER_MODEL" >> /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.
+# Use $NUM_GPU workers unless overriden by $TS_NUM_WORKERS
+CMD ["TOTAL=$(nvidia-smi", "--list-gpus","|","wc","-l)","&&", "TS_NUM_WORKERS=${TS_NUM_WORKERS:-$TOTAL}","&&", "sed","-i","\"s/DEFAULT_WORKERS_PER_MODEL/$TS_NUM_WORKERS/g\"","/home/model-server/config.properties", "&&", \
+    "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_model_garden/model_oss/llava/handler.py

@@ -0,0 +1,174 @@
+"""Customer handler for LLava 1.5 OSS model.
+
+The code is based on here: https://github.com/haotian-liu/LLaVA
+handler based on:
+https://github.com/haotian-liu/LLaVA/blob/main/llava/eval/run_llava.py
+There are two supported variant:
+1. liuhaotian/llava-v1.5-13b: 13B params
+2. liuhaotian/llava-v1.5-7b: 7B params
+"""
+
+import os
+import re
+from typing import Any, Dict, List
+
+from llava import constants as llava_constants
+from llava import conversation
+from llava import mm_utils
+from llava.model import builder
+import model_handler_setup
+import torch
+from ts.torch_handler import base_handler
+
+from util import constants
+from util import image_format_converter
+
+
+DEFAULT_MODEL_ID = "liuhaotian/llava-v1.5-7b"
+
+
+class LlavaHandler(base_handler.BaseHandler):
+  """Custom handler for LLava model."""
+
+  def initialize(self, context: Any):
+    """Initializes model, tokenizer, and other components."""
+    self.map_location = model_handler_setup.get_map_location(context=context)
+    self.device = model_handler_setup.get_model_device(
+        map_location=self.map_location, context=context
+    )
+    self.manifest = context.manifest
+    self.model_id = model_handler_setup.get_model_id(
+        default_model_id=DEFAULT_MODEL_ID
+    )
+
+    # Allows 4bit and 8bit quantiziation using BnB nf4.
+    precision = os.environ.get("PRECISION_MODE")
+    load_8bit = precision == constants.PRECISION_MODE_8
+    load_4bit = precision == constants.PRECISION_MODE_4
+
+    self.tokenizer, self.model, self.image_processor, self.context_len = (
+        builder.load_pretrained_model(
+            model_path=self.model_id,
+            model_base=None,
+            model_name=mm_utils.get_model_name_from_path(self.model_id),
+            load_8bit=load_8bit,
+            load_4bit=load_4bit,
+        )
+    )
+
+  def preprocess(self, data: List[Dict[str, Any]]) -> Any:
+    """Runs the preprocessing to tokenize image and the prompt."""
+    if len(data) > 1:
+      raise ValueError(
+          "LLava original repo currently does not support batch inference."
+          " https://github.com/haotian-liu/LLaVA/issues/754"
+      )
+    data = data[0]
+    prompt, base64_image = data["prompt"], data["base64_image"]
+
+    # Adds proper image token to the prompt.
+    image_token_se = (
+        llava_constants.DEFAULT_IM_START_TOKEN
+        + llava_constants.DEFAULT_IMAGE_TOKEN
+        + llava_constants.DEFAULT_IM_END_TOKEN
+    )
+    if llava_constants.IMAGE_PLACEHOLDER in prompt:
+      if self.model.config.mm_use_im_start_end:
+        prompt = re.sub(
+            llava_constants.IMAGE_PLACEHOLDER, image_token_se, prompt
+        )
+      else:
+        prompt = re.sub(
+            llava_constants.IMAGE_PLACEHOLDER,
+            llava_constants.DEFAULT_IMAGE_TOKEN,
+            prompt,
+        )
+    else:
+      if self.model.config.mm_use_im_start_end:
+        prompt = image_token_se + "\n" + prompt
+      else:
+        prompt = llava_constants.DEFAULT_IMAGE_TOKEN + "\n" + prompt
+
+    # Formats the prompt as a conversation to be fed to the model.
+    conv = conversation.conv_llava_v1.copy()
+    conv.append_message(role=conv.roles[0], message=prompt)
+    conv.append_message(role=conv.roles[1], message=None)
+    prompt = conv.get_prompt()
+
+    # Tokenizes the prompt that includes special image token as well.
+    input_ids = (
+        mm_utils.tokenizer_image_token(
+            prompt=prompt,
+            tokenizer=self.tokenizer,
+            image_token_index=llava_constants.IMAGE_TOKEN_INDEX,
+            return_tensors="pt",
+        )
+        .unsqueeze(0)
+        .to(self.device)
+    )
+
+    images = [
+        image_format_converter.base64_to_image(image_str=base64_image).convert(
+            "RGB"
+        )
+    ]
+    # Gets the image embedding.
+    images_tensor = mm_utils.process_images(
+        images=images,
+        image_processor=self.image_processor,
+        model_cfg=self.model.config,
+    ).to(self.device, dtype=torch.float16)
+
+    self.stop_str = conversation.conv_llava_v1.sep2
+    self.keywords = [self.stop_str]
+
+    return input_ids, images_tensor
+
+  def inference(
+      self, input_ids: List[torch.Tensor], images_tensor: torch.Tensor
+  ) -> List[torch.Tensor]:
+    """Runs the inference."""
+    stopping_criteria = mm_utils.KeywordsStoppingCriteria(
+        keywords=self.keywords, tokenizer=self.tokenizer, input_ids=input_ids
+    )
+
+    with torch.inference_mode():
+      output_ids = self.model.generate(
+          input_ids=input_ids,
+          images=images_tensor,
+          do_sample=False,
+          temperature=0,
+          top_p=None,
+          num_beams=1,
+          max_new_tokens=512,
+          use_cache=True,
+          stopping_criteria=[stopping_criteria],
+      )
+
+    return output_ids
+
+  def postprocess(
+      self, output_ids: List[torch.Tensor], input_token_len: int
+  ) -> List[str]:
+    """Runs the postprocessing to convert token ids to string."""
+    outputs = self.tokenizer.batch_decode(
+        output_ids[:, input_token_len:], skip_special_tokens=True
+    )[0]
+    outputs = outputs.strip()
+    if outputs.endswith(self.stop_str):
+      outputs = outputs[: -len(self.stop_str)]
+    outputs = outputs.strip()
+
+    return [outputs]
+
+  def handle(self, data: List[Dict[str, Any]], context: Any) -> List[str]:
+    """Handles an incoming request by passing it through `preprocess`, `inference`, and `postprocess`."""
+    input_ids, images_tensor = self.preprocess(data=data)
+    model_output = self.inference(
+        input_ids=input_ids, images_tensor=images_tensor
+    )
+
+    input_token_len = input_ids.shape[1]
+    return self.postprocess(
+        output_ids=model_output, input_token_len=input_token_len
+    )

community-content/vertex_model_garden/model_oss/llava/model_handler_setup.py

@@ -0,0 +1,77 @@
+"""Common utility functions for setting up and initializing the model and the handler."""
+
+import logging
+import os
+from typing import Any
+
+import torch
+
+from util import constants
+from util import fileutils
+
+
+def get_model_id(default_model_id: str) -> str:
+  """Gets a model id or a local model path.
+
+  Args:
+    default_model_id: Default model id for the corresponding model set in the
+      handler.
+
+  Returns:
+    str: model id or a local model path.
+  """
+  # The model id can be either:
+  # 1) a huggingface model card id, like "Salesforce/blip", or
+  # 2) a GCS path to the model files, like "gs://foo/bar".
+  # If it's a model card id, the model will be loaded from huggingface.
+  model_id = (
+      default_model_id
+      if os.environ.get("MODEL_ID") is None
+      else os.environ["MODEL_ID"]
+  )
+
+  # Else it will be downloaded from GCS to local first.
+  # Since the transformers from_pretrained API can't read from GCS.
+  if model_id.startswith(constants.GCS_URI_PREFIX):
+    gcs_path = model_id[len(constants.GCS_URI_PREFIX) :]
+    local_model_dir = os.path.join(constants.LOCAL_MODEL_DIR, gcs_path)
+    logging.info("Download %s to %s", model_id, local_model_dir)
+    fileutils.download_gcs_dir_to_local(model_id, local_model_dir)
+    model_id = local_model_dir
+
+  return model_id
+
+
+def get_map_location(context: Any) -> str:
+  """Gets model map location.
+
+  Args:
+    context: Torchserve worker context.
+
+  Returns:
+    str: Mapping location.
+  """
+  properties = context.system_properties
+  return (
+      "cuda"
+      if torch.cuda.is_available() and properties.get("gpu_id") is not None
+      else "cpu"
+  )
+
+
+def get_model_device(map_location: str, context: Any) -> torch.device:
+  """Gets model accelerator device.
+
+  Args:
+    map_location: Model map location.
+    context: TorchServe worker context.
+
+  Returns:
+    torch.Device: Device to load the model into.
+  """
+  properties = context.system_properties
+  return torch.device(
+      map_location + ":" + str(properties.get("gpu_id"))
+      if torch.cuda.is_available() and properties.get("gpu_id") is not None
+      else map_location
+  )

notebooks/community/CODEOWNERS

@@ -111,3 +111,4 @@
 /notebooks/community/model_garden/model_garden_pytorch_wizard_coder.ipynb @KCFindstr
 /notebooks/community/model_registry/get_started_with_vertex_ai_deployer.ipynb angelmontero@ @inardini
 /notebooks/community/model_garden/model_garden_pytorch_wizard_lm.ipynb @KCFindstr
+/notebooks/community/model_garden/model_garden_pytorch_llava.ipynb @py4