Added handler.py to manage the model’s endpoints

handler.py

@@ -0,0 +1,39 @@
+from typing import Dict, Any
+import torch
+from torch.cuda.amp import autocast
+from diffusers import StableDiffusionPipeline
+import base64
+from io import BytesIO
+
+# Setting the device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+class EndpointHandler():
+
+    def __init__(self, path=""):
+        # Load the model
+        self.pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float32)
+        self.pipe = self.pipe.to(device)
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
+        """
+        Args:
+            data (dict): Includes the input data for inference.
+            
+        Return:
+            dict: Base64 encoded image.
+        """
+        inputs = data.get("inputs")  # Getting the inputs from the data dictionary
+
+        # Run inference pipeline
+        with autocast():
+            output = self.pipe(inputs, guidance_scale=7.5)
+            image = output['images'][0]  # Accessing the 'images' key in the output
+
+        # Encoding image as base 64
+        buffered = BytesIO()
+        image.save(buffered, format="PNG")
+        img_str = base64.b64encode(buffered.getvalue())
+
+        # Returning the base64 image as a dictionary
+        return {"image": img_str.decode()}