fix handler

handler.py

@@ -1,3 +1,4 @@
+# handler.py  (replace your __init__)
 import io, base64
 from PIL import Image
 import torch
@@ -5,31 +6,53 @@ from diffusers import StableDiffusionXLInpaintPipeline
 
 class EndpointHandler:
     def __init__(self, path="."):
-        self.pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
-            "andro-flock/LUSTIFY-SDXL-NSFW-checkpoint-v2-0-INPAINTING",
-            torch_dtype=torch.float16,
-            variant="fp16"
-        ).to("cuda")
+        model_id = "andro-flock/LUSTIFY-SDXL-NSFW-checkpoint-v2-0-INPAINTING"
 
-    def _to_pil(self, data):
-        if isinstance(data, str):
-            data = base64.b64decode(data)
-        return Image.open(io.BytesIO(data))
+        # Try pure fp16 first (no variant); fall back if the GPU/dtype mismatches
+        try:
+            self.pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+                model_id,
+                torch_dtype=torch.float16,   # no variant arg
+                use_safetensors=True
+            ).to("cuda")
+        except Exception:
+            try:
+                self.pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+                    model_id,
+                    torch_dtype=torch.bfloat16,
+                    use_safetensors=True
+                ).to("cuda")
+            except Exception:
+                self.pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+                    model_id,
+                    torch_dtype=torch.float32,
+                    use_safetensors=True
+                ).to("cuda")
+
+    def _to_pil(self, b64_or_bytes, mode):
+        if isinstance(b64_or_bytes, str):
+            b = base64.b64decode(b64_or_bytes)
+        else:
+            b = b64_or_bytes
+        return Image.open(io.BytesIO(b)).convert(mode)
 
     def __call__(self, data):
-        prompt = data.get("prompt", "")
-        init_img = self._to_pil(data["image"]).convert("RGB")
-        mask_img = self._to_pil(data["mask"]).convert("L")
+        prompt   = data.get("prompt", "")
+        init_img = self._to_pil(data["image"], "RGB")
+        mask_img = self._to_pil(data["mask"],  "L")   # white=repaint, black=keep
+
+        steps    = int(data.get("num_inference_steps", 30))
+        guidance = float(data.get("guidance_scale", 7.0))
+        strength = float(data.get("strength", 0.85))
 
-        result = self.pipe(
+        out = self.pipe(
             prompt=prompt,
             image=init_img,
             mask_image=mask_img,
-            num_inference_steps=int(data.get("num_inference_steps", 30)),
-            guidance_scale=float(data.get("guidance_scale", 7.0)),
-            strength=float(data.get("strength", 0.85))
+            num_inference_steps=steps,
+            guidance_scale=guidance,
+            strength=strength
         ).images[0]
 
-        buf = io.BytesIO()
-        result.save(buf, format="PNG")
+        buf = io.BytesIO(); out.save(buf, format="PNG")
         return {"image_base64": base64.b64encode(buf.getvalue()).decode()}