update handler

handler.py

@@ -11,8 +11,9 @@ class EndpointHandler():
         # Preload all the elements you are going to need at inference.
         # pseudo:
         # self.model= load_model(path)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
-        self.model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
+        self.model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined").to(self.device)
         self.depth_pipe = pipeline("depth-estimation", model="depth-anything/Depth-Anything-V2-Small-hf")
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
@@ -23,14 +24,18 @@ class EndpointHandler():
       Return:
             A :obj:`list` | `dict`: will be serialized and returned
         """
-        if "image" not in data or "text" not in data:
+        if "inputs" not in data:
+            return [{"error": "Missing 'inputs' key"}]
+
+        inputs_data = data["inputs"]
+        if "image" not in inputs_data or "text" not in inputs_data:
             return [{"error": "Missing 'image' or 'text' key in input data"}]
 
         try:
             # Decode base64 image
-            image = self.decode_image(data["image"])
-            prompts = data["text"].split(",")
-
+            image = self.decode_image(inputs_data["image"])
+            prompts = inputs_data["text"]
+        
             # Preprocess input
             inputs = self.processor(
                 text=prompts,