Skip to content

Commit debafe6

Browse files
fix(mm): vllm model calculates its own size
1 parent ffb5f6c commit debafe6

File tree

2 files changed

+9
-0
lines changed

2 files changed

+9
-0
lines changed

invokeai/backend/llava_onevision_model.py

+7
Original file line numberDiff line numberDiff line change
@@ -47,3 +47,10 @@ def run(self, prompt: str, images: list[Image], device: torch.device, dtype: tor
4747

4848
def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
4949
self._vllm_model.to(device=device, dtype=dtype)
50+
51+
def calc_size(self) -> int:
52+
"""Get size of the model in memory in bytes."""
53+
# HACK(ryand): Fix this issue with circular imports.
54+
from invokeai.backend.model_manager.load.model_util import calc_module_size
55+
56+
return calc_module_size(self._vllm_model)

invokeai/backend/model_manager/load/model_util.py

+2
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
from invokeai.backend.image_util.grounding_dino.grounding_dino_pipeline import GroundingDinoPipeline
1616
from invokeai.backend.image_util.segment_anything.segment_anything_pipeline import SegmentAnythingPipeline
1717
from invokeai.backend.ip_adapter.ip_adapter import IPAdapter
18+
from invokeai.backend.llava_onevision_model import LlavaOnevisionModel
1819
from invokeai.backend.model_manager.taxonomy import AnyModel
1920
from invokeai.backend.onnx.onnx_runtime import IAIOnnxRuntimeModel
2021
from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
@@ -50,6 +51,7 @@ def calc_model_size_by_data(logger: logging.Logger, model: AnyModel) -> int:
5051
SegmentAnythingPipeline,
5152
DepthAnythingPipeline,
5253
SigLipPipeline,
54+
LlavaOnevisionModel,
5355
),
5456
):
5557
return model.calc_size()

0 commit comments

Comments
 (0)