Add handler.py for inference

handler.py

@@ -0,0 +1,17 @@
+from diffusers import DiffusionPipeline
+import torch
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        print("Loading pipeline...")
+        self.pipe = DiffusionPipeline.from_pretrained(
+            path,
+            torch_dtype=torch.float16,
+            revision="fp16",
+            use_safetensors=True
+        ).to("cuda")
+
+    def __call__(self, data):
+        inputs = data.pop("inputs", data)
+        prompt = inputs if isinstance(inputs, str) else inputs[0]
+        return self.pipe(prompt).images[0]