a

handler.py

@@ -6,15 +6,22 @@ import numpy as np
 from PIL import Image, ImageOps
 from diffusers import StableDiffusionXLInpaintPipeline
 
+
 # ==========================================================
-# 🧠 EndpointHandler — Hugging Face Inference Entrypoint
+# 🧠 EndpointHandler — main entry for Hugging Face Endpoint
 # ==========================================================
 class EndpointHandler:
-    def __init__(self):
+    def __init__(self, model_dir: str = None):
+        """
+        Hugging Face automatically passes model_dir when starting the endpoint.
+        We don't need to use it, but we must accept it in the signature to avoid
+        a TypeError on initialization.
+        """
         print("[HANDLER] 🚀 Initializing model...")
+
         model_id = "stabilityai/stable-diffusion-xl-base-1.0"
 
-        # Load SDXL inpainting model (supports masks + expansion)
+        # Load model
         self.pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
             model_id,
             torch_dtype=torch.float16,
@@ -24,18 +31,18 @@ class EndpointHandler:
         print("[HANDLER] ✅ Model loaded successfully on cuda")
 
     # -------------------------------------------------------
-    # Main inference function called by Hugging Face Endpoint
+    # Called automatically for each request
     # -------------------------------------------------------
     def __call__(self, data: dict) -> dict:
         print("[HANDLER] 📩 Received request")
 
         try:
-            # Hugging Face sends { "inputs": {...} }
+            # Handle both { "inputs": {...} } and {...}
             inputs = data.get("inputs", data)
             if not isinstance(inputs, dict):
                 raise ValueError("Invalid input payload")
 
-            # Extract fields
+            # Extract user parameters
             b64_image = inputs.get("image")
             top = int(inputs.get("top", 0))
             bottom = int(inputs.get("bottom", 0))
@@ -47,42 +54,42 @@ class EndpointHandler:
             guidance = float(inputs.get("guidance_scale", 7.0))
             seed = int(inputs.get("seed", 42))
 
-            print(f"[HANDLER] Params → top={top}, bottom={bottom}, left={left}, right={right}, prompt='{prompt}'")
+            print(
+                f"[HANDLER] Params → top={top}, bottom={bottom}, left={left}, right={right}, prompt='{prompt}'"
+            )
 
-            # Decode image
+            # Decode base64 → PIL
             image_bytes = base64.b64decode(b64_image)
             image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
             width, height = image.size
             print(f"[HANDLER] Original image size: {width}x{height}")
 
-            # -------------------------------------------------
-            # Expand the canvas with transparent borders
-            # -------------------------------------------------
+            # Create expanded canvas
             new_w = width + left + right
             new_h = height + top + bottom
             canvas = Image.new("RGB", (new_w, new_h), (128, 128, 128))
             canvas.paste(image, (left, top))
             print(f"[HANDLER] Canvas created: {canvas.size}")
 
-            # Create a mask: white for new areas, black for existing
+            # Create mask (white = new area)
             mask = Image.new("L", (new_w, new_h), color=255)
-            mask_draw = ImageOps.expand(Image.new("L", (width, height), color=0), (left, top, right, bottom), fill=255)
+            mask_draw = ImageOps.expand(
+                Image.new("L", (width, height), color=0),
+                (left, top, right, bottom),
+                fill=255,
+            )
             mask.paste(mask_draw, (0, 0))
             print(f"[HANDLER] Mask created: {mask.size}")
 
-            # -------------------------------------------------
-            # Convert PIL to NumPy arrays (important fix)
-            # -------------------------------------------------
+            # Convert to NumPy arrays (diffusers requires .shape)
             canvas_np = np.array(canvas.convert("RGB"))
             mask_np = np.array(mask.convert("L"))
-
             print(f"[HANDLER] Converted to NumPy → canvas={canvas_np.shape}, mask={mask_np.shape}")
 
-            # -------------------------------------------------
-            # Diffusion process
-            # -------------------------------------------------
+            # Run diffusion
             print("[HANDLER] 🚀 Running diffusion process...")
             generator = torch.Generator(device="cuda").manual_seed(seed)
+
             result = self.pipe(
                 prompt=prompt,
                 negative_prompt=negative_prompt,
@@ -95,14 +102,12 @@ class EndpointHandler:
 
             print("[HANDLER] ✅ Diffusion complete")
 
-            # -------------------------------------------------
-            # Encode result
-            # -------------------------------------------------
-            buffered = io.BytesIO()
-            result.save(buffered, format="PNG")
-            img_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
-
+            # Encode output as base64
+            buffer = io.BytesIO()
+            result.save(buffer, format="PNG")
+            img_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
             print("[HANDLER] ✅ Returning base64 image")
+
             return {"image_base64": img_b64}
 
         except Exception as e: