[Feat/Fix] Refactoring Llava models into single file

examples/usage/llava/http_llama3_llava_test.py

@@ -22,11 +22,7 @@
 import requests
 
 from llava.conversation import (
-    default_conversation,
-    conv_templates,
-    SeparatorStyle,
     conv_llava_llama_3,
-    conv_qwen,
 )
 
 
@@ -43,7 +39,8 @@ async def test_concurrent(args):
 
     prompt = "<image>\nPlease generate caption towards this image."
     conv_template = copy.deepcopy(conv_llava_llama_3)
-    conv_template.append_message(role="user", message=prompt)
+    conv_template.append_message(role=conv_template.roles[0], message=prompt)
+    conv_template.append_message(role=conv_template.roles[1], message=None)
     prompt_with_template = conv_template.get_prompt()
     response = []
     for i in range(1):
@@ -74,7 +71,8 @@ def test_streaming(args):
     url = f"{args.host}:{args.port}"
     prompt = "<image>\nPlease generate caption towards this image."
     conv_template = copy.deepcopy(conv_llava_llama_3)
-    conv_template.append_message(role="user", message=prompt)
+    conv_template.append_message(role=conv_template.roles[0], message=prompt)
+    conv_template.append_message(role=conv_template.roles[1], message=None)
     prompt_with_template = conv_template.get_prompt()
     pload = {
         "text": prompt_with_template,

examples/usage/llava/http_qwen_llava_test.py

@@ -22,11 +22,7 @@
 import requests
 
 from llava.conversation import (
-    default_conversation,
-    conv_templates,
-    SeparatorStyle,
-    conv_llava_llama_3,
-    conv_qwen,
+    conv_qwen
 )
 
 
@@ -43,7 +39,8 @@ async def test_concurrent(args):
 
     prompt = "<image>\nPlease generate caption towards this image."
     conv_template = copy.deepcopy(conv_qwen)
-    conv_template.append_message(role="user", message=prompt)
+    conv_template.append_message(role=conv_template.roles[0], message=prompt)
+    conv_template.append_message(role=conv_template.roles[1], message=None)
     prompt_with_template = conv_template.get_prompt()
     response = []
     for i in range(1):
@@ -74,7 +71,8 @@ def test_streaming(args):
     url = f"{args.host}:{args.port}"
     prompt = "<image>\nPlease generate caption towards this image."
     conv_template = copy.deepcopy(conv_qwen)
-    conv_template.append_message(role="user", message=prompt)
+    conv_template.append_message(role=conv_template.roles[0], message=prompt)
+    conv_template.append_message(role=conv_template.roles[1], message=None)
     prompt_with_template = conv_template.get_prompt()
     pload = {
         "text": prompt_with_template,
@@ -113,5 +111,5 @@ def test_streaming(args):
     parser.add_argument("--host", type=str, default="http://127.0.0.1")
     parser.add_argument("--port", type=int, default=30000)
     args = parser.parse_args()
-    # asyncio.run(test_concurrent(args))
+    asyncio.run(test_concurrent(args))
     test_streaming(args)

python/sglang/srt/managers/router/model_runner.py

@@ -421,7 +421,12 @@ def import_model_classes():
         if not ispkg:
             module = importlib.import_module(name)
             if hasattr(module, "EntryClass"):
-                model_arch_name_to_cls[module.EntryClass.__name__] = module.EntryClass
+                entry = module.EntryClass
+                if isinstance(entry, list): # To support multiple model classes in one module
+                    for cls in entry:
+                        model_arch_name_to_cls[cls.__name__] = cls
+                else:
+                    model_arch_name_to_cls[entry.__name__] = entry
     return model_arch_name_to_cls
 
 

python/sglang/srt/models/llava.py

@@ -5,7 +5,7 @@
 import numpy as np
 import torch
 from torch import nn
-from transformers import CLIPVisionModel, LlavaConfig
+from transformers import CLIPVisionModel, CLIPVisionConfig, LlavaConfig, Qwen2Config, MistralConfig 
 from transformers.models.llava.modeling_llava import LlavaMultiModalProjector
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -18,6 +18,8 @@
     unpad_image_shape,
 )
 from sglang.srt.models.llama2 import LlamaForCausalLM
+from sglang.srt.models.qwen2 import Qwen2ForCausalLM
+from sglang.srt.models.mistral import MistralForCausalLM
 
 
 class LlavaLlamaForCausalLM(nn.Module):
@@ -287,8 +289,101 @@ def num_patches_per_side(self):
         return self.image_size // self.patch_size
 
 
-first_call = True
+class LlavaQwenForCausalLM(LlavaLlamaForCausalLM):
+    def __init__(
+        self,
+        config: LlavaConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+    ) -> None:
+        super().__init__(config, quant_config=quant_config)
+        self.config = config
+        self.vision_tower = None
+        if getattr(self.config, "vision_config", None) is None:
+            self.config.vision_config = CLIPVisionConfig(self.config.mm_vision_tower)
 
+        if getattr(self.config, "text_config", None) is None:
+            self.config.text_config = Qwen2Config(self.config._name_or_path)
+
+        self.config.vision_config.hidden_size = config.mm_hidden_size
+        self.config.text_config.hidden_size = config.hidden_size
+
+        if getattr(self.config, "projector_hidden_act", None) is None:
+            self.config.projector_hidden_act = "gelu"
+
+        if getattr(self.config, "image_token_index", None) is None:
+            self.config.image_token_index = 151646
+
+        self.multi_modal_projector = LlavaMultiModalProjector(config)
+        self.language_model = Qwen2ForCausalLM(config, quant_config=quant_config)
+        if "unpad" in getattr(config, "mm_patch_merge_type", ""):
+            self.language_model.model.image_newline = nn.Parameter(
+                torch.empty(config.text_config.hidden_size, dtype=torch.float16)
+            )
+
+    def pad_input_ids(self, input_ids, pad_value, pt_shape=None, image_size=None):
+        new_image_feature_len = self.image_feature_len
+        # now only support spatial_unpad + anyres
+        if self.mm_patch_merge_type.startswith("spatial"):
+            height = width = self.num_patches_per_side
+            if pt_shape[0] > 1:
+                if self.image_aspect_ratio == "anyres":
+                    num_patch_width, num_patch_height = get_anyres_image_grid_shape(
+                        image_size,
+                        self.image_grid_pinpoints,
+                        self.vision_tower.config.image_size,
+                    )
+                if "unpad" in self.mm_patch_merge_type:
+                    h = num_patch_height * height
+                    w = num_patch_width * width
+                    new_h, new_w = unpad_image_shape(h, w, image_size)
+                    new_image_feature_len += new_h * (new_w + 1)
+
+        pad_ids = pad_value * (
+            (new_image_feature_len + len(pad_value)) // len(pad_value)
+        )
+        offset = input_ids.index(self.config.image_token_index)
+        # old_len + pad_len - 1, because we need to remove image_token_id
+        new_input_ids = (
+            input_ids[:offset]
+            + pad_ids[:new_image_feature_len]
+            + input_ids[offset + 1 :]
+        )
+        return new_input_ids, offset
+
+
+class LlavaMistralForCausalLM(LlavaLlamaForCausalLM):
+    def __init__(
+        self,
+        config: LlavaConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+    ) -> None:
+        super().__init__(config, quant_config=quant_config)
+        self.config = config
+        self.vision_tower = None
+        if getattr(self.config, "vision_config", None) is None:
+            self.config.vision_config = CLIPVisionConfig(self.config.mm_vision_tower)
+
+        if getattr(self.config, "text_config", None) is None:
+            self.config.text_config = MistralConfig(self.config._name_or_path)
+
+        self.config.vision_config.hidden_size = config.mm_hidden_size
+        self.config.text_config.hidden_size = config.hidden_size
+
+        if getattr(self.config, "projector_hidden_act", None) is None:
+            self.config.projector_hidden_act = "gelu"
+
+        if getattr(self.config, "image_token_index", None) is None:
+            self.config.image_token_index = 32000
+
+        self.multi_modal_projector = LlavaMultiModalProjector(config)
+        self.language_model = MistralForCausalLM(config, quant_config=quant_config)
+        if "unpad" in getattr(config, "mm_patch_merge_type", ""):
+            self.language_model.model.image_newline = nn.Parameter(
+                torch.empty(config.text_config.hidden_size, dtype=torch.float16)
+            )
+
+
+first_call = True
 
 def clip_vision_embed_forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
     batch_size = pixel_values.shape[0]
@@ -319,4 +414,8 @@ def monkey_path_clip_vision_embed_forward():
     )
 
 
-EntryClass = LlavaLlamaForCausalLM
+EntryClass = [
+    LlavaLlamaForCausalLM, 
+    LlavaQwenForCausalLM,
+    LlavaMistralForCausalLM
+]