Update gradio_utils/utils.py

gradio_utils/utils.py

@@ -111,10 +111,13 @@ def plot_results(support_img, query_img, support_kp, support_w, query_kp, query_
 @spaces.GPU(duration=30)
 def estimate(model, data):
     model.cuda()
+    data['img_s'] = [s.cuda() for s in data['img_s']]
+    data['img_q'] = data['img_q'].cuda()
+    data['target_s'] = [s.cuda() for s in data['target_s']]
+    data['target_weight_s'] = [s.cuda() for s in data['target_weight_s']]
     with torch.no_grad():
         return model(**data)
         
-@spaces.GPU(duration=30)
 def process(query_img, state,
             cfg_path='configs/test/1shot_split1.py',
             checkpoint_path='ckpt/1shot_split1.pth'):
@@ -159,10 +162,10 @@ def process(query_img, state,
         torch.tensor(target_weight_s).float()[None])
 
     data = {
-        'img_s': [support_img.to(device)],
-        'img_q': q_img.to(device),
-        'target_s': [target_s.to(device)],
-        'target_weight_s': [target_weight_s.to(device)],
+        'img_s': [support_img],
+        'img_q': q_img,
+        'target_s': [target_s],
+        'target_weight_s': [target_weight_s],
         'target_q': None,
         'target_weight_q': None,
         'return_loss': False,