Update app.py

app.py

@@ -21,14 +21,22 @@ def load_image(img_path):
     
     return d_img
 
+# import time
 def predict(image):
     global model
     trans = torchvision.transforms.Compose([Normalize(0.5, 0.5), ToTensor()])
 
     """Run a single prediction on the model"""
     img = load_image(image)
-    img_tensor = trans(img).unsqueeze(0).cuda()
-    iq = model(img_tensor).cpu().detach().numpy().tolist()[0]
+    # t = time.time()
+    if is_gpu:
+        img_tensor = trans(img).unsqueeze(0).cuda()
+        iq = model(img_tensor).cpu().detach().numpy().tolist()[0]
+        # print('GPU ', time.time() - t)
+    else:
+        img_tensor = trans(img).unsqueeze(0)
+        iq = model(img_tensor).detach().numpy().tolist()[0]
+        # print('CPU Time: ', time.time() - t)
 
     return "The image quality of the image is: {}".format(round(iq, 4))
 
@@ -40,8 +48,13 @@ parser.add_argument('--backbone', dest='backbone', type=str, default='vit_base_p
 parser.add_argument('--mal_num', dest='mal_num', type=int, default=3, help='The number of the MAL modules.')
 config = parser.parse_args()
 
-model = MoNet.MoNet(config).cuda()
-model.load_state_dict(torch.load('best_model.pkl'))
+is_gpu = torch.cuda.is_available()
+if is_gpu:
+    model = MoNet.MoNet(config, is_gpu=is_gpu).cuda()
+    model.load_state_dict(torch.load('best_model.pkl'))
+else:
+    model = MoNet.MoNet(config, is_gpu=is_gpu)
+    model.load_state_dict(torch.load('best_model.pkl', map_location="cpu"))
 model.eval()
 
 interface = gr.Interface(fn=predict, inputs="image", outputs="text")