Update app.py

app.py

@@ -3,9 +3,9 @@ from PIL import Image
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Set default device to CUDA for GPU acceleration
-device = 'cuda' if torch.cuda.is_available() else "cpu"
-# torch.set_default_device("cuda")
+# # Set default device to CUDA for GPU acceleration
+# device = 'cuda' if torch.cuda.is_available() else "cpu"
+torch.set_default_device("cuda")
 
 # Initialize the model and tokenizer
 model = AutoModelForCausalLM.from_pretrained("ManishThota/Sparrow",  torch_dtype=torch.float16, 
@@ -44,7 +44,7 @@ def predict_answer(image, question, max_tokens):
     text = f"A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: <image>\n{question}? ASSISTANT:"
     image = image.convert("RGB")
     
-    input_ids = tokenizer(text, return_tensors='pt').input_ids
+    input_ids = tokenizer(text, return_tensors='pt').input_ids.to("cuda:0", torch.float16)
     image_tensor = model.image_preprocess(image)
     
     #Generate the answer