Update app.py

app.py

@@ -2,8 +2,13 @@ import requests
 from PIL import Image
 
 import torch
+
 from transformers import AutoProcessor, LlavaForConditionalGeneration
 
+=
+
+
+
 model_id = "llava-hf/llava-1.5-7b-hf"
 model = LlavaForConditionalGeneration.from_pretrained(
     model_id, 
@@ -11,6 +16,9 @@ model = LlavaForConditionalGeneration.from_pretrained(
     low_cpu_mem_usage=True, 
 ).to(0)
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
 processor = AutoProcessor.from_pretrained(model_id)
 
 # Define a chat history and use `apply_chat_template` to get correctly formatted prompt