model: support nvidia/LocateAnything-3B (#28958)
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
This commit is contained in:
co-authored by
Xinyuan Tong
parent
5169df70f6
commit
473a278dd1
@@ -21,6 +21,7 @@ from sglang.srt.configs.laguna import LagunaConfig
|
||||
from sglang.srt.configs.lfm2 import Lfm2Config
|
||||
from sglang.srt.configs.lfm2_moe import Lfm2MoeConfig
|
||||
from sglang.srt.configs.lfm2_vl import Lfm2VlConfig
|
||||
from sglang.srt.configs.locate_anything import LocateAnythingConfig
|
||||
from sglang.srt.configs.longcat_flash import LongcatFlashConfig
|
||||
from sglang.srt.configs.minicpmv4_6 import MiniCPMV4_6Config, MiniCPMV4_6VisionConfig
|
||||
from sglang.srt.configs.nano_nemotron_vl import (
|
||||
@@ -71,6 +72,7 @@ __all__ = [
|
||||
"Lfm2Config",
|
||||
"Lfm2MoeConfig",
|
||||
"Lfm2VlConfig",
|
||||
"LocateAnythingConfig",
|
||||
"MiniCPMV4_6Config",
|
||||
"MiniCPMV4_6VisionConfig",
|
||||
"NemotronHConfig",
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Adapted from https://huggingface.co/nvidia/LocateAnything-3B/blob/main/configuration_locateanything.py
|
||||
"""Config for nvidia/LocateAnything-3B.
|
||||
|
||||
LocateAnything is a multimodal grounding/detection model composed of a MoonViT
|
||||
vision encoder, an InternVL-style ``mlp1`` projector, and a Qwen2 language model
|
||||
backbone. The config is a composite that wraps a ``MoonViTConfig`` (vision) and a
|
||||
``Qwen2Config`` (text) plus the special token ids used for the grounding grammar
|
||||
(``<box>``/``<ref>``/coordinate tokens).
|
||||
"""
|
||||
|
||||
from typing import Optional, Union
|
||||
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.models.qwen2 import Qwen2Config
|
||||
|
||||
from sglang.srt.configs.kimi_vl_moonvit import MoonViTConfig
|
||||
|
||||
|
||||
class LocateAnythingConfig(PretrainedConfig):
|
||||
model_type = "locateanything"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vision_config: Optional[Union[dict, MoonViTConfig]] = None,
|
||||
text_config: Optional[Union[dict, Qwen2Config]] = None,
|
||||
image_token_index: int = 151665,
|
||||
box_start_token_id: int = 151668,
|
||||
box_end_token_id: int = 151669,
|
||||
ref_start_token_id: int = 151672,
|
||||
ref_end_token_id: int = 151673,
|
||||
coord_start_token_id: int = 151677,
|
||||
coord_end_token_id: int = 152677,
|
||||
none_token_id: int = 4064,
|
||||
mlp_connector_layers: int = 2,
|
||||
**kwargs,
|
||||
):
|
||||
if vision_config is None:
|
||||
vision_config = MoonViTConfig()
|
||||
elif isinstance(vision_config, dict):
|
||||
vision_config = MoonViTConfig(**vision_config)
|
||||
self.vision_config = vision_config
|
||||
|
||||
if text_config is None:
|
||||
text_config = Qwen2Config()
|
||||
elif isinstance(text_config, dict):
|
||||
text_config = Qwen2Config(**text_config)
|
||||
self.text_config = text_config
|
||||
|
||||
self.image_token_index = image_token_index
|
||||
self.box_start_token_id = box_start_token_id
|
||||
self.box_end_token_id = box_end_token_id
|
||||
# ref_*_token_id and mlp_connector_layers are kept for round-trip
|
||||
# fidelity with the HF config; the box-grammar processor reads the box /
|
||||
# coord / none ids, and the projector hardcodes its 2-layer structure.
|
||||
self.ref_start_token_id = ref_start_token_id
|
||||
self.ref_end_token_id = ref_end_token_id
|
||||
self.coord_start_token_id = coord_start_token_id
|
||||
self.coord_end_token_id = coord_end_token_id
|
||||
self.none_token_id = none_token_id
|
||||
self.mlp_connector_layers = mlp_connector_layers
|
||||
|
||||
super().__init__(**kwargs)
|
||||
@@ -1693,6 +1693,7 @@ multimodal_model_archs = [
|
||||
"Qwen3ASRForConditionalGeneration",
|
||||
"Qwen3OmniMoeForConditionalGeneration",
|
||||
"KimiVLForConditionalGeneration",
|
||||
"LocateAnythingForConditionalGeneration",
|
||||
"InternVLChatModel",
|
||||
"InternS1ForConditionalGeneration",
|
||||
"InternS1ProForConditionalGeneration",
|
||||
|
||||
@@ -1326,6 +1326,22 @@ def _get_length(value):
|
||||
return None
|
||||
|
||||
|
||||
def _is_rank2_grid(value):
|
||||
"""True if `value` is a rank-2 grid ([N, dims]) suitable for per-row prod.
|
||||
|
||||
Tensors/arrays must have ndim == 2; nested lists/tuples must have each row
|
||||
be a sequence. Anything flat (1-D / scalars) is rejected so callers fall
|
||||
back to a simple split instead of mis-collapsing it with prod(dim=-1).
|
||||
"""
|
||||
if isinstance(value, (torch.Tensor, np.ndarray)):
|
||||
return value.ndim == 2
|
||||
if isinstance(value, (list, tuple)):
|
||||
return len(value) > 0 and all(
|
||||
isinstance(row, (list, tuple, torch.Tensor, np.ndarray)) for row in value
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def _slice_value(value, start, end):
|
||||
if isinstance(value, torch.Tensor):
|
||||
return value[start:end]
|
||||
@@ -1490,7 +1506,15 @@ def get_new_expanded_mm_items(original_mm_items):
|
||||
num_items = len(item.offsets)
|
||||
|
||||
if item.is_image():
|
||||
# MoonViT-style models (e.g. LocateAnything) carry per-image
|
||||
# grids under `image_grid_hws` ([h, w]) rather than
|
||||
# `image_grid_thw` ([t, h, w]); both encode dim-0 patch counts
|
||||
# via prod over the last axis, so accept either key. (Use an
|
||||
# explicit None check, not `a or b`: the value is a multi-element
|
||||
# tensor whose truthiness is ambiguous.)
|
||||
image_grid_thw = item.model_specific_data.get("image_grid_thw")
|
||||
if image_grid_thw is None:
|
||||
image_grid_thw = item.model_specific_data.get("image_grid_hws")
|
||||
grid_len = _get_length(image_grid_thw)
|
||||
if image_grid_thw is None or grid_len != num_items:
|
||||
# No grid info — fall back to simple split by feature dim-0
|
||||
@@ -1498,6 +1522,18 @@ def get_new_expanded_mm_items(original_mm_items):
|
||||
expanded_mm_items.append(item)
|
||||
continue
|
||||
|
||||
# The grid must be rank-2 ([N, dims]) so `prod` over the last
|
||||
# axis yields one patch count per image. A flat 1-D grid (e.g.
|
||||
# `tensor([h, w])` with num_items==2) would pass the length check
|
||||
# above but `prod(dim=-1)` collapses it to a scalar and mis-splits.
|
||||
# The HF processor always emits rank-2, so this only guards the
|
||||
# degenerate case — fall back to simple split rather than corrupt
|
||||
# the slice boundaries.
|
||||
if not _is_rank2_grid(image_grid_thw):
|
||||
if not _try_simple_split(item, num_items, expanded_mm_items):
|
||||
expanded_mm_items.append(item)
|
||||
continue
|
||||
|
||||
if isinstance(image_grid_thw, torch.Tensor):
|
||||
patches_per_item = (
|
||||
torch.prod(image_grid_thw, dim=-1).long().tolist()
|
||||
|
||||
@@ -0,0 +1,403 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Adapted from https://huggingface.co/nvidia/LocateAnything-3B/blob/main/modeling_locateanything.py
|
||||
# and from vllm-project/vllm PR #44182.
|
||||
"""Inference-only LocateAnything-3B model for SGLang.
|
||||
|
||||
LocateAnything-3B is a multimodal grounding/detection model:
|
||||
|
||||
* MoonViT vision encoder (reused unchanged from Kimi-VL)
|
||||
* An InternVL-style ``mlp1`` projector (LayerNorm applied AFTER the 2x2 patch
|
||||
merge, i.e. over ``hidden_size * merge_h * merge_w``)
|
||||
* A Qwen2 language-model backbone
|
||||
|
||||
The model emits structured grounding outputs such as
|
||||
``<ref>object</ref><box>...</box>`` when special tokens are preserved
|
||||
(``skip_special_tokens=False``). An optional constrained-decoding logit
|
||||
processor (:class:`LocateAnythingBoxGrammarLogitProcessor`) restricts the tokens
|
||||
emitted inside a ``<box>...</box>`` block to a valid ``none`` / point / bbox
|
||||
pattern.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from sglang.srt.configs.kimi_vl_moonvit import MoonViTConfig
|
||||
from sglang.srt.configs.locate_anything import LocateAnythingConfig
|
||||
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
||||
from sglang.srt.managers.mm_utils import (
|
||||
MultiModalityDataPaddingPatternMultimodalTokens,
|
||||
general_mm_embed_routine,
|
||||
)
|
||||
from sglang.srt.managers.schedule_batch import (
|
||||
Modality,
|
||||
MultimodalDataItem,
|
||||
MultimodalInputs,
|
||||
)
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||
from sglang.srt.model_loader.weight_utils import default_weight_loader
|
||||
from sglang.srt.models.kimi_vl_moonvit import MoonVitPretrainedModel
|
||||
from sglang.srt.models.qwen2 import Qwen2ForCausalLM
|
||||
from sglang.srt.sampling.custom_logit_processor import CustomLogitProcessor
|
||||
from sglang.srt.utils import add_prefix
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LocateAnythingMultiModalProjector(nn.Module):
|
||||
"""InternVL-style ``mlp1`` projector.
|
||||
|
||||
Unlike Kimi-VL's projector (which LayerNorms the per-patch features over
|
||||
``hidden_size`` *before* the 2x2 merge), LocateAnything merges first and then
|
||||
LayerNorms over the merged width ``hidden_size * merge_h * merge_w``.
|
||||
|
||||
HF checkpoint layout (``mlp1`` Sequential):
|
||||
mlp1.0 = LayerNorm(merged_size)
|
||||
mlp1.1 = Linear(merged_size, text_hidden)
|
||||
mlp1.2 = GELU
|
||||
mlp1.3 = Linear(text_hidden, text_hidden)
|
||||
"""
|
||||
|
||||
def __init__(self, config: LocateAnythingConfig):
|
||||
super().__init__()
|
||||
|
||||
merge = config.vision_config.merge_kernel_size
|
||||
self.merged_size = config.vision_config.hidden_size * merge[0] * merge[1]
|
||||
text_hidden = config.text_config.hidden_size
|
||||
|
||||
self.pre_norm = nn.LayerNorm(self.merged_size, eps=1e-5)
|
||||
self.linear_1 = nn.Linear(self.merged_size, text_hidden, bias=True)
|
||||
# Plain (exact, erf-based) GELU to match the HF checkpoint's nn.GELU().
|
||||
self.act = nn.GELU()
|
||||
self.linear_2 = nn.Linear(text_hidden, text_hidden, bias=True)
|
||||
|
||||
def forward(self, image_features: torch.Tensor) -> torch.Tensor:
|
||||
# MoonViT's patch_merger yields per-image tensors of shape
|
||||
# (num_merged_tokens, merge_h * merge_w, hidden_size); concatenated and
|
||||
# flattened to (num_merged_tokens, merged_size) the 4 sub-patches sit
|
||||
# contiguously per token, matching the trained LayerNorm(merged_size).
|
||||
# reshape (not view) since the concatenated input may be non-contiguous.
|
||||
hidden_states = image_features.reshape(-1, self.merged_size)
|
||||
hidden_states = self.pre_norm(hidden_states)
|
||||
hidden_states = self.linear_1(hidden_states)
|
||||
hidden_states = self.act(hidden_states)
|
||||
hidden_states = self.linear_2(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class LocateAnythingForConditionalGeneration(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
config: LocateAnythingConfig,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
prefix: str = "",
|
||||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
assert isinstance(config.vision_config, MoonViTConfig)
|
||||
|
||||
self.vision_tower = MoonVitPretrainedModel(config.vision_config)
|
||||
self.multi_modal_projector = LocateAnythingMultiModalProjector(config)
|
||||
self.quant_config = quant_config
|
||||
|
||||
self.language_model = Qwen2ForCausalLM(
|
||||
config=config.text_config,
|
||||
quant_config=quant_config,
|
||||
prefix=add_prefix("language_model", prefix),
|
||||
)
|
||||
|
||||
def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
|
||||
pixel_values = (
|
||||
torch.cat([item.feature for item in items], dim=0)
|
||||
.type(self.vision_tower.dtype)
|
||||
.to(self.vision_tower.device)
|
||||
)
|
||||
|
||||
# Already-projected embeddings (e.g. precomputed) pass through.
|
||||
if (
|
||||
pixel_values.dim() == 2
|
||||
and pixel_values.shape[-1] == self.config.text_config.hidden_size
|
||||
):
|
||||
return pixel_values
|
||||
|
||||
# image_grid_hws may arrive as numpy arrays from the HF image processor;
|
||||
# coerce each to a tensor before concatenating.
|
||||
image_grid_hws = torch.cat(
|
||||
[torch.as_tensor(item.image_grid_hws) for item in items], dim=0
|
||||
).to(self.vision_tower.device)
|
||||
image_features = self.vision_tower(pixel_values, image_grid_hws)
|
||||
assert isinstance(image_features, list)
|
||||
return self.multi_modal_projector(torch.cat(image_features))
|
||||
|
||||
def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs):
|
||||
pattern = MultiModalityDataPaddingPatternMultimodalTokens()
|
||||
return pattern.pad_input_tokens(input_ids, mm_inputs)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
get_embedding: bool = False,
|
||||
):
|
||||
hidden_states = general_mm_embed_routine(
|
||||
input_ids=input_ids,
|
||||
forward_batch=forward_batch,
|
||||
language_model=self.language_model,
|
||||
data_embedding_funcs={
|
||||
Modality.IMAGE: self.get_image_feature,
|
||||
},
|
||||
positions=positions,
|
||||
)
|
||||
return hidden_states
|
||||
|
||||
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
|
||||
# Remap HF checkpoint prefixes onto SGLang submodule names.
|
||||
prefix_mapping = {
|
||||
"vision_model.": "vision_tower.",
|
||||
"mlp1.0.": "multi_modal_projector.pre_norm.",
|
||||
"mlp1.1.": "multi_modal_projector.linear_1.",
|
||||
"mlp1.3.": "multi_modal_projector.linear_2.",
|
||||
}
|
||||
|
||||
# Qwen2 packs qkv / gate-up; apply the same shard mapping for the LM part.
|
||||
stacked_params_mapping = [
|
||||
(".qkv_proj", ".q_proj", "q"),
|
||||
(".qkv_proj", ".k_proj", "k"),
|
||||
(".qkv_proj", ".v_proj", "v"),
|
||||
(".gate_up_proj", ".gate_proj", 0),
|
||||
(".gate_up_proj", ".up_proj", 1),
|
||||
]
|
||||
|
||||
tie_word_embeddings = getattr(
|
||||
self.config.text_config, "tie_word_embeddings", False
|
||||
)
|
||||
|
||||
params_dict = dict(self.named_parameters())
|
||||
loaded_params: Set[str] = set()
|
||||
|
||||
for name, loaded_weight in weights:
|
||||
for src, dst in prefix_mapping.items():
|
||||
if name.startswith(src):
|
||||
name = dst + name[len(src) :]
|
||||
break
|
||||
|
||||
if "rotary_emb.inv_freq" in name:
|
||||
continue
|
||||
if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
|
||||
continue
|
||||
# Under tied embeddings the checkpoint's lm_head duplicates the input
|
||||
# embedding and has no separate destination.
|
||||
if tie_word_embeddings and name.startswith("language_model.lm_head."):
|
||||
continue
|
||||
|
||||
is_vision_weight = name.startswith("vision_tower.") or name.startswith(
|
||||
"multi_modal_projector."
|
||||
)
|
||||
|
||||
if is_vision_weight:
|
||||
if name.endswith(".bias") and name not in params_dict:
|
||||
continue
|
||||
if name not in params_dict:
|
||||
logger.warning(f"Parameter {name} not found in params_dict")
|
||||
continue
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(param, loaded_weight)
|
||||
loaded_params.add(name)
|
||||
continue
|
||||
|
||||
# Language-model weights: apply Qwen2 stacked shard mapping.
|
||||
for param_name, weight_name, shard_id in stacked_params_mapping:
|
||||
if weight_name not in name:
|
||||
continue
|
||||
mapped = name.replace(weight_name, param_name)
|
||||
if mapped.endswith(".bias") and mapped not in params_dict:
|
||||
continue
|
||||
if mapped not in params_dict:
|
||||
continue
|
||||
param = params_dict[mapped]
|
||||
param.weight_loader(param, loaded_weight, shard_id)
|
||||
loaded_params.add(mapped)
|
||||
break
|
||||
else:
|
||||
if name.endswith(".bias") and name not in params_dict:
|
||||
continue
|
||||
if name not in params_dict:
|
||||
logger.warning(f"Parameter {name} not found in params_dict")
|
||||
continue
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(param, loaded_weight)
|
||||
loaded_params.add(name)
|
||||
|
||||
# Reconcile: warn about any model parameter that never received a weight,
|
||||
# so a partial/mismatched checkpoint is visible in the logs rather than
|
||||
# silently serving garbage. Tied lm_head shares embed_tokens' storage and
|
||||
# is loaded via it, so it is expected to be absent here.
|
||||
missing = set(params_dict.keys()) - loaded_params
|
||||
if tie_word_embeddings:
|
||||
missing = {
|
||||
n for n in missing if not n.startswith("language_model.lm_head.")
|
||||
}
|
||||
if missing:
|
||||
logger.warning(
|
||||
f"LocateAnything: {len(missing)} parameters did not receive "
|
||||
f"weights, e.g. {sorted(missing)[:10]}"
|
||||
)
|
||||
|
||||
return loaded_params
|
||||
|
||||
|
||||
class LocateAnythingBoxGrammarLogitProcessor(CustomLogitProcessor):
|
||||
"""Constrained decoding for LocateAnything ``<box>...</box>`` blocks.
|
||||
|
||||
Outside an open box the logits are untouched. Inside an open box (a
|
||||
``box_start`` with no matching ``box_end`` yet) the next token is restricted
|
||||
so that the box body is one of:
|
||||
|
||||
* ``none`` -> ``[none]``
|
||||
* a 2-coordinate point -> ``[c, c]``
|
||||
* a 4-coordinate bounding box -> ``[c, c, c, c]``
|
||||
|
||||
where ``c`` is any token in ``[coord_start_token_id, coord_end_token_id]``.
|
||||
|
||||
Token ids are read per-request from ``custom_param_list[i]`` (keys
|
||||
``box_start_token_id``, ``box_end_token_id``, ``coord_start_token_id``,
|
||||
``coord_end_token_id``, ``none_token_id``) so the processor stays generic.
|
||||
The ``__req__`` entry supplies the generated-so-far token ids.
|
||||
|
||||
This processor is **opt-in**: it is never attached server-side, and the
|
||||
server must be started with ``--enable-custom-logit-processor`` (off by
|
||||
default) or the tokenizer rejects the request. A client enables it by
|
||||
passing both the serialized processor and the matching token ids.
|
||||
:meth:`build_sampling_params` wires both from a
|
||||
:class:`LocateAnythingConfig` so callers don't hand-build the id dict.
|
||||
|
||||
The two pieces live in **different** request fields, so do NOT spread them
|
||||
both into ``sampling_params``: ``custom_logit_processor`` is a top-level
|
||||
:class:`~sglang.srt.managers.io_struct.GenerateReqInput` field, while
|
||||
``custom_params`` is a :class:`SamplingParams` field. (Spreading both into
|
||||
``sampling_params`` raises ``TypeError: Unexpected keyword argument
|
||||
'custom_logit_processor'`` because ``SamplingParams`` is a strict
|
||||
``msgspec.Struct``.) Wire them like the OpenAI ``to_sampling_params`` path::
|
||||
|
||||
from sglang.srt.managers.io_struct import GenerateReqInput
|
||||
from sglang.srt.models.locate_anything import (
|
||||
LocateAnythingBoxGrammarLogitProcessor,
|
||||
)
|
||||
|
||||
extra = LocateAnythingBoxGrammarLogitProcessor.build_sampling_params(config)
|
||||
req = GenerateReqInput(
|
||||
text=prompt,
|
||||
image_data=image,
|
||||
sampling_params={
|
||||
"max_new_tokens": 8192,
|
||||
"custom_params": extra["custom_params"],
|
||||
},
|
||||
custom_logit_processor=extra["custom_logit_processor"],
|
||||
)
|
||||
|
||||
Passing the processor without ``custom_params`` (or vice versa) silently
|
||||
no-ops — both must be present together.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def build_sampling_params(cls, config: "LocateAnythingConfig") -> Dict[str, Any]:
|
||||
"""Build the two request fields needed to enable constrained decoding.
|
||||
|
||||
Returns a dict with ``custom_logit_processor`` (the serialized
|
||||
processor) and ``custom_params`` (the box/coord/none token ids read from
|
||||
``config``). These go to **different** request fields — put
|
||||
``custom_params`` inside ``sampling_params`` and pass
|
||||
``custom_logit_processor`` as a top-level ``GenerateReqInput`` field
|
||||
(see the class docstring). The server also needs
|
||||
``--enable-custom-logit-processor``.
|
||||
"""
|
||||
return {
|
||||
"custom_logit_processor": cls.to_str(),
|
||||
"custom_params": {
|
||||
"box_start_token_id": config.box_start_token_id,
|
||||
"box_end_token_id": config.box_end_token_id,
|
||||
"coord_start_token_id": config.coord_start_token_id,
|
||||
"coord_end_token_id": config.coord_end_token_id,
|
||||
"none_token_id": config.none_token_id,
|
||||
},
|
||||
}
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
custom_param_list: Optional[List[Dict[str, Any]]] = None,
|
||||
) -> torch.Tensor:
|
||||
if not custom_param_list:
|
||||
return logits
|
||||
|
||||
neg_inf = float("-inf")
|
||||
for batch_idx, params in enumerate(custom_param_list):
|
||||
if not params:
|
||||
continue
|
||||
req = params.get("__req__")
|
||||
if req is None:
|
||||
continue
|
||||
|
||||
box_start = params.get("box_start_token_id")
|
||||
box_end = params.get("box_end_token_id")
|
||||
coord_start = params.get("coord_start_token_id")
|
||||
coord_end = params.get("coord_end_token_id")
|
||||
none_id = params.get("none_token_id")
|
||||
if None in (box_start, box_end, coord_start, coord_end, none_id):
|
||||
continue
|
||||
|
||||
# Only the generated tokens are scanned (not origin_input_ids),
|
||||
# which avoids an O(prompt_len) reverse scan per decode step over the
|
||||
# long <IMG_CONTEXT> run. Assumes the prompt contains no *unclosed*
|
||||
# <box>: a closed <box>...</box> in a few-shot / multi-turn prompt is
|
||||
# harmless (last_open finds no open box here), but an unclosed <box>
|
||||
# left dangling in the prompt would not be constrained.
|
||||
output_ids = list(req.output_ids)
|
||||
|
||||
# Find the last box_start; if a box_end follows it, no box is open.
|
||||
try:
|
||||
last_open = len(output_ids) - 1 - output_ids[::-1].index(box_start)
|
||||
except ValueError:
|
||||
continue # no box opened yet
|
||||
body = output_ids[last_open + 1 :]
|
||||
if box_end in body:
|
||||
continue # last box already closed
|
||||
|
||||
num_coords = sum(1 for t in body if coord_start <= t <= coord_end)
|
||||
has_none = none_id in body
|
||||
|
||||
# Determine which token classes are allowed next. The coordinate
|
||||
# range is contiguous, so it is masked as a slice rather than an
|
||||
# enumerated set (the range can span ~1000 ids per decode step).
|
||||
allow_coords = False
|
||||
allow_scalars: Set[int] = set()
|
||||
if has_none:
|
||||
allow_scalars = {box_end}
|
||||
elif num_coords == 0:
|
||||
allow_coords, allow_scalars = True, {none_id}
|
||||
elif num_coords in (1, 3):
|
||||
allow_coords = True
|
||||
elif num_coords == 2:
|
||||
allow_coords, allow_scalars = True, {box_end}
|
||||
else: # >= 4 coords -> must close
|
||||
allow_scalars = {box_end}
|
||||
|
||||
mask = torch.full_like(logits[batch_idx], neg_inf)
|
||||
if allow_coords:
|
||||
mask[coord_start : coord_end + 1] = logits[
|
||||
batch_idx, coord_start : coord_end + 1
|
||||
]
|
||||
for tok in allow_scalars:
|
||||
mask[tok] = logits[batch_idx, tok]
|
||||
logits[batch_idx] = mask
|
||||
|
||||
return logits
|
||||
|
||||
|
||||
EntryClass = [LocateAnythingForConditionalGeneration]
|
||||
@@ -0,0 +1,56 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
import re
|
||||
from typing import Dict, List, Union
|
||||
|
||||
from sglang.srt.managers.schedule_batch import MultimodalProcessorOutput
|
||||
from sglang.srt.models.locate_anything import LocateAnythingForConditionalGeneration
|
||||
from sglang.srt.multimodal.processors.base_processor import (
|
||||
BaseMultimodalProcessor as SGLangBaseProcessor,
|
||||
)
|
||||
from sglang.srt.multimodal.processors.base_processor import (
|
||||
MultimodalSpecialTokens,
|
||||
)
|
||||
|
||||
|
||||
# Compatible with LocateAnythingForConditionalGeneration
|
||||
class LocateAnythingImageProcessor(SGLangBaseProcessor):
|
||||
models = [LocateAnythingForConditionalGeneration]
|
||||
# The LocateAnything HF processor is remote-code and does not support tensor inputs.
|
||||
gpu_image_decode = False
|
||||
|
||||
def __init__(self, hf_config, server_args, _processor, *args, **kwargs):
|
||||
super().__init__(hf_config, server_args, _processor, *args, **kwargs)
|
||||
# The model's chat template emits numbered ``<image-N>`` placeholders.
|
||||
# The HF LocateAnythingProcessor expands each into
|
||||
# ``<img>`` + N×``<IMG_CONTEXT>`` + ``</img>`` and only the
|
||||
# ``<IMG_CONTEXT>`` (id 151665) run carries vision embeddings, so the
|
||||
# offset/embedding token id is image_token_index while the prompt-level
|
||||
# placeholder we split on is ``<image-N>``.
|
||||
self.mm_tokens = MultimodalSpecialTokens(
|
||||
image_token_id=hf_config.image_token_index,
|
||||
image_token_regex=re.compile(r"<image-\d+>"),
|
||||
).build(_processor)
|
||||
|
||||
async def process_mm_data_async(
|
||||
self,
|
||||
image_data: List[Union[str, bytes, Dict]],
|
||||
input_text,
|
||||
request_obj,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
base_output = await self.load_mm_data(
|
||||
prompt=input_text,
|
||||
image_data=image_data,
|
||||
multimodal_tokens=self.mm_tokens,
|
||||
)
|
||||
|
||||
mm_items, input_ids, _ = self.process_and_combine_mm_data(
|
||||
base_output, self.mm_tokens
|
||||
)
|
||||
|
||||
return MultimodalProcessorOutput(
|
||||
input_ids=input_ids.tolist(),
|
||||
mm_items=mm_items,
|
||||
im_token_id=self.mm_tokens.image_token_id,
|
||||
)
|
||||
@@ -39,6 +39,7 @@ from sglang.srt.configs import (
|
||||
KimiLinearConfig,
|
||||
KimiVLConfig,
|
||||
LagunaConfig,
|
||||
LocateAnythingConfig,
|
||||
LongcatFlashConfig,
|
||||
MiniCPMV4_6Config,
|
||||
MiniCPMV4_6VisionConfig,
|
||||
@@ -84,6 +85,7 @@ _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
|
||||
DeepseekVL2Config,
|
||||
MultiModalityConfig,
|
||||
KimiVLConfig,
|
||||
LocateAnythingConfig,
|
||||
InternVLChatConfig,
|
||||
LagunaConfig,
|
||||
Step3VLConfig,
|
||||
|
||||
Reference in New Issue
Block a user