update

app.py

@@ -3,6 +3,7 @@ import torch
 from PIL import Image
 from threading import Thread
 from transformers import AutoProcessor, AutoModelForVision2Seq, TextIteratorStreamer
+from transformers.image_utils import load_image
 import spaces
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -26,46 +27,45 @@ def respond(
     temperature,
     top_p,
 ):
+    text = message["text"]
+    files = message["files"]
+    all_images = []
+    current_message_images = []
+
     messages = [{"role": "system", "content": system_message}]
 
     print(message)
     print(history)
 
-    messages.extend(history)
-
-    images = []
-    if message["files"]:
-        pil_image = Image.open(message["files"][0]).convert("RGB")
-        images.append(pil_image)
-
-    current_user_message = {"role": "user", "content": message["text"]}
-    messages.append(current_user_message)
-
+    current_message_images = [load_image(image) for image in files]
+    messages.append({
+        "role": "user",
+        "content": [
+            *[{"type": "image", "image": image} for image in current_message_images],
+            {"type": "text", "text": text},
+        ],
+    })
     prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    inputs = processor(text=prompt, images=images, return_tensors="pt").to(device, torch_dtype)
-
-    streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
+    inputs = processor(
+        text=[prompt],
+        images=current_message_images if current_message_images else None,
+        return_tensors="pt",
+        padding=True,
+    ).to(device, torch_dtype)
 
-    generation_kwargs = dict(
-        **inputs,
-        streamer=streamer,
-        max_new_tokens=max_tokens,
-        do_sample=True,
-        temperature=temperature,
-        top_p=top_p,
-    )
+    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
 
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
 
-    response = ""
+    buffer = ""
     for new_text in streamer:
-        response += new_text
-        yield response
+        buffer += new_text
+        yield buffer
 
 demo = gr.ChatInterface(
     respond,
-    type='messages',
     multimodal=True,
     additional_inputs=[
         gr.Textbox(value="You are a helpful and friendly multimodal assistant. You can analyze images and answer questions about them.", label="System message"),