Create handler.py

handler.py

@@ -0,0 +1,47 @@
+
+# handler.py
+import os
+import io
+import base64
+import torch
+from diffusers import DiffusionPipeline
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        # The default container mounts your repo at /repository
+        model_dir = path or "/repository"
+        # Load your SDXL pipeline in fp16, no device_map, no offloading
+        self.pipe = DiffusionPipeline.from_pretrained(
+            model_dir,
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+        ).to("cuda")
+        self.pipe.set_progress_bar_config(disable=True)
+
+    def __call__(self, data: dict):
+        # Accept either {"inputs": "..."} or {"prompt": "..."} + optional "parameters"
+        prompt = data.get("inputs") or data.get("prompt") or ""
+        params = data.get("parameters") or {}
+
+        width  = int(params.get("width", 768))
+        height = int(params.get("height", 768))
+        steps  = int(params.get("num_inference_steps", 25))
+        guidance = float(params.get("guidance_scale", 7.0))
+        negative = params.get("negative_prompt")
+        seed = params.get("seed")
+        generator = (torch.Generator(device="cuda").manual_seed(int(seed))
+                     if seed is not None else None)
+
+        image = self.pipe(
+            prompt=prompt,
+            negative_prompt=negative,
+            width=width,
+            height=height,
+            num_inference_steps=steps,
+            guidance_scale=guidance,
+            generator=generator,
+        ).images[0]
+
+        buf = io.BytesIO()
+        image.save(buf, format="PNG")
+        return {"image_base64": base64.b64encode(buf.getvalue()).decode("utf-8")}