Update gradio_utils/utils.py

gradio_utils/utils.py

@@ -106,7 +106,22 @@ def plot_results(support_img, query_img, support_kp, support_w, query_kp, query_
         plt.axis('off')  # command for hiding the axis.
         return plt
 
-@spaces.GPU
+
+
+@spaces.GPU(duration=30)
+def estimate(model, data):
+    with torch.no_grad():
+        return model(data)
+
+
+# Custom JSON encoder to handle non-serializable objects
+class CustomEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        return super().default(obj)
+        
+
 def process(query_img, state,
             cfg_path='configs/test/1shot_split1.py',
             checkpoint_path='ckpt/1shot_split1.pth'):
@@ -185,8 +200,9 @@ def process(query_img, state,
         wrap_fp16_model(model)
     load_checkpoint(model, checkpoint_path, map_location='cpu')
     model.eval().to(device)
-    with torch.no_grad():
-        outputs = model(**data)
+    
+    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)