Update app.py

app.py

@@ -2,6 +2,19 @@ import gradio as gr
 from gradio_client import Client
 import os
 import json
+from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration
+import torch
+from PIL import Image
+import requests
+import spaces
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+processor = LlavaNextProcessor.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf")
+
+model = LlavaNextForConditionalGeneration.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf", torch_dtype=torch.float16, low_cpu_mem_usage=True) 
+model.to("device")
+
 
 def postprocess_kosmos_out(result):
   token = ""
@@ -84,16 +97,42 @@ def generate_caption(image_path, caption_bool):
 
   kosmos_caption = generate_caption_kosmos(image_path, caption_bool)
   fuyu_caption = generate_caption_fuyu(image_path, caption_bool)
+  llava_caption = generate_caption_llava(image_path, caption_bool)
 
-  return kosmos_caption, fuyu_caption
+  return kosmos_caption, fuyu_caption, llava_caption
 
 
 def generate_answers(image_path, question):
 
   kosmos_answer = generate_answer_kosmos(image_path, question)
   fuyu_answer = generate_answer_fuyu(image_path, question)
+  llava_answer = generate_answer_llava(image_path, question)
+
+  return kosmos_answer, fuyu_answer, llava_answer
+
+
+@spaces.GPU
+def generate_caption_llava(image_path, caption_bool):
+    if caption_bool:
+        text_prompt =f"[INST]  \nCaption this image in detail in objective manner.[/INST]"
+    else:
+        text_prompt =f"[INST]  \nCaption this image briefly in objective manner. [/INST]"
+    
+    inputs = processor(prompt, Image.open(image_path), return_tensors="pt").to(device)
+
+    # autoregressively complete prompt
+    output = model.generate(**inputs, max_new_tokens=100)
+
+    return processor.decode(output[0], skip_special_tokens=True)["generated_text"][len(text_prompt):]
+
+@spaces.GPU
+def generate_answer_llava(image_path, question):
+    text_prompt =f"[INST]  \n{question} [/INST]"
+    inputs = processor(prompt, Image.open(image_path), return_tensors="pt").to(device)
+    output = model.generate(**inputs, max_new_tokens=100)
+    return processor.decode(output[0], skip_special_tokens=True)["generated_text"][len(text_prompt):]
+
 
-  return kosmos_answer, fuyu_answer
 
 title = "# Comparing Vision Language Models"
 
@@ -116,14 +155,14 @@ with gr.Blocks(css=css) as demo:
       with gr.Tab("Visual Question Answering"):
         with gr.Column():
             input_image = gr.Image(label = "Input Image", type="filepath")
-            question = gr.Textbox(label = "question")
+            question = gr.Textbox(label = "Question")
             run_button = gr.Button("Answer")
         with gr.Column():
             answer_kosmos = gr.Textbox(label="Answer generated by KOSMOS-2")
             answer_fuyu = gr.Textbox(label="Answer generated by Fuyu-8B")
-
+            answer_llava = gr.Textbox(label="Answer generated by LLaVA-NeXT")
         outputs_answer = [
-            answer_kosmos, answer_fuyu
+            answer_kosmos, answer_fuyu, answer_llava
         ]
 
         gr.Examples(
@@ -149,7 +188,7 @@ with gr.Blocks(css=css) as demo:
             caption_kosmos = gr.Textbox(label="Caption generated by KOSMOS-2")
             caption_fuyu = gr.Textbox(label="Caption generated by Fuyu-8B")
 
-        outputs_caption = [caption_kosmos, caption_fuyu]
+        outputs_caption = [caption_kosmos, caption_fuyu, caption_llava]
 
         gr.Examples(
             examples = [["./cat.png", True], ["./cat.png", False]],