From 1598e5d3551fd4d02d7e9de17a3fb8fab6c38c9a Mon Sep 17 00:00:00 2001 From: hiyouga Date: Mon, 4 Nov 2024 08:27:20 +0000 Subject: [PATCH] add image input type Former-commit-id: ffa39ba3db0dbfd375cdf20b9f3cbecd359be1a1 --- src/llamafactory/data/mm_plugin.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/llamafactory/data/mm_plugin.py b/src/llamafactory/data/mm_plugin.py index f6748883..6a174838 100644 --- a/src/llamafactory/data/mm_plugin.py +++ b/src/llamafactory/data/mm_plugin.py @@ -30,7 +30,7 @@ if TYPE_CHECKING: path: Optional[str] bytes: Optional[bytes] - ImageInput = Union[str, EncodedImage, ImageObject] + ImageInput = Union[str, bytes, EncodedImage, ImageObject] VideoInput = str @@ -104,6 +104,8 @@ class BasePlugin: for image in images: if isinstance(image, str): image = Image.open(image) + elif isinstance(image, bytes): + image = Image.open(BytesIO(image)) elif isinstance(image, dict): if image["bytes"] is not None: image = Image.open(BytesIO(image["bytes"]))