Update app.py

app.py

@@ -22,17 +22,7 @@ def load_image(img_path):
     return d_img
 
 def predict(image):
-    parser = argparse.ArgumentParser()
-    # model related
-    parser.add_argument('--backbone', dest='backbone', type=str, default='vit_base_patch8_224',
-                        help='The backbone for MoNet.')
-    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'))
-    model.eval()
-
+    global model
     trans = torchvision.transforms.Compose([Normalize(0.5, 0.5), ToTensor()])
 
     """Run a single prediction on the model"""
@@ -42,7 +32,17 @@ def predict(image):
 
     return "The image quality of the image is: {}".format(round(iq, 4))
 
-os.system("wget -O best_model.pkl https://huggingface.co/Zevin2023/MoC-IQA/resolve/main/Koniq10K_570908.pkl")
+# os.system("wget https://huggingface.co/Zevin2023/MoC-IQA/resolve/main/Koniq10K_570908.pkl")
+
+parser = argparse.ArgumentParser()
+# model related
+parser.add_argument('--backbone', dest='backbone', type=str, default='vit_base_patch8_224', help='The backbone for MoNet.')
+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'))
+model.eval()
 
 interface = gr.Interface(fn=predict, inputs="image", outputs="text")
 interface.launch()