Update flask_app.py

flask_app.py

@@ -147,26 +147,29 @@ def img_text():
     
     sentences = data['sentences']
     print(sentences)
-    device = torch.device('cpu')
     
     image_size = 384
-    image = load_demo_image(image_size=image_size, device=device, img_url = image_url)
+    image = load_demo_image(image_size=image_size,device=device, img_url = image_url)
 
     
     model = blip_itm(pretrained=model_url, image_size=image_size, vit='base')
     model.eval()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     #model = model.to(device='cpu')
     #model = model.to('cuda')
 
     #caption = 'a cute kitten with orange color'
     #print('text: %s' %sentences)
     score = []
+
     for sentence in sentences:
-        itm_output = model(image,sentence,match_head='itm')
-        #itm_score = torch.nn.functional.softmax(itm_output,dim=1)[:,1]
-        itm_score = torch.nn.functional.softmax(itm_output, dim=1)[:, 1].item()  # Convert tensor to Python float
-        print('The image and text is matched with a probability of %.4f'%itm_score)
-        score.append(itm_score)
+        with torch.no_grad():
+            itm_output = model(image, sentence, match_head='itm')
+            itm_score = torch.nn.functional.softmax(itm_output, dim=1)[:, 1].item()
+            score.append(itm_score)
+
+        del itm_output
+        torch.cuda.empty_cache()  # Only needed if running on CUDA
     
     del model, image        
     return jsonify(scores=score)