added error handling

handler.py

@@ -41,18 +41,28 @@ class EndpointHandler:
 
         image = data["inputs"]
         inputs = self.processor(image, return_tensors="pt")
-        with torch.no_grad():
-            outputs = self.model(**inputs)
 
-        print(subprocess.run(["nvidia-smi"]))
-        
-        output = outputs.reconstruction.data.squeeze().float().cpu().clamp_(0, 1).numpy()
-        output = np.moveaxis(output, source=0, destination=-1)
-        output = (output * 255.0).round().astype(np.uint8)
+        try:
+            with torch.no_grad():
+                outputs = self.model(**inputs)
+
+            print(subprocess.run(["nvidia-smi"]))
+
+            output = outputs.reconstruction.data.squeeze().float().cpu().clamp_(0, 1).numpy()
+            output = np.moveaxis(output, source=0, destination=-1)
+            output = (output * 255.0).round().astype(np.uint8)
+
+            img = Image.fromarray(output)
+            buffered = BytesIO()
+            img.save(buffered, format="JPEG")
+            img_str = base64.b64encode(buffered.getvalue())
+
+            return img_str.decode()
 
-        img = Image.fromarray(output)
-        buffered = BytesIO()
-        img.save(buffered, format="JPEG")
-        img_str = base64.b64encode(buffered.getvalue())
+        except Exception as e:
+            logger.error(str(e))
+            del inputs
+            gc.collect()
+            torch.cuda.empty_cache()
 
-        return img_str.decode()
+            return {"error": str(e)}