Update app.py

app.py

@@ -64,7 +64,7 @@ class StoppingCriteriaSub(StoppingCriteria):
 
         return False
 
-device = torch.device("cpu" if torch.cuda.is_available() else "cpu")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 def get_blip_config(model="base"):
     config = dict()
@@ -129,7 +129,7 @@ tar_img_feats = torch.cat(tar_img_feats, dim=0)
 
 class Chat:
 
-    def __init__(self, model, transform, dataframe, tar_img_feats, device='cpu', stopping_criteria=None):
+    def __init__(self, model, transform, dataframe, tar_img_feats, device='cuda', stopping_criteria=None):
         self.device = device
         self.model = model
         self.transform = transform
@@ -393,7 +393,7 @@ def answer_formatter_node(question, context):
 CURR_CONTEXT = ''
 CURR_SESSION_KEY = generate_session_key()
 
-# @spaces.GPU
+@spaces.GPU
 def get_answer(image=[], message='', sessionID='abc123'):
     global CURR_CONTEXT
     global CURR_SESSION_KEY
@@ -401,7 +401,7 @@ def get_answer(image=[], message='', sessionID='abc123'):
     if image is not None:
         try:
             # Process the image and message here
-            device = torch.device("cpu" if torch.cuda.is_available() else "cpu")
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
             chat = Chat(model,transform,df,tar_img_feats, device)
             chat.encode_image(image)
             data = chat.ask()
@@ -470,7 +470,7 @@ def handle_message(data):
                     image_bytes = base64.b64decode(image_bytes)
                 image = Image.open(BytesIO(image_bytes))
                 image_array = np.array(image)
-                device = torch.device("cpu" if torch.cuda.is_available() else "cpu")
+                device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
                 chat = Chat(model, transform, df, tar_img_feats, device)
                 chat.encode_image(image_array)
                 context = chat.ask()
@@ -545,7 +545,7 @@ def handle_message(data):
 # @spaces.GPU
 def respond_to_user(image, message):
     # Process the image and message here
-    device = torch.device("cpu" if torch.cuda.is_available() else "cpu")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     chat = Chat(model,transform,df,tar_img_feats, device)
     chat.encode_image(image)
     data = chat.ask()