瀏覽代碼

refactor: update sglang version checks to use packaging.version for comparison

myhloli 3 月之前
父節點
當前提交
62d1ef184e
共有 2 個文件被更改,包括 6 次插入4 次删除
  1. 3 2
      mineru/model/vlm_sglang_model/image_processor.py
  2. 3 2
      mineru/model/vlm_sglang_model/model.py

+ 3 - 2
mineru/model/vlm_sglang_model/image_processor.py

@@ -6,7 +6,8 @@ from typing import List, Optional, Union
 import numpy as np
 
 from sglang.version import __version__ as sglang_version
-if sglang_version >= "0.4.9":
+from packaging import version
+if version.parse(sglang_version) >= version.parse("0.4.9"):
     # sglang >= 0.4.9
     from sglang.srt.multimodal.processors.base_processor import (
         BaseMultimodalProcessor as BaseProcessor,
@@ -184,7 +185,7 @@ class Mineru2ImageProcessor(BaseProcessor):
             elif request_obj.modalities[0] == "video":
                 modality = Modality.VIDEO
 
-        if sglang_version >= "0.4.9.post3":
+        if version.parse(sglang_version) >= version.parse("0.4.9.post3"):
             # sglang >= 0.4.9.post3
             return {
                 "mm_items": [

+ 3 - 2
mineru/model/vlm_sglang_model/model.py

@@ -7,7 +7,8 @@ import torch
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 
 from sglang.version import __version__ as sglang_version
-if sglang_version >= "0.4.9":
+from packaging import version
+if version.parse(sglang_version) >= version.parse("0.4.9"):
     # sglang >= 0.4.9
     from sglang.srt.multimodal.mm_utils import (
             get_anyres_image_grid_shape,
@@ -234,7 +235,7 @@ class Mineru2QwenForCausalLM(nn.Module):
             if need_vision.any():
                 bs = forward_batch.batch_size
 
-                if sglang_version >= "0.4.9.post3":
+                if version.parse(sglang_version) >= version.parse("0.4.9.post3"):
                     # sglang >= 0.4.9.post3
                     pixel_values = flatten_nested_list(
                         [[item.feature for item in image_inputs[i].mm_items] for i in range(bs) if need_vision[i]]