From f2edacb02d729b8a9febf4a2284586ca84d38137 Mon Sep 17 00:00:00 2001
From: hiyouga <467089858@qq.com>
Date: Thu, 6 Jun 2024 17:29:19 +0800
Subject: [PATCH] fix base64 image read #4061

Former-commit-id: 67aa78cde0d89e9f4b5640bce5fc506ef73ff7b6
---
 src/llamafactory/api/chat.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

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")