Update gradio_utils/utils.py

gradio_utils/utils.py

@@ -110,6 +110,7 @@ def plot_results(support_img, query_img, support_kp, support_w, query_kp, query_
 
 @spaces.GPU(duration=30)
 def estimate(model, data):
+    model.cuda()
     with torch.no_grad():
         return model(data)
 
@@ -166,17 +167,17 @@ 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,
         'img_metas': [{'sample_skeleton': [state['skeleton']],
                        'query_skeleton': state['skeleton'],
-                       'sample_joints_3d': [kp_src_3d.to(device)],
-                       'query_joints_3d': kp_src_3d.to(device),
+                       'sample_joints_3d': [kp_src_3d],
+                       'query_joints_3d': kp_src_3d,
                        'sample_center': [kp_src_tensor.mean(dim=0)],
                        'query_center': kp_src_tensor.mean(dim=0),
                        'sample_scale': [
@@ -199,10 +200,24 @@ def process(query_img, state,
     if fp16_cfg is not None:
         wrap_fp16_model(model)
     load_checkpoint(model, checkpoint_path, map_location='cpu')
-    model.eval().to(device)
+	
+	data["img_s"] = data["img_s"][0].cpu().numpy().tolist()
+	data["img_q"] = data["img_q"].cpu().numpy().tolist()
+    data['target_weight_s'][0] = data['target_weight_s'][0].cpu().numpy().tolist()
+    data['target_s'][0] = data['target_s'][0].cpu().numpy().tolist()
+    
+    data['img_metas'][0]['sample_joints_3d'][0] = data['img_metas'][0]['sample_joints_3d'][0].cpu().tolist()
+    data['img_metas'][0]['query_joints_3d'] = data['img_metas'][0]['query_joints_3d'].cpu().tolist()
+    data['img_metas'][0]['sample_center'][0] = data['img_metas'][0]['sample_center'][0].cpu().tolist()
+    data['img_metas'][0]['query_center'] = data['img_metas'][0]['query_center'].cpu().tolist()
+    data['img_metas'][0]['sample_scale'][0] = data['img_metas'][0]['sample_scale'][0].cpu().tolist()
+    data['img_metas'][0]['query_scale'] = data['img_metas'][0]['query_scale'].cpu().tolist()
+	
+    model.eval()
     
     str_data = json.dumps(data, cls=CustomEncoder)
     outputs = estimate(model, str_data)
+	
     # visualize results
     vis_s_weight = target_weight_s[0]
     vis_s_image = support_img[0].detach().cpu().numpy().transpose(1, 2, 0)