fix handler

handler.py

@@ -1,5 +1,5 @@
-# HANDLER v5 — SDXL txt2img + inpaint, no 'variant' usage, robust inputs
-import os, io, json, base64
+# HANDLER v6 — SDXL txt2img + inpaint, supports image_url/mask_url, guards UNet channels
+import os, io, json, base64, requests
 from typing import Any, Dict
 from PIL import Image
 import torch
@@ -10,47 +10,53 @@ MODEL_ID = os.getenv("MODEL_ID", "andro-flock/LUSTIFY-SDXL-NSFW-checkpoint-v2-0-
 
 class EndpointHandler:
     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
+        print("HANDLER v6: init start")
+        token = os.getenv("HF_TOKEN")
         local_dir = snapshot_download(MODEL_ID, token=token)
-        print(f"HANDLER v5: snapshot at {local_dir}")
+        print(f"HANDLER v6: 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 to load txt2img
                 try:
-                    self.pipe_txt2img = StableDiffusionXLPipeline.from_pretrained(
+                    p = StableDiffusionXLPipeline.from_pretrained(
                         local_dir, torch_dtype=dtype, use_safetensors=True
                     ).to(self.device)
-                    print(f"HANDLER v5: txt2img OK ({dtype})")
+                    # Keep txt2img ONLY if UNet is 4-ch (proper base)
+                    if getattr(p.unet.config, "in_channels", 4) == 4:
+                        self.pipe_txt2img = p
+                        print(f"HANDLER v6: txt2img OK ({dtype}, in_ch=4)")
+                    else:
+                        print("HANDLER v6: txt2img UNet in_ch != 4; disabling txt2img for this repo")
+                        try:
+                            p.to("cpu"); del p
+                        except Exception:
+                            pass
+                        self.pipe_txt2img = None
                 except Exception as e:
                     self.pipe_txt2img = None
-                    print(f"HANDLER v5: txt2img failed on {dtype}: {e}")
+                    print(f"HANDLER v6: 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
+                print(f"HANDLER v6: inpaint OK ({dtype}, in_ch={getattr(self.pipe_inpaint.unet.config, 'in_channels', 'NA')})")
+                break
             except Exception as e:
                 last_err = e
                 self.pipe_txt2img = None
                 self.pipe_inpaint = None
-                print(f"HANDLER v5: inpaint failed on {dtype}: {e}")
+                print(f"HANDLER v6: 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:
@@ -58,11 +64,10 @@ class EndpointHandler:
         except Exception:
             pass
 
-        print("HANDLER v5: ready")
+        print("HANDLER v6: 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):
@@ -74,12 +79,20 @@ class EndpointHandler:
                 return inner
         return data
 
+    def _fetch_url_bytes(self, url: str) -> bytes:
+        r = requests.get(url, timeout=60)
+        r.raise_for_status()
+        return r.content
+
     def _to_pil(self, payload: Any, mode: str) -> Image.Image:
-        # Accept pure base64 bytes OR data URLs
+        # Accept: bytes, base64, or data URL, or HTTP(S) URL
         if isinstance(payload, str):
-            if payload.startswith("data:"):
-                payload = payload.split(",", 1)[1]
-            payload = base64.b64decode(payload)
+            if payload.startswith("http://") or payload.startswith("https://"):
+                payload = self._fetch_url_bytes(payload)
+            else:
+                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):
@@ -111,12 +124,23 @@ class EndpointHandler:
             except Exception:
                 generator = None
 
-        # --------- decide mode ---------
-        # (1) txt2img path: no init image provided
-        if "image" not in data and "init_image" not in data:
+        # Normalize keys for images/masks
+        # Accept: image / init_image / image_url ; mask / mask_url
+        init_img_payload = None
+        if "image" in data:
+            init_img_payload = data["image"]
+        elif "init_image" in data:
+            init_img_payload = data["init_image"]
+        elif "image_url" in data:
+            init_img_payload = data["image_url"]
+
+        mask_payload = data.get("mask") or data.get("mask_url")
+
+        # --------- choose mode ---------
+        if init_img_payload is None:
+            # txt2img mode (only if we truly have a 4-ch UNet)
             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)
 
@@ -131,9 +155,9 @@ class EndpointHandler:
                     generator=generator,
                 ).images[0]
             else:
-                # Fallback: synthesize from blank canvas with inpaint
+                # Fallback: blank-canvas inpaint (works with 9-ch UNet)
                 canvas = Image.new("RGB", (width, height), (255, 255, 255))
-                mask = Image.new("L", (width, height), 255)  # edit-all
+                mask = Image.new("L", (width, height), 255)
                 image = self.pipe_inpaint(
                     prompt=prompt,
                     image=canvas,
@@ -143,17 +167,14 @@ class EndpointHandler:
                     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")
+            # inpaint mode
+            init_img = self._to_pil(init_img_payload, "RGB")
 
-            if "mask" in data:
-                mask_img = self._to_pil(data["mask"], "L")
+            if mask_payload is not None:
+                mask_img = self._to_pil(mask_payload, "L").resize(init_img.size, Image.NEAREST)
             else:
-                # default to "edit-all" if mask omitted
-                mask_img = Image.new("L", init_img.size, 255)
+                mask_img = Image.new("L", init_img.size, 255)  # edit-all default
 
             strength = self._float(data, "strength", 0.85)