fix handler

handler.py

@@ -1,33 +1,175 @@
-# HANDLER_MARKER_2025-08-29
-import io, base64
+# HANDLER v5 — SDXL txt2img + inpaint, no 'variant' usage, robust inputs
+import os, io, json, base64
+from typing import Any, Dict
 from PIL import Image
 import torch
-from diffusers import StableDiffusionXLInpaintPipeline
+from diffusers import StableDiffusionXLPipeline, StableDiffusionXLInpaintPipeline
+from huggingface_hub import snapshot_download
+
+MODEL_ID = os.getenv("MODEL_ID", "andro-flock/LUSTIFY-SDXL-NSFW-checkpoint-v2-0-INPAINTING")
 
 class EndpointHandler:
-    def __init__(self, path="."):
-        print("HANDLER_MARKER_2025-08-29: loading WITHOUT variant arg")
-        model_id = "andro-flock/LUSTIFY-SDXL-NSFW-checkpoint-v2-0-INPAINTING"
-        # NO variant= here
-        self.pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
-            model_id, torch_dtype=torch.float16, use_safetensors=True
-        ).to("cuda")
-        self.pipe.enable_attention_slicing()
-
-    def _to_pil(self, data, mode):
-        if isinstance(data, str):
-            import base64, io
-            data = base64.b64decode(data)
-        return Image.open(io.BytesIO(data)).convert(mode)
-
-    def __call__(self, data):
-        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))
-        out = self.pipe(prompt=prompt, image=init_img, mask_image=mask_img,
-                        num_inference_steps=steps, guidance_scale=guidance, strength=strength).images[0]
-        buf = io.BytesIO(); out.save(buf, format="PNG")
-        return {"image_base64": base64.b64encode(buf.getvalue()).decode()}
+    def __init__(self, path: str = "."):
+        print("HANDLER v5: init start")
+        token = os.getenv("HF_TOKEN")  # optional, for gated/private repos
+        # Download repo locally first to avoid variant resolution issues
+        local_dir = snapshot_download(MODEL_ID, token=token)
+        print(f"HANDLER v5: snapshot at {local_dir}")
+
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.pipe_txt2img = None
+        self.pipe_inpaint = None
+        last_err = None
+
+        # Try fp16 → bf16 → fp32
+        for dtype in (torch.float16, torch.bfloat16, torch.float32):
+            try:
+                # Try to load txt2img (some inpaint repos still work for txt2img)
+                try:
+                    self.pipe_txt2img = StableDiffusionXLPipeline.from_pretrained(
+                        local_dir, torch_dtype=dtype, use_safetensors=True
+                    ).to(self.device)
+                    print(f"HANDLER v5: txt2img OK ({dtype})")
+                except Exception as e:
+                    self.pipe_txt2img = None
+                    print(f"HANDLER v5: txt2img failed on {dtype}: {e}")
+
+                # Load inpaint (required)
+                self.pipe_inpaint = StableDiffusionXLInpaintPipeline.from_pretrained(
+                    local_dir, torch_dtype=dtype, use_safetensors=True
+                ).to(self.device)
+                print(f"HANDLER v5: inpaint OK ({dtype})")
+
+                break  # success on this dtype
+            except Exception as e:
+                last_err = e
+                self.pipe_txt2img = None
+                self.pipe_inpaint = None
+                print(f"HANDLER v5: inpaint failed on {dtype}: {e}")
+
+        if self.pipe_inpaint is None:
+            raise RuntimeError(f"Failed to load pipelines: {last_err}")
+
+        # Light memory tweaks
+        try:
+            self.pipe_inpaint.enable_attention_slicing()
+            if self.pipe_txt2img:
+                self.pipe_txt2img.enable_attention_slicing()
+        except Exception:
+            pass
+
+        print("HANDLER v5: ready")
+
+    # ---------- helpers ----------
+    def _unwrap(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        # Accept {"inputs": {...}} or raw dict
+        if "inputs" in data:
+            inner = data["inputs"]
+            if isinstance(inner, str):
+                try:
+                    return json.loads(inner)
+                except Exception:
+                    return {}
+            if isinstance(inner, dict):
+                return inner
+        return data
+
+    def _to_pil(self, payload: Any, mode: str) -> Image.Image:
+        # Accept pure base64 bytes OR data URLs
+        if isinstance(payload, str):
+            if payload.startswith("data:"):
+                payload = payload.split(",", 1)[1]
+            payload = base64.b64decode(payload)
+        return Image.open(io.BytesIO(payload)).convert(mode)
+
+    def _int(self, data, key, default):
+        try:
+            return int(data.get(key, default))
+        except Exception:
+            return default
+
+    def _float(self, data, key, default):
+        try:
+            return float(data.get(key, default))
+        except Exception:
+            return default
+
+    # ---------- main entry ----------
+    def __call__(self, data: Dict[str, Any]):
+        data = self._unwrap(data)
+
+        prompt = data.get("prompt", "")
+        negative_prompt = data.get("negative_prompt", None)
+        steps = self._int(data, "num_inference_steps", 30)
+        guidance = self._float(data, "guidance_scale", 7.0)
+        seed = data.get("seed", None)
+
+        generator = None
+        if seed is not None:
+            try:
+                generator = torch.Generator(device=self.device).manual_seed(int(seed))
+            except Exception:
+                generator = None
+
+        # --------- decide mode ---------
+        # (1) txt2img path: no init image provided
+        if "image" not in data and "init_image" not in data:
+            width = self._int(data, "width", 1024)
+            height = self._int(data, "height", 1024)
+            # SDXL likes multiples of 8
+            width = max(64, (width // 8) * 8)
+            height = max(64, (height // 8) * 8)
+
+            if self.pipe_txt2img is not None:
+                image = self.pipe_txt2img(
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    width=width,
+                    height=height,
+                    num_inference_steps=steps,
+                    guidance_scale=guidance,
+                    generator=generator,
+                ).images[0]
+            else:
+                # Fallback: synthesize from blank canvas with inpaint
+                canvas = Image.new("RGB", (width, height), (255, 255, 255))
+                mask = Image.new("L", (width, height), 255)  # edit-all
+                image = self.pipe_inpaint(
+                    prompt=prompt,
+                    image=canvas,
+                    mask_image=mask,
+                    negative_prompt=negative_prompt,
+                    num_inference_steps=steps,
+                    guidance_scale=guidance,
+                    generator=generator,
+                ).images[0]
+
+        # (2) inpaint path: init image (and optional mask)
+        else:
+            init_key = "image" if "image" in data else "init_image"
+            init_img = self._to_pil(data[init_key], "RGB")
+
+            if "mask" in data:
+                mask_img = self._to_pil(data["mask"], "L")
+            else:
+                # default to "edit-all" if mask omitted
+                mask_img = Image.new("L", init_img.size, 255)
+
+            strength = self._float(data, "strength", 0.85)
+
+            image = self.pipe_inpaint(
+                prompt=prompt,
+                image=init_img,
+                mask_image=mask_img,
+                negative_prompt=negative_prompt,
+                num_inference_steps=steps,
+                guidance_scale=guidance,
+                strength=strength,
+                generator=generator,
+            ).images[0]
+
+        # Return PNG as base64
+        buf = io.BytesIO()
+        image.save(buf, format="PNG")
+        out_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+        return {"image_base64": out_b64}