Add Llava model to examples (#2576)

Summary:
Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #2576

Pull Request resolved: #2576

Reviewed By: cccclai

Differential Revision: D55268288

Pulled By: iseeyuan

fbshipit-source-id: 4574fdfec46594161d0e7f34d2c9187bb1edb1d5

Author: Martin Yuan

.ci/scripts/gather_test_models.py

@@ -13,7 +13,6 @@
 from examples.models import MODEL_NAME_TO_MODEL
 from examples.xnnpack import MODEL_NAME_TO_OPTIONS
 
-
 DEFAULT_RUNNERS = {
     "linux": "linux.2xlarge",
     "macos": "macos-m1-stable",
@@ -24,6 +23,7 @@
         "w2l": "linux.12xlarge",
         "ic4": "linux.12xlarge",
         "resnet50": "linux.12xlarge",
+        "llava_encoder": "linux.4xlarge",
         # This one causes timeout on smaller runner, the root cause is unclear (T161064121)
         "dl3": "linux.12xlarge",
         "emformer_join": "linux.12xlarge",
@@ -83,7 +83,17 @@ def model_should_run_on_event(model: str, event: str) -> bool:
     We put higher priority and fast models to pull request and rest to push.
     """
     if event == "pull_request":
-        return model in ["add", "ic3", "mv2", "mv3", "resnet18", "vit"]
+        return model in ["add", "ic3", "mv2", "mv3", "resnet18", "vit", "llava_encoder"]
+    return True
+
+
+def model_should_run_on_target_os(model: str, target_os: str) -> bool:
+    """
+    A helper function to decide whether a model should be tested on a target os (linux/macos).
+    For example, a big model can be disabled in macos due to the limited macos resources.
+    """
+    if target_os == "macos":
+        return model not in ["llava_encoder"]
     return True
 
 
@@ -119,6 +129,9 @@ def export_models_for_ci() -> dict[str, dict]:
         if not model_should_run_on_event(name, event):
             continue
 
+        if not model_should_run_on_target_os(name, target_os):
+            continue
+
         if backend == "xnnpack":
             if name not in MODEL_NAME_TO_OPTIONS:
                 continue

.ci/scripts/test.sh

@@ -67,6 +67,10 @@ test_model() {
     run_portable_executor_runner
     rm "./${MODEL_NAME}.pte"
   fi
+  if [[ "${MODEL_NAME}" == "llava_encoder" ]]; then
+    # Install requirements for llava
+    bash examples/models/llava_encoder/install_requirements.sh
+  fi
   # python3 -m examples.portable.scripts.export --model_name="llama2" should works too
   "${PYTHON_EXECUTABLE}" -m examples.portable.scripts.export --model_name="${MODEL_NAME}"
   run_portable_executor_runner

.gitmodules

@@ -62,3 +62,6 @@
 [submodule "kernels/optimized/third-party/eigen"]
 	path = kernels/optimized/third-party/eigen
 	url = https://gitlab.com/libeigen/eigen.git
+[submodule "examples/third-party/LLaVA"]
+	path = examples/third-party/LLaVA
+	url = https://github.com/haotian-liu/LLaVA.git

examples/models/__init__.py

@@ -26,6 +26,7 @@
     "ic4": ("inception_v4", "InceptionV4Model"),
     "resnet18": ("resnet", "ResNet18Model"),
     "resnet50": ("resnet", "ResNet50Model"),
+    "llava_encoder": ("llava_encoder", "LlavaModel"),
 }
 
 __all__ = [

examples/models/llava_encoder/README.md

@@ -0,0 +1,20 @@
+## Summary
+In this example, we initiate the process of running multi modality through ExecuTorch.
+- Demonstrate how to export the image encoder model in the [LLava](https://github.com/haotian-liu/LLaVA) multimodal model.
+- Provide TODO steps on how to use the exported .pte file and the existing [exported Llama2 model](https://github.com/pytorch/executorch/tree/main/examples/models/llama2), to build the multimodal pipeline.
+
+## Instructions
+Note that this folder does not host the pretrained LLava model.
+- To have Llava available, follow the [Install instructions](https://github.com/haotian-liu/LLaVA?tab=readme-ov-file#install) in the LLava github. Follow the licence in the specific repo when using L
+- Since the pytorch model version may not be updated, `cd executorch`, run `./install_requirements.sh`.
+- If there is numpy compatibility issue, run `pip install bitsandbytes -I`.
+- Alternatively, run `examples/models/llava_encoder/install_requirements.sh`, to replace the steps above.
+- Run `python3 -m examples.portable.scripts.export --model_name="llava_encoder"`. The llava_encoder.pte file will be generated.
+- Run `./cmake-out/executor_runner --model_path ./llava_encoder.pte` to verify the exported model with ExecuTorch runtime with portable kernels. Note that the portable kernels are not performance optimized. Please refer to other examples like those in llama2 folder for optimization.
+
+## TODO
+- Write the pipeline in cpp
+  - Have image and text prompts as inputs.
+  - Call image processing functions to preprocess the image tensor.
+  - Load the llava_encoder.pte model, run it using the image tensor.
+  - The output of the encoder can be combined with the prompt, as inputs to the llama model. Call functions in llama_runner.cpp to run the llama model and get outputs. The ExecuTorch end to end flow for the llama model is located at `examples/models/llama2`.

examples/models/llava_encoder/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .model import LlavaModel
+
+__all__ = [
+    LlavaModel,
+]

examples/models/llava_encoder/install_requirements.sh

@@ -0,0 +1,22 @@
+#!/bin/bash
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# install llava from the submodule
+pip install --force-reinstall -e examples/third-party/LLaVA
+
+# not included in the pip install package, but needed in llava
+pip install protobuf
+
+# bitsandbytes depends on numpy 1.x, which is not compatible with numpy 2.x.
+# Reinstall bitsandbytes to make it compatible.
+pip install bitsandbytes -I
+
+# The deps of llava can have different versions than deps of ExecuTorch.
+# For example, torch version required from llava is older than ExecuTorch.
+# To make both work, recover ExecuTorch's original dependencies by rerunning
+# the install_requirements.sh.
+./install_requirements.sh

examples/models/llava_encoder/model.py

@@ -0,0 +1,52 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+
+from examples.models.model_base import EagerModelBase
+from llava.eval.run_llava import load_images, process_images
+from llava.mm_utils import get_model_name_from_path
+
+from llava.model.builder import load_pretrained_model
+from torch import nn
+
+
+class EncoderModel(nn.Module):
+    def __init__(self, llava_model):
+        super().__init__()
+        self.model_ = llava_model
+
+    def forward(self, images_tensor):
+        features = self.model_.get_model().get_vision_tower()(images_tensor)
+        features = self.model_.get_model().mm_projector(features)
+        return features
+
+
+class LlavaModel(EagerModelBase):
+    def __init__(self):
+        model_path = "liuhaotian/llava-v1.5-7b"
+        tokenizer, self.model_, self.image_processor_, context_len = (
+            load_pretrained_model(
+                model_path=model_path,
+                model_base=None,
+                model_name=get_model_name_from_path(model_path),
+            )
+        )
+        self.device = "cpu"
+        self.dtype = torch.float32
+        self.model_.to(device=self.device, dtype=self.dtype)
+
+    def get_eager_model(self):
+        model = EncoderModel(self.model_)
+        return model
+
+    def get_example_inputs(self):
+        image_file = "https://llava-vl.github.io/static/images/view.jpg"
+        images = load_images([image_file])
+        images_tensor = process_images(
+            images, self.image_processor_, self.model_.config
+        ).to(self.model_.device)
+        return (images_tensor,)