Update app.py

app.py

@@ -1,19 +1,35 @@
-import torch
-from transformers import LlavaProcessor, LlavaForConditionalGeneration
-import gradio as gr
+import requests
 from PIL import Image
 
-# Load LLaVA model and processor
-model_id = "ICTNLP/llava-mini-llama-3.1-8b"
-processor = LlavaProcessor.from_pretrained(model_id)
-model = LlavaForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
+import torch
+from transformers import AutoProcessor, LlavaForConditionalGeneration
+
+model_id = "llava-hf/llava-1.5-7b-hf"
+model = LlavaForConditionalGeneration.from_pretrained(
+    model_id, 
+    torch_dtype=torch.float16, 
+    low_cpu_mem_usage=True, 
+).to(0)
+
+processor = AutoProcessor.from_pretrained(model_id)
+
+# Define a chat history and use `apply_chat_template` to get correctly formatted prompt
+# Each value in "content" has to be a list of dicts with types ("text", "image") 
+conversation = [
+    {
+
+      "role": "user",
+      "content": [
+          {"type": "text", "text": "What are these?"},
+          {"type": "image"},
+        ],
+    },
+]
+prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
 
-# Function to generate captions
-def generate_caption(image, prompt="Describe this image."):
-    inputs = processor(images=image, text=prompt, return_tensors="pt").to("cuda")
-    output = model.generate(**inputs, max_new_tokens=50)
-    return processor.batch_decode(output, skip_special_tokens=True)[0]
+image_file = "http://images.cocodataset.org/val2017/000000039769.jpg"
+raw_image = Image.open(requests.get(image_file, stream=True).raw)
+inputs = processor(images=raw_image, text=prompt, return_tensors='pt').to(0, torch.float16)
 
-# Gradio UI
-demo = gr.Interface(fn=generate_caption, inputs=[gr.Image(type="pil"), gr.Textbox()], outputs="text")
-demo.launch()
+output = model.generate(**inputs, max_new_tokens=200, do_sample=False)
+print(processor.decode(output[0][2:], skip_special_tokens=True))