To Chat Interface

app.py

@@ -10,7 +10,7 @@ from mistral_inference.transformer import Transformer
 from mistral_inference.generate import generate
 
 from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
-from mistral_common.protocol.instruct.messages import UserMessage, TextChunk, ImageURLChunk
+from mistral_common.protocol.instruct.messages import UserMessage, AssistantMessage, TextChunk, ImageURLChunk
 from mistral_common.protocol.instruct.request import ChatCompletionRequest
 
 models_path = Path.home().joinpath('pixtral', 'Pixtral')
@@ -28,10 +28,23 @@ def image_to_base64(image_path):
         encoded_string = base64.b64encode(img.read()).decode('utf-8')
     return f"data:image/jpeg;base64,{encoded_string}"
 
-@spaces.GPU(duration=30)
-def run_inference(image_url, prompt):
-    base64 = image_to_base64(image_url)
-    completion_request = ChatCompletionRequest(messages=[UserMessage(content=[ImageURLChunk(image_url=base64), TextChunk(text=prompt)])])
+@spaces.GPU(duration=60)
+def run_inference(message, history):
+    ## may work
+    messages = []
+    images = []
+    for couple in history:
+        if type(couple[0]) is tuple:
+            images += couple[0]
+        elif couple[0][1]:
+            messages.append(UserMessage(content = [ImageURLChunk(image_url=image_to_base64(path)) for path in images]+[TextChunk(text=couple[0][1])]))
+            messages.append(AssistantMessage(content = couple[1]))
+            images = []
+    ##
+        
+    messages.append(UserMessage(content = [ImageURLChunk(image_url=image_to_base64(file["path"])) for file in message["files"]]+[TextChunk(text=message["text"])]))
+    
+    completion_request = ChatCompletionRequest(messages=messages)
     
     encoded = tokenizer.encode_chat_completion(completion_request)
     
@@ -40,26 +53,7 @@ def run_inference(image_url, prompt):
     
     out_tokens, _ = generate([tokens], model, images=[images], max_tokens=512, temperature=0.45, eos_id=tokenizer.instruct_tokenizer.tokenizer.eos_id)
     result = tokenizer.decode(out_tokens[0])
-    return [[prompt, result]]
-
-with gr.Blocks() as demo:
-    with gr.Row():
-        image_box = gr.Image(type="filepath")
-   
-        chatbot = gr.Chatbot(
-            scale = 2,
-            height=750
-        )
-    text_box = gr.Textbox(
-            placeholder="Enter your text and press enter, or upload an image.",
-            container=False,
-        )
-
-
-    btn = gr.Button("Submit")
-    clicked = btn.click(run_inference,
-                        [image_box,text_box],
-                        chatbot
-                        )
+    return result
 
+demo = gr.ChatInterface(fn=run_inference, title="Pixtral 12B", multimodal=True)
 demo.queue().launch()