diff --git a/src/llamafactory/api/chat.py b/src/llamafactory/api/chat.py index 712b6940..50892a54 100644 --- a/src/llamafactory/api/chat.py +++ b/src/llamafactory/api/chat.py @@ -1,3 +1,5 @@ +import base64 +import io import json import os import uuid @@ -83,9 +85,12 @@ def _process_request( input_messages.append({"role": ROLE_MAPPING[message.role], "content": input_item.text}) else: image_url = input_item.image_url.url - if os.path.isfile(image_url): + if image_url.startswith("data:image"): # base64 image + image_data = base64.b64decode(image_url.split(",", maxsplit=1)[1]) + image_path = io.BytesIO(image_data) + elif os.path.isfile(image_url): # local file image_path = open(image_url, "rb") - else: + else: # web uri image_path = requests.get(image_url, stream=True).raw image = Image.open(image_path).convert("RGB")