diff --git a/docs/basic_usage/deepseek_ocr.md b/docs/basic_usage/deepseek_ocr.md new file mode 100644 index 000000000..6f62713eb --- /dev/null +++ b/docs/basic_usage/deepseek_ocr.md @@ -0,0 +1,54 @@ +# DeepSeek OCR (OCR-1 / OCR-2) + +DeepSeek OCR models are multimodal (image + text) models for OCR and document understanding. + +## Launch server + +```shell +python -m sglang.launch_server \ + --model-path deepseek-ai/DeepSeek-OCR-2 \ + --trust-remote-code \ + --host 0.0.0.0 \ + --port 30000 +``` + +> You can replace `deepseek-ai/DeepSeek-OCR-2` with `deepseek-ai/DeepSeek-OCR`. + +## Prompt examples + +Recommended prompts from the model card: + +``` + +<|grounding|>Convert the document to markdown. +``` + +``` + +Free OCR. +``` + +## OpenAI-compatible request example + +```python +import requests + +url = "http://localhost:30000/v1/chat/completions" + +data = { + "model": "deepseek-ai/DeepSeek-OCR-2", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "\n<|grounding|>Convert the document to markdown."}, + {"type": "image_url", "image_url": {"url": "https://example.com/your_image.jpg"}}, + ], + } + ], + "max_tokens": 512, +} + +response = requests.post(url, json=data) +print(response.text) +``` diff --git a/docs/index.rst b/docs/index.rst index b823bd2b7..5e11ed15c 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -33,6 +33,7 @@ Its core features include: basic_usage/ollama_api.md basic_usage/offline_engine_api.ipynb basic_usage/native_api.ipynb + basic_usage/deepseek_ocr.md basic_usage/sampling_params.md basic_usage/popular_model_usage.rst diff --git a/docs/supported_models/multimodal_language_models.md b/docs/supported_models/multimodal_language_models.md index 9f645ab95..bc6330640 100644 --- a/docs/supported_models/multimodal_language_models.md +++ b/docs/supported_models/multimodal_language_models.md @@ -30,6 +30,7 @@ in the GitHub search bar. |----------------------------|--------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|-------| | **Qwen-VL** | `Qwen/Qwen3-VL-235B-A22B-Instruct` | Alibaba's vision-language extension of Qwen; for example, Qwen2.5-VL (7B and larger variants) can analyze and converse about image content. | | | **DeepSeek-VL2** | `deepseek-ai/deepseek-vl2` | Vision-language variant of DeepSeek (with a dedicated image processor), enabling advanced multimodal reasoning on image and text inputs. | | +| **DeepSeek-OCR / OCR-2** | `deepseek-ai/DeepSeek-OCR-2` | OCR-focused DeepSeek models for document understanding and text extraction. | Use `--trust-remote-code`. | | **Janus-Pro** (1B, 7B) | `deepseek-ai/Janus-Pro-7B` | DeepSeek's open-source multimodal model capable of both image understanding and generation. Janus-Pro employs a decoupled architecture for separate visual encoding paths, enhancing performance in both tasks. | | | **MiniCPM-V / MiniCPM-o** | `openbmb/MiniCPM-V-2_6` | MiniCPM-V (2.6, ~8B) supports image inputs, and MiniCPM-o adds audio/video; these multimodal LLMs are optimized for end-side deployment on mobile/edge devices. | | | **Llama 3.2 Vision** (11B) | `meta-llama/Llama-3.2-11B-Vision-Instruct` | Vision-enabled variant of Llama 3 (11B) that accepts image inputs for visual question answering and other multimodal tasks. | | diff --git a/python/sglang/srt/configs/deepseek_ocr.py b/python/sglang/srt/configs/deepseek_ocr.py index b1f2488d3..1677423d1 100644 --- a/python/sglang/srt/configs/deepseek_ocr.py +++ b/python/sglang/srt/configs/deepseek_ocr.py @@ -196,6 +196,7 @@ class DeepseekOCRProcessor(ProcessorMixin): sft_format: str = "deepseek", mask_prompt: bool = True, ignore_id: int = -100, + ocr2_mode: bool = False, **kwargs, ): @@ -243,6 +244,7 @@ class DeepseekOCRProcessor(ProcessorMixin): self.sft_format = sft_format self.mask_prompt = mask_prompt self.ignore_id = ignore_id + self.ocr2_mode = ocr2_mode super().__init__( tokenizer, @@ -359,6 +361,13 @@ class DeepseekOCRProcessor(ProcessorMixin): target_ids = torch.LongTensor(masked_tokenized_str) + has_images = len(images_list) > 0 + has_local_crops = False + if len(images_spatial_crop) > 0: + has_local_crops = any( + crop[0] > 1 or crop[1] > 1 for crop in images_spatial_crop + ) + if len(images_list) == 0: images = torch.zeros((1, 3, self.image_size, self.image_size)) else: @@ -376,6 +385,8 @@ class DeepseekOCRProcessor(ProcessorMixin): images_seq_mask=images_seq_mask, images_spatial_crop=images_spatial_crop, ) + prepare.has_images = has_images + prepare.has_local_crops = has_local_crops return prepare @@ -481,15 +492,27 @@ class DeepseekOCRProcessor(ProcessorMixin): (self.base_size // self.patch_size) / self.downsample_ratio ) - tokenized_image = ( - [self.image_token_id] * num_queries_base + [self.image_token_id] - ) * num_queries_base - tokenized_image += [self.image_token_id] - if num_width_tiles > 1 or num_height_tiles > 1: - tokenized_image += ( - [self.image_token_id] * (num_queries * num_width_tiles) - + [self.image_token_id] - ) * (num_queries * num_height_tiles) + if self.ocr2_mode: + tokenized_image = [] + if num_width_tiles > 1 or num_height_tiles > 1: + tokenized_image += [self.image_token_id] * ( + num_queries * num_width_tiles * num_queries * num_height_tiles + ) + tokenized_image += [self.image_token_id] * ( + num_queries_base * num_queries_base + ) + # One extra token for the view separator. + tokenized_image += [self.image_token_id] + else: + tokenized_image = ( + [self.image_token_id] * num_queries_base + [self.image_token_id] + ) * num_queries_base + tokenized_image += [self.image_token_id] + if num_width_tiles > 1 or num_height_tiles > 1: + tokenized_image += ( + [self.image_token_id] * (num_queries * num_width_tiles) + + [self.image_token_id] + ) * (num_queries * num_height_tiles) tokenized_str += tokenized_image images_seq_mask += [True] * len(tokenized_image) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 68d0388ce..edb183d84 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1050,7 +1050,12 @@ def _get_and_verify_dtype( ) -> torch.dtype: # NOTE: getattr(config, "torch_dtype", torch.float32) is not correct # because config.torch_dtype can be None. - config_dtype = getattr(config, "dtype", None) + if isinstance(config, dict): + config_dtype = config.get("dtype", None) or config.get("torch_dtype", None) + model_type = config.get("model_type", "") + else: + config_dtype = getattr(config, "dtype", None) + model_type = getattr(config, "model_type", "") if isinstance(config_dtype, str): config_dtype = _STR_DTYPE_TO_TORCH_DTYPE.get(config_dtype, None) if config_dtype is None: @@ -1060,11 +1065,11 @@ def _get_and_verify_dtype( dtype = dtype.lower() if dtype == "auto": if config_dtype == torch.float32: - if config.model_type.startswith("gemma"): - if config.model_type == "gemma": + if model_type.startswith("gemma"): + if model_type == "gemma": gemma_version = "" else: - gemma_version = config.model_type[5] + gemma_version = model_type[5] logger.info( f"For Gemma {gemma_version}, we downcast float32 to bfloat16 instead " "of float16 by default. Please specify `dtype` if you " diff --git a/python/sglang/srt/model_loader/utils.py b/python/sglang/srt/model_loader/utils.py index f6cabe6db..8f38e20ad 100644 --- a/python/sglang/srt/model_loader/utils.py +++ b/python/sglang/srt/model_loader/utils.py @@ -59,14 +59,18 @@ def resolve_transformers_arch(model_config: ModelConfig, architectures: list[str ) model_module = auto_modules["AutoModel"] if model_config.model_impl == ModelImpl.TRANSFORMERS: - if not model_module.is_backend_compatible(): + if hasattr(model_module, "is_backend_compatible") and ( + not model_module.is_backend_compatible() + ): raise ValueError( f"The Transformers implementation of {arch} is not " "compatible with SGLang." ) architectures[i] = "TransformersForCausalLM" if model_config.model_impl == ModelImpl.AUTO: - if not model_module.is_backend_compatible(): + if hasattr(model_module, "is_backend_compatible") and ( + not model_module.is_backend_compatible() + ): raise ValueError( f"{arch} has no SGlang implementation and the Transformers " "implementation is not compatible with SGLang." diff --git a/python/sglang/srt/models/deepseek_ocr.py b/python/sglang/srt/models/deepseek_ocr.py index fca372a18..b3506606c 100644 --- a/python/sglang/srt/models/deepseek_ocr.py +++ b/python/sglang/srt/models/deepseek_ocr.py @@ -24,6 +24,7 @@ from typing import Iterable, List, Optional, Set, Tuple, Type, TypeAlias, Union import torch import torch.nn.functional as F +import transformers from torch import Tensor, nn from transformers.models.vitdet.modeling_vitdet import get_rel_pos @@ -702,6 +703,7 @@ class ImageEncoderViT(nn.Module): rel_pos_zero_init: bool = True, window_size: int = 0, global_attn_indexes: Tuple[int, ...] = (), + net_3_out_channels: int = 1024, ) -> None: """ Args: @@ -776,7 +778,7 @@ class ImageEncoderViT(nn.Module): self.net_2 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1, bias=False) self.net_3 = nn.Conv2d( - 512, 1024, kernel_size=3, stride=2, padding=1, bias=False + 512, net_3_out_channels, kernel_size=3, stride=2, padding=1, bias=False ) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -800,6 +802,7 @@ def _build_sam( encoder_num_heads, encoder_global_attn_indexes, checkpoint=None, + net_3_out_channels: int = 1024, ): prompt_embed_dim = 256 image_size = 1024 @@ -817,6 +820,7 @@ def _build_sam( global_attn_indexes=encoder_global_attn_indexes, window_size=14, out_chans=prompt_embed_dim, + net_3_out_channels=net_3_out_channels, ) image_encoder.eval() if checkpoint is not None: @@ -828,13 +832,14 @@ def _build_sam( return image_encoder -def build_sam_vit_b(checkpoint=None): +def build_sam_vit_b(checkpoint=None, net_3_out_channels: int = 1024): return _build_sam( encoder_embed_dim=768, encoder_depth=12, encoder_num_heads=12, encoder_global_attn_indexes=[2, 5, 8, 11], checkpoint=checkpoint, + net_3_out_channels=net_3_out_channels, ) @@ -1146,6 +1151,257 @@ def build_clip_l(): ) +class CustomQwen2Decoder(nn.Module): + """Qwen2 decoder with mixed causal masking for OCR2 vision encoder.""" + + def __init__( + self, + decoder_layer: int = 24, + max_position_embeddings: int = 131072, + hidden_dimension: int = 896, + num_attention_heads: int = 14, + num_key_value_heads: int = 2, + intermediate_size: int = 4864, + vocab_size: int = 151936, + attn_implementation: str = "sdpa", + rms_norm_eps: float = 1e-6, + rope_theta: float = 1000000.0, + attention_dropout: float = 0.0, + hidden_act: str = "silu", + initializer_range: float = 0.02, + ): + super().__init__() + if attn_implementation == "flash_attention_2": + raise ValueError( + "CustomQwen2Decoder does not support flash_attention_2; " + "use sdpa or eager." + ) + + Qwen2Model = getattr(transformers.models.qwen2.modeling_qwen2, "Qwen2Model") + Qwen2Config = getattr(transformers, "Qwen2Config") + + config = Qwen2Config( + hidden_size=hidden_dimension, + num_hidden_layers=decoder_layer, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + intermediate_size=intermediate_size, + max_position_embeddings=max_position_embeddings, + vocab_size=vocab_size, + rms_norm_eps=rms_norm_eps, + rope_theta=rope_theta, + attention_dropout=attention_dropout, + hidden_act=hidden_act, + initializer_range=initializer_range, + _attn_implementation=attn_implementation, + ) + + self.model = self._create_custom_model(Qwen2Model, config) + del self.model.embed_tokens + + def _create_custom_model(self, Qwen2Model, config): + class CustomQwen2ModelInner(Qwen2Model): + def forward( + self, + input_ids=None, + attention_mask=None, + position_ids=None, + past_key_values=None, + inputs_embeds=None, + token_type_ids=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + cache_position=None, + ): + self._current_token_type_ids = token_type_ids + return super().forward( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + def _update_causal_mask( + self, + attention_mask, + input_tensor, + cache_position, + past_key_values, + output_attentions, + ): + dtype, device = input_tensor.dtype, input_tensor.device + min_dtype = torch.finfo(dtype).min + batch_size, sequence_length = ( + input_tensor.shape[0], + input_tensor.shape[1], + ) + + token_type_ids = getattr(self, "_current_token_type_ids", None) + if token_type_ids is None: + return super()._update_causal_mask( + attention_mask, + input_tensor, + cache_position, + past_key_values, + output_attentions, + ) + + causal_mask = self._create_custom_4d_mask( + sequence_length=sequence_length, + dtype=dtype, + device=device, + batch_size=batch_size, + token_type_ids=token_type_ids, + ) + + if attention_mask is not None and attention_mask.dim() == 2: + padding_mask = attention_mask[:, None, None, :].to(dtype=dtype) + padding_mask = (1.0 - padding_mask) * min_dtype + causal_mask = causal_mask + padding_mask + + return causal_mask + + def _create_custom_4d_mask( + self, + sequence_length, + dtype, + device, + batch_size, + token_type_ids, + ): + min_dtype = torch.finfo(dtype).min + masks = [] + for b in range(batch_size): + mask = torch.full( + (sequence_length, sequence_length), + fill_value=min_dtype, + dtype=dtype, + device=device, + ) + + type_ids = token_type_ids[b] + image_positions = (type_ids == 0).nonzero(as_tuple=True)[0] + text_positions = (type_ids == 1).nonzero(as_tuple=True)[0] + + if len(image_positions) > 0: + mask[image_positions[:, None], image_positions] = 0.0 + + for i, text_pos in enumerate(text_positions): + if len(image_positions) > 0: + mask[text_pos, image_positions] = 0.0 + mask[text_pos, text_positions[: i + 1]] = 0.0 + + masks.append(mask) + + mask = torch.stack(masks, dim=0).unsqueeze(1) + return mask + + return CustomQwen2ModelInner(config) + + def forward(self, inputs_embeds, token_type_ids, attention_mask=None, **kwargs): + return self.model( + inputs_embeds=inputs_embeds, + token_type_ids=token_type_ids, + attention_mask=attention_mask, + **kwargs, + ) + + +class Qwen2Decoder2Encoder(nn.Module): + """Decoder-as-encoder for OCR2 vision tokens.""" + + def __init__( + self, + decoder_layer: int, + hidden_dimension: int, + num_attention_heads: int, + num_key_value_heads: int, + intermediate_size: int, + max_query: int, + ): + super().__init__() + self.model = CustomQwen2Decoder( + decoder_layer=decoder_layer, + hidden_dimension=hidden_dimension, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + intermediate_size=intermediate_size, + attn_implementation="sdpa", + ) + + self.query_768 = nn.Embedding(144, hidden_dimension) + self.query_1024 = nn.Embedding(256, hidden_dimension) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x.flatten(2).transpose(1, 2) + bs, n_query, _ = x.shape + + if n_query == 144: + param_img = self.query_768.weight + elif n_query == 256: + param_img = self.query_1024.weight + else: + base = ( + self.query_1024.weight + if n_query > self.query_768.num_embeddings + else self.query_768.weight + ) + param_img = ( + F.interpolate( + base.T.unsqueeze(0), + size=n_query, + mode="linear", + align_corners=False, + ) + .squeeze(0) + .T + ) + + batch_query_imgs = param_img.unsqueeze(0).expand(bs, -1, -1) + x_combined = torch.cat([x, batch_query_imgs], dim=1) + token_type_ids = torch.cat( + [ + torch.zeros(bs, n_query, dtype=torch.long, device=x.device), + torch.ones(bs, n_query, dtype=torch.long, device=x.device), + ], + dim=1, + ) + y = self.model(x_combined, token_type_ids)[0] + y = y[:, n_query:, :] + return y + + +def build_qwen2_decoder_as_encoder( + decoder_layer: int = 24, + hidden_dimension: int = 896, + num_attention_heads: int = 14, + num_key_value_heads: int = 2, + intermediate_size: int = 4864, + max_query: int = 400, + checkpoint=None, +): + decoder_as_encoder = Qwen2Decoder2Encoder( + decoder_layer=decoder_layer, + hidden_dimension=hidden_dimension, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + intermediate_size=intermediate_size, + max_query=max_query, + ) + if checkpoint is not None: + state_dict = torch.load(checkpoint) + decoder_as_encoder.load_state_dict(state_dict, strict=True) + return decoder_as_encoder + + class DeepseekOCRForCausalLM(nn.Module): def __init__( self, @@ -1161,8 +1417,12 @@ class DeepseekOCRForCausalLM(nn.Module): self.vision_config = config.vision_config self.projector_config = config.projector_config self.text_config = config.text_config - - n_embed = 1280 + self.is_ocr2 = ( + str(getattr(self.vision_config, "model_name", "")).lower() + == "deepencoderv2" + or getattr(self.projector_config, "input_dim", None) == 896 + ) + n_embed = getattr(self.projector_config, "n_embed", 1280) self.tile_tag = config.tile_tag self.global_view_pos = config.global_view_pos @@ -1171,48 +1431,136 @@ class DeepseekOCRForCausalLM(nn.Module): embed_std = 1 / torch.sqrt(torch.tensor(n_embed, dtype=torch.float32)) if self.tile_tag == "2D": # <|view_separator|>, <|\n|> - self.image_newline = nn.Parameter(torch.randn(n_embed) * embed_std) self.view_seperator = nn.Parameter(torch.randn(n_embed) * embed_std) + if not self.is_ocr2: + self.image_newline = nn.Parameter(torch.randn(n_embed) * embed_std) else: raise ValueError( f"Only 2D tile_tag is supported currently, got: {self.tile_tag}" ) - if self.text_config.topk_method == "noaux_tc": - self.model = DeepseekV3ForCausalLM( - config=config.text_config, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "language"), - ) - elif not self.text_config.use_mla: + if not self.is_ocr2: + if self.text_config.topk_method == "noaux_tc": + self.model = DeepseekV3ForCausalLM( + config=config.text_config, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "language"), + ) + elif not self.text_config.use_mla: + self.model = DeepseekForCausalLM( + config=config.text_config, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "language"), + ) + else: + self.model = DeepseekV2ForCausalLM( + config=config.text_config, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "language"), + ) + else: + # OCR2 language_config uses non-MLA attention (qk_* dims are 0). + # Use the non-MLA Deepseek model to avoid MLA-specific assumptions. self.model = DeepseekForCausalLM( config=config.text_config, quant_config=quant_config, prefix=maybe_prefix(prefix, "language"), ) + + if not self.is_ocr2: + self.sam_model = build_sam_vit_b() + self.vision_model = build_clip_l() else: - self.model = DeepseekV2ForCausalLM( - config=config.text_config, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "language"), + projector_input_dim = getattr(self.projector_config, "input_dim", 896) + self.sam_model = build_sam_vit_b(net_3_out_channels=projector_input_dim) + self.qwen2_model = build_qwen2_decoder_as_encoder( + hidden_dimension=projector_input_dim ) - self.sam_model = build_sam_vit_b() - self.vision_model = build_clip_l() - n_embed = 1280 self.projector = MlpProjector( - projector_type="linear", - input_dim=2048, + projector_type=self.projector_config.projector_type, + input_dim=self.projector_config.input_dim, n_embed=n_embed, + depth=self.projector_config.depth, + mlp_ratio=self.projector_config.mlp_ratio, + downsample_ratio=self.projector_config.downsample_ratio, ) + @staticmethod + def _collect_mm_flag( + items: List[MultimodalDataItem], flag_name: str + ) -> Optional[List[bool]]: + values = [] + for item in items: + value = getattr(item, flag_name, None) + if value is None: + return None + values.append(bool(value)) + return values + + def _encode_ocr2_features(self, images: torch.Tensor) -> torch.Tensor: + features = self.sam_model(images) + features = self.qwen2_model(features) + features = self.projector(features) + return features.view(-1, features.shape[-1]) + + def _encode_ocr1_features(self, images: torch.Tensor) -> torch.Tensor: + features_1 = self.sam_model(images) + features_2 = self.vision_model(images, features_1) + features = torch.cat( + ( + features_2[:, 1:], + features_1.flatten(2).permute(0, 2, 1), + ), + dim=-1, + ) + return self.projector(features) + + def _format_ocr1_global_features(self, features: torch.Tensor) -> torch.Tensor: + _, hw, n_dim = features.shape + h = w = int(hw**0.5) + features = features.view(h, w, n_dim) + features = torch.cat( + [features, self.image_newline[None, None, :].expand(h, 1, n_dim)], + dim=1, + ) + return features.view(-1, n_dim) + + def _format_ocr1_local_features( + self, features: torch.Tensor, crop_shape: torch.Tensor + ) -> torch.Tensor: + _, hw2, n_dim2 = features.shape + h2 = w2 = int(hw2**0.5) + width_crop_num, height_crop_num = int(crop_shape[0]), int(crop_shape[1]) + features = ( + features.view(height_crop_num, width_crop_num, h2, w2, n_dim2) + .permute(0, 2, 1, 3, 4) + .reshape(height_crop_num * h2, width_crop_num * w2, n_dim2) + ) + features = torch.cat( + [ + features, + self.image_newline[None, None, :].expand( + height_crop_num * h2, 1, n_dim2 + ), + ], + dim=1, + ) + return features.view(-1, n_dim2) + def _parse_and_validate_image_input(self, **kwargs: object): pixel_values = kwargs.pop("pixel_values", None) images_spatial_crop = kwargs.pop("images_spatial_crop", None) images_crop = kwargs.pop("images_crop", None) + has_images = kwargs.pop("has_images", None) - if pixel_values is None or torch.sum(pixel_values).item() == 0: + if pixel_values is None: + return None + if has_images is not None: + if not has_images: + return None + elif torch.sum(pixel_values).item() == 0: return None if pixel_values is not None: @@ -1241,6 +1589,7 @@ class DeepseekOCRForCausalLM(nn.Module): pixel_values: torch.Tensor, images_crop: torch.Tensor, images_spatial_crop: torch.Tensor, + has_local_crops: Optional[List[bool]] = None, ) -> NestedTensors: # Pixel_values (global view): [n_image, batch_size, 3, height, width] @@ -1250,108 +1599,61 @@ class DeepseekOCRForCausalLM(nn.Module): images_in_this_batch = [] + if not self.is_ocr2: + with torch.no_grad(): + for jdx in range(images_spatial_crop.size(0)): + patches = images_crop[jdx][0].to(torch.bfloat16) + image_ori = pixel_values[jdx] + crop_shape = images_spatial_crop[jdx][0] + use_local_crops = ( + has_local_crops[jdx] + if has_local_crops is not None + else torch.sum(patches).item() != 0 + ) + + global_features = self._encode_ocr1_features(image_ori) + global_features = self._format_ocr1_global_features(global_features) + + if use_local_crops: + local_features = self._encode_ocr1_features(patches) + local_features = self._format_ocr1_local_features( + local_features, crop_shape + ) + global_local_features = torch.cat( + [ + local_features, + global_features, + self.view_seperator[None, :], + ], + dim=0, + ) + else: + global_local_features = torch.cat( + [global_features, self.view_seperator[None, :]], dim=0 + ) + + images_in_this_batch.append(global_local_features) + + return images_in_this_batch + with torch.no_grad(): for jdx in range(images_spatial_crop.size(0)): patches = images_crop[jdx][0].to(torch.bfloat16) image_ori = pixel_values[jdx] - crop_shape = images_spatial_crop[jdx][0] - - if torch.sum(patches).item() != 0: - local_features_1 = self.sam_model(patches) - local_features_2 = self.vision_model(patches, local_features_1) - - local_features = torch.cat( - ( - local_features_2[:, 1:], - local_features_1.flatten(2).permute(0, 2, 1), - ), - dim=-1, - ) - local_features = self.projector(local_features) - - global_features_1 = self.sam_model(image_ori) - global_features_2 = self.vision_model(image_ori, global_features_1) - global_features = torch.cat( - ( - global_features_2[:, 1:], - global_features_1.flatten(2).permute(0, 2, 1), - ), - dim=-1, - ) - global_features = self.projector(global_features) - - _, hw, n_dim = global_features.shape - h = w = int(hw**0.5) - - _2, hw2, n_dim2 = local_features.shape - h2 = w2 = int(hw2**0.5) - - width_crop_num, height_crop_num = int(crop_shape[0]), int( - crop_shape[1] - ) - - global_features = global_features.view(h, w, n_dim) - - global_features = torch.cat( - [ - global_features, - self.image_newline[None, None, :].expand(h, 1, n_dim), - ], - dim=1, - ) - - global_features = global_features.view(-1, n_dim) - - local_features = ( - local_features.view( - height_crop_num, width_crop_num, h2, w2, n_dim2 - ) - .permute(0, 2, 1, 3, 4) - .reshape(height_crop_num * h2, width_crop_num * w2, n_dim2) - ) - local_features = torch.cat( - [ - local_features, - self.image_newline[None, None, :].expand( - height_crop_num * h2, 1, n_dim2 - ), - ], - dim=1, - ) - local_features = local_features.view(-1, n_dim2) + use_local_crops = ( + has_local_crops[jdx] + if has_local_crops is not None + else torch.sum(patches).item() != 0 + ) + global_features = self._encode_ocr2_features(image_ori) + if use_local_crops: + local_features = self._encode_ocr2_features(patches) global_local_features = torch.cat( [local_features, global_features, self.view_seperator[None, :]], dim=0, ) - else: - global_features_1 = self.sam_model(image_ori) - global_features_2 = self.vision_model(image_ori, global_features_1) - global_features = torch.cat( - ( - global_features_2[:, 1:], - global_features_1.flatten(2).permute(0, 2, 1), - ), - dim=-1, - ) - global_features = self.projector(global_features) - - _, hw, n_dim = global_features.shape - h = w = int(hw**0.5) - - global_features = global_features.view(h, w, n_dim) - - global_features = torch.cat( - [ - global_features, - self.image_newline[None, None, :].expand(h, 1, n_dim), - ], - dim=1, - ) - - global_features = global_features.view(-1, n_dim) - global_local_features = torch.cat( [global_features, self.view_seperator[None, :]], dim=0 ) @@ -1361,8 +1663,14 @@ class DeepseekOCRForCausalLM(nn.Module): return images_in_this_batch def _process_image_input(self, mm_items: List[MultimodalDataItem]) -> torch.Tensor: + target_dtype = ( + next(self.sam_model.parameters()).dtype + if self.is_ocr2 + else self.vision_model.dtype + ) + has_local_crops = self._collect_mm_flag(mm_items, "has_local_crops") pixel_values = torch.stack([item.feature for item in mm_items], dim=0).type( - self.vision_model.dtype + target_dtype ) images_crop = ( @@ -1383,10 +1691,9 @@ class DeepseekOCRForCausalLM(nn.Module): pixel_values=pixel_values, images_crop=images_crop, images_spatial_crop=images_spatial_crop, + has_local_crops=has_local_crops, ) - vision_features = torch.cat(vision_feature_lists, dim=0).type( - self.vision_model.dtype - ) + vision_features = torch.cat(vision_feature_lists, dim=0).type(target_dtype) return vision_features @@ -1458,6 +1765,7 @@ class DeepseekOCRForCausalLM(nn.Module): for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue + is_qwen2_weight = "qwen2_model." in name if name == "lm_head.weight": name = "model.lm_head.weight" elif name.startswith("model."): @@ -1465,6 +1773,7 @@ class DeepseekOCRForCausalLM(nn.Module): "image_newline" in name or ".projector" in name or "vision_model" in name + or "qwen2_model" in name or "sam_model" in name or "view_seperator" in name ): @@ -1472,11 +1781,32 @@ class DeepseekOCRForCausalLM(nn.Module): elif not ( ".projector" in name or "vision_model" in name + or "qwen2_model" in name or "sam_model" in name or "image_newline" in name ): name = name.replace("model.", "model.model.") + if is_qwen2_weight: + target_name = name + if target_name not in params_dict: + if ".model.model." in target_name: + alt_name = target_name.replace(".model.model.", ".model.") + else: + alt_name = target_name.replace(".model.", ".model.model.", 1) + if alt_name in params_dict: + target_name = alt_name + if target_name.endswith(".bias") and target_name not in params_dict: + continue + if target_name in params_dict: + param = params_dict[target_name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(target_name) + continue + for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue diff --git a/python/sglang/srt/multimodal/processors/deepseek_ocr.py b/python/sglang/srt/multimodal/processors/deepseek_ocr.py index 8f0d583be..9b9002d8d 100644 --- a/python/sglang/srt/multimodal/processors/deepseek_ocr.py +++ b/python/sglang/srt/multimodal/processors/deepseek_ocr.py @@ -12,6 +12,14 @@ class DeepseekOCRProcessor(BaseMultimodalProcessor): def __init__(self, hf_config, server_args, _processor, *args, **kwargs): _processor.image_size = 640 + _processor.ocr2_mode = ( + str( + getattr(getattr(hf_config, "vision_config", None), "model_name", "") + ).lower() + == "deepencoderv2" + or getattr(getattr(hf_config, "projector_config", None), "input_dim", None) + == 896 + ) super().__init__(hf_config, server_args, _processor, *args, **kwargs) self.mm_tokens = MultimodalSpecialTokens( image_token="", image_token_id=self._processor.image_token_id diff --git a/python/sglang/srt/utils/hf_transformers_utils.py b/python/sglang/srt/utils/hf_transformers_utils.py index c78d75c6a..cc1534c56 100644 --- a/python/sglang/srt/utils/hf_transformers_utils.py +++ b/python/sglang/srt/utils/hf_transformers_utils.py @@ -230,11 +230,13 @@ def _load_mistral_large_3_for_causal_LM( def _is_deepseek_ocr_model(config: PretrainedConfig) -> bool: # TODO: Remove this workaround related when AutoConfig correctly identifies deepseek-ocr. # Hugging Face's AutoConfig currently misidentifies it as deepseekvl2. - return ( - getattr(config, "auto_map", None) is not None - and config.auto_map.get("AutoModel") - == "modeling_deepseekocr.DeepseekOCRForCausalLM" - ) + auto_map = getattr(config, "auto_map", None) or {} + return auto_map.get("AutoModel") == "modeling_deepseekocr.DeepseekOCRForCausalLM" + + +def _is_deepseek_ocr2_model(config: PretrainedConfig) -> bool: + auto_map = getattr(config, "auto_map", None) or {} + return auto_map.get("AutoModel") == "modeling_deepseekocr2.DeepseekOCR2ForCausalLM" def _override_deepseek_ocr_v_head_dim(config: DeepseekVLV2Config) -> None: @@ -248,6 +250,30 @@ def _override_deepseek_ocr_v_head_dim(config: DeepseekVLV2Config) -> None: ) +def _override_v_head_dim_if_zero(config: PretrainedConfig, patch: int = 128) -> None: + text_config = getattr(config, "text_config", None) + language_config = getattr(config, "language_config", None) + target = text_config or language_config + if target is None: + return + if getattr(target, "v_head_dim", None) == 0: + setattr(target, "v_head_dim", patch) + logger.warning( + f"Overriding v_head_dim from 0 to {patch} to avoid potential issues." + ) + + +def _ensure_llama_flash_attention2_compat() -> None: + """Ensure LlamaFlashAttention2 symbol exists for remote code compatibility.""" + try: + from transformers.models.llama import modeling_llama + except Exception: + return + if not hasattr(modeling_llama, "LlamaFlashAttention2"): + if hasattr(modeling_llama, "LlamaAttention"): + modeling_llama.LlamaFlashAttention2 = modeling_llama.LlamaAttention + + @lru_cache_frozenset(maxsize=32) def get_config( model: str, @@ -274,6 +300,7 @@ def get_config( model, trust_remote_code=trust_remote_code, revision=revision, **kwargs ) else: + _ensure_llama_flash_attention2_compat() try: config = AutoConfig.from_pretrained( model, trust_remote_code=trust_remote_code, revision=revision, **kwargs @@ -308,20 +335,38 @@ def get_config( text_config = get_hf_text_config(config=config) if isinstance(model, str) and text_config is not None: - for key, val in text_config.__dict__.items(): - if not hasattr(config, key) and getattr(text_config, key, None) is not None: + items = ( + text_config.items() + if hasattr(text_config, "items") + else vars(text_config).items() + ) + for key, val in items: + if not hasattr(config, key) and val is not None: setattr(config, key, val) - if config.model_type in _CONFIG_REGISTRY: + if _is_deepseek_ocr2_model(config): + _override_v_head_dim_if_zero(config) + # Temporary hack for load deepseek-ocr2 + config.model_type = "deepseek-ocr" + config.update({"architectures": ["DeepseekOCRForCausalLM"]}) + config = DeepseekVLV2Config.from_pretrained(model, revision=revision) + _override_v_head_dim_if_zero(config) + config.update({"architectures": ["DeepseekOCRForCausalLM"]}) + setattr(config, "_name_or_path", model) + elif config.model_type in _CONFIG_REGISTRY: model_type = config.model_type if model_type == "deepseek_vl_v2": - if _is_deepseek_ocr_model(config): + if _is_deepseek_ocr_model(config) or _is_deepseek_ocr2_model(config): model_type = "deepseek-ocr" config_class = _CONFIG_REGISTRY[model_type] config = config_class.from_pretrained(model, revision=revision) if _is_deepseek_ocr_model(config): _override_deepseek_ocr_v_head_dim(config) + config.update({"architectures": ["DeepseekOCRForCausalLM"]}) + elif _is_deepseek_ocr2_model(config): + _override_v_head_dim_if_zero(config) + config.update({"architectures": ["DeepseekOCRForCausalLM"]}) # NOTE(HandH1998): Qwen2VL requires `_name_or_path` attribute in `config`. setattr(config, "_name_or_path", model) @@ -536,6 +581,7 @@ def get_processor( **kwargs, ) else: + _ensure_llama_flash_attention2_compat() config = AutoConfig.from_pretrained( tokenizer_name, trust_remote_code=trust_remote_code, @@ -545,6 +591,12 @@ def get_processor( if _is_deepseek_ocr_model(config): # Temporary hack for load deepseek-ocr config.model_type = "deepseek-ocr" + config.update({"architectures": ["DeepseekOCRForCausalLM"]}) + elif _is_deepseek_ocr2_model(config): + # Temporary hack for load deepseek-ocr2 + config.model_type = "deepseek-ocr" + config.update({"architectures": ["DeepseekOCRForCausalLM"]}) + _override_v_head_dim_if_zero(config) # fix: for Qwen2-VL and Sarashina2Vision models, inject default 'size' if not provided. if config.model_type in {"qwen2_vl", "sarashina2_vision"}: