Update app.py

app.py

@@ -1,11 +1,11 @@
-from threading import Thread
-import torch
 from PIL import Image
 import gradio as gr
 import spaces
-from transformers import AutoModel, AutoTokenizer, TextIteratorStreamer
 import os
-import time
+from huggingface_hub import hf_hub_download
+import base64
+from llama_cpp import Llama
+from llama_cpp.llama_chat_format import MoondreamChatHandler
 
 
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
@@ -27,24 +27,34 @@ CSS = """
 }
 """
 
-model = AutoModel.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.float16,
-    trust_remote_code=True
-).to(0)
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
-model.eval()
+chat_handler = MoondreamChatHandler.from_pretrained(
+  repo_id="openbmb/MiniCPM-Llama3-V-2_5-gguf",
+  filename="*mmproj*",
+)
 
+llm = Llama.from_pretrained(
+  repo_id="openbmb/MiniCPM-Llama3-V-2_5-gguf",
+  filename="ggml-model-Q5_K_M.gguf",
+  chat_handler=chat_handler,
+  n_ctx=2048, # n_ctx should be increased to accommodate the image embedding
+)
 
 
 @spaces.GPU(queue=False)
 def stream_chat(message, history: list, temperature: float, max_new_tokens: int):
     print(f'message is - {message}')
     print(f'history is - {history}')
-    conversation = [] 
+    messages = []
+
     if message["files"]:
         image = Image.open(message["files"][-1]).convert('RGB')
-        conversation.append({"role": "user", "content": message['text']})
+        messages.append({
+            "role": "user", 
+            "content": [
+                {"type": "text", "text": message['text']},
+                {"type": "image_url", "image_url":{"url": image}}
+            ]
+        })
     else:
         if len(history) == 0:
             raise gr.Error("Please upload an image first.")
@@ -53,25 +63,39 @@ def stream_chat(message, history: list, temperature: float, max_new_tokens: int)
             image = Image.open(history[0][0][0])
             for prompt, answer in history:
                 if answer is None:
-                    conversation.extend([{"role": "user", "content": prompt},{"role": "assistant", "content": ""}])
+                    messages.extend([{
+                        "role": "user", 
+                        "content": [
+                            {"type": "text", "text": prompt},
+                            {"type": "image_url", "image_url": {"url": image}}
+                        ]
+                    },{
+                        "role": "assistant", 
+                        "content": ""
+                    }])
                 else:
-                    conversation.extend([{"role": "user", "content": prompt}, {"role": "assistant", "content": answer}])
-            conversation.append({"role": "user", "content": message['text']})    
-    print(f"Conversation is -\n{conversation}")
-
-    generate_kwargs = dict(
-        image=image,
-        msgs=conversation,
-        max_new_tokens=max_new_tokens,
+                    messages.extend([{
+                        "role": "user", 
+                        "content": [
+                            {"type": "text", "text": prompt},
+                            {"type": "image_url", "image_url": {"url": image}}
+                        ]
+                    }, {
+                        "role": "assistant", 
+                        "content": answer
+                    }])
+            messages.append({"role": "user", "content": message['text']})    
+    print(f"Messages is -\n{messages}")
+
+
+    response = llm.create_chat_completion(
+        messages = messages,
         temperature=temperature,
-        sampling=True,
-        tokenizer=tokenizer,
+        max_tokens=max_new_tokens,
+        stream=True
     )
-    if temperature == 0:
-        generate_kwargs["sampling"] = False
 
-    response = model.chat(**generate_kwargs)
-    return response
+    return response["choices"][0]["text"]
 
 
 chatbot = gr.Chatbot(height=450)