p04798526
/
LLaMA-Factory-310P3
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix base64 image read #4061

This commit is contained in:
hiyouga 2024-06-06 17:29:19 +08:00
parent 53eb2de75e
commit 67aa78cde0
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
src/llamafactory/api/chat.py
View File

@ -1,3 +1,5 @@
import base64
import io
import json import json
import os import os
import uuid import uuid
@ -83,9 +85,12 @@ def _process_request(
input_messages.append({"role": ROLE_MAPPING[message.role], "content": input_item.text}) input_messages.append({"role": ROLE_MAPPING[message.role], "content": input_item.text})
else: else:
image_url = input_item.image_url.url 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") image_path = open(image_url, "rb")
else: else: # web uri
image_path = requests.get(image_url, stream=True).raw image_path = requests.get(image_url, stream=True).raw
image = Image.open(image_path).convert("RGB") image = Image.open(image_path).convert("RGB")