Update app.py

app.py

@@ -103,6 +103,17 @@ def predict(radio, dict, word_mask, prompt=""):
             mask = dict["mask"].convert("RGB").resize((imgRes, imgRes))
     elif(radio == "type what to keep"):
         img = transform(dict["image"]).squeeze(0)
+        
+        #-----New Lines-----
+        if torch.cuda.is_available():
+            img.cuda()
+            print ("yes, CUDA is available here !! ")
+        
+        model = model.to(torch.device(device))
+        img = img.to(torch.device(device))
+        prompt = labels.to(torch.device(device))
+        #------------------
+        
         word_masks = [word_mask]
         with torch.no_grad(): 
             #torch.cuda.amp.autocast(): #
@@ -118,6 +129,17 @@ def predict(radio, dict, word_mask, prompt=""):
         os.remove(filename)
     else:
         img = transform(dict["image"]).unsqueeze(0)
+        
+        #-----New Lines-----
+        if torch.cuda.is_available():
+            img.cuda()
+            print ("yes, CUDA is available here !! ")
+        
+        model = model.to(torch.device(device))
+        img = img.to(torch.device(device))
+        prompt = labels.to(torch.device(device))
+        #------------------
+        
         word_masks = [word_mask]
         #with torch.cuda.amp.autocast(): #
         with torch.no_grad():