support image input in api #3971 #4061

2024-06-06 02:29:55 +08:00
parent dc4a00dd63
commit 946f601136
4 changed files with 49 additions and 8 deletions
--- a/src/llamafactory/api/protocol.py
+++ b/src/llamafactory/api/protocol.py
@@ -56,9 +56,19 @@ class FunctionCall(BaseModel):
    function: Function


+class ImageURL(BaseModel):
+    url: str
+
+
+class MultimodalInputItem(BaseModel):
+    type: Literal["text", "image_url"]
+    text: Optional[str] = None
+    image_url: Optional[ImageURL] = None
+
+
 class ChatMessage(BaseModel):
    role: Role
-    content: Optional[str] = None
+    content: Optional[Union[str, List[MultimodalInputItem]]] = None
    tool_calls: Optional[List[FunctionCall]] = None