Update app.py

app.py

@@ -2,35 +2,33 @@ import time
 from threading import Thread
 
 import gradio as gr
+import spaces
 import torch
 from PIL import Image
-from transformers import AutoProcessor, LlavaForConditionalGeneration
+from transformers import AutoProcessor, AutoModelForCausalLM
 from transformers import TextIteratorStreamer
 
-import spaces
-
-
 PLACEHOLDER = """
 <div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
-   <img src="https://cdn-uploads.huggingface.co/production/uploads/64ccdc322e592905f922a06e/DDIW0kbWmdOQWwy4XMhwX.png" style="width: 80%; max-width: 550px; height: auto; opacity: 0.55;  "> 
-   <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">LLaVA-Llama-3-8B</h1>
-   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Llava-Llama-3-8b is a LLaVA model fine-tuned from Meta-Llama-3-8B-Instruct and CLIP-ViT-Large-patch14-336 with ShareGPT4V-PT and InternVL-SFT by XTuner</p>
+   <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">microsoft/Phi-3-vision-128k-instruct</h1>
 </div>
 """
+user_prompt = '<|user|>\n'
+assistant_prompt = '<|assistant|>\n'
+prompt_suffix = "<|end|>\n"
 
+model_id = "microsoft/Phi-3-vision-128k-instruct"
 
-model_id = "xtuner/llava-llama-3-8b-v1_1-transformers"
-
-processor = AutoProcessor.from_pretrained(model_id)
+processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
 
-model = LlavaForConditionalGeneration.from_pretrained(
+model = AutoModelForCausalLM.from_pretrained(
     model_id,
     torch_dtype=torch.float16,
     low_cpu_mem_usage=True,
+    trust_remote_code=True,
 )
 
 model.to("cuda:0")
-model.generation_config.eos_token_id = 128009
 
 
 @spaces.GPU
@@ -51,15 +49,15 @@ def bot_streaming(message, history):
     try:
         if image is None:
             # Handle the case where image is None
-            gr.Error("You need to upload an image for LLaVA to work.")
+            gr.Error("You need to upload an image for Phi-3-vision to work.")
     except NameError:
         # Handle the case where 'image' is not defined at all
-        gr.Error("You need to upload an image for LLaVA to work.")
+        gr.Error("You need to upload an image for Phi-3-vision to work.")
 
-    prompt = f"<|start_header_id|>user<|end_header_id|>\n\n<image>\n{message['text']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+    prompt = f"{message['text']}<|image_1|>\nCan you convert the table to markdown format?{prompt_suffix}{assistant_prompt}"
     # print(f"prompt: {prompt}")
     image = Image.open(image)
-    inputs = processor(prompt, image, return_tensors='pt').to(0, torch.float16)
+    inputs = processor(prompt, [image], return_tensors='pt').to(0, torch.float16)
 
     streamer = TextIteratorStreamer(processor, **{"skip_special_tokens": False, "skip_prompt": True})
     generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024, do_sample=False)
@@ -67,9 +65,6 @@ def bot_streaming(message, history):
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
 
-    text_prompt = f"<|start_header_id|>user<|end_header_id|>\n\n{message['text']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-    # print(f"text_prompt: {text_prompt}")
-
     buffer = ""
     time.sleep(0.5)
     for new_text in streamer:
@@ -78,7 +73,6 @@ def bot_streaming(message, history):
             new_text = new_text.split("<|eot_id|>")[0]
         buffer += new_text
 
-        # generated_text_without_prompt = buffer[len(text_prompt):]
         generated_text_without_prompt = buffer
         # print(generated_text_without_prompt)
         time.sleep(0.06)
@@ -86,19 +80,20 @@ def bot_streaming(message, history):
         yield generated_text_without_prompt
 
 
-chatbot=gr.Chatbot(placeholder=PLACEHOLDER,scale=1)
-chat_input = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter message or upload file...", show_label=False)
+chatbot = gr.Chatbot(placeholder=PLACEHOLDER, scale=1)
+chat_input = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter message or upload file...",
+                                  show_label=False)
 with gr.Blocks(fill_height=True, ) as demo:
     gr.ChatInterface(
-    fn=bot_streaming,
-    title="LLaVA Llama-3-8B",
-    examples=[{"text": "What is on the flower?", "files": ["./bee.jpg"]},
-              {"text": "How to make this pastry?", "files": ["./baklava.png"]}],
-    description="Try [LLaVA Llama-3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers). Upload an image and start chatting about it, or simply try one of the examples below. If you don't upload an image, you will receive an error.",
-    stop_btn="Stop Generation",
-    multimodal=True,
-    textbox=chat_input,
-    chatbot=chatbot,
+        fn=bot_streaming,
+        title="Phi-3 Vision 128k Instruct",
+        examples=[{"text": "What is on the flower?", "files": ["./bee.jpg"]},
+                  {"text": "How to make this pastry?", "files": ["./baklava.png"]}],
+        description="Try [microsoft/Phi-3-vision-128k-instruct](https://huggingface.co/microsoft/Phi-3-vision-128k-instruct). Upload an image and start chatting about it, or simply try one of the examples below. If you don't upload an image, you will receive an error.",
+        stop_btn="Stop Generation",
+        multimodal=True,
+        textbox=chat_input,
+        chatbot=chatbot,
     )
 
 demo.queue(api_open=False)