Update handler.py

handler.py

@@ -29,18 +29,25 @@ class EndpointHandler:
     def __call__(self, data: Any):
         image = data["inputs"]
         inputs = self.processor(image, return_tensors="pt")
-        with torch.no_grad():
-            outputs = self.model(**inputs)
 
-        print(subprocess.run(["nvidia-smi"]))
+        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()
         
-        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()
+        except Exception as e:
+            logger.error(str(e))
+            del inputs
+            torch.cuda.empty_cache()