hf hub added

app.py

@@ -18,14 +18,15 @@ cfg.MODEL.CKPT = model_path
 net, _ = build_model(cfg)
 net = net.eval()
 vgg16 = models.vgg16(pretrained=True).features.eval()
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 
 
-def load_checkpoints_from_ckpt(ckpt_path):
-    checkpoints = torch.load(ckpt_path, map_location=torch.device('cuda'))
+def load_checkpoints_from_ckpt(ckpt_path, device):
+    checkpoints = torch.load(ckpt_path, map_location=device)
     net.load_state_dict(checkpoints["ifr"])
 
 
-load_checkpoints_from_ckpt(cfg.MODEL.CKPT)
+load_checkpoints_from_ckpt(cfg.MODEL.CKPT, device)
 
 
 def filter_removal(img):