From 3da0d026d897e3ba848039951c77e35b0f02bf47 Mon Sep 17 00:00:00 2001 From: chaunceyjiang Date: Fri, 1 Nov 2024 21:18:18 +0800 Subject: [PATCH] [Misc] Multi-Modality Support for Loading Local image Files FIX #8730 Signed-off-by: chaunceyjiang --- vllm/multimodal/utils.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/vllm/multimodal/utils.py b/vllm/multimodal/utils.py index 3c801464383ad..01f409cb1822a 100644 --- a/vllm/multimodal/utils.py +++ b/vllm/multimodal/utils.py @@ -18,13 +18,21 @@ cached_get_tokenizer = lru_cache(get_tokenizer) -def _load_image_from_bytes(b: bytes): +def _load_image_from_bytes(b: bytes) -> Image.Image: image = Image.open(BytesIO(b)) image.load() return image -def _load_image_from_data_url(image_url: str): +def _load_image_from_file(image_url: str) -> Image.Image: + # Only split once and assume the second part is the image path + _, image_path = image_url.split("file://", 1) + image = Image.open(image_path) + image.load() + return image + + +def _load_image_from_data_url(image_url: str) -> Image.Image: # Only split once and assume the second part is the base64 encoded image _, image_base64 = image_url.split(",", 1) return load_image_from_base64(image_base64) @@ -43,9 +51,11 @@ def fetch_image(image_url: str, *, image_mode: str = "RGB") -> Image.Image: elif image_url.startswith('data:image'): image = _load_image_from_data_url(image_url) + elif image_url.startswith('file://'): + image = _load_image_from_file(image_url) else: raise ValueError("Invalid 'image_url': A valid 'image_url' must start " - "with either 'data:image' or 'http'.") + "with either 'data:image', 'file://' or 'http'.") return image.convert(image_mode) @@ -65,9 +75,11 @@ async def async_fetch_image(image_url: str, elif image_url.startswith('data:image'): image = _load_image_from_data_url(image_url) + elif image_url.startswith('file://'): + image = _load_image_from_file(image_url) else: raise ValueError("Invalid 'image_url': A valid 'image_url' must start " - "with either 'data:image' or 'http'.") + "with either 'data:image', 'file://' or 'http'.") return image.convert(image_mode)